CERBERUS/ollama
1
0
Fork 0
mirror of https://github.com/ollama/ollama.git synced 2026-10-12 14:01:39 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

feat: allow ten web searches per response (#18602)

Browse source
This commit is contained in:
Parth Sareen 2026-09-28 14:21:25 -07:00 • committed by GitHub
parent cc4069396f
commit 05a7a91dae
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
4 changed files with 134 additions and 108 deletions
Download patch file Download diff file Expand all files Collapse all files

16
middleware/anthropic.go
View file

@ -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(&currentToolCall)
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(&currentToolCall)
maxLoopToolUseID := loopServerToolUseID(w.inner.id, maxWebSearchLoops+1)
maxLoopToolUseID := loopServerToolUseID(w.inner.id, maxLoops+1)
serverContent = append(serverContent,
anthropic.ContentBlock{
Type: "server_tool_use",

200
middleware/anthropic_test.go
View file

@ -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)
}
})
}
}

24
middleware/responses_web_search_test.go
View file

@ -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) {

2
middleware/web_search.go
View file

@ -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.