diff --git a/middleware/anthropic.go b/middleware/anthropic.go index b9485e413..60cbcf987 100644 --- a/middleware/anthropic.go +++ b/middleware/anthropic.go @@ -219,7 +219,15 @@ func (w *WebSearchAnthropicWriter) Write(data []byte) (int, error) { } func (w *WebSearchAnthropicWriter) runWebSearchLoop(ctx context.Context, initialResponse api.ChatResponse, initialToolCall api.ToolCall, initialUsage anthropic.Usage) (anthropic.MessagesResponse, *webSearchLoopError) { - followUpMessages := make([]api.Message, 0, len(w.chatReq.Messages)+maxWebSearchLoops*2) + // Omitted or non-positive max_uses retains the server default. + maxLoops := maxWebSearchLoops + for _, tool := range w.req.Tools { + if strings.HasPrefix(tool.Type, "web_search") && tool.MaxUses > 0 { + maxLoops = min(maxLoops, tool.MaxUses) + } + } + + followUpMessages := make([]api.Message, 0, len(w.chatReq.Messages)+maxLoops*2) followUpMessages = append(followUpMessages, w.chatReq.Messages...) followUpTools := append(api.Tools(nil), w.chatReq.Tools...) @@ -229,7 +237,7 @@ func (w *WebSearchAnthropicWriter) runWebSearchLoop(ctx context.Context, initial "tool_call", anthropic.TraceToolCall(initialToolCall), "messages", len(followUpMessages), "tools", len(followUpTools), - "max_loops", maxWebSearchLoops, + "max_loops", maxLoops, ) currentResponse := initialResponse @@ -237,7 +245,7 @@ func (w *WebSearchAnthropicWriter) runWebSearchLoop(ctx context.Context, initial var serverContent []anthropic.ContentBlock - for loop := 1; loop <= maxWebSearchLoops; loop++ { + for loop := 1; loop <= maxLoops; loop++ { query := extractQueryFromToolCall(¤tToolCall) logutil.TraceContext(ctx, "anthropic middleware: web_search loop iteration", "loop", loop, @@ -340,7 +348,7 @@ func (w *WebSearchAnthropicWriter) runWebSearchLoop(ctx context.Context, initial } maxLoopQuery := extractQueryFromToolCall(¤tToolCall) - maxLoopToolUseID := loopServerToolUseID(w.inner.id, maxWebSearchLoops+1) + maxLoopToolUseID := loopServerToolUseID(w.inner.id, maxLoops+1) serverContent = append(serverContent, anthropic.ContentBlock{ Type: "server_tool_use", diff --git a/middleware/anthropic_test.go b/middleware/anthropic_test.go index b17214fbf..373e0715e 100644 --- a/middleware/anthropic_test.go +++ b/middleware/anthropic_test.go @@ -3,6 +3,7 @@ package middleware import ( "bytes" "encoding/json" + "fmt" "io" "net/http" "net/http/httptest" @@ -2588,112 +2589,129 @@ func TestWebSearchMultiIterationLoop(t *testing.T) { } func TestWebSearchLoopMaxLimit(t *testing.T) { - gin.SetMode(gin.TestMode) - enableCloudForTest(t) + for _, test := range []struct { + name string + maxUses int + wantCalls int + }{ + {name: "default", wantCalls: 10}, + {name: "stricter caller limit", maxUses: 2, wantCalls: 2}, + {name: "caller cannot raise cap", maxUses: 20, wantCalls: 10}, + } { + t.Run(test.name, func(t *testing.T) { + gin.SetMode(gin.TestMode) + enableCloudForTest(t) - followupCall := 0 - followupServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { - followupCall++ - resp := api.ChatResponse{ - Model: "test-model", - Message: api.Message{ - Role: "assistant", - ToolCalls: []api.ToolCall{ - { - ID: "call_ws_loop_limit", - Function: api.ToolCallFunction{ - Name: "web_search", - Arguments: makeArgs("query", "loop query next"), + followupCall := 0 + followupServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + followupCall++ + resp := api.ChatResponse{ + Model: "test-model", + Message: api.Message{ + Role: "assistant", + ToolCalls: []api.ToolCall{ + { + ID: "call_ws_loop_limit", + Function: api.ToolCallFunction{ + Name: "web_search", + Arguments: makeArgs("query", "loop query next"), + }, + }, }, }, - }, - }, - Done: true, - DoneReason: "stop", - Metrics: api.Metrics{PromptEvalCount: 7, EvalCount: 2}, - } - _ = json.NewEncoder(w).Encode(resp) - })) - defer followupServer.Close() - t.Setenv("OLLAMA_HOST", followupServer.URL) + Done: true, + DoneReason: "stop", + Metrics: api.Metrics{PromptEvalCount: 7, EvalCount: 2}, + } + _ = json.NewEncoder(w).Encode(resp) + })) + defer followupServer.Close() + t.Setenv("OLLAMA_HOST", followupServer.URL) - searchServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { - resp := anthropic.OllamaWebSearchResponse{ - Results: []anthropic.OllamaWebSearchResult{ - {Title: "Result", URL: "https://example.com", Content: "content"}, - }, - } - _ = json.NewEncoder(w).Encode(resp) - })) - defer searchServer.Close() - originalEndpoint := anthropic.WebSearchEndpoint - anthropic.WebSearchEndpoint = searchServer.URL - defer func() { anthropic.WebSearchEndpoint = originalEndpoint }() + searchCalls := 0 + searchServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + searchCalls++ + resp := anthropic.OllamaWebSearchResponse{ + Results: []anthropic.OllamaWebSearchResult{ + {Title: "Result", URL: "https://example.com", Content: "content"}, + }, + } + _ = json.NewEncoder(w).Encode(resp) + })) + defer searchServer.Close() + originalEndpoint := anthropic.WebSearchEndpoint + anthropic.WebSearchEndpoint = searchServer.URL + defer func() { anthropic.WebSearchEndpoint = originalEndpoint }() - router := gin.New() - router.Use(AnthropicMessagesMiddleware()) - router.POST("/v1/messages", func(c *gin.Context) { - resp := api.ChatResponse{ - Model: "test-model", - Message: api.Message{ - Role: "assistant", - ToolCalls: []api.ToolCall{ - { - ID: "call_ws_initial", - Function: api.ToolCallFunction{ - Name: "web_search", - Arguments: makeArgs("query", "loop query 1"), + router := gin.New() + router.Use(AnthropicMessagesMiddleware()) + router.POST("/v1/messages", func(c *gin.Context) { + resp := api.ChatResponse{ + Model: "test-model", + Message: api.Message{ + Role: "assistant", + ToolCalls: []api.ToolCall{ + { + ID: "call_ws_initial", + Function: api.ToolCallFunction{ + Name: "web_search", + Arguments: makeArgs("query", "loop query 1"), + }, + }, }, }, - }, - }, - Done: true, - DoneReason: "stop", - Metrics: api.Metrics{PromptEvalCount: 5, EvalCount: 1}, - } - data, _ := json.Marshal(resp) - c.Writer.WriteHeader(http.StatusOK) - _, _ = c.Writer.Write(data) - }) + Done: true, + DoneReason: "stop", + Metrics: api.Metrics{PromptEvalCount: 5, EvalCount: 1}, + } + data, _ := json.Marshal(resp) + c.Writer.WriteHeader(http.StatusOK) + _, _ = c.Writer.Write(data) + }) - body := `{ + body := fmt.Sprintf(`{ "model":"test-model:cloud", "max_tokens":100, "messages":[{"role":"user","content":"keep searching"}], - "tools":[{"type":"web_search_20250305","name":"web_search"}] - }` - req, _ := http.NewRequest(http.MethodPost, "/v1/messages", strings.NewReader(body)) - req.Header.Set("Content-Type", "application/json") + "tools":[{"type":"web_search_20250305","name":"web_search","max_uses":%d}] + }`, test.maxUses) + req, _ := http.NewRequest(http.MethodPost, "/v1/messages", strings.NewReader(body)) + req.Header.Set("Content-Type", "application/json") - resp := httptest.NewRecorder() - router.ServeHTTP(resp, req) + resp := httptest.NewRecorder() + router.ServeHTTP(resp, req) - if resp.Code != http.StatusOK { - t.Fatalf("expected 200, got %d: %s", resp.Code, resp.Body.String()) - } - if followupCall != 3 { - t.Fatalf("expected 3 followup calls before max loop error, got %d", followupCall) - } + if resp.Code != http.StatusOK { + t.Fatalf("expected 200, got %d: %s", resp.Code, resp.Body.String()) + } + if searchCalls != test.wantCalls { + t.Fatalf("searches = %d, want %d", searchCalls, test.wantCalls) + } + if followupCall != test.wantCalls { + t.Fatalf("expected %d followup calls before max loop error, got %d", test.wantCalls, followupCall) + } - var result anthropic.MessagesResponse - if err := json.Unmarshal(resp.Body.Bytes(), &result); err != nil { - t.Fatalf("unmarshal error: %v", err) - } + var result anthropic.MessagesResponse + if err := json.Unmarshal(resp.Body.Bytes(), &result); err != nil { + t.Fatalf("unmarshal error: %v", err) + } - last := result.Content[len(result.Content)-1] - if last.Type != "web_search_tool_result" { - t.Fatalf("expected last block web_search_tool_result, got %q", last.Type) - } - contentJSON, _ := json.Marshal(last.Content) - var errContent anthropic.WebSearchToolResultError - if err := json.Unmarshal(contentJSON, &errContent); err != nil { - t.Fatalf("failed to parse web search error content: %v", err) - } - if errContent.ErrorCode != "max_uses_exceeded" { - t.Fatalf("expected max_uses_exceeded error, got %q", errContent.ErrorCode) - } - if result.StopReason != "end_turn" { - t.Fatalf("expected end_turn, got %q", result.StopReason) + last := result.Content[len(result.Content)-1] + if last.Type != "web_search_tool_result" { + t.Fatalf("expected last block web_search_tool_result, got %q", last.Type) + } + contentJSON, _ := json.Marshal(last.Content) + var errContent anthropic.WebSearchToolResultError + if err := json.Unmarshal(contentJSON, &errContent); err != nil { + t.Fatalf("failed to parse web search error content: %v", err) + } + if errContent.ErrorCode != "max_uses_exceeded" { + t.Fatalf("expected max_uses_exceeded error, got %q", errContent.ErrorCode) + } + if result.StopReason != "end_turn" { + t.Fatalf("expected end_turn, got %q", result.StopReason) + } + }) } } diff --git a/middleware/responses_web_search_test.go b/middleware/responses_web_search_test.go index 3d5198f41..92f57ddd6 100644 --- a/middleware/responses_web_search_test.go +++ b/middleware/responses_web_search_test.go @@ -728,7 +728,7 @@ func TestWebSearchResponsesWriterFinalizesAtSearchLimit(t *testing.T) { } } response := api.ChatResponse{Done: true, Message: api.Message{Role: "assistant"}, Metrics: api.Metrics{PromptEvalCount: 5, PromptEvalCachedCount: testIntPtr(2), EvalCount: 3}} - if searches < maxWebSearchLoops { + if searches < 10 { if !reflect.DeepEqual(tools, originalTools) { t.Fatalf("tools removed before search limit: %#v", tools) } @@ -767,8 +767,8 @@ func TestWebSearchResponsesWriterFinalizesAtSearchLimit(t *testing.T) { if _, err := writer.Write(data); err != nil { t.Fatal(err) } - if searches != maxWebSearchLoops || followUps != maxWebSearchLoops { - t.Fatalf("searches=%d follow-ups=%d, want %d each", searches, followUps, maxWebSearchLoops) + if searches != 10 || followUps != 10 { + t.Fatalf("searches=%d follow-ups=%d, want 10 each", searches, followUps) } if !reflect.DeepEqual(chat.Tools, originalTools) { t.Fatalf("original request tools mutated: %#v", chat.Tools) @@ -793,15 +793,15 @@ func TestWebSearchResponsesWriterFinalizesAtSearchLimit(t *testing.T) { } else if err := json.Unmarshal(recorder.Body.Bytes(), &response); err != nil { t.Fatal(err) } - if recorder.Code != http.StatusOK || response.Status != "completed" || len(response.Output) != maxWebSearchLoops+1 { + if recorder.Code != http.StatusOK || response.Status != "completed" || len(response.Output) != 11 { t.Fatalf("unexpected final response: %s", body) } - for _, item := range response.Output[:maxWebSearchLoops] { + for _, item := range response.Output[:10] { if item.Type != "web_search_call" || item.Status != "completed" { t.Fatalf("search output lost: %#v", item) } } - final := response.Output[maxWebSearchLoops] + final := response.Output[10] if test.clientTool { if final.Type != "function_call" || final.Name != "get_weather" || final.CallID != "call_weather" { t.Fatalf("client tool output lost: %#v", final) @@ -809,8 +809,8 @@ func TestWebSearchResponsesWriterFinalizesAtSearchLimit(t *testing.T) { } else if final.Type != "message" || len(final.Content) != 1 || final.Content[0].Text != "Rain expected, based on the available results." { t.Fatalf("final answer lost: %#v", final) } - if response.Usage == nil || response.Usage.InputTokens != 20 || response.Usage.OutputTokens != 12 || response.Usage.InputTokensDetails.CachedTokens != 8 { - t.Fatalf("usage = %#v, want all four model responses", response.Usage) + if response.Usage == nil || response.Usage.InputTokens != 55 || response.Usage.OutputTokens != 33 || response.Usage.InputTokensDetails.CachedTokens != 22 { + t.Fatalf("usage = %#v, want all eleven model responses", response.Usage) } }) } @@ -824,7 +824,7 @@ func TestWebSearchResponsesWriterFinalizationFailure(t *testing.T) { status int message string }{ - {name: "model requests another search", status: http.StatusBadGateway, message: "web_search exceeded the maximum"}, + {name: "model requests another search", status: http.StatusBadGateway, message: "web_search exceeded the maximum of 10 calls"}, {name: "model request fails", err: api.StatusError{StatusCode: http.StatusServiceUnavailable, ErrorMessage: "model unavailable"}, status: http.StatusServiceUnavailable, message: "model unavailable"}, } { t.Run(fmt.Sprintf("%s/stream=%t", test.name, stream), func(t *testing.T) { @@ -844,7 +844,7 @@ func TestWebSearchResponsesWriterFinalizationFailure(t *testing.T) { }, followUpChat: func(context.Context, []api.Message, api.Tools) (api.ChatResponse, error) { followUps++ - if followUps == maxWebSearchLoops && test.err != nil { + if followUps == 10 && test.err != nil { return api.ChatResponse{}, test.err } return api.ChatResponse{Done: true, Message: api.Message{Role: "assistant", ToolCalls: []api.ToolCall{call}}}, nil @@ -855,8 +855,8 @@ func TestWebSearchResponsesWriterFinalizationFailure(t *testing.T) { if _, err := writer.Write(data); err != nil { t.Fatal(err) } - if searches != maxWebSearchLoops || followUps != maxWebSearchLoops { - t.Fatalf("searches=%d follow-ups=%d, want %d each", searches, followUps, maxWebSearchLoops) + if searches != 10 || followUps != 10 { + t.Fatalf("searches=%d follow-ups=%d, want 10 each", searches, followUps) } body := recorder.Body.String() if !strings.Contains(body, test.message) { diff --git a/middleware/web_search.go b/middleware/web_search.go index 70b5ac85b..d2aa2ac3f 100644 --- a/middleware/web_search.go +++ b/middleware/web_search.go @@ -6,7 +6,7 @@ import ( "github.com/ollama/ollama/api" ) -const maxWebSearchLoops = 3 +const maxWebSearchLoops = 10 // doFollowUpChat sends a non-streaming /api/chat request with the accumulated // messages and tools so the model can continue after a web search result.