Cyrene_v2/plugins/tools/io/web_search.go

package io

import (
	"context"
	"encoding/json"
	"fmt"
	"log/slog"
	"net/http"
	"net/url"
	"strings"
	"time"

	"git.yeij.top/AskaEth/Cyrene_v2/pkg/plugins/sdk"
)

// Compile-time assertion that webSearchTool implements sdk.Tool.
var _ sdk.Tool = (*webSearchTool)(nil)

// webSearchTool queries a configured SearXNG-compatible (or custom) search
// endpoint and returns ranked results, capped in size and count.
type webSearchTool struct {
	client       *http.Client
	endpoint     string
	allowedHosts []string
	maxBytes     int64
	maxResults   int
	log          *slog.Logger
	// timeout 是本次搜索的超时覆盖（0 = 沿用共享客户端自己的超时）。它只约束
	// web-search 一次请求，不牵动 http / web-fetch 共用的那个客户端（WS1-A）。
	timeout time.Duration
}

func newWebSearchTool(t *tools) *webSearchTool {
	// 端点来自**运维配置**，不是模型输入（模型的 query 只会进查询串，主机名恒为
	// 配置里那一个）：因此把它并入本工具**自己**的 SSRF 白名单。否则本机兜底端点
	// http://127.0.0.1:38471 会被 SSRF 守卫按回环地址拒掉——那正是「端点配了却
	// 仍然用不了」的另一种形态（SearXNG 还会 308 跳到 /search，重定向那一跳同样
	// 要过守卫，所以加宽后的名单也必须交给客户端，见 tools.clientFor）。
	//
	// 刻意只并入 web-search：http / web-fetch 的 URL 由模型给出，它们共用
	// t.allowedHosts，把回环塞进**那份**名单会让模型能伸手去摸本机服务。
	allowedHosts := withSearchEndpointHost(t.allowedHosts, t.searchURL)
	return &webSearchTool{
		client:       t.clientFor(allowedHosts),
		endpoint:     t.searchURL,
		allowedHosts: allowedHosts,
		maxBytes:     t.maxHTTPBytes,
		maxResults:   t.maxResults,
		log:          t.log,
		timeout:      t.searchTimeout,
	}
}

// Name implements sdk.Tool.
func (s *webSearchTool) Name() string { return "web-search" }

// Description implements sdk.Tool.
func (s *webSearchTool) Description() string {
	return "Search the web via a configured SearXNG or custom search endpoint and return ranked results."
}

// Parameters implements sdk.Tool.
func (s *webSearchTool) Parameters() (map[string]any, error) {
	return map[string]any{
		"type": "object",
		"properties": map[string]any{
			"query": map[string]any{"type": "string", "description": "Search query."},
			"count": map[string]any{"type": "integer", "description": "Max results to return."},
		},
		"required": []any{"query"},
	}, nil
}

// Execute implements sdk.Tool.
func (s *webSearchTool) Execute(ctx context.Context, args any) (any, error) {
	m, ok := args.(map[string]any)
	if !ok {
		return nil, fmt.Errorf("io: web-search: args must be an object")
	}
	query := argString(m, "query")
	if strings.TrimSpace(query) == "" {
		return nil, fmt.Errorf("io: web-search: query is required")
	}
	if strings.TrimSpace(s.endpoint) == "" {
		return nil, fmt.Errorf("io: web-search: no search endpoint configured (set a SearXNG or custom URL)")
	}
	// 每次调用各自的超时预算：ctx 在此派生，取消随本次调用一起结束。
	if s.timeout > 0 {
		var cancel context.CancelFunc
		ctx, cancel = context.WithTimeout(ctx, s.timeout)
		defer cancel()
	}
	target, err := buildSearchURL(s.endpoint, query)
	if err != nil {
		return nil, err
	}
	if _, err := validateRequestURL(ctx, target, s.allowedHosts); err != nil {
		return nil, err
	}
	if s.client == nil {
		return nil, fmt.Errorf("io: web-search: no underlying client")
	}
	req, err := http.NewRequestWithContext(ctx, http.MethodGet, target, nil)
	if err != nil {
		return nil, fmt.Errorf("io: web-search: build request: %w", err)
	}
	resp, err := s.client.Do(req)
	if err != nil {
		return nil, fmt.Errorf("io: web-search: do: %w", err)
	}
	defer resp.Body.Close()
	if resp.StatusCode != http.StatusOK {
		return nil, fmt.Errorf("io: web-search: endpoint returned status %d", resp.StatusCode)
	}
	data, truncated, err := readLimited(resp.Body, s.maxBytes)
	if err != nil {
		return nil, fmt.Errorf("io: web-search: read: %w", err)
	}
	if truncated {
		return nil, fmt.Errorf("io: web-search: response exceeds the %d-byte limit", s.maxBytes)
	}
	results, err := parseSearchResults(data)
	if err != nil {
		return nil, err
	}
	n := s.maxResults
	if v, ok := argInt(m, "count"); ok && v > 0 {
		n = v
	}
	if len(results) > n {
		results = results[:n]
	}
	return map[string]any{
		"query":     query,
		"count":     len(results),
		"results":   results,
		"truncated": truncated,
	}, nil
}

// buildSearchURL expands an endpoint against a query. It supports a "{query}"
// placeholder, otherwise appends q= and format=json for a SearXNG JSON API.
func buildSearchURL(endpoint, query string) (string, error) {
	if strings.TrimSpace(endpoint) == "" {
		return "", fmt.Errorf("io: web-search: no search endpoint configured")
	}
	q := url.QueryEscape(query)
	if strings.Contains(endpoint, "{query}") {
		return strings.ReplaceAll(endpoint, "{query}", q), nil
	}
	sep := "?"
	if strings.Contains(endpoint, "?") {
		sep = "&"
	}
	return endpoint + sep + "q=" + q + "&format=json", nil
}

// withSearchEndpointHost 把**配置里给定的**搜索端点主机并入 web-search 自己的
// SSRF 白名单（返回新切片，不改动入参切片——组合根可能把它交给别的工具）。
// 端点解析不出主机、或主机已在名单里时原样返回。
//
// 理由见 newWebSearchTool：本机兜底端点必然是回环地址（127.0.0.1:38471），而
// SSRF 守卫默认拒绝回环；端点是运维配置而非模型输入，因此对它放行**不会**给模型
// 打开新的目标面——工具能去的主机恒等于配置里的那一个。
func withSearchEndpointHost(allowed []string, endpoint string) []string {
	host := searchEndpointHost(endpoint)
	if host == "" || isHostAllowed(host, allowed) {
		return allowed
	}
	out := make([]string, 0, len(allowed)+1)
	out = append(out, allowed...)
	return append(out, host)
}

// searchEndpointHost 取出搜索端点的主机名（去端口、去 IPv6 方括号）；解析不出来
// 或没有主机时返回空串。
func searchEndpointHost(endpoint string) string {
	u, err := url.Parse(strings.TrimSpace(endpoint))
	if err != nil {
		return ""
	}
	return hostOnly(u.Host)
}

// searchResult is one entry of a SearXNG-compatible results payload.
type searchResult struct {
	Title   string `json:"title"`
	URL     string `json:"url"`
	Content string `json:"content"`
	Snippet string `json:"snippet"`
}

// parseSearchResults decodes a {"results":[...]} payload, falling back to the
// snippet field when content is absent.
func parseSearchResults(data []byte) ([]searchResult, error) {
	var payload struct {
		Results []searchResult `json:"results"`
	}
	if err := json.Unmarshal(data, &payload); err != nil {
		return nil, fmt.Errorf("io: web-search: decode results: %w", err)
	}
	for i := range payload.Results {
		if payload.Results[i].Content == "" {
			payload.Results[i].Content = payload.Results[i].Snippet
		}
	}
	return payload.Results, nil
}