feat: vector similarity tool filtering via Gemini embeddings

At startup, embeddings are computed for all MCP tools using
gemini-embedding-001 and cached in memory. Each request embeds
the user query and selects top-K most relevant tools via cosine
similarity, reducing LLM context size without keyword maintenance.

Falls back to all tools if embeddings unavailable or top_k=0.
Configured via models.embedding and tool_filter.top_k in config.yaml.

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>

Author: Alexey Panfilov

CLAUDE.md

@@ -52,6 +52,7 @@ flowchart LR
 - Connects on startup (`initialize` → `tools/list`). Supports JSON and SSE responses.
 - Per-server tool filtering: `allowTools` (allowlist) checked after `denyTools` (blocklist)
 - Auth via generic `headers` map (Claude Desktop format)
+- **Vector tool filtering**: `EnableEmbeddings(apiKey, model, topK)` + `EmbedTools(ctx)` at startup computes Gemini embeddings for all tools; `LLMToolsForQuery(ctx, query)` returns top-K most relevant tools per request via cosine similarity. Falls back to all tools if embeddings unavailable or query is empty.
 
 **`internal/agent`** — agentic loop.
 - `Process(ctx, chatID, llm.Message, onToolCall)` — prepends `"Current date and time: ..."` to the system prompt on every call using `time.Now()` (respects `TZ` env var set in Docker)
@@ -68,7 +69,7 @@ flowchart LR
 | File | Purpose |
 |---|---|
 | `.env` | Secrets: `TELEGRAM_BOT_TOKEN`, `DEEPSEEK_API_KEY`, `GEMINI_API_KEY`, `TELEGRAM_OWNER_CHAT_ID`, `TZ` (default `Europe/Belgrade`) — auto-loaded by Docker Compose from project root |
-| `config/config.yaml` | Models (all require `base_url`), routing, Telegram IDs — `${ENV_VAR}` substitution |
+| `config/config.yaml` | Models (all require `base_url`; `embedding` model is exception — no `base_url`/`max_tokens`), routing, tool_filter, Telegram IDs — `${ENV_VAR}` substitution |
 | `config/mcp.json` | MCP servers in Claude Desktop format — `allowTools`, `denyTools` per server |
 | `config/system_prompt.md` | System prompt injected on every LLM request |
 
@@ -81,6 +82,15 @@ Paths are hardcoded in `main.go` as `config/config.yaml`, `config/system_prompt.
 3. **Primary** (DeepSeek Chat) — default
 4. **Fallback** (Gemini 3.1 Flash Lite) — primary returns 5xx/429/network error
 
+### Tool filtering (vector similarity)
+
+Configured via `tool_filter.top_k` in `config.yaml` and `models.embedding` (Gemini `gemini-embedding-001`).
+
+- At startup: embeddings computed for all tools (`name + ": " + description`) and cached in memory
+- Per request: user message embedded → cosine similarity → top-K tools sent to LLM
+- Fallback to all tools if: embeddings not ready, `top_k=0`, `top_k >= total tools`, or embed API error
+- `top_k: 0` disables filtering entirely
+
 ### SQLite schema notes
 
 Flag columns: `is_reset`, `is_compacted`, `is_summary`, `parts` (JSON). Queries always filter by `id > lastResetID`. `GetHistory` returns last 30 non-compacted messages. `ALTER TABLE ADD COLUMN parts` runs at startup for migration of existing DBs.

README.md

@@ -8,7 +8,7 @@ A lightweight Telegram bot that acts as a personal AI assistant. Written in Go 
 - **Image support** — send a photo (with or without caption) and it's routed automatically to the vision model
 - **Forwarded messages** — forward any message (text, photo, link) to the bot, then ask your question; messages arriving within 2 s are batched automatically
 - **Link extraction** — hidden hyperlinks (`text_link` entities) in forwarded messages are surfaced as plain URLs for the LLM
-- **MCP tool support** — connects to any MCP-compatible server (HTTP/SSE), same `mcp.json` format as Claude Desktop; per-server `allowTools`/`denyTools` filtering
+- **MCP tool support** — connects to any MCP-compatible server (HTTP/SSE), same `mcp.json` format as Claude Desktop; per-server `allowTools`/`denyTools` filtering; vector similarity filtering selects only the most relevant tools per request
 - **Persistent memory** — SQLite-backed conversation history with automatic session management
 - **Context compaction** — auto-summarises old history to stay within token limits
 - **Rich formatting** — Markdown converted to Telegram HTML; responses ≥ 4096 chars sent as `response.md`
@@ -20,7 +20,7 @@ A lightweight Telegram bot that acts as a personal AI assistant. Written in Go 
 - Go 1.24+ (or Docker)
 - [Telegram Bot Token](https://t.me/BotFather)
 - [DeepSeek API key](https://platform.deepseek.com)
-- Gemini API key (optional — for fallback and image support)
+- Gemini API key (optional — for fallback, image support, and tool filtering)
 
 ## Quick start (NAS / Pi / server)
 
@@ -131,6 +131,18 @@ models:
     api_key: ${GEMINI_API_KEY}
     max_tokens: 4096
     base_url: https://generativelanguage.googleapis.com/v1beta/openai/
+  embedding:
+    provider: gemini
+    model: gemini-embedding-001
+    api_key: ${GEMINI_API_KEY}
+
+routing:
+  default: default
+  fallback: flash_lite
+  compaction_model: default
+
+tool_filter:
+  top_k: 20   # top-K tools selected per request via vector similarity; 0 = disabled
 ```
 
 ### `config/mcp.json`
@@ -197,6 +209,7 @@ flowchart TD
     Router["LLM Router"]
     MCP["MCP Client"]
     Store[("SQLite")]
+    Emb["Gemini Embedding\ngemini-embedding-001"]
 
     subgraph LLMs ["LLM Providers"]
         DS["DeepSeek Chat"]
@@ -215,7 +228,9 @@ flowchart TD
     Handler --> Agent
     Agent <--> Store
     Agent --> Router
-    Agent <-->|"tools/call"| MCP
+    Agent -->|"query embed\ntop-K filter"| MCP
+    MCP <-->|"embed tools\nat startup"| Emb
+    MCP <-->|"tools/call"| Servers
     Router --> DS
     Router --> DSR
     Router --> GL

cmd/agent/main.go

@@ -98,6 +98,13 @@ func main() {
 		initCtx, cancel := context.WithTimeout(context.Background(), 15*time.Second)
 		mcpClient.Initialize(initCtx)
 		cancel()
+
+		if cfg.ToolFilter.TopK > 0 && cfg.Models.Embedding.APIKey != "" {
+			mcpClient.EnableEmbeddings(cfg.Models.Embedding.APIKey, cfg.Models.Embedding.Model, cfg.ToolFilter.TopK)
+			embedCtx, embedCancel := context.WithTimeout(context.Background(), 60*time.Second)
+			mcpClient.EmbedTools(embedCtx)
+			embedCancel()
+		}
 	}
 
 	ag := agent.New(router, s, mcpClient, compacter, string(sysPromptBytes), logger)

config/config.yaml

@@ -29,8 +29,15 @@ models:
     api_key: ${GEMINI_API_KEY}
     max_tokens: 4096
     base_url: https://generativelanguage.googleapis.com/v1beta/openai/
+  embedding:
+    provider: gemini
+    model: gemini-embedding-001
+    api_key: ${GEMINI_API_KEY}
 
 routing:
   default: default
   fallback: flash_lite
   compaction_model: default
+
+tool_filter:
+  top_k: 20

internal/agent/agent.go

@@ -4,6 +4,7 @@ import (
 	"context"
 	"fmt"
 	"log/slog"
+	"strings"
 	"time"
 
 	"telegram-agent/internal/llm"
@@ -47,7 +48,7 @@ func (a *Agent) Process(ctx context.Context, chatID int64, userMsg llm.Message,
 
 	var tools []llm.Tool
 	if a.mcp != nil {
-		tools = a.mcp.LLMTools()
+		tools = a.mcp.LLMToolsForQuery(ctx, messageText(userMsg))
 	}
 
 	for i := 0; i < maxToolIterations; i++ {
@@ -131,3 +132,17 @@ func (a *Agent) ListTools() []ToolInfo {
 	}
 	return result
 }
+
+// messageText extracts plain text from a message (handles both Content and Parts).
+func messageText(msg llm.Message) string {
+	if msg.Content != "" {
+		return msg.Content
+	}
+	var sb strings.Builder
+	for _, p := range msg.Parts {
+		if p.Type == "text" {
+			sb.WriteString(p.Text)
+		}
+	}
+	return sb.String()
+}

internal/config/config.go

@@ -9,9 +9,14 @@ import (
 )
 
 type Config struct {
-	Telegram TelegramConfig `yaml:"telegram"`
-	Models   ModelsConfig   `yaml:"models"`
-	Routing  RoutingConfig  `yaml:"routing"`
+	Telegram   TelegramConfig   `yaml:"telegram"`
+	Models     ModelsConfig     `yaml:"models"`
+	Routing    RoutingConfig    `yaml:"routing"`
+	ToolFilter ToolFilterConfig `yaml:"tool_filter"`
+}
+
+type ToolFilterConfig struct {
+	TopK int `yaml:"top_k"` // 0 = disabled
 }
 
 type MCPServerConfig struct {
@@ -40,6 +45,7 @@ type ModelsConfig struct {
 	Reasoner   ModelConfig `yaml:"reasoner"`
 	FlashLite  ModelConfig `yaml:"flash_lite"`
 	Multimodal ModelConfig `yaml:"multimodal"`
+	Embedding  ModelConfig `yaml:"embedding"`
 }
 
 type RoutingConfig struct {

internal/mcp/client.go

@@ -9,6 +9,7 @@ import (
 	"io"
 	"log/slog"
 	"net/http"
+	"sort"
 	"strings"
 	"sync/atomic"
 	"time"
@@ -19,10 +20,14 @@ import (
 
 // Client manages connections to multiple MCP servers.
 type Client struct {
-	servers     map[string]*server
-	tools       []Tool
-	toolServers map[string]string // tool name → server name
-	logger      *slog.Logger
+	servers         map[string]*server
+	tools           []Tool
+	toolServers     map[string]string // tool name → server name
+	logger          *slog.Logger
+	embeddingAPIKey   string
+	embeddingModel    string
+	topK              int
+	embeddingsReady   bool
 }
 
 type server struct {
@@ -42,6 +47,7 @@ type Tool struct {
 	Description string
 	InputSchema json.RawMessage
 	ServerName  string
+	Embedding   []float32
 }
 
 func NewClient(configs map[string]config.MCPServerConfig, logger *slog.Logger) *Client {
@@ -102,6 +108,77 @@ func (c *Client) Initialize(ctx context.Context) {
 	}
 }
 
+// EnableEmbeddings configures vector-based tool filtering.
+// Must be called before EmbedTools.
+func (c *Client) EnableEmbeddings(apiKey, model string, topK int) {
+	c.embeddingAPIKey = apiKey
+	c.embeddingModel = model
+	c.topK = topK
+}
+
+// EmbedTools computes and caches embeddings for all tools.
+// Should be called once after Initialize.
+func (c *Client) EmbedTools(ctx context.Context) {
+	if c.embeddingAPIKey == "" {
+		return
+	}
+	ok := 0
+	for i := range c.tools {
+		text := c.tools[i].Name + ": " + c.tools[i].Description
+		emb, err := embed(ctx, c.embeddingAPIKey, c.embeddingModel, text)
+		if err != nil {
+			c.logger.Warn("failed to embed tool", "tool", c.tools[i].Name, "err", err)
+			continue
+		}
+		c.tools[i].Embedding = emb
+		ok++
+	}
+	if ok == len(c.tools) {
+		c.embeddingsReady = true
+		c.logger.Info("tool embeddings ready", "tools", ok)
+	} else {
+		c.logger.Warn("some tool embeddings failed, disabling filtering", "ok", ok, "total", len(c.tools))
+	}
+}
+
+// LLMToolsForQuery returns the top-K most relevant tools for the given query.
+// Falls back to all tools if embeddings are not ready or topK >= total tools.
+func (c *Client) LLMToolsForQuery(ctx context.Context, query string) []llm.Tool {
+	if !c.embeddingsReady || c.topK <= 0 || c.topK >= len(c.tools) || query == "" {
+		return c.LLMTools()
+	}
+
+	queryEmb, err := embed(ctx, c.embeddingAPIKey, c.embeddingModel, query)
+	if err != nil {
+		c.logger.Warn("query embedding failed, using all tools", "err", err)
+		return c.LLMTools()
+	}
+
+	type scored struct {
+		tool  Tool
+		score float64
+	}
+	candidates := make([]scored, len(c.tools))
+	for i, t := range c.tools {
+		candidates[i] = scored{t, cosineSimilarity(queryEmb, t.Embedding)}
+	}
+	sort.Slice(candidates, func(i, j int) bool {
+		return candidates[i].score > candidates[j].score
+	})
+
+	result := make([]llm.Tool, 0, c.topK)
+	for i := 0; i < c.topK; i++ {
+		t := candidates[i].tool
+		result = append(result, llm.Tool{
+			Name:        t.Name,
+			Description: t.Description,
+			InputSchema: t.InputSchema,
+		})
+	}
+	c.logger.Debug("tool filter applied", "query_len", len(query), "selected", len(result), "total", len(c.tools))
+	return result
+}
+
 // LLMTools returns tools in the format expected by the LLM provider.
 func (c *Client) LLMTools() []llm.Tool {
 	result := make([]llm.Tool, 0, len(c.tools))

internal/mcp/embeddings.go

@@ -0,0 +1,74 @@
+package mcp
+
+import (
+	"bytes"
+	"context"
+	"encoding/json"
+	"fmt"
+	"io"
+	"math"
+	"net/http"
+	"time"
+)
+
+const geminiEmbeddingBaseURL = "https://generativelanguage.googleapis.com/v1beta/models/%s:embedContent"
+
+func embed(ctx context.Context, apiKey, model, text string) ([]float32, error) {
+	reqBody := map[string]any{
+		"model": "models/" + model,
+		"content": map[string]any{
+			"parts": []map[string]any{{"text": text}},
+		},
+	}
+	body, err := json.Marshal(reqBody)
+	if err != nil {
+		return nil, err
+	}
+
+	url := fmt.Sprintf(geminiEmbeddingBaseURL, model)
+	req, err := http.NewRequestWithContext(ctx, "POST", url, bytes.NewReader(body))
+	if err != nil {
+		return nil, err
+	}
+	req.Header.Set("Content-Type", "application/json")
+	req.Header.Set("x-goog-api-key", apiKey)
+
+	client := &http.Client{Timeout: 10 * time.Second}
+	resp, err := client.Do(req)
+	if err != nil {
+		return nil, err
+	}
+	defer resp.Body.Close()
+
+	if resp.StatusCode != 200 {
+		bodyBytes, _ := io.ReadAll(resp.Body)
+		return nil, fmt.Errorf("embedding API returned %d: %s", resp.StatusCode, string(bodyBytes))
+	}
+
+	var result struct {
+		Embedding struct {
+			Values []float32 `json:"values"`
+		} `json:"embedding"`
+	}
+	if err := json.NewDecoder(resp.Body).Decode(&result); err != nil {
+		return nil, err
+	}
+	return result.Embedding.Values, nil
+}
+
+func cosineSimilarity(a, b []float32) float64 {
+	if len(a) != len(b) || len(a) == 0 {
+		return 0
+	}
+	var dot, normA, normB float64
+	for i := range a {
+		dot += float64(a[i]) * float64(b[i])
+		normA += float64(a[i]) * float64(a[i])
+		normB += float64(b[i]) * float64(b[i])
+	}
+	denom := math.Sqrt(normA) * math.Sqrt(normB)
+	if denom == 0 {
+		return 0
+	}
+	return dot / denom
+}