feat: issue299 add timer for OpenAi service call

fidelity · Mar 4, 2024 · d88512e · d88512e
1 parent 1e9c5ab
commit d88512e
Show file tree

Hide file tree

Showing 2 changed files with 28 additions and 21 deletions.
diff --git a/internal/investigators/common.go b/internal/investigators/common.go
@@ -157,10 +157,33 @@ func appendSolution(problem *problem.Problem, solutions interface{}, commands in
 	}
 }
 
-func getAiSuggestion(issue string) string {
-	return fmt.Sprintf(ai.DefaultPrompt, issue)
+func callAiKnowledge(ctx context.Context, wg *sync.WaitGroup, problem *problem.Problem,
+	input *problem.DetectorCreationInput, prompt string) {
+	defer wg.Done()
+	knowledge, err := input.AiClient.GetCompletion(ctx, fmt.Sprintf(ai.KnowledgePrompt, prompt))
+	if err != nil {
+		return
+	}
+	problem.AiKnowledge.Append(strings.Split(knowledge, "\n\n")...)
+}
+
+func callAiSuggestion(ctx context.Context, wg *sync.WaitGroup, problem *problem.Problem,
+	input *problem.DetectorCreationInput, prompt string) {
+	defer wg.Done()
+	suggestions, err := input.AiClient.GetCompletion(ctx, fmt.Sprintf(ai.DefaultPrompt, prompt))
+	if err != nil {
+		return
+	}
+	problem.AiSuggestions.Append(strings.Split(suggestions, "\n")...)
 }
 
-func getAiKnowledge(issue string) string {
-	return fmt.Sprintf(ai.KnowledgePrompt, issue)
+func getAiResults(ctx context.Context, problem *problem.Problem,
+	input *problem.DetectorCreationInput, issue string, knowledge string) {
+	var wg sync.WaitGroup
+	t1 := time.Now()
+	wg.Add(2)
+	go callAiSuggestion(ctx, &wg, problem, input, issue)
+	go callAiKnowledge(ctx, &wg, problem, input, knowledge)
+	wg.Wait()
+	log.SWithContext(ctx).Infof("Calling OpenAi for %v.", time.Since(t1))
 }
diff --git a/internal/investigators/pod_pending_investigator.go b/internal/investigators/pod_pending_investigator.go
@@ -122,11 +122,7 @@ func PodNotRunningInvestigator(ctx context.Context, wg *sync.WaitGroup, problem
 		} else {
 			solutions, commands = getPendingPodUnknownSolution(ctx, pod)
 		}
-		var wg2 sync.WaitGroup
-		wg2.Add(2)
-		callAiSuggestion(ctx, &wg2, problem, input, "PodFailedScheduled, "+failSchedule)
-		callAiKnowledge(ctx, &wg2, problem, input, "PodFailedScheduled")
-		wg2.Wait()
+		getAiResults(ctx, problem, input, "PodFailedScheduled, "+failSchedule, "PodFailedScheduled")
 	} else if len(failMount) > 0 {
 		commands = appendSeq(commands, GetSolutionsByTemplate(ctx, KubeDescribePoCmd, pod, true)[0])
 		commands = appendSeq(commands, GetSolutionsByTemplate(ctx, GetEventsCmd, pod, true)[0])
@@ -153,18 +149,6 @@ func PodNotRunningInvestigator(ctx context.Context, wg *sync.WaitGroup, problem
 	appendSolution(problem, solutions, commands)
 }
 
-func callAiKnowledge(ctx context.Context, wg *sync.WaitGroup, problem *problem.Problem, input *problem.DetectorCreationInput, prompt string) {
-	defer wg.Done()
-	knowledge, _ := input.AiClient.GetCompletion(ctx, getAiKnowledge(prompt))
-	problem.AiKnowledge.Append(strings.Split(knowledge, "\n")...)
-}
-
-func callAiSuggestion(ctx context.Context, wg *sync.WaitGroup, problem *problem.Problem, input *problem.DetectorCreationInput, prompt string) {
-	defer wg.Done()
-	suggestions, _ := input.AiClient.GetCompletion(ctx, getAiSuggestion(prompt))
-	problem.AiSuggestions.Append(strings.Split(suggestions, "\n")...)
-}
-
 func appendPVSolution(ctx context.Context, po v1.Pod, solutions []string, solution string) ([]string, []string) {
 	addSolutions := GetSolutionsByTemplate(ctx, solution, po, true)
 	solutions = append(solutions, addSolutions...)