Add web search capabilities

2026-08-28 22:11:38 +00:00 · 2026-07-17 10:49:24 +02:00
parent 0543283b8a
commit 7fd7818142
6 changed files with 344 additions and 52 deletions
@@ -17,6 +17,8 @@ var ErrModelNotFound = errors.New("model not found or deprecated")

 const maxFileNotFoundRetries = 3

+const maxPauseTurnContinuations = 5
+
 const defaultMaxTokens = 1000

 const mcpUnsupportedSentinel = "format not currently supported by the Anthropic API"
@@ -65,9 +67,10 @@ func (b *Bot) getAnthropicResponse(ctx context.Context, chatID int64, messages [
 		params.Thinking = thinking
 	}

+	var tools []anthropic.BetaToolUnionParam
+
 	if len(b.config.MCPServers) > 0 {
 		mcpServers := make([]anthropic.BetaRequestMCPServerURLDefinitionParam, 0, len(b.config.MCPServers))
-		tools := make([]anthropic.BetaToolUnionParam, 0, len(b.config.MCPServers))
 		for _, s := range b.config.MCPServers {
 			srv := anthropic.BetaRequestMCPServerURLDefinitionParam{
 				Name: s.Name,
@@ -95,32 +98,99 @@ func (b *Bot) getAnthropicResponse(ctx context.Context, chatID int64, messages [
 			tools = append(tools, anthropic.BetaToolUnionParam{OfMCPToolset: toolset})
 		}
 		params.MCPServers = mcpServers
-		params.Tools = tools
 		params.Betas = append(params.Betas, anthropic.AnthropicBetaMCPClient2025_11_20)
 	}

-	for attempt := 0; attempt < maxFileNotFoundRetries; attempt++ {
-		joined, streamErr := b.streamMessages(ctx, params, onSegment)
-		if streamErr == nil {
-			return joined, nil
-		}
-		var apiErr *anthropic.Error
-		if !errors.As(streamErr, &apiErr) || apiErr.StatusCode != http.StatusNotFound {
-			return "", fmt.Errorf("error creating Anthropic message: %w", streamErr)
-		}
-		missingFileID := extractMissingFileID(streamErr)
-		if missingFileID == "" {
-			return "", fmt.Errorf("%w: %s", ErrModelNotFound, b.config.Model)
-		}
-		ErrorLogger.Printf("[%s] self-heal: stripping dead file_id %s from chat %d (attempt %d/%d)",
-			b.config.ID, missingFileID, chatID, attempt+1, maxFileNotFoundRetries)
-		b.stripDeadFileIDFromMemory(chatID, missingFileID)
-		if _, cleanupErr := b.markFilesPendingCleanup(ctx, chatID, []string{missingFileID}); cleanupErr != nil {
-			ErrorLogger.Printf("[%s] mark files pending cleanup: %v", b.config.ID, cleanupErr)
-		}
-		params.Messages = b.prepareContextMessages(b.getOrCreateChatMemory(chatID))
+	tools = append(tools, webSearchTools(b.config.WebSearch)...)
+
+	if len(tools) > 0 {
+		params.Tools = tools
 	}
-	return "", fmt.Errorf("max self-heal retries (%d) exceeded: too many file_ids gone from anthropic", maxFileNotFoundRetries)
+
+	var fullText strings.Builder
+	var lastMsg anthropic.BetaMessage
+	fileRetries, pauseContinuations := 0, 0
+	for {
+		joined, msg, streamErr := b.streamMessages(ctx, params, onSegment)
+		if streamErr != nil {
+			var apiErr *anthropic.Error
+			if !errors.As(streamErr, &apiErr) || apiErr.StatusCode != http.StatusNotFound {
+				return "", fmt.Errorf("error creating Anthropic message: %w", streamErr)
+			}
+			missingFileID := extractMissingFileID(streamErr)
+			if missingFileID == "" {
+				return "", fmt.Errorf("%w: %s", ErrModelNotFound, b.config.Model)
+			}
+			fileRetries++
+			if fileRetries > maxFileNotFoundRetries {
+				return "", fmt.Errorf("max self-heal retries (%d) exceeded: too many file_ids gone from anthropic", maxFileNotFoundRetries)
+			}
+			ErrorLogger.Printf("[%s] self-heal: stripping dead file_id %s from chat %d (attempt %d/%d)",
+				b.config.ID, missingFileID, chatID, fileRetries, maxFileNotFoundRetries)
+			b.stripDeadFileIDFromMemory(chatID, missingFileID)
+			if _, cleanupErr := b.markFilesPendingCleanup(ctx, chatID, []string{missingFileID}); cleanupErr != nil {
+				ErrorLogger.Printf("[%s] mark files pending cleanup: %v", b.config.ID, cleanupErr)
+			}
+			params.Messages = b.prepareContextMessages(b.getOrCreateChatMemory(chatID))
+			continue
+		}
+
+		lastMsg = msg
+		if joined != "" {
+			if fullText.Len() > 0 {
+				fullText.WriteString("\n\n")
+			}
+			fullText.WriteString(joined)
+		}
+
+		if msg.StopReason == anthropic.BetaStopReasonPauseTurn {
+			pauseContinuations++
+			if pauseContinuations > maxPauseTurnContinuations {
+				ErrorLogger.Printf("[%s] pause_turn continuations exceeded (%d); returning partial answer",
+					b.config.ID, maxPauseTurnContinuations)
+				break
+			}
+			params.Messages = append(params.Messages, msg.ToParam())
+			continue
+		}
+		break
+	}
+
+	if fullText.Len() == 0 {
+		return "", emptyStreamError(string(lastMsg.StopReason),
+			lastMsg.Usage.OutputTokensDetails.ThinkingTokens, params.MaxTokens)
+	}
+	return fullText.String(), nil
+}
+
+func webSearchTools(cfg *WebSearchConfig) []anthropic.BetaToolUnionParam {
+	if cfg == nil {
+		return nil
+	}
+	search := &anthropic.BetaWebSearchTool20250305Param{
+		AllowedDomains: cfg.AllowedDomains,
+		BlockedDomains: cfg.BlockedDomains,
+	}
+	if cfg.MaxUses > 0 {
+		search.MaxUses = param.NewOpt(int64(cfg.MaxUses))
+	}
+	tools := []anthropic.BetaToolUnionParam{{OfWebSearchTool20250305: search}}
+
+	if cfg.Fetch {
+		fetch := &anthropic.BetaWebFetchTool20250910Param{
+			AllowedDomains: cfg.AllowedDomains,
+			BlockedDomains: cfg.BlockedDomains,
+			Citations:      anthropic.BetaCitationsConfigParam{Enabled: param.NewOpt(true)},
+		}
+		if cfg.MaxUses > 0 {
+			fetch.MaxUses = param.NewOpt(int64(cfg.MaxUses))
+		}
+		if cfg.MaxContentTokens > 0 {
+			fetch.MaxContentTokens = param.NewOpt(int64(cfg.MaxContentTokens))
+		}
+		tools = append(tools, anthropic.BetaToolUnionParam{OfWebFetchTool20250910: fetch})
+	}
+	return tools
 }

 func buildUserContext(username, firstName, lastName string, isPremium bool, languageCode string, messageTime int) string {
@@ -179,7 +249,7 @@ func thinkingParamFromConfig(mode, display string) (anthropic.BetaThinkingConfig
 	}
 }

-func (b *Bot) streamMessages(ctx context.Context, params anthropic.BetaMessageNewParams, onSegment func(string) error) (string, error) {
+func (b *Bot) streamMessages(ctx context.Context, params anthropic.BetaMessageNewParams, onSegment func(string) error) (string, anthropic.BetaMessage, error) {
 	stream := b.anthropicClient.Beta.Messages.NewStreaming(ctx, params)
 	defer func() {
 		if err := stream.Close(); err != nil {
@@ -188,6 +258,7 @@ func (b *Bot) streamMessages(ctx context.Context, params anthropic.BetaMessageNe
 	}()

 	var (
+		message                                           anthropic.BetaMessage
 		allSegments                                       []string
 		currentKind                                       string
 		currentText                                       strings.Builder
@@ -197,31 +268,24 @@ func (b *Bot) streamMessages(ctx context.Context, params anthropic.BetaMessageNe
 		currentTResultUseID, currentTResultServer         string
 		currentTResultIsError                             bool
 		currentTResultContent                             string
+		currentServerToolName, currentServerToolID        string
+		currentServerResult                               string
 		mcpCalls                                          = map[string]mcpCall{}
-		startInputTokens                                  int64
-		finalUsage                                        anthropic.BetaMessageDeltaUsage
-		stopReason                                        string
 	)

 	for stream.Next() {
 		e := stream.Current()
+		if accErr := message.Accumulate(e); accErr != nil {
+			ErrorLogger.Printf("[stream] accumulate failed: %v", accErr)
+		}
 		switch e.Type {
-		case "message_start":
-			startInputTokens = e.AsMessageStart().Message.Usage.InputTokens
-
-		case "message_delta":
-			md := e.AsMessageDelta()
-			finalUsage = md.Usage
-			if md.Delta.StopReason != "" {
-				stopReason = string(md.Delta.StopReason)
-			}
-
 		case "content_block_start":
 			cbs := e.AsContentBlockStart()
 			currentKind = cbs.ContentBlock.Type
 			currentText.Reset()
 			currentThinking.Reset()
 			currentInputJSON.Reset()
+			currentServerResult = ""
 			switch currentKind {
 			case "mcp_tool_use":
 				currentTUseName = cbs.ContentBlock.Name
@@ -232,6 +296,11 @@ func (b *Bot) streamMessages(ctx context.Context, params anthropic.BetaMessageNe
 				currentTResultServer = cbs.ContentBlock.ServerName
 				currentTResultIsError = cbs.ContentBlock.IsError
 				currentTResultContent = cbs.ContentBlock.JSON.Content.Raw()
+			case "server_tool_use":
+				currentServerToolName = cbs.ContentBlock.Name
+				currentServerToolID = cbs.ContentBlock.ID
+			case "web_search_tool_result", "web_fetch_tool_result":
+				currentServerResult = cbs.ContentBlock.JSON.Content.Raw()
 			}

 		case "content_block_delta":
@@ -246,7 +315,7 @@ func (b *Bot) streamMessages(ctx context.Context, params anthropic.BetaMessageNe
 					currentThinking.WriteString(cbd.Delta.Thinking)
 				}
 			case "input_json_delta":
-				if currentKind == "mcp_tool_use" {
+				if currentKind == "mcp_tool_use" || currentKind == "server_tool_use" {
 					currentInputJSON.WriteString(cbd.Delta.PartialJSON)
 				}
 			}
@@ -285,6 +354,15 @@ func (b *Bot) streamMessages(ctx context.Context, params anthropic.BetaMessageNe
 						"(total=%d): server=%q tool=%q input=%s tool_use_id=%q",
 						b.config.ID, total, call.server, call.name, call.input, currentTResultUseID)
 				}
+			case "server_tool_use":
+				InfoLogger.Printf("[web] %s id=%q input=%s",
+					currentServerToolName, currentServerToolID, currentInputJSON.String())
+			case "web_search_tool_result", "web_fetch_tool_result":
+				preview := currentServerResult
+				if len(preview) > 500 {
+					preview = preview[:500] + "...(truncated)"
+				}
+				InfoLogger.Printf("[web] %s content=%s", currentKind, preview)
 			case "thinking", "redacted_thinking":
 				if summary := strings.TrimSpace(currentThinking.String()); summary != "" {
 					if len(summary) > 500 {
@@ -296,7 +374,7 @@ func (b *Bot) streamMessages(ctx context.Context, params anthropic.BetaMessageNe
 				}
 			default:
 				if currentKind != "" {
-					InfoLogger.Printf("[mcp] block type=%q (unhandled)", currentKind)
+					InfoLogger.Printf("[stream] block type=%q (unhandled)", currentKind)
 				}
 			}
 			currentKind = ""
@@ -304,27 +382,21 @@ func (b *Bot) streamMessages(ctx context.Context, params anthropic.BetaMessageNe
 	}

 	if err := stream.Err(); err != nil {
-		return "", err
+		return "", message, err
 	}

-	if stopReason != "" || finalUsage.OutputTokens > 0 {
-		inputTokens := finalUsage.InputTokens
-		if inputTokens == 0 {
-			inputTokens = startInputTokens
-		}
+	stopReason := string(message.StopReason)
+	if stopReason != "" || message.Usage.OutputTokens > 0 {
 		InfoLogger.Printf("[usage] model=%s in=%d out=%d thinking=%d stop=%s",
-			params.Model, inputTokens, finalUsage.OutputTokens,
-			finalUsage.OutputTokensDetails.ThinkingTokens, stopReason)
-		if stopReason == "max_tokens" {
+			params.Model, message.Usage.InputTokens, message.Usage.OutputTokens,
+			message.Usage.OutputTokensDetails.ThinkingTokens, stopReason)
+		if message.StopReason == anthropic.BetaStopReasonMaxTokens {
 			ErrorLogger.Printf("[usage] response truncated at max_tokens=%d - raise max_tokens (thinking counts toward it)",
 				params.MaxTokens)
 		}
 	}

-	if len(allSegments) == 0 {
-		return "", emptyStreamError(stopReason, finalUsage.OutputTokensDetails.ThinkingTokens, params.MaxTokens)
-	}
-	return strings.Join(allSegments, "\n\n"), nil
+	return strings.Join(allSegments, "\n\n"), message, nil
 }

 func emptyStreamError(stopReason string, thinkingTokens, maxTokens int64) error {