From 3adb59b44d36934c52991ac182ae2f79517090d4 Mon Sep 17 00:00:00 2001 From: instax-dutta Date: Thu, 12 Feb 2026 18:04:46 +0530 Subject: [PATCH] feat: add Ollama search tools and update LLM providers --- config.example.json | 14 ++- pkg/agent/loop.go | 12 ++- pkg/config/config.go | 21 ++++ pkg/tools/web.go | 246 ++++++++++++++++++++++++++++++++++++++++++- 4 files changed, 283 insertions(+), 10 deletions(-) diff --git a/config.example.json b/config.example.json index 12dc47316..35219177e 100644 --- a/config.example.json +++ b/config.example.json @@ -12,7 +12,9 @@ "telegram": { "enabled": false, "token": "YOUR_TELEGRAM_BOT_TOKEN", - "allow_from": ["YOUR_USER_ID"] + "allow_from": [ + "YOUR_USER_ID" + ] }, "discord": { "enabled": false, @@ -76,6 +78,10 @@ "api_key": "", "api_base": "" }, + "nvidia": { + "api_key": "nvapi-xxx", + "api_base": "" + }, "vllm": { "api_key": "", "api_base": "" @@ -86,6 +92,10 @@ "search": { "api_key": "YOUR_BRAVE_API_KEY", "max_results": 5 + }, + "ollama": { + "api_key": "YOUR_OLLAMA_API_KEY", + "max_results": 5 } } }, @@ -93,4 +103,4 @@ "host": "0.0.0.0", "port": 18790 } -} +} \ No newline at end of file diff --git a/pkg/agent/loop.go b/pkg/agent/loop.go index cc14ceaf0..d2eac1f4f 100644 --- a/pkg/agent/loop.go +++ b/pkg/agent/loop.go @@ -61,9 +61,15 @@ func NewAgentLoop(cfg *config.Config, msgBus *bus.MessageBus, provider providers toolsRegistry.Register(&tools.ListDirTool{}) toolsRegistry.Register(tools.NewExecTool(workspace)) - braveAPIKey := cfg.Tools.Web.Search.APIKey - toolsRegistry.Register(tools.NewWebSearchTool(braveAPIKey, cfg.Tools.Web.Search.MaxResults)) - toolsRegistry.Register(tools.NewWebFetchTool(50000)) + ollamaAPIKey := cfg.Tools.Web.Ollama.APIKey + if ollamaAPIKey != "" { + toolsRegistry.Register(tools.NewOllamaSearchTool(ollamaAPIKey, cfg.Tools.Web.Ollama.MaxResults)) + toolsRegistry.Register(tools.NewOllamaFetchTool(ollamaAPIKey, 50000)) + } else { + braveAPIKey := cfg.Tools.Web.Search.APIKey + toolsRegistry.Register(tools.NewWebSearchTool(braveAPIKey, cfg.Tools.Web.Search.MaxResults)) + toolsRegistry.Register(tools.NewWebFetchTool(50000)) + } // Register message tool messageTool := tools.NewMessageTool() diff --git a/pkg/config/config.go b/pkg/config/config.go index cd81dfa4d..48e6e9f56 100644 --- a/pkg/config/config.go +++ b/pkg/config/config.go @@ -105,6 +105,7 @@ type ProvidersConfig struct { Zhipu ProviderConfig `json:"zhipu"` VLLM ProviderConfig `json:"vllm"` Gemini ProviderConfig `json:"gemini"` + Nvidia ProviderConfig `json:"nvidia"` } type ProviderConfig struct { @@ -125,6 +126,12 @@ type WebSearchConfig struct { type WebToolsConfig struct { Search WebSearchConfig `json:"search"` + Ollama OllamaConfig `json:"ollama"` +} + +type OllamaConfig struct { + APIKey string `json:"api_key" env:"PICOCLAW_TOOLS_WEB_OLLAMA_API_KEY"` + MaxResults int `json:"max_results" env:"PICOCLAW_TOOLS_WEB_OLLAMA_MAX_RESULTS"` } type ToolsConfig struct { @@ -200,6 +207,7 @@ func DefaultConfig() *Config { Zhipu: ProviderConfig{}, VLLM: ProviderConfig{}, Gemini: ProviderConfig{}, + Nvidia: ProviderConfig{}, }, Gateway: GatewayConfig{ Host: "0.0.0.0", @@ -211,6 +219,10 @@ func DefaultConfig() *Config { APIKey: "", MaxResults: 5, }, + Ollama: OllamaConfig{ + APIKey: "", + MaxResults: 5, + }, }, }, } @@ -285,6 +297,9 @@ func (c *Config) GetAPIKey() string { if c.Providers.VLLM.APIKey != "" { return c.Providers.VLLM.APIKey } + if c.Providers.Nvidia.APIKey != "" { + return c.Providers.Nvidia.APIKey + } return "" } @@ -303,6 +318,12 @@ func (c *Config) GetAPIBase() string { if c.Providers.VLLM.APIKey != "" && c.Providers.VLLM.APIBase != "" { return c.Providers.VLLM.APIBase } + if c.Providers.Nvidia.APIKey != "" { + if c.Providers.Nvidia.APIBase != "" { + return c.Providers.Nvidia.APIBase + } + return "https://integrate.api.nvidia.com/v1" + } return "" } diff --git a/pkg/tools/web.go b/pkg/tools/web.go index 3a3596865..d2995e104 100644 --- a/pkg/tools/web.go +++ b/pkg/tools/web.go @@ -1,12 +1,12 @@ package tools import ( + "bytes" "context" "encoding/json" "fmt" "io" "net/http" - "net/url" "regexp" "strings" "time" @@ -16,6 +16,240 @@ const ( userAgent = "Mozilla/5.0 (compatible; picoclaw/1.0)" ) +// --- Ollama Search Tool --- + +type OllamaSearchTool struct { + apiKey string + maxResults int +} + +func NewOllamaSearchTool(apiKey string, maxResults int) *OllamaSearchTool { + if maxResults <= 0 || maxResults > 10 { + maxResults = 5 + } + return &OllamaSearchTool{ + apiKey: apiKey, + maxResults: maxResults, + } +} + +func (t *OllamaSearchTool) Name() string { + return "web_search" +} + +func (t *OllamaSearchTool) Description() string { + return "Search the web for current information using Ollama. Returns titles, URLs, and snippets." +} + +func (t *OllamaSearchTool) Parameters() map[string]interface{} { + return map[string]interface{}{ + "type": "object", + "properties": map[string]interface{}{ + "query": map[string]interface{}{ + "type": "string", + "description": "Search query", + }, + "count": map[string]interface{}{ + "type": "integer", + "description": "Number of results (1-10)", + "minimum": 1.0, + "maximum": 10.0, + }, + }, + "required": []string{"query"}, + } +} + +func (t *OllamaSearchTool) Execute(ctx context.Context, args map[string]interface{}) (string, error) { + query, ok := args["query"].(string) + if !ok { + return "", fmt.Errorf("query is required") + } + + count := t.maxResults + if c, ok := args["count"].(float64); ok { + if int(c) > 0 && int(c) <= 10 { + count = int(c) + } + } + + requestBody := map[string]interface{}{ + "query": query, + "max_results": count, + } + jsonData, err := json.Marshal(requestBody) + if err != nil { + return "", fmt.Errorf("failed to marshal request: %w", err) + } + + req, err := http.NewRequestWithContext(ctx, "POST", "https://ollama.com/api/web_search", bytes.NewReader(jsonData)) + if err != nil { + return "", fmt.Errorf("failed to create request: %w", err) + } + + req.Header.Set("Content-Type", "application/json") + if t.apiKey != "" { + req.Header.Set("Authorization", "Bearer "+t.apiKey) + } else { + return "Error: OLLAMA_API_KEY not configured", nil + } + + client := &http.Client{Timeout: 30 * time.Second} + resp, err := client.Do(req) + if err != nil { + return "", fmt.Errorf("request failed: %w", err) + } + defer resp.Body.Close() + + body, err := io.ReadAll(resp.Body) + if err != nil { + return "", fmt.Errorf("failed to read response: %w", err) + } + + if resp.StatusCode != http.StatusOK { + return fmt.Sprintf("Error: Ollama API returned %d: %s", resp.StatusCode, string(body)), nil + } + + var searchResp struct { + Results []struct { + Title string `json:"title"` + URL string `json:"url"` + Content string `json:"content"` + } `json:"results"` + } + + if err := json.Unmarshal(body, &searchResp); err != nil { + return "", fmt.Errorf("failed to parse response: %w", err) + } + + if len(searchResp.Results) == 0 { + return fmt.Sprintf("No results for: %s", query), nil + } + + var lines []string + lines = append(lines, fmt.Sprintf("Results for: %s", query)) + for i, item := range searchResp.Results { + lines = append(lines, fmt.Sprintf("%d. %s\n %s", i+1, item.Title, item.URL)) + if item.Content != "" { + lines = append(lines, fmt.Sprintf(" %s", item.Content)) + } + } + + return strings.Join(lines, "\n"), nil +} + +// --- Ollama Fetch Tool --- + +type OllamaFetchTool struct { + apiKey string + maxChars int +} + +func NewOllamaFetchTool(apiKey string, maxChars int) *OllamaFetchTool { + if maxChars <= 0 { + maxChars = 50000 + } + return &OllamaFetchTool{ + apiKey: apiKey, + maxChars: maxChars, + } +} + +func (t *OllamaFetchTool) Name() string { + return "web_fetch" +} + +func (t *OllamaFetchTool) Description() string { + return "Fetch a URL and extract readable content using Ollama's Web Fetch API." +} + +func (t *OllamaFetchTool) Parameters() map[string]interface{} { + return map[string]interface{}{ + "type": "object", + "properties": map[string]interface{}{ + "url": map[string]interface{}{ + "type": "string", + "description": "URL to fetch", + }, + }, + "required": []string{"url"}, + } +} + +func (t *OllamaFetchTool) Execute(ctx context.Context, args map[string]interface{}) (string, error) { + urlStr, ok := args["url"].(string) + if !ok { + return "", fmt.Errorf("url is required") + } + + requestBody := map[string]interface{}{ + "url": urlStr, + } + jsonData, err := json.Marshal(requestBody) + if err != nil { + return "", fmt.Errorf("failed to marshal request: %w", err) + } + + req, err := http.NewRequestWithContext(ctx, "POST", "https://ollama.com/api/web_fetch", bytes.NewReader(jsonData)) + if err != nil { + return "", fmt.Errorf("failed to create request: %w", err) + } + + req.Header.Set("Content-Type", "application/json") + if t.apiKey != "" { + req.Header.Set("Authorization", "Bearer "+t.apiKey) + } else { + return "Error: OLLAMA_API_KEY not configured", nil + } + + client := &http.Client{Timeout: 60 * time.Second} + resp, err := client.Do(req) + if err != nil { + return "", fmt.Errorf("request failed: %w", err) + } + defer resp.Body.Close() + + body, err := io.ReadAll(resp.Body) + if err != nil { + return "", fmt.Errorf("failed to read response: %w", err) + } + + if resp.StatusCode != http.StatusOK { + return fmt.Sprintf("Error: Ollama API returned %d: %s", resp.StatusCode, string(body)), nil + } + + var fetchResp struct { + Title string `json:"title"` + Content string `json:"content"` + Links []string `json:"links"` + } + + if err := json.Unmarshal(body, &fetchResp); err != nil { + return "", fmt.Errorf("failed to parse response: %w", err) + } + + text := fetchResp.Content + if len(text) > t.maxChars { + text = text[:t.maxChars] + } + + result := map[string]interface{}{ + "url": urlStr, + "title": fetchResp.Title, + "status": resp.StatusCode, + "extractor": "ollama", + "truncated": len(fetchResp.Content) > t.maxChars, + "length": len(text), + "text": text, + "links": fetchResp.Links, + } + + resultJSON, _ := json.MarshalIndent(result, "", " ") + return string(resultJSON), nil +} + +// --- Original Brave Search Tool --- + type WebSearchTool struct { apiKey string maxResults int @@ -32,11 +266,11 @@ func NewWebSearchTool(apiKey string, maxResults int) *WebSearchTool { } func (t *WebSearchTool) Name() string { - return "web_search" + return "web_search_brave" } func (t *WebSearchTool) Description() string { - return "Search the web for current information. Returns titles, URLs, and snippets from search results." + return "Search the web for current information using Brave Search. Returns titles, URLs, and snippets." } func (t *WebSearchTool) Parameters() map[string]interface{} { @@ -132,6 +366,8 @@ func (t *WebSearchTool) Execute(ctx context.Context, args map[string]interface{} return strings.Join(lines, "\n"), nil } +// --- Original Web Fetch Tool --- + type WebFetchTool struct { maxChars int } @@ -146,11 +382,11 @@ func NewWebFetchTool(maxChars int) *WebFetchTool { } func (t *WebFetchTool) Name() string { - return "web_fetch" + return "web_fetch_raw" } func (t *WebFetchTool) Description() string { - return "Fetch a URL and extract readable content (HTML to text). Use this to get weather info, news, articles, or any web content." + return "Fetch a URL and extract readable content (HTML to text) directly. Use if Ollama fetch fails." } func (t *WebFetchTool) Parameters() map[string]interface{} {