Fetch total costs from OpenRouter after query

Author: ahumenberger

evaluate/metrics/assessment.go

@@ -59,6 +59,12 @@ var (
 	AssessmentKeyTokenInput = RegisterAssessmentKey("token-input")
 	// AssessmentKeyTokenOutput collects the number of output token.
 	AssessmentKeyTokenOutput = RegisterAssessmentKey("token-output")
+	// AssessmentKeyNativeTokenInput collects the number of input token.
+	AssessmentKeyNativeTokenInput = RegisterAssessmentKey("native-token-input")
+	// AssessmentKeyNativeTokenOutput collects the number of output token.
+	AssessmentKeyNativeTokenOutput = RegisterAssessmentKey("native-token-output")
+	// AssessmentKeyCostsTokenActual collects the number of output token.
+	AssessmentKeyCostsTokenActual = RegisterAssessmentKey("costs-total-actual")
 )
 
 // Assessments holds a collection of numerical assessment metrics.

evaluate/metrics/assessment_test.go

@@ -84,7 +84,7 @@ func TestAssessmentString(t *testing.T) {
 
 		Assessment: NewAssessments(),
 
-		ExpectedString: "coverage=0, files-executed=0, files-executed-maximum-reachable=0, generate-tests-for-file-character-count=0, processing-time=0, response-character-count=0, response-no-error=0, response-no-excess=0, response-with-code=0, tests-passing=0, token-input=0, token-output=0",
+		ExpectedString: "costs-total-actual=0, coverage=0, files-executed=0, files-executed-maximum-reachable=0, generate-tests-for-file-character-count=0, native-token-input=0, native-token-output=0, processing-time=0, response-character-count=0, response-no-error=0, response-no-excess=0, response-with-code=0, tests-passing=0, token-input=0, token-output=0",
 	})
 
 	validate(t, &testCase{
@@ -105,7 +105,7 @@ func TestAssessmentString(t *testing.T) {
 			AssessmentKeyTokenOutput:                        456,
 		},
 
-		ExpectedString: "coverage=1, files-executed=2, files-executed-maximum-reachable=2, generate-tests-for-file-character-count=50, processing-time=200, response-character-count=100, response-no-error=3, response-no-excess=4, response-with-code=5, tests-passing=7, token-input=123, token-output=456",
+		ExpectedString: "costs-total-actual=0, coverage=1, files-executed=2, files-executed-maximum-reachable=2, generate-tests-for-file-character-count=50, native-token-input=0, native-token-output=0, processing-time=200, response-character-count=100, response-no-error=3, response-no-excess=4, response-with-code=5, tests-passing=7, token-input=123, token-output=456",
 	})
 }
 

evaluate/report/csv_test.go

@@ -24,7 +24,7 @@ func TestNewEvaluationFile(t *testing.T) {
 	require.NoError(t, err)
 
 	expectedEvaluationFileContent := bytesutil.StringTrimIndentations(`
-		model-id,language,repository,case,task,run,coverage,files-executed,files-executed-maximum-reachable,generate-tests-for-file-character-count,processing-time,response-character-count,response-no-error,response-no-excess,response-with-code,tests-passing,token-input,token-output
+		model-id,language,repository,case,task,run,costs-total-actual,coverage,files-executed,files-executed-maximum-reachable,generate-tests-for-file-character-count,native-token-input,native-token-output,processing-time,response-character-count,response-no-error,response-no-excess,response-with-code,tests-passing,token-input,token-output
 	`)
 
 	assert.Equal(t, expectedEvaluationFileContent, string(actualEvaluationFileContent))
@@ -65,8 +65,8 @@ func TestWriteEvaluationRecord(t *testing.T) {
 		},
 
 		ExpectedCSV: `
-			model-id,language,repository,case,task,run,coverage,files-executed,files-executed-maximum-reachable,generate-tests-for-file-character-count,processing-time,response-character-count,response-no-error,response-no-excess,response-with-code,tests-passing,token-input,token-output
-			mocked-model,golang,golang/plain,plain.go,write-tests,1,0,0,0,0,0,0,0,0,0,0,0,0
+			model-id,language,repository,case,task,run,costs-total-actual,coverage,files-executed,files-executed-maximum-reachable,generate-tests-for-file-character-count,native-token-input,native-token-output,processing-time,response-character-count,response-no-error,response-no-excess,response-with-code,tests-passing,token-input,token-output
+			mocked-model,golang,golang/plain,plain.go,write-tests,1,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0
 		`,
 	})
 	validate(t, &testCase{
@@ -90,9 +90,9 @@ func TestWriteEvaluationRecord(t *testing.T) {
 		},
 
 		ExpectedCSV: `
-			model-id,language,repository,case,task,run,coverage,files-executed,files-executed-maximum-reachable,generate-tests-for-file-character-count,processing-time,response-character-count,response-no-error,response-no-excess,response-with-code,tests-passing,token-input,token-output
-			mocked-model,golang,golang/plain,plain.go,write-tests,1,0,1,1,0,0,0,1,0,0,0,0,0
-			mocked-model,golang,golang/plain,plain.go,write-tests-symflower-fix,1,10,1,1,0,0,0,1,0,0,0,0,0
+			model-id,language,repository,case,task,run,costs-total-actual,coverage,files-executed,files-executed-maximum-reachable,generate-tests-for-file-character-count,native-token-input,native-token-output,processing-time,response-character-count,response-no-error,response-no-excess,response-with-code,tests-passing,token-input,token-output
+			mocked-model,golang,golang/plain,plain.go,write-tests,1,0,0,1,1,0,0,0,0,0,1,0,0,0,0,0
+			mocked-model,golang,golang/plain,plain.go,write-tests-symflower-fix,1,0,10,1,1,0,0,0,0,0,1,0,0,0,0,0
 		`,
 	})
 }

model/llm/llm.go

@@ -339,7 +339,11 @@ func (m *Model) query(logger *log.Logger, request string) (queryResult *provider
 				return err
 			}
 			duration = time.Since(start)
-			logger.Info("model responded", "model", m.ID(), "id", id, "duration", duration.Milliseconds(), "response-id", queryResult.ResponseID, "token-input", queryResult.Usage.PromptTokens, "token-output", queryResult.Usage.CompletionTokens, "response", string(bytesutil.PrefixLines([]byte(queryResult.Message), []byte("\t"))))
+			totalCosts := float64(-1)
+			if queryResult.GenerationInfo != nil {
+				totalCosts = queryResult.GenerationInfo.TotalCost
+			}
+			logger.Info("model responded", "model", m.ID(), "id", id, "duration", duration.Milliseconds(), "response-id", queryResult.ResponseID, "costs-total", totalCosts, "token-input", queryResult.Usage.PromptTokens, "token-output", queryResult.Usage.CompletionTokens, "response", string(bytesutil.PrefixLines([]byte(queryResult.Message), []byte("\t"))))
 
 			return nil
 		},
@@ -496,6 +500,11 @@ func handleQueryResult(queryResult *provider.QueryResult, filePathAbsolute strin
 	assessment[metrics.AssessmentKeyGenerateTestsForFileCharacterCount] = float64(len(sourceFileContent))
 	assessment[metrics.AssessmentKeyTokenInput] = float64(queryResult.Usage.PromptTokens)
 	assessment[metrics.AssessmentKeyTokenOutput] = float64(queryResult.Usage.CompletionTokens)
+	if queryResult.GenerationInfo != nil {
+		assessment[metrics.AssessmentKeyNativeTokenInput] = float64(queryResult.GenerationInfo.NativeTokensPrompt)
+		assessment[metrics.AssessmentKeyNativeTokenOutput] = float64(queryResult.GenerationInfo.NativeTokensCompletion)
+		assessment[metrics.AssessmentKeyCostsTokenActual] = queryResult.GenerationInfo.TotalCost
+	}
 
 	if err := os.MkdirAll(filepath.Dir(filePathAbsolute), 0755); err != nil {
 		return nil, pkgerrors.WithStack(err)

provider/openrouter/openrouter.go

@@ -138,7 +138,17 @@ var _ provider.Query = (*Provider)(nil)
 
 // Query queries the provider with the given model name.
 func (p *Provider) Query(ctx context.Context, model model.Model, promptText string) (result *provider.QueryResult, err error) {
-	return openaiapi.QueryOpenAIAPIModel(ctx, p.client(), model.ModelIDWithoutProvider(), model.Attributes(), promptText)
+	queryResult, err := openaiapi.QueryOpenAIAPIModel(ctx, p.client(), model.ModelIDWithoutProvider(), model.Attributes(), promptText)
+	if err != nil {
+		return nil, pkgerrors.WithStack(err)
+	}
+
+	queryResult.GenerationInfo, err = p.fetchGenerationInfo(queryResult.ResponseID)
+	if err != nil {
+		return nil, pkgerrors.WithStack(err)
+	}
+
+	return queryResult, nil
 }
 
 // client returns a new client with the current configuration.
@@ -148,3 +158,54 @@ func (p *Provider) client() (client *openai.Client) {
 
 	return openai.NewClientWithConfig(config)
 }
+
+func (p *Provider) fetchGenerationInfo(generationID string) (generationInfo *provider.GenerationInfo, err error) {
+	request, err := http.NewRequest("GET", "https://openrouter.ai/api/v1/generation?id="+generationID, nil)
+	if err != nil {
+		return nil, pkgerrors.WithStack(err)
+	}
+	request.Header.Set("Accept", "application/json")
+	request.Header.Set("Authorization", "Bearer "+p.token)
+
+	client := &http.Client{}
+	var responseBody []byte
+	if err := retry.Do( // Query available models with a retry logic cause "openrouter.ai" has failed us in the past.
+		func() error {
+			response, err := client.Do(request)
+			if err != nil {
+				return pkgerrors.WithStack(err)
+			}
+			defer func() {
+				if e := response.Body.Close(); e != nil {
+					err = errors.Join(err, pkgerrors.WithStack(e))
+				}
+			}()
+
+			if response.StatusCode != http.StatusOK {
+				return pkgerrors.Errorf("received status code %d when querying provider models", response.StatusCode)
+			}
+
+			responseBody, err = io.ReadAll(response.Body)
+			if err != nil {
+				return pkgerrors.WithStack(err)
+			}
+
+			return nil
+		},
+		retry.Attempts(3),
+		retry.Delay(5*time.Second),
+		retry.DelayType(retry.BackOffDelay),
+		retry.LastErrorOnly(true),
+	); err != nil {
+		return nil, err
+	}
+
+	var dataResponse struct {
+		provider.GenerationInfo `json:"data"`
+	}
+	if err := json.Unmarshal(responseBody, &dataResponse); err != nil {
+		return nil, err
+	}
+
+	return &dataResponse.GenerationInfo, nil
+}

provider/provider.go

@@ -54,6 +54,20 @@ type QueryResult struct {
 	Duration time.Duration
 	// Usage holds the usage metrics of the query.
 	Usage openai.Usage
+	// GenerationInfo holds information about a generation.
+	GenerationInfo *GenerationInfo
+}
+
+// GenerationInfo holds information about a generation.
+// See https://openrouter.ai/docs/api-reference/overview#querying-cost-and-stats for more details.
+type GenerationInfo struct {
+	ID                     string  `json:"id"`
+	TotalCost              float64 `json:"total_cost"`
+	TokensPrompt           int     `json:"tokens_prompt"`
+	TokensCompletion       int     `json:"tokens_completion"`
+	NativeTokensPrompt     int     `json:"native_tokens_prompt"`
+	NativeTokensCompletion int     `json:"native_tokens_completion"`
+	NativeTokensReasoning  int     `json:"native_tokens_reasoning"`
 }
 
 // Query is a provider that allows to query a model directly.