feat: implement semantic search for memory blocks using configurable …

.claude/settings.json

@@ -50,7 +50,26 @@
       "Bash(git check-ignore *)",
       "Bash(git pull *)",
       "Bash(curl -s -o /dev/null -w \"%{http_code}\" https://ai.antonijevic.rs/static/villa-demo.html)",
-      "Bash(curl -s https://openrouter.ai/api/v1/models)"
+      "Bash(curl -s https://openrouter.ai/api/v1/models)",
+      "Bash(sqlite3 my_agent_data.db \"PRAGMA table_info\\(memory_blocks\\);\")",
+      "Read(//private/tmp/**)",
+      "Bash(/Users/ivanantonijevic/development/bt/adk/mate/.venv/bin/python -c \"import hashlib; print\\('sha256 ok:', hashlib.sha256\\(b'x'\\).hexdigest\\(\\)[:8]\\)\")",
+      "Bash(/Users/ivanantonijevic/.pyenv/versions/3.11.11/bin/python -c \"import _blake2\")",
+      "Bash(echo \"--- exit: $?\")",
+      "Bash(/Users/ivanantonijevic/.pyenv/versions/3.11.11/bin/python -c \"import ssl; print\\(ssl.OPENSSL_VERSION\\)\")",
+      "Read(//opt/homebrew/opt/**)",
+      "Bash(brew list *)",
+      "Read(//opt/homebrew/Cellar/**)",
+      "Read(//opt/homebrew/lib/**)",
+      "Bash(brew info *)",
+      "Bash(brew install *)",
+      "Bash(/Users/ivanantonijevic/development/bt/adk/mate/.venv/bin/python -c \"import hashlib; print\\('blake2b ok:', hashlib.blake2b\\(b'x'\\).hexdigest\\(\\)[:8]\\); print\\('sha256 ok:', hashlib.sha256\\(b'x'\\).hexdigest\\(\\)[:8]\\)\")",
+      "Bash(pkill -f \"auth_server.py\")",
+      "Bash(pkill -f \"adk_main.py\")",
+      "Bash(lsof -ti:8000 -ti:8001)",
+      "Bash(docker info *)",
+      "Bash(/Users/ivanantonijevic/development/bt/adk/mate/.venv/bin/python -c ' *)",
+      "Bash(otool -L /Users/ivanantonijevic/.pyenv/versions/3.11.11/lib/python3.11/lib-dynload/_blake2.cpython-311-darwin.so)"
     ]
   }
 }

.env.example

@@ -79,6 +79,14 @@ OPENROUTER_API_KEY=
 # Ollama (local model hosting)
 OLLAMA_API_BASE=http://localhost:11434
 
+# Embedding model for semantic search over memory blocks (litellm format).
+# Any litellm-supported embedding model works, e.g.:
+#   gemini/gemini-embedding-001 (default, uses GOOGLE_API_KEY)
+#   text-embedding-3-small (OpenAI)
+#   mistral/mistral-embed
+# When no key/model is available, search falls back to keyword (LIKE) matching.
+# EMBEDDING_MODEL=gemini/gemini-embedding-001
+
 # Alternative Local API Servers (defaults are used if not set)
 # LM_STUDIO_BASE_URL=http://localhost:1234/v1
 # LLAMACPP_BASE_URL=http://localhost:8080/v1

CLAUDE.md

@@ -96,6 +96,7 @@ Follow PEP 8 with type hints on all function signatures. Use f-strings for forma
 | `TRUSTED_PROXY_HOSTS` | Proxy hosts trusted for `X-Forwarded-*` headers |
 | `ADK_HOST` / `ADK_PORT` | ADK server address (default `127.0.0.1:8001`) |
 | `ARTIFACT_SERVICE` | `local_folder`, `supabase`, or `s3` |
+| `EMBEDDING_MODEL` | Embedding model for memory block semantic search (litellm format, default `gemini/gemini-embedding-001`) |
 | `RATE_LIMIT_ENABLED` | Enable per-user/agent/project budgets |
 | `OTEL_TRACING_ENABLED` | OpenTelemetry distributed tracing (ADK runtime) |
 | `LANGSMITH_TRACING` / `LANGSMITH_API_KEY` | LangSmith run tracing (langgraph runtime) |

documents/DYNAMIC_MEMORY_INSTRUCTIONS.md

@@ -28,6 +28,9 @@ LAZY LOADING PROTOCOL:
 - IF the user asks for "visualization", "frontend data", or "smart object":
 - THEN call `list_shared_blocks(label="smart_object_output_format_json")`.
 - AND use that schema to format your response.
+- IF you need information but do NOT know the exact block label:
+- THEN call `search_shared_blocks(query="<what you are looking for>")` (semantic search by meaning).
+- AND call `get_shared_block(block_id="<label>")` to read the full content of a relevant result.
 
 MEMORY UPDATE PROTOCOL:
 1.  **User Memory (`human_current_user`)**:
@@ -90,6 +93,8 @@ Use a **Namespace Strategy** to separate shared instructions from agent-specific
 
 For large or rarely used instructions (like complex JSON schemas), do NOT load them at startup. Instead, load them only when the user asks for them. Use a distinct prefix like `smart_object_output_format_` to differentiate them from core `system_instruction_` blocks.
 
+When the agent does not know the exact label, it can use `search_shared_blocks(query="...", top_k=5)` — semantic search that ranks blocks by meaning (embeddings via the `EMBEDDING_MODEL` env var, default `gemini/gemini-embedding-001`). Results include a short `value_preview` plus a `search_type` field (`semantic`, or `keyword` when embeddings are unavailable and the tool falls back to LIKE matching); the agent then reads the full block with `get_shared_block`.
+
 **Example Block 4: Smart Object JSON Schema**
 - **Label**: `smart_object_output_format_json`
 - **Value**:

shared/sql/migrations/mysql/V024__memory_block_embeddings.sql

@@ -0,0 +1,7 @@
+-- Migration: Embedding columns on memory_blocks for semantic search
+-- Version: V024
+-- Database: MySQL
+
+ALTER TABLE memory_blocks ADD COLUMN embedding TEXT;
+ALTER TABLE memory_blocks ADD COLUMN embedding_model VARCHAR(200);
+ALTER TABLE memory_blocks ADD COLUMN embedding_hash VARCHAR(64);

shared/sql/migrations/postgresql/V024__memory_block_embeddings.sql

@@ -0,0 +1,7 @@
+-- Migration: Embedding columns on memory_blocks for semantic search
+-- Version: V024
+-- Database: PostgreSQL
+
+ALTER TABLE memory_blocks ADD COLUMN IF NOT EXISTS embedding TEXT;
+ALTER TABLE memory_blocks ADD COLUMN IF NOT EXISTS embedding_model VARCHAR(200);
+ALTER TABLE memory_blocks ADD COLUMN IF NOT EXISTS embedding_hash VARCHAR(64);

shared/sql/migrations/sqlite/V024__memory_block_embeddings.sql

@@ -0,0 +1,7 @@
+-- Migration: Embedding columns on memory_blocks for semantic search
+-- Version: V024
+-- Database: SQLite
+
+ALTER TABLE memory_blocks ADD COLUMN embedding TEXT;
+ALTER TABLE memory_blocks ADD COLUMN embedding_model TEXT;
+ALTER TABLE memory_blocks ADD COLUMN embedding_hash TEXT;

shared/test/test_memory_blocks_semantic_search.py

@@ -0,0 +1,223 @@
+#!/usr/bin/env python3
+"""
+Unit tests for semantic search over memory blocks:
+embedding_service helpers, MemoryBlocksService embedding/backfill,
+and the search_shared_blocks agent tool (semantic + keyword fallback).
+"""
+
+import json
+import unittest
+from unittest.mock import Mock, patch
+import sys
+import os
+
+sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
+
+from sqlalchemy import create_engine
+from sqlalchemy.orm import sessionmaker
+from sqlalchemy.pool import StaticPool
+
+from shared.utils import embedding_service
+from shared.utils.embedding_service import cosine_similarity, embedding_hash_for_text, embedding_text_for_block
+from shared.utils.memory_blocks_service import MemoryBlocksService
+from shared.utils.models import Base, MemoryBlock
+
+
+class FakeDbClient:
+    """Minimal db client exposing get_session() over an in-memory SQLite DB."""
+
+    def __init__(self):
+        self.engine = create_engine(
+            "sqlite://",
+            connect_args={"check_same_thread": False},
+            poolclass=StaticPool,
+        )
+        Base.metadata.create_all(self.engine)
+        self._factory = sessionmaker(bind=self.engine)
+
+    def get_session(self):
+        return self._factory()
+
+
+class TestCosineSimilarity(unittest.TestCase):
+
+    def test_identical_vectors(self):
+        self.assertAlmostEqual(cosine_similarity([1.0, 2.0], [1.0, 2.0]), 1.0)
+
+    def test_orthogonal_vectors(self):
+        self.assertAlmostEqual(cosine_similarity([1.0, 0.0], [0.0, 1.0]), 0.0)
+
+    def test_zero_vector(self):
+        self.assertEqual(cosine_similarity([0.0, 0.0], [1.0, 2.0]), 0.0)
+
+    def test_mismatched_or_empty(self):
+        self.assertEqual(cosine_similarity([1.0], [1.0, 2.0]), 0.0)
+        self.assertEqual(cosine_similarity([], []), 0.0)
+
+
+class TestEmbedTexts(unittest.TestCase):
+
+    def test_returns_none_on_provider_error(self):
+        with patch("litellm.embedding", side_effect=Exception("no api key")):
+            self.assertIsNone(embedding_service.embed_texts(["hello"]))
+
+    def test_returns_vectors_in_input_order(self):
+        response = Mock()
+        response.data = [
+            {"index": 1, "embedding": [0.0, 1.0]},
+            {"index": 0, "embedding": [1.0, 0.0]},
+        ]
+        with patch("litellm.embedding", return_value=response):
+            vectors = embedding_service.embed_texts(["a", "b"])
+        self.assertEqual(vectors, [[1.0, 0.0], [0.0, 1.0]])
+
+    def test_empty_input(self):
+        self.assertEqual(embedding_service.embed_texts([]), [])
+
+
+class TestServiceEmbeddingOnWrite(unittest.TestCase):
+
+    def setUp(self):
+        self.service = MemoryBlocksService(FakeDbClient())
+
+    def _get_row(self, label):
+        session = self.service.db_client.get_session()
+        try:
+            return session.query(MemoryBlock).filter(MemoryBlock.label == label).first()
+        finally:
+            session.close()
+
+    def test_create_block_stores_embedding(self):
+        with patch("shared.utils.embedding_service.embed_texts", return_value=[[1.0, 0.0]]):
+            result = self.service.create_block(project_id=1, label="greeting", value="hello world")
+        self.assertEqual(result["status"], "success")
+        row = self._get_row("greeting")
+        self.assertEqual(json.loads(row.embedding), [1.0, 0.0])
+        self.assertEqual(row.embedding_model, embedding_service.get_embedding_model())
+        expected_hash = embedding_hash_for_text(embedding_text_for_block("greeting", None, "hello world"))
+        self.assertEqual(row.embedding_hash, expected_hash)
+
+    def test_create_block_succeeds_when_embedding_fails(self):
+        with patch("shared.utils.embedding_service.embed_texts", return_value=None):
+            result = self.service.create_block(project_id=1, label="no_embed", value="text")
+        self.assertEqual(result["status"], "success")
+        row = self._get_row("no_embed")
+        self.assertIsNone(row.embedding)
+
+    def test_modify_block_refreshes_embedding(self):
+        with patch("shared.utils.embedding_service.embed_texts", return_value=[[1.0, 0.0]]):
+            self.service.create_block(project_id=1, label="doc", value="old")
+        with patch("shared.utils.embedding_service.embed_texts", return_value=[[0.0, 1.0]]):
+            result = self.service.modify_block(project_id=1, block_id="doc", value="new")
+        self.assertEqual(result["status"], "success")
+        row = self._get_row("doc")
+        self.assertEqual(json.loads(row.embedding), [0.0, 1.0])
+        expected_hash = embedding_hash_for_text(embedding_text_for_block("doc", None, "new"))
+        self.assertEqual(row.embedding_hash, expected_hash)
+
+
+class TestSemanticSearch(unittest.TestCase):
+
+    def setUp(self):
+        self.service = MemoryBlocksService(FakeDbClient())
+        # Blocks created without embeddings — exercises lazy backfill at search time.
+        with patch("shared.utils.embedding_service.embed_texts", return_value=None):
+            self.service.create_block(project_id=1, label="alpha", value="about cats")
+            self.service.create_block(project_id=1, label="beta", value="about dogs")
+            self.service.create_block(project_id=1, label="gamma", value="about kittens")
+            self.service.create_block(project_id=2, label="other_project", value="about cats")
+
+    def test_returns_none_when_query_embedding_unavailable(self):
+        with patch("shared.utils.embedding_service.embed_texts", return_value=None):
+            self.assertIsNone(self.service.semantic_search_blocks(project_id=1, query="cats"))
+
+    def test_backfills_and_ranks_by_similarity(self):
+        def fake_embed(texts):
+            if len(texts) == 1:  # the query
+                return [[1.0, 0.0]]
+            # backfill batch for alpha, beta, gamma (creation order)
+            return [[1.0, 0.0], [0.0, 1.0], [0.7, 0.7]]
+
+        with patch("shared.utils.embedding_service.embed_texts", side_effect=fake_embed):
+            result = self.service.semantic_search_blocks(project_id=1, query="cats", top_k=2)
+
+        self.assertEqual(result["status"], "success")
+        labels = [b["label"] for b in result["blocks"]]
+        self.assertEqual(labels, ["alpha", "gamma"])
+        self.assertGreater(result["blocks"][0]["score"], result["blocks"][1]["score"])
+        self.assertNotIn("value", result["blocks"][0])
+        self.assertEqual(result["blocks"][0]["value_preview"], "about cats")
+
+        # Embeddings persisted: a second search only embeds the query.
+        calls = []
+
+        def count_embed(texts):
+            calls.append(texts)
+            return [[1.0, 0.0]]
+
+        with patch("shared.utils.embedding_service.embed_texts", side_effect=count_embed):
+            result2 = self.service.semantic_search_blocks(project_id=1, query="cats", top_k=2)
+        self.assertEqual(result2["status"], "success")
+        self.assertEqual(len(calls), 1)
+
+    def test_scoped_to_project(self):
+        def fake_embed(texts):
+            return [[1.0, 0.0]] * len(texts)
+
+        with patch("shared.utils.embedding_service.embed_texts", side_effect=fake_embed):
+            result = self.service.semantic_search_blocks(project_id=2, query="cats", top_k=10)
+        self.assertEqual([b["label"] for b in result["blocks"]], ["other_project"])
+
+
+class TestSearchSharedBlocksTool(unittest.TestCase):
+
+    def test_tool_is_registered(self):
+        from shared.utils.tools.memory_blocks_tools import create_memory_blocks_tools_from_config
+        tools = create_memory_blocks_tools_from_config({"project_id": 1})
+        names = [t.__name__ for t in tools]
+        self.assertIn("search_shared_blocks", names)
+
+    def test_semantic_path(self):
+        mock_service = Mock()
+        mock_service.semantic_search_blocks.return_value = {
+            "status": "success",
+            "blocks": [{"label": "alpha", "value_preview": "x", "score": 0.9}],
+            "block_count": 1,
+        }
+        from shared.utils.tools.memory_blocks_tools import create_memory_blocks_tools_from_config
+        with patch("shared.utils.tools.memory_blocks_tools._get_service", return_value=mock_service):
+            tools = create_memory_blocks_tools_from_config({"project_id": 1})
+            tool = next(t for t in tools if t.__name__ == "search_shared_blocks")
+            result = tool(query="cats")
+        self.assertEqual(result["search_type"], "semantic")
+        self.assertEqual(result["blocks"][0]["label"], "alpha")
+
+    def test_keyword_fallback(self):
+        mock_service = Mock()
+        mock_service.semantic_search_blocks.return_value = None
+        mock_service.list_blocks.return_value = {
+            "status": "success",
+            "blocks": [{"label": "alpha", "value": "v" * 300, "description": None}],
+            "block_count": 1,
+        }
+        from shared.utils.tools.memory_blocks_tools import create_memory_blocks_tools_from_config
+        with patch("shared.utils.tools.memory_blocks_tools._get_service", return_value=mock_service):
+            tools = create_memory_blocks_tools_from_config({"project_id": 1})
+            tool = next(t for t in tools if t.__name__ == "search_shared_blocks")
+            result = tool(query="cats")
+        self.assertEqual(result["status"], "success")
+        self.assertEqual(result["search_type"], "keyword")
+        self.assertEqual(len(result["blocks"]), 1)  # deduped across label/value search
+        self.assertEqual(len(result["blocks"][0]["value_preview"]), 200)
+        self.assertNotIn("value", result["blocks"][0])
+
+    def test_empty_query(self):
+        from shared.utils.tools.memory_blocks_tools import create_memory_blocks_tools_from_config
+        tools = create_memory_blocks_tools_from_config({"project_id": 1})
+        tool = next(t for t in tools if t.__name__ == "search_shared_blocks")
+        result = tool(query="  ")
+        self.assertEqual(result["status"], "error")
+
+
+if __name__ == "__main__":
+    unittest.main()

shared/utils/embedding_service.py

@@ -0,0 +1,74 @@
+"""
+Embedding helpers for semantic search over memory blocks.
+
+Model is configurable via the EMBEDDING_MODEL env var (any litellm-supported
+embedding model). Vectors are stored as JSON text on memory_blocks and ranked
+with pure-Python cosine similarity — portable across sqlite/postgresql/mysql.
+"""
+
+import hashlib
+import logging
+import math
+import os
+from typing import List, Optional
+
+logger = logging.getLogger(__name__)
+
+DEFAULT_EMBEDDING_MODEL = "gemini/gemini-embedding-001"
+
+# Keep the embedded text bounded; embedding APIs have input limits and
+# block values can hold whole crawled pages.
+MAX_EMBED_CHARS = 8000
+
+
+def get_embedding_model() -> str:
+    """Embedding model name, configurable via EMBEDDING_MODEL env var."""
+    return os.getenv("EMBEDDING_MODEL", DEFAULT_EMBEDDING_MODEL)
+
+
+def embedding_text_for_block(label: str, description: Optional[str], value: str) -> str:
+    """Canonical text embedded for a memory block."""
+    return f"{label}\n{description or ''}\n{(value or '')[:MAX_EMBED_CHARS]}"
+
+
+def embedding_hash_for_text(text: str) -> str:
+    """Stable hash used to detect stale embeddings after value changes."""
+    return hashlib.sha256(text.encode("utf-8")).hexdigest()
+
+
+def embed_texts(texts: List[str]) -> Optional[List[List[float]]]:
+    """Embed texts via litellm. Returns None on any failure (caller falls back)."""
+    if not texts:
+        return []
+    try:
+        import litellm
+
+        # litellm's Gemini provider reads GEMINI_API_KEY; MATE configures GOOGLE_API_KEY.
+        if not os.getenv("GEMINI_API_KEY") and os.getenv("GOOGLE_API_KEY"):
+            os.environ["GEMINI_API_KEY"] = os.environ["GOOGLE_API_KEY"]
+
+        response = litellm.embedding(model=get_embedding_model(), input=texts)
+        vectors: List[Optional[List[float]]] = [None] * len(texts)
+        for item in response.data:
+            index = item["index"] if isinstance(item, dict) else item.index
+            vector = item["embedding"] if isinstance(item, dict) else item.embedding
+            vectors[index] = list(vector)
+        if any(v is None for v in vectors):
+            logger.warning("Embedding response missing vectors for some inputs")
+            return None
+        return vectors  # type: ignore[return-value]
+    except Exception as e:
+        logger.warning(f"Embedding failed (model={get_embedding_model()}): {e}")
+        return None
+
+
+def cosine_similarity(a: List[float], b: List[float]) -> float:
+    """Cosine similarity in pure Python; 0.0 for mismatched or zero vectors."""
+    if not a or not b or len(a) != len(b):
+        return 0.0
+    dot = sum(x * y for x, y in zip(a, b))
+    norm_a = math.sqrt(sum(x * x for x in a))
+    norm_b = math.sqrt(sum(y * y for y in b))
+    if norm_a == 0.0 or norm_b == 0.0:
+        return 0.0
+    return dot / (norm_a * norm_b)