Query timeout for LLM queries to be able to abort long queries automatically

Closes #443

Author: bauersimon

cmd/eval-dev-quality/cmd/evaluate.go

@@ -60,7 +60,9 @@ type Evaluate struct {
 	// ProviderUrls holds all custom inference endpoint urls for the providers.
 	ProviderUrls map[string]string `long:"urls" description:"Custom OpenAI API compatible inference endpoints (of the form '$provider:$url,...'). Use '$provider=custom-$name' to manually register a custom OpenAI API endpoint provider. Note that the models of a custom OpenAI API endpoint provider must be declared explicitly using the '--model' option. When using the environment variable, separate multiple definitions with ','." env:"PROVIDER_URL" env-delim:","`
 	// QueryAttempts holds the number of query attempts to perform when a model request errors in the process of solving a task.
-	QueryAttempts uint `long:"attempts" description:"Number of query attempts to perform when a model request errors in the process of solving a task." default:"3"`
+	QueryAttempts uint `long:"query-attempts" description:"Number of query attempts to perform when a model request errors in the process of solving a task." default:"3"`
+	// QueryTimeout holds the timeout for model requests.
+	QueryTimeout uint `long:"query-timeout" description:"Timeout of a model query in seconds. ("0" to disable)" default:"1200"`
 
 	// Repositories determines which repository should be used for the evaluation, or empty if all repositories should be used.
 	Repositories []string `long:"repository" description:"Evaluate with this repository. By default all repositories are used."`
@@ -170,6 +172,7 @@ func (command *Evaluate) Initialize(args []string) (evaluationContext *evaluate.
 			command.logger.Panicf("number of configured query attempts must be greater than zero")
 		}
 		evaluationContext.QueryAttempts = command.QueryAttempts
+		evaluationContext.QueryTimeout = command.QueryTimeout
 
 		if command.ExecutionTimeout == 0 {
 			command.logger.Panicf("execution timeout for compilation and tests must be greater than zero")

evaluate/evaluate.go

@@ -30,6 +30,8 @@ type Context struct {
 	ProviderForModel map[evalmodel.Model]provider.Provider
 	// QueryAttempts holds the number of query attempts to perform when a model request errors in the process of solving a task.
 	QueryAttempts uint
+	// QueryTimeout holds the timeout for model queries in seconds.
+	QueryTimeout uint
 
 	// RepositoryPaths determines which relative repository paths should be used for the evaluation, or empty if all repositories should be used.
 	RepositoryPaths []string
@@ -130,8 +132,9 @@ func Evaluate(ctx *Context) {
 						modelSucceededBasicChecksOfLanguage[model] = map[evallanguage.Language]bool{}
 					}
 
-					if r, ok := model.(evalmodel.SetQueryAttempts); ok {
+					if r, ok := model.(evalmodel.SetQueryHandling); ok {
 						r.SetQueryAttempts(ctx.QueryAttempts)
+						r.SetQueryTimeout(ctx.QueryTimeout)
 					}
 
 					for _, taskIdentifier := range temporaryRepository.Configuration().Tasks {

model/llm/llm.go

@@ -35,6 +35,8 @@ type Model struct {
 	attributes map[string]string
 	// queryAttempts holds the number of query attempts to perform when a model request errors in the process of solving a task.
 	queryAttempts uint
+	// queryTimeout holds the timeout for model requests in seconds.
+	queryTimeout uint
 
 	// metaInformation holds a model meta information.
 	metaInformation *model.MetaInformation
@@ -47,6 +49,7 @@ func NewModel(provider provider.Query, modelIDWithAttributes string) (llmModel *
 		provider: provider,
 
 		queryAttempts: 1,
+		queryTimeout:  0,
 	}
 	llmModel.modelID, llmModel.attributes = model.ParseModelID(modelIDWithAttributes)
 
@@ -61,6 +64,7 @@ func NewModelWithMetaInformation(provider provider.Query, modelIdentifier string
 		modelID:  modelIdentifier,
 
 		queryAttempts: 1,
+		queryTimeout:  0,
 
 		metaInformation: metaInformation,
 	}
@@ -333,10 +337,19 @@ func (m *Model) query(logger *log.Logger, request string) (queryResult *provider
 	if err := retry.Do(
 		func() error {
 			logger.Info("querying model", "model", m.ID(), "query-id", id, "prompt", string(bytesutil.PrefixLines([]byte(request), []byte("\t"))))
+			ctx := context.Background()
+			if m.queryTimeout > 0 {
+				c, cancel := context.WithTimeoutCause(ctx, time.Second*time.Duration(m.queryTimeout), pkgerrors.Errorf("request query timeout (%d seconds)", m.queryTimeout))
+				defer cancel()
+				ctx = c
+			}
+
 			start := time.Now()
-			queryResult, err = m.provider.Query(context.Background(), m, request)
+			queryResult, err = m.provider.Query(ctx, m, request)
 			if err != nil {
 				return err
+			} else if ctx.Err() != nil {
+				return context.Cause(ctx)
 			}
 			duration = time.Since(start)
 			totalCosts := float64(-1)
@@ -520,9 +533,14 @@ func handleQueryResult(queryResult *provider.QueryResult, filePathAbsolute strin
 	return assessment, nil
 }
 
-var _ model.SetQueryAttempts = (*Model)(nil)
+var _ model.SetQueryHandling = (*Model)(nil)
 
 // SetQueryAttempts sets the number of query attempts to perform when a model request errors in the process of solving a task.
 func (m *Model) SetQueryAttempts(queryAttempts uint) {
 	m.queryAttempts = queryAttempts
 }
+
+// SetQueryTimeout sets the timeout for model requests in seconds.
+func (m *Model) SetQueryTimeout(timeout uint) {
+	m.queryTimeout = timeout
+}

model/llm/llm_test.go

@@ -173,6 +173,7 @@ func TestModelQuery(t *testing.T) {
 		SetupMock func(mockedProvider *providertesting.MockQuery)
 
 		QueryAttempts uint
+		QueryTimeout  uint
 		Request       string
 
 		ExpectedResponse *provider.QueryResult
@@ -196,6 +197,7 @@ func TestModelQuery(t *testing.T) {
 			}
 			llm := NewModel(mock, "some-model")
 			llm.SetQueryAttempts(tc.QueryAttempts)
+			llm.SetQueryTimeout(tc.QueryTimeout)
 
 			queryResult, actualError := llm.query(logger, tc.Request)
 
@@ -277,6 +279,39 @@ func TestModelQuery(t *testing.T) {
 
 		ValidateLogs: assertAllIDsMatch,
 	})
+
+	validate(t, &testCase{
+		Name: "Timeout",
+		SetupMock: func(mockedProvider *providertesting.MockQuery) {
+			queryResult := &provider.QueryResult{
+				Message: "test response",
+			}
+			mockedProvider.On("Query", mock.Anything, mock.Anything, "test request").Return(queryResult, nil).After(time.Second * 2)
+		},
+		QueryAttempts: 1,
+		QueryTimeout:  1,
+		Request:       "test request",
+		ExpectedError: "request query timeout",
+	})
+
+	validate(t, &testCase{
+		Name: "Multiple Timeouts",
+		SetupMock: func(mockedProvider *providertesting.MockQuery) {
+			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)
+		},
+		QueryAttempts: 2,
+		QueryTimeout:  1,
+		Request:       "test request",
+		ExpectedError: "request query timeout",
+
+		ValidateLogs: func(t *testing.T, logs string) {
+			assert.Equal(t, 2, strings.Count(logs, "querying model"))
+		},
+	})
 }
 
 func TestModelRepairSourceCodeFile(t *testing.T) {

model/model.go

@@ -82,8 +82,10 @@ type Context struct {
 	Logger *log.Logger
 }
 
-// SetQueryAttempts defines a model that can set the number of query attempts when a model request errors in the process of solving a task.
-type SetQueryAttempts interface {
+// SetQueryHandling defines a model that can configure how API queries are handled.
+type SetQueryHandling interface {
 	// SetQueryAttempts sets the number of query attempts to perform when a model request errors in the process of solving a task.
 	SetQueryAttempts(attempts uint)
+	// SetQueryTimeout sets the timeout for model requests in seconds.
+	SetQueryTimeout(timeout uint)
 }