update tests and LLM latency score

livepeer · Jan 7, 2025 · ddb3a7c · ddb3a7c
1 parent 6a31695
commit ddb3a7c
Show file tree

Hide file tree

Showing 3 changed files with 9 additions and 3 deletions.
diff --git a/core/ai_test.go b/core/ai_test.go
@@ -652,7 +652,10 @@ func (a *stubAIWorker) SegmentAnything2(ctx context.Context, req worker.GenSegme
 }
 
 func (a *stubAIWorker) LLM(ctx context.Context, req worker.GenLLMJSONRequestBody) (interface{}, error) {
-	return &worker.LLMResponse{Response: "response tokens", TokensUsed: 10}, nil
+	var choices []worker.LLMChoice
+	choices = append(choices, worker.LLMChoice{Delta: worker.LLMMessage{Content: "choice1", Role: "assistant"}, Index: 0})
+	tokensUsed := worker.LLMTokenUsage{PromptTokens: 40, CompletionTokens: 10, TotalTokens: 50}
+	return &worker.LLMResponse{Choices: choices, Created: 1, Model: "llm_model", TokensUsed: tokensUsed}, nil
 }
 
 func (a *stubAIWorker) ImageToText(ctx context.Context, req worker.GenImageToTextMultipartRequestBody) (*worker.ImageToTextResponse, error) {

diff --git a/server/ai_process.go b/server/ai_process.go
@@ -1228,7 +1228,7 @@ func handleNonStreamingResponse(ctx context.Context, body io.ReadCloser, sess *A
 	}
 
 	took := time.Since(start)
-	sess.LatencyScore = CalculateLLMLatencyScore(took, res.TokensUsed)
+	sess.LatencyScore = CalculateLLMLatencyScore(took, res.TokensUsed.TotalTokens)
 
 	if monitor.Enabled {
 		var pricePerAIUnit float64

diff --git a/server/ai_worker_test.go b/server/ai_worker_test.go
@@ -605,7 +605,10 @@ func (a *stubAIWorker) LLM(ctx context.Context, req worker.GenLLMJSONRequestBody
 	if a.Err != nil {
 		return nil, a.Err
 	} else {
-		return &worker.LLMResponse{Response: "output tokens", TokensUsed: 10}, nil
+		var choices []worker.LLMChoice
+		choices = append(choices, worker.LLMChoice{Delta: worker.LLMMessage{Content: "choice1", Role: "assistant"}, Index: 0})
+		tokensUsed := worker.LLMTokenUsage{PromptTokens: 40, CompletionTokens: 10, TotalTokens: 50}
+		return &worker.LLMResponse{Choices: choices, Created: 1, Model: "llm_model", TokensUsed: tokensUsed}, nil
 	}
 }