opencode/picoclaw
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

feat: add Ollama search tools and update LLM providers

Browse source
This commit is contained in:
instax-dutta 2026-02-12 18:04:46 +05:30
parent 8968d58fed
commit 3adb59b44d
4 changed files with 283 additions and 10 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

14
config.example.json
View file

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

12
pkg/agent/loop.go
View file

@ -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()

21
pkg/config/config.go
View file

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

246
pkg/tools/web.go
View file

@ -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{} {