286 lines
8.7 KiB
Go
286 lines
8.7 KiB
Go
package tools
|
||
|
||
import (
|
||
"bytes"
|
||
"context"
|
||
"encoding/json"
|
||
"fmt"
|
||
"io"
|
||
"net/http"
|
||
"regexp"
|
||
"strings"
|
||
"sync"
|
||
"time"
|
||
)
|
||
|
||
// SearchProvider 是网络搜索能力的注入端口（U9b Batch1：替代 PI Tavily 扩展）。
|
||
// 端点与密钥由 Core 装配；测试注入 fake。
|
||
type SearchProvider interface {
|
||
Search(ctx context.Context, req SearchRequest) (SearchResponse, error)
|
||
}
|
||
|
||
// SearchRequest 是一次搜索的参数（与扩展参数面一致）。
|
||
type SearchRequest struct {
|
||
Query string
|
||
Depth string // basic | advanced
|
||
MaxResults int // 1..10，0 取默认 5
|
||
}
|
||
|
||
// SearchResponse 是 Tavily 兼容响应形状的最小集。
|
||
type SearchResponse struct {
|
||
Answer string
|
||
Results []SearchResult
|
||
}
|
||
|
||
// SearchResult 是单条检索结果。
|
||
type SearchResult struct {
|
||
Title string `json:"title"`
|
||
Content string `json:"content"`
|
||
Score float64 `json:"score"`
|
||
}
|
||
|
||
const (
|
||
searchDefaultMaxResults = 5
|
||
searchTimeout = 25 * time.Second
|
||
)
|
||
|
||
// TavilyProvider 以密钥池方式调用 Tavily Search API：
|
||
// 429/401/403 时轮换下一把密钥重试（对齐扩展 TavilyKeyPool 语义）。
|
||
type TavilyProvider struct {
|
||
endpoint string
|
||
keys []string
|
||
client *http.Client
|
||
|
||
mu sync.Mutex
|
||
next int
|
||
}
|
||
|
||
// NewTavilyProvider 创建提供者；endpoint 空则官方默认；keys 为空则 Search 报错
|
||
// （由 Core 在装配期决定是否注册搜索工具，而非静默降级）。
|
||
func NewTavilyProvider(endpoint string, keys []string) *TavilyProvider {
|
||
if strings.TrimSpace(endpoint) == "" {
|
||
endpoint = "https://api.tavily.com/search"
|
||
}
|
||
return &TavilyProvider{
|
||
endpoint: endpoint,
|
||
keys: keys,
|
||
client: &http.Client{Timeout: searchTimeout},
|
||
}
|
||
}
|
||
|
||
func (p *TavilyProvider) currentKey() (string, int, bool) {
|
||
p.mu.Lock()
|
||
defer p.mu.Unlock()
|
||
if len(p.keys) == 0 {
|
||
return "", 0, false
|
||
}
|
||
idx := p.next % len(p.keys)
|
||
return p.keys[idx], idx, true
|
||
}
|
||
|
||
func (p *TavilyProvider) rotate(idx int) {
|
||
p.mu.Lock()
|
||
p.next = (idx + 1) % len(p.keys)
|
||
p.mu.Unlock()
|
||
}
|
||
|
||
// rotatable 报告该 HTTP 状态是否应轮换密钥重试。
|
||
func rotatable(status int) bool { return status == 429 || status == 401 || status == 403 }
|
||
|
||
func (p *TavilyProvider) Search(ctx context.Context, req SearchRequest) (SearchResponse, error) {
|
||
depth := "basic"
|
||
if req.Depth == "advanced" {
|
||
depth = "advanced"
|
||
}
|
||
max := req.MaxResults
|
||
if max < 1 {
|
||
max = searchDefaultMaxResults
|
||
}
|
||
if max > 10 {
|
||
max = 10
|
||
}
|
||
|
||
total := len(p.keys)
|
||
if total == 0 {
|
||
return SearchResponse{}, fmt.Errorf("search: no API key configured")
|
||
}
|
||
var lastErr error
|
||
for attempt := 0; attempt < total; attempt++ {
|
||
key, idx, ok := p.currentKey()
|
||
if !ok {
|
||
return SearchResponse{}, fmt.Errorf("search: no API key configured")
|
||
}
|
||
payload, _ := json.Marshal(map[string]any{
|
||
"api_key": key,
|
||
"query": req.Query,
|
||
"search_depth": depth,
|
||
"include_answer": true,
|
||
"max_results": max,
|
||
})
|
||
httpReq, err := http.NewRequestWithContext(ctx, http.MethodPost, p.endpoint, bytes.NewReader(payload))
|
||
if err != nil {
|
||
return SearchResponse{}, err
|
||
}
|
||
httpReq.Header.Set("Content-Type", "application/json")
|
||
httpReq.Header.Set("User-Agent", "agent-native-search/1.0")
|
||
|
||
resp, err := p.client.Do(httpReq)
|
||
if err != nil {
|
||
lastErr = err
|
||
p.rotate(idx) // 网络失败也换下一把（与扩展一致）
|
||
continue
|
||
}
|
||
body, _ := io.ReadAll(io.LimitReader(resp.Body, 1<<20))
|
||
resp.Body.Close()
|
||
|
||
if resp.StatusCode/100 != 2 {
|
||
lastErr = fmt.Errorf("search: HTTP %d: %s", resp.StatusCode, strings.TrimSpace(string(body)))
|
||
if rotatable(resp.StatusCode) {
|
||
p.rotate(idx)
|
||
continue
|
||
}
|
||
return SearchResponse{}, lastErr
|
||
}
|
||
|
||
var decoded struct {
|
||
Answer string `json:"answer"`
|
||
Results []SearchResult `json:"results"`
|
||
}
|
||
if err := json.Unmarshal(body, &decoded); err != nil {
|
||
return SearchResponse{}, fmt.Errorf("search: decode response: %w", err)
|
||
}
|
||
return SearchResponse{Answer: decoded.Answer, Results: decoded.Results}, nil
|
||
}
|
||
if lastErr == nil {
|
||
lastErr = fmt.Errorf("search: no key attempt succeeded")
|
||
}
|
||
return SearchResponse{}, fmt.Errorf("search: all %d key(s) failed: %w", total, lastErr)
|
||
}
|
||
|
||
// tavilySearchTool 与 deploy/pi-agent/extensions/tavily-search.mjs 的
|
||
// 工具面逐字段对齐（名称/描述/参数/输出 Markdown/失败降级文案）。
|
||
type tavilySearchTool struct{ provider SearchProvider }
|
||
|
||
func (t *tavilySearchTool) Name() string { return "tavily_search" }
|
||
|
||
func (t *tavilySearchTool) Description() string {
|
||
return "Search the live internet for real-time information, competitor intelligence, regulatory news, product specs, or market data via Tavily (with automatic API key pool rotation)."
|
||
}
|
||
|
||
func (t *tavilySearchTool) Parameters() json.RawMessage {
|
||
return json.RawMessage(`{
|
||
"type": "object",
|
||
"properties": {
|
||
"query": {"type": "string", "description": "The search query keywords in Chinese or English"},
|
||
"search_depth": {"type": "string", "enum": ["basic", "advanced"], "description": "Search depth: 'basic' (fast summary) or 'advanced' (deep analysis, default 'basic')"},
|
||
"max_results": {"type": "number", "description": "Maximum number of search results to retrieve (1 to 10, default 5)"}
|
||
},
|
||
"required": ["query"]
|
||
}`)
|
||
}
|
||
|
||
func (t *tavilySearchTool) Execute(ctx context.Context, args json.RawMessage) (Result, error) {
|
||
var params struct {
|
||
Query string `json:"query"`
|
||
SearchDepth string `json:"search_depth"`
|
||
MaxResults float64 `json:"max_results"`
|
||
}
|
||
if err := json.Unmarshal(args, ¶ms); err != nil {
|
||
return Result{}, err
|
||
}
|
||
query := strings.TrimSpace(params.Query)
|
||
if query == "" {
|
||
return Result{Content: "Error: query cannot be empty.", IsError: true}, nil
|
||
}
|
||
|
||
resp, err := t.provider.Search(ctx, SearchRequest{
|
||
Query: query,
|
||
Depth: params.SearchDepth,
|
||
MaxResults: int(params.MaxResults),
|
||
})
|
||
if err != nil {
|
||
// 失败对模型可见并指引回退（与扩展降级文案一致）；标错误利于分类。
|
||
return Result{
|
||
Content: fmt.Sprintf("[Tavily Search Failed]: %v\n(Agent should rely on workspace materials or local knowledge base as fallback.)", err),
|
||
IsError: true,
|
||
}, nil
|
||
}
|
||
return Result{Content: formatSearchResults(resp, query)}, nil
|
||
}
|
||
|
||
// webSearchTool 是 web_search 别名（扩展同时注册两个名字）。
|
||
type webSearchTool struct{ inner *tavilySearchTool }
|
||
|
||
func (t *webSearchTool) Name() string { return "web_search" }
|
||
func (t *webSearchTool) Description() string {
|
||
return "Alias for tavily_search. Searches the live internet for information."
|
||
}
|
||
func (t *webSearchTool) Parameters() json.RawMessage {
|
||
return json.RawMessage(`{
|
||
"type": "object",
|
||
"properties": {"query": {"type": "string", "description": "The search query keywords"}},
|
||
"required": ["query"]
|
||
}`)
|
||
}
|
||
|
||
func (t *webSearchTool) Execute(ctx context.Context, args json.RawMessage) (Result, error) {
|
||
return t.inner.Execute(ctx, args)
|
||
}
|
||
|
||
// RegisterSearch 把搜索工具注册进 registry（provider 非空才调用；
|
||
// reasoning profile 不注册）。
|
||
func RegisterSearch(r *Registry, provider SearchProvider) {
|
||
inner := &tavilySearchTool{provider: provider}
|
||
r.Register(inner)
|
||
r.Register(&webSearchTool{inner: inner})
|
||
}
|
||
|
||
var (
|
||
urlPattern = regexp.MustCompile(`http[s]?://[^\s]+`)
|
||
sensitivePattern = regexp.MustCompile(`避孕|计划生育|性腺|性器官|阴道|人流`)
|
||
)
|
||
|
||
// sanitizeSearchText 对齐扩展的脱敏：敏感词→[医疗]、URL→#（结果不外泄链接）。
|
||
func sanitizeSearchText(s string) string {
|
||
if s == "" {
|
||
return ""
|
||
}
|
||
s = sensitivePattern.ReplaceAllString(s, "[医疗]")
|
||
return urlPattern.ReplaceAllString(s, "#")
|
||
}
|
||
|
||
// formatSearchResults 渲染与扩展 formatTavilyResults 一致的 Markdown。
|
||
func formatSearchResults(resp SearchResponse, query string) string {
|
||
var b strings.Builder
|
||
fmt.Fprintf(&b, "## 🌐 Tavily 实时网络搜索结果: %q\n\n", query)
|
||
if strings.TrimSpace(resp.Answer) != "" {
|
||
fmt.Fprintf(&b, "> **智能总结**: %s\n\n", sanitizeSearchText(strings.TrimSpace(resp.Answer)))
|
||
}
|
||
if len(resp.Results) > 0 {
|
||
fmt.Fprintf(&b, "### 检索到的参考条目 (%d 条):\n\n", len(resp.Results))
|
||
for i, item := range resp.Results {
|
||
title := strings.TrimSpace(item.Title)
|
||
if title == "" {
|
||
title = "无标题网页"
|
||
}
|
||
title = sanitizeSearchText(title)
|
||
snippet := "（无摘要）"
|
||
if c := strings.TrimSpace(strings.Join(strings.Fields(item.Content), " ")); c != "" {
|
||
snippet = c
|
||
}
|
||
if len(snippet) > 300 {
|
||
snippet = snippet[:300]
|
||
}
|
||
fmt.Fprintf(&b, "%d. **%s**\n", i+1, title)
|
||
fmt.Fprintf(&b, " - **摘要**: %s\n", snippet)
|
||
if item.Score != 0 {
|
||
fmt.Fprintf(&b, " - **相关度评分**: %.3f\n", item.Score)
|
||
}
|
||
b.WriteString("\n")
|
||
}
|
||
} else {
|
||
fmt.Fprintf(&b, "*未检索到与 %q 相关的网页结果。*\n\n", query)
|
||
}
|
||
return b.String()
|
||
}
|