Cyrene_v2/plugins/tools/io/web_search.go

package io

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

	"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
}

func newWebSearchTool(t *tools) *webSearchTool {
	return &webSearchTool{
		client:       t.client,
		endpoint:     t.searchURL,
		allowedHosts: t.allowedHosts,
		maxBytes:     t.maxHTTPBytes,
		maxResults:   t.maxResults,
		log:          t.log,
	}
}

// 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)")
	}
	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
}

// 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
}