ollama source for Momentry Core verification

x/tools/websearch.go

@@ -0,0 +1,180 @@
+package tools
+
+import (
+	"bytes"
+	"context"
+	"encoding/json"
+	"errors"
+	"fmt"
+	"io"
+	"net/http"
+	"net/url"
+	"strconv"
+	"strings"
+	"time"
+
+	"github.com/ollama/ollama/api"
+	"github.com/ollama/ollama/auth"
+	internalcloud "github.com/ollama/ollama/internal/cloud"
+)
+
+const (
+	webSearchAPI     = "https://ollama.com/api/web_search"
+	webSearchTimeout = 15 * time.Second
+)
+
+// ErrWebSearchAuthRequired is returned when web search requires authentication
+var ErrWebSearchAuthRequired = errors.New("web search requires authentication")
+
+// WebSearchTool implements web search using Ollama's hosted API.
+type WebSearchTool struct{}
+
+// Name returns the tool name.
+func (w *WebSearchTool) Name() string {
+	return "web_search"
+}
+
+// Description returns a description of the tool.
+func (w *WebSearchTool) Description() string {
+	return "Search the web for current information. Use this when you need up-to-date information that may not be in your training data."
+}
+
+// Schema returns the tool's parameter schema.
+func (w *WebSearchTool) Schema() api.ToolFunction {
+	props := api.NewToolPropertiesMap()
+	props.Set("query", api.ToolProperty{
+		Type:        api.PropertyType{"string"},
+		Description: "The search query to look up on the web",
+	})
+	return api.ToolFunction{
+		Name:        w.Name(),
+		Description: w.Description(),
+		Parameters: api.ToolFunctionParameters{
+			Type:       "object",
+			Properties: props,
+			Required:   []string{"query"},
+		},
+	}
+}
+
+// webSearchRequest is the request body for the web search API.
+type webSearchRequest struct {
+	Query      string `json:"query"`
+	MaxResults int    `json:"max_results,omitempty"`
+}
+
+// webSearchResponse is the response from the web search API.
+type webSearchResponse struct {
+	Results []webSearchResult `json:"results"`
+}
+
+// webSearchResult is a single search result.
+type webSearchResult struct {
+	Title   string `json:"title"`
+	URL     string `json:"url"`
+	Content string `json:"content"`
+}
+
+// Execute performs the web search.
+// Uses Ollama key signing for authentication - this makes requests via ollama.com API.
+func (w *WebSearchTool) Execute(args map[string]any) (string, error) {
+	if internalcloud.Disabled() {
+		return "", errors.New(internalcloud.DisabledError("web search is unavailable"))
+	}
+
+	query, ok := args["query"].(string)
+	if !ok || query == "" {
+		return "", fmt.Errorf("query parameter is required")
+	}
+
+	// Prepare request
+	reqBody := webSearchRequest{
+		Query:      query,
+		MaxResults: 5,
+	}
+
+	jsonBody, err := json.Marshal(reqBody)
+	if err != nil {
+		return "", fmt.Errorf("marshaling request: %w", err)
+	}
+
+	// Parse URL and add timestamp for signing
+	searchURL, err := url.Parse(webSearchAPI)
+	if err != nil {
+		return "", fmt.Errorf("parsing search URL: %w", err)
+	}
+
+	q := searchURL.Query()
+	q.Add("ts", strconv.FormatInt(time.Now().Unix(), 10))
+	searchURL.RawQuery = q.Encode()
+
+	// Sign the request using Ollama key (~/.ollama/id_ed25519)
+	// This authenticates with ollama.com using the local signing key
+	ctx := context.Background()
+	data := fmt.Appendf(nil, "%s,%s", http.MethodPost, searchURL.RequestURI())
+	signature, err := auth.Sign(ctx, data)
+	if err != nil {
+		return "", fmt.Errorf("signing request: %w", err)
+	}
+
+	req, err := http.NewRequestWithContext(ctx, http.MethodPost, searchURL.String(), bytes.NewBuffer(jsonBody))
+	if err != nil {
+		return "", fmt.Errorf("creating request: %w", err)
+	}
+
+	req.Header.Set("Content-Type", "application/json")
+	if signature != "" {
+		req.Header.Set("Authorization", fmt.Sprintf("Bearer %s", signature))
+	}
+
+	// Send request
+	client := &http.Client{Timeout: webSearchTimeout}
+	resp, err := client.Do(req)
+	if err != nil {
+		return "", fmt.Errorf("sending request: %w", err)
+	}
+	defer resp.Body.Close()
+
+	body, err := io.ReadAll(resp.Body)
+	if err != nil {
+		return "", fmt.Errorf("reading response: %w", err)
+	}
+
+	if resp.StatusCode == http.StatusUnauthorized {
+		return "", ErrWebSearchAuthRequired
+	}
+	if resp.StatusCode != http.StatusOK {
+		return "", fmt.Errorf("web search API returned status %d: %s", resp.StatusCode, string(body))
+	}
+
+	// Parse response
+	var searchResp webSearchResponse
+	if err := json.Unmarshal(body, &searchResp); err != nil {
+		return "", fmt.Errorf("parsing response: %w", err)
+	}
+
+	// Format results
+	if len(searchResp.Results) == 0 {
+		return "No results found for query: " + query, nil
+	}
+
+	var sb strings.Builder
+	sb.WriteString(fmt.Sprintf("Search results for: %s\n\n", query))
+
+	for i, result := range searchResp.Results {
+		sb.WriteString(fmt.Sprintf("%d. %s\n", i+1, result.Title))
+		sb.WriteString(fmt.Sprintf("   URL: %s\n", result.URL))
+		if result.Content != "" {
+			// Truncate long content (UTF-8 safe)
+			content := result.Content
+			runes := []rune(content)
+			if len(runes) > 300 {
+				content = string(runes[:300]) + "..."
+			}
+			sb.WriteString(fmt.Sprintf("   %s\n", content))
+		}
+		sb.WriteString("\n")
+	}
+
+	return sb.String(), nil
+}