refactor, Make logger available within the provider query

bauersimon · bauersimon · commit f6e28031b563 · 2025-04-24T08:37:16.000+02:00
diff --git a/evaluate/evaluate_test.go b/evaluate/evaluate_test.go
@@ -181,7 +181,7 @@ func TestEvaluate(t *testing.T) {
 					},
 				}
 				// Set up mocks, when test is running.
-				mockedQuery.On("Query", mock.Anything, mock.Anything, mock.Anything).Return(queryResult1, nil).Once().After(10 * time.Millisecond) // Simulate a model response delay because our internal safety measures trigger when a query is done in 0 milliseconds.
+				mockedQuery.On("Query", mock.Anything, mock.Anything, mock.Anything, mock.Anything).Return(queryResult1, nil).Once().After(10 * time.Millisecond) // Simulate a model response delay because our internal safety measures trigger when a query is done in 0 milliseconds.
 
 				queryResult2 := &provider.QueryResult{
 					Message: "",
@@ -192,7 +192,7 @@ func TestEvaluate(t *testing.T) {
 					},
 				}
 				// Set up mocks, when test is running.
-				mockedQuery.On("Query", mock.Anything, mock.Anything, mock.Anything).Return(queryResult2, nil).Once().After(10 * time.Millisecond) // Simulate a model response delay because our internal safety measures trigger when a query is done in 0 milliseconds.
+				mockedQuery.On("Query", mock.Anything, mock.Anything, mock.Anything, mock.Anything).Return(queryResult2, nil).Once().After(10 * time.Millisecond) // Simulate a model response delay because our internal safety measures trigger when a query is done in 0 milliseconds.
 			},
 			After: func(t *testing.T, logger *log.Logger, resultPath string) {
 				mockedQuery.AssertNumberOfCalls(t, "Query", 2)
@@ -296,7 +296,7 @@ func TestEvaluate(t *testing.T) {
 
 				Before: func(t *testing.T, logger *log.Logger, resultPath string) {
 					// Set up mocks, when test is running.
-					mockedQuery.On("Query", mock.Anything, mock.Anything, mock.Anything).Return(nil, ErrEmptyResponseFromModel)
+					mockedQuery.On("Query", mock.Anything, mock.Anything, mock.Anything, mock.Anything).Return(nil, ErrEmptyResponseFromModel)
 				},
 				After: func(t *testing.T, logger *log.Logger, resultPath string) {
 					mockedQuery.AssertNumberOfCalls(t, "Query", 2)
@@ -380,10 +380,10 @@ func TestEvaluate(t *testing.T) {
 						Message: "model-response",
 					}
 					// Set up mocks, when test is running.
-					mockedQuery.On("Query", mock.Anything, mock.Anything, mock.Anything).Return(nil, ErrEmptyResponseFromModel).Once()
-					mockedQuery.On("Query", mock.Anything, mock.Anything, mock.Anything).Return(queryResult, nil).Once().After(10 * time.Millisecond) // Simulate a model response delay because our internal safety measures trigger when a query is done in 0 milliseconds.
-					mockedQuery.On("Query", mock.Anything, mock.Anything, mock.Anything).Return(nil, ErrEmptyResponseFromModel).Once()
-					mockedQuery.On("Query", mock.Anything, mock.Anything, mock.Anything).Return(queryResult, nil).Once().After(10 * time.Millisecond) // Simulate a model response delay because our internal safety measures trigger when a query is done in 0 milliseconds.
+					mockedQuery.On("Query", mock.Anything, mock.Anything, mock.Anything, mock.Anything).Return(nil, ErrEmptyResponseFromModel).Once()
+					mockedQuery.On("Query", mock.Anything, mock.Anything, mock.Anything, mock.Anything).Return(queryResult, nil).Once().After(10 * time.Millisecond) // Simulate a model response delay because our internal safety measures trigger when a query is done in 0 milliseconds.
+					mockedQuery.On("Query", mock.Anything, mock.Anything, mock.Anything, mock.Anything).Return(nil, ErrEmptyResponseFromModel).Once()
+					mockedQuery.On("Query", mock.Anything, mock.Anything, mock.Anything, mock.Anything).Return(queryResult, nil).Once().After(10 * time.Millisecond) // Simulate a model response delay because our internal safety measures trigger when a query is done in 0 milliseconds.
 				},
 				After: func(t *testing.T, logger *log.Logger, resultPath string) {
 					mockedQuery.AssertNumberOfCalls(t, "Query", 4)
@@ -486,7 +486,7 @@ func TestEvaluate(t *testing.T) {
 						},
 					}
 					// Set up mocks, when test is running.
-					mockedQuery.On("Query", mock.Anything, mock.Anything, mock.Anything).Return(queryResult, nil).After(10 * time.Millisecond) // Simulate a model response delay because our internal safety measures trigger when a query is done in 0 milliseconds.
+					mockedQuery.On("Query", mock.Anything, mock.Anything, mock.Anything, mock.Anything).Return(queryResult, nil).After(10 * time.Millisecond) // Simulate a model response delay because our internal safety measures trigger when a query is done in 0 milliseconds.
 				},
 				After: func(t *testing.T, logger *log.Logger, resultPath string) {
 					mockedQuery.AssertNumberOfCalls(t, "Query", 2)
diff --git a/model/llm/llm.go b/model/llm/llm.go
@@ -345,7 +345,7 @@ func (m *Model) query(logger *log.Logger, request string) (queryResult *provider
 			}
 
 			start := time.Now()
-			queryResult, err = m.provider.Query(ctx, m, request)
+			queryResult, err = m.provider.Query(ctx, logger, m, request)
 			if err != nil {
 				return err
 			} else if ctx.Err() != nil {
diff --git a/model/llm/llm_test.go b/model/llm/llm_test.go
@@ -120,7 +120,7 @@ func TestModelGenerateTestsForFile(t *testing.T) {
 				` + "```" + `
 			`),
 			}
-			mockedProvider.On("Query", mock.Anything, mock.Anything, promptMessage).Return(queryResult, nil)
+			mockedProvider.On("Query", mock.Anything, mock.Anything, mock.Anything, promptMessage).Return(queryResult, nil)
 		},
 
 		Language:          &golang.Language{},
@@ -151,7 +151,7 @@ func TestModelGenerateTestsForFile(t *testing.T) {
 					TotalCost: 0.123456789,
 				},
 			}
-			mockedProvider.On("Query", mock.Anything, mock.Anything, promptMessage).Return(queryResult, nil)
+			mockedProvider.On("Query", mock.Anything, mock.Anything, mock.Anything, promptMessage).Return(queryResult, nil)
 		},
 
 		Language:          &golang.Language{},
@@ -242,7 +242,7 @@ func TestModelQuery(t *testing.T) {
 			queryResult := &provider.QueryResult{
 				Message: "test response",
 			}
-			mockedProvider.On("Query", mock.Anything, mock.Anything, "test request").Return(queryResult, nil)
+			mockedProvider.On("Query", mock.Anything, mock.Anything, mock.Anything, "test request").Return(queryResult, nil)
 		},
 		QueryAttempts: 1,
 		Request:       "test request",
@@ -256,7 +256,7 @@ func TestModelQuery(t *testing.T) {
 	validate(t, &testCase{
 		Name: "Failed query no retry",
 		SetupMock: func(mockedProvider *providertesting.MockQuery) {
-			mockedProvider.On("Query", mock.Anything, mock.Anything, "test request").Return(nil, assert.AnError)
+			mockedProvider.On("Query", mock.Anything, mock.Anything, mock.Anything, "test request").Return(nil, assert.AnError)
 		},
 		QueryAttempts: 1,
 		Request:       "test request",
@@ -266,8 +266,8 @@ func TestModelQuery(t *testing.T) {
 	validate(t, &testCase{
 		Name: "Failed query with retry",
 		SetupMock: func(mockedProvider *providertesting.MockQuery) {
-			mockedProvider.On("Query", mock.Anything, mock.Anything, "test request").Return(nil, assert.AnError).Once()
-			mockedProvider.On("Query", mock.Anything, mock.Anything, "test request").Return(&provider.QueryResult{
+			mockedProvider.On("Query", mock.Anything, mock.Anything, mock.Anything, "test request").Return(nil, assert.AnError).Once()
+			mockedProvider.On("Query", mock.Anything, mock.Anything, mock.Anything, "test request").Return(&provider.QueryResult{
 				Message: "test response",
 			}, nil).Once()
 		},
@@ -286,7 +286,7 @@ func TestModelQuery(t *testing.T) {
 			queryResult := &provider.QueryResult{
 				Message: "test response",
 			}
-			mockedProvider.On("Query", mock.Anything, mock.Anything, "test request").Return(queryResult, nil).After(time.Second * 2)
+			mockedProvider.On("Query", mock.Anything, mock.Anything, mock.Anything, "test request").Return(queryResult, nil).After(time.Second * 2)
 		},
 		QueryAttempts: 1,
 		QueryTimeout:  1,
@@ -300,8 +300,8 @@ func TestModelQuery(t *testing.T) {
 			queryResult := &provider.QueryResult{
 				Message: "test response",
 			}
-			mockedProvider.On("Query", mock.Anything, mock.Anything, "test request").Return(queryResult, nil).After(time.Second * 2)
-			mockedProvider.On("Query", mock.Anything, mock.Anything, "test request").Return(queryResult, nil).After(time.Second * 2)
+			mockedProvider.On("Query", mock.Anything, mock.Anything, mock.Anything, "test request").Return(queryResult, nil).After(time.Second * 2)
+			mockedProvider.On("Query", mock.Anything, mock.Anything, mock.Anything, "test request").Return(queryResult, nil).After(time.Second * 2)
 		},
 		QueryAttempts: 2,
 		QueryTimeout:  1,
@@ -392,7 +392,7 @@ func TestModelRepairSourceCodeFile(t *testing.T) {
 						` + "```" + `
 					`),
 				}
-				mockedProvider.On("Query", mock.Anything, mock.Anything, mock.Anything).Return(queryResult, nil)
+				mockedProvider.On("Query", mock.Anything, mock.Anything, mock.Anything, mock.Anything).Return(queryResult, nil)
 			},
 
 			Language:       &golang.Language{},
@@ -446,7 +446,7 @@ func TestModelRepairSourceCodeFile(t *testing.T) {
 						` + "```" + `
 					`),
 				}
-				mockedProvider.On("Query", mock.Anything, mock.Anything, mock.Anything).Return(queryResult, nil)
+				mockedProvider.On("Query", mock.Anything, mock.Anything, mock.Anything, mock.Anything).Return(queryResult, nil)
 			},
 
 			Language:       &java.Language{},
@@ -876,7 +876,7 @@ func TestModelTranspile(t *testing.T) {
 				queryResult := &provider.QueryResult{
 					Message: "```\n" + transpiledFileContent + "```\n",
 				}
-				mockedProvider.On("Query", mock.Anything, mock.Anything, mock.Anything).Return(queryResult, nil)
+				mockedProvider.On("Query", mock.Anything, mock.Anything, mock.Anything, mock.Anything).Return(queryResult, nil)
 			},
 
 			Language:       &golang.Language{},
@@ -929,7 +929,7 @@ func TestModelTranspile(t *testing.T) {
 				queryResult := &provider.QueryResult{
 					Message: "```\n" + transpiledFileContent + "```\n",
 				}
-				mockedProvider.On("Query", mock.Anything, mock.Anything, mock.Anything).Return(queryResult, nil)
+				mockedProvider.On("Query", mock.Anything, mock.Anything, mock.Anything, mock.Anything).Return(queryResult, nil)
 			},
 
 			Language:       &java.Language{},
@@ -1032,7 +1032,7 @@ func TestModelMigrate(t *testing.T) {
 			queryResult := &provider.QueryResult{
 				Message: "```\n" + migratedTestFile + "```\n",
 			}
-			mockedProvider.On("Query", mock.Anything, mock.Anything, mock.Anything).Return(queryResult, nil)
+			mockedProvider.On("Query", mock.Anything, mock.Anything, mock.Anything, mock.Anything).Return(queryResult, nil)
 		},
 
 		Language: &java.Language{},
diff --git a/provider/ollama/ollama.go b/provider/ollama/ollama.go
@@ -81,7 +81,7 @@ func (p *Provider) Models() (models []model.Model, err error) {
 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) {
+func (p *Provider) Query(ctx context.Context, logger *log.Logger, model model.Model, promptText string) (result *provider.QueryResult, err error) {
 	return openaiapi.QueryOpenAIAPIModel(ctx, p.client(), model.ModelIDWithoutProvider(), model.Attributes(), promptText)
 }
 
diff --git a/provider/openai-api/openai.go b/provider/openai-api/openai.go
@@ -59,7 +59,7 @@ func (p *Provider) SetToken(token string) {
 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) {
+func (p *Provider) Query(ctx context.Context, logger *log.Logger, model model.Model, promptText string) (result *provider.QueryResult, err error) {
 	return QueryOpenAIAPIModel(ctx, p.client(), model.ModelIDWithoutProvider(), model.Attributes(), promptText)
 }
 
diff --git a/provider/openrouter/openrouter.go b/provider/openrouter/openrouter.go
@@ -137,13 +137,13 @@ func (p *Provider) SetToken(token string) {
 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) {
+func (p *Provider) Query(ctx context.Context, logger *log.Logger, model model.Model, promptText string) (result *provider.QueryResult, err error) {
 	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)
+	queryResult.GenerationInfo, err = p.fetchGenerationInfo(logger, queryResult.ResponseID)
 	if err != nil {
 		return nil, pkgerrors.WithStack(err)
 	}
@@ -159,7 +159,7 @@ func (p *Provider) client() (client *openai.Client) {
 	return openai.NewClientWithConfig(config)
 }
 
-func (p *Provider) fetchGenerationInfo(generationID string) (generationInfo *provider.GenerationInfo, err error) {
+func (p *Provider) fetchGenerationInfo(logger *log.Logger, 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)
diff --git a/provider/provider.go b/provider/provider.go
@@ -73,7 +73,7 @@ type GenerationInfo struct {
 // Query is a provider that allows to query a model directly.
 type Query interface {
 	// Query queries the provider with the given model name.
-	Query(ctx context.Context, model model.Model, promptText string) (result *QueryResult, err error)
+	Query(ctx context.Context, logger *log.Logger, model model.Model, promptText string) (result *QueryResult, err error)
 }
 
 // Service is a provider that requires background services.
diff --git a/provider/testing/Query_mock_gen.go b/provider/testing/Query_mock_gen.go

Original file line number	Diff line number	Diff line change
`@@ -345,7 +345,7 @@ func (m Model) query(logger log.Logger, request string) (queryResult *provider`
`345`	`345`	`}`
`346`	`346`
`347`	`347`	`start := time.Now()`
`348`		`- queryResult, err = m.provider.Query(ctx, m, request)`
	`348`	`+ queryResult, err = m.provider.Query(ctx, logger, m, request)`
`349`	`349`	`if err != nil {`
`350`	`350`	`return err`
`351`	`351`	`} else if ctx.Err() != nil {`
Original file line number	Diff line number	Diff line change
`@@ -81,7 +81,7 @@ func (p *Provider) Models() (models []model.Model, err error) {`
`81`	`81`	`var _ provider.Query = (*Provider)(nil)`
`82`	`82`
`83`	`83`	`// Query queries the provider with the given model name.`
`84`		`-func (p Provider) Query(ctx context.Context, model model.Model, promptText string) (result provider.QueryResult, err error) {`
	`84`	`+func (p Provider) Query(ctx context.Context, logger log.Logger, model model.Model, promptText string) (result *provider.QueryResult, err error) {`
`85`	`85`	`return openaiapi.QueryOpenAIAPIModel(ctx, p.client(), model.ModelIDWithoutProvider(), model.Attributes(), promptText)`
`86`	`86`	`}`
`87`	`87`
Original file line number	Diff line number	Diff line change
`@@ -59,7 +59,7 @@ func (p *Provider) SetToken(token string) {`
`59`	`59`	`var _ provider.Query = (*Provider)(nil)`
`60`	`60`
`61`	`61`	`// Query queries the provider with the given model name.`
`62`		`-func (p Provider) Query(ctx context.Context, model model.Model, promptText string) (result provider.QueryResult, err error) {`
	`62`	`+func (p Provider) Query(ctx context.Context, logger log.Logger, model model.Model, promptText string) (result *provider.QueryResult, err error) {`
`63`	`63`	`return QueryOpenAIAPIModel(ctx, p.client(), model.ModelIDWithoutProvider(), model.Attributes(), promptText)`
`64`	`64`	`}`
`65`	`65`
Original file line number	Diff line number	Diff line change
`@@ -73,7 +73,7 @@ type GenerationInfo struct {`
`73`	`73`	`// Query is a provider that allows to query a model directly.`
`74`	`74`	`type Query interface {`
`75`	`75`	`// Query queries the provider with the given model name.`
`76`		`- Query(ctx context.Context, model model.Model, promptText string) (result *QueryResult, err error)`
	`76`	`+ Query(ctx context.Context, logger log.Logger, model model.Model, promptText string) (result QueryResult, err error)`
`77`	`77`	`}`
`78`	`78`
`79`	`79`	`// Service is a provider that requires background services.`