diff --git a/config/settings.yaml b/config/settings.yaml index 6a726a4..83dd254 100644 --- a/config/settings.yaml +++ b/config/settings.yaml @@ -9,6 +9,8 @@ embedder: provider: ollama model: nomic-embed-text dimensions: 768 + max_chars: 4000 + chunk_overlap: 200 ollama_base_url: http://ollama:11434 mcp: diff --git a/src/api/main.py b/src/api/main.py index 79dc760..31a54c0 100644 --- a/src/api/main.py +++ b/src/api/main.py @@ -30,6 +30,7 @@ ) from ..db.queries import get_memory_stats from ..embedder import create_embedding +from ..embedding_regeneration import regenerate_memory_embedding from ..extractors.entities import extract_entities from ..extractors.tagger import get_tagger from ..version import get_version @@ -105,6 +106,10 @@ class SearchRequest(BaseModel): captured_by: Optional[List[str]] = None +class EmbeddingRegenerationRequest(BaseModel): + force: bool = False + + @app.get("/") async def root(): """Root endpoint.""" @@ -210,6 +215,31 @@ async def get_memory(memory_id: str): return memory +@app.post("/memories/{memory_id}/regenerate-embedding", response_model=dict) +async def regenerate_embedding_endpoint( + memory_id: str, + request: EmbeddingRegenerationRequest, +): + """Regenerate a failed embedding or force replacement during model migration.""" + import uuid + + try: + parsed_id = uuid.UUID(memory_id) + except ValueError as exc: + raise HTTPException(status_code=400, detail=f"Invalid UUID: {memory_id}") from exc + + try: + return regenerate_memory_embedding(parsed_id, force=request.force) + except KeyError as exc: + raise HTTPException(status_code=404, detail="Memory not found") from exc + except Exception as exc: + logger.warning("embedding regeneration failed: %s", type(exc).__name__) + raise HTTPException( + status_code=503, + detail="Embedding provider unavailable", + ) from exc + + @app.post("/memories/search", response_model=List[MemoryResponse]) async def search_memories_endpoint(search: SearchRequest): """Search memories by semantic content and structured filters.""" diff --git a/src/db/attribution.py b/src/db/attribution.py index b13afd2..e261d18 100644 --- a/src/db/attribution.py +++ b/src/db/attribution.py @@ -164,6 +164,42 @@ def get_memory_by_id(memory_id: uuid.UUID) -> Optional[Dict[str, Any]]: return _decode_memory(row) if row else None +def get_memory_embedding_target(memory_id: uuid.UUID) -> Optional[Dict[str, Any]]: + """Return content and whether a memory already has an embedding.""" + with get_db_cursor() as cursor: + cursor.execute( + """ + SELECT content, embedding IS NOT NULL AS has_embedding + FROM memory + WHERE id = %s + """, + (memory_id,), + ) + row = cursor.fetchone() + return dict(row) if row else None + + +def update_memory_embedding( + memory_id: uuid.UUID, + embedding: List[float], + *, + force: bool = False, +) -> bool: + """Set a memory embedding, overwriting an existing vector only when forced.""" + with get_db_cursor() as cursor: + cursor.execute( + """ + UPDATE memory + SET embedding = %s + WHERE id = %s + AND (%s OR embedding IS NULL) + RETURNING id + """, + (embedding, memory_id, force), + ) + return cursor.fetchone() is not None + + def get_recent_memories( limit: int = 50, offset: int = 0, diff --git a/src/db/migrations/015_embedding_dim_change.sql b/src/db/migrations/015_embedding_dim_change.sql index 7ceaf9b..6b74ab5 100644 --- a/src/db/migrations/015_embedding_dim_change.sql +++ b/src/db/migrations/015_embedding_dim_change.sql @@ -6,14 +6,17 @@ -- local Ollama service with nomic-embed-text (768 dimensions). The embedding -- column must match so semantic search can store and compare vectors. -- --- This migration is safe on a fresh database that has not yet stored real --- embeddings. On an installation with previously-stored 1536-dim vectors the --- column would need to be dropped and re-created (loss of prior embeddings); --- that has been intentionally avoided here. +-- The legacy memory table is installed from schema.sql rather than the v2 +-- migration ledger, so a fresh migration-only database may not have it yet. +-- Existing installations still receive the dimension change and index rebuild. +DO $$ +BEGIN + IF to_regclass('public.memory') IS NOT NULL THEN + ALTER TABLE memory ALTER COLUMN embedding TYPE vector(768); -ALTER TABLE memory ALTER COLUMN embedding TYPE vector(768); - --- The HNSW index on embedding needs to be rebuilt to match the new dimension. -DROP INDEX IF EXISTS idx_memory_embedding; -CREATE INDEX IF NOT EXISTS idx_memory_embedding - ON memory USING hnsw (embedding vector_cosine_ops); + -- The HNSW index on embedding needs to be rebuilt to match the new dimension. + DROP INDEX IF EXISTS idx_memory_embedding; + CREATE INDEX IF NOT EXISTS idx_memory_embedding + ON memory USING hnsw (embedding vector_cosine_ops); + END IF; +END $$; diff --git a/src/embedder/__init__.py b/src/embedder/__init__.py index d4e6b1d..2eb6cfb 100644 --- a/src/embedder/__init__.py +++ b/src/embedder/__init__.py @@ -29,6 +29,7 @@ dimensions: 768 ``` """ +import math import os from typing import ClassVar, List, Optional from abc import ABC, abstractmethod @@ -37,6 +38,10 @@ import yaml +DEFAULT_MAX_CHARS = 4000 +DEFAULT_CHUNK_OVERLAP = 200 + + class EmbedderConfig: """Configuration for embedder with multiple provider support.""" @@ -63,6 +68,20 @@ def __init__(self, config_path: str = None): # Dimensions self.dimensions = embedder_cfg.get('dimensions', 768) + + # Chunking guardrail for providers with smaller context windows + self.max_chars = int(os.environ.get( + 'EMBEDDING_MAX_CHARS', + embedder_cfg.get('max_chars', DEFAULT_MAX_CHARS) + )) + self.chunk_overlap = int(os.environ.get( + 'EMBEDDING_CHUNK_OVERLAP', + embedder_cfg.get('chunk_overlap', DEFAULT_CHUNK_OVERLAP) + )) + if self.max_chars <= 0: + raise ValueError("embedder.max_chars must be greater than zero") + if self.chunk_overlap < 0 or self.chunk_overlap >= self.max_chars: + raise ValueError("embedder.chunk_overlap must be >= 0 and < max_chars") # OpenRouter (default) self.openrouter_api_key = os.environ.get( @@ -458,14 +477,63 @@ def get_embedder(config_path: str = None) -> BaseEmbedder: return _embedder +def _chunk_text(text: str, max_chars: int, overlap: int) -> List[str]: + """Split text into bounded overlapping chunks.""" + if len(text) <= max_chars: + return [text] + + chunks = [] + start = 0 + while start < len(text): + end = min(start + max_chars, len(text)) + chunks.append(text[start:end]) + if end == len(text): + break + start = end - overlap + return chunks + + +def _pool_embeddings(embeddings: List[List[float]]) -> List[float]: + """Mean-pool chunk vectors and L2-normalize the result.""" + if not embeddings or not embeddings[0]: + raise ValueError("embedding provider returned an empty vector") + + dimensions = len(embeddings[0]) + if any(len(embedding) != dimensions for embedding in embeddings): + raise ValueError("embedding provider returned inconsistent dimensions") + + pooled = [ + sum(embedding[index] for embedding in embeddings) / len(embeddings) + for index in range(dimensions) + ] + norm = math.sqrt(sum(value * value for value in pooled)) + if norm == 0: + raise ValueError("embedding provider returned only zero vectors") + return [value / norm for value in pooled] + + def create_embedding(text: str) -> List[float]: - """Convenience function to create an embedding.""" - return get_embedder().embed(text) + """Convenience function to create an embedding with long-input chunking.""" + embedder = get_embedder() + config = getattr(embedder, 'config', None) + if config is None: + config = EmbedderConfig.get_instance() + max_chars = getattr(config, 'max_chars', DEFAULT_MAX_CHARS) + overlap = getattr(config, 'chunk_overlap', DEFAULT_CHUNK_OVERLAP) + chunks = _chunk_text(text, max_chars, overlap) + + if len(chunks) == 1: + return embedder.embed(chunks[0]) + + # Embed chunks individually so provider failures are never hidden by a + # provider-specific batch fallback (notably Ollama's zero-vector fallback). + embeddings = [embedder.embed(chunk) for chunk in chunks] + return _pool_embeddings(embeddings) def create_embeddings(texts: List[str]) -> List[List[float]]: """Convenience function to create multiple embeddings.""" - return get_embedder().embed_batch(texts) + return [create_embedding(text) for text in texts] # For backward compatibility diff --git a/src/embedding_regeneration.py b/src/embedding_regeneration.py new file mode 100644 index 0000000..746c9b9 --- /dev/null +++ b/src/embedding_regeneration.py @@ -0,0 +1,46 @@ +"""Shared embedding regeneration workflow for REST and MCP surfaces.""" +from __future__ import annotations + +import uuid +from typing import Any, Dict + +from .db.attribution import get_memory_embedding_target, update_memory_embedding +from .embedder import create_embedding + + +def regenerate_memory_embedding( + memory_id: uuid.UUID, + *, + force: bool = False, +) -> Dict[str, Any]: + """Regenerate one memory embedding without changing memory identity or timestamps. + + By default an existing embedding is left untouched. ``force=True`` is intended + for embedding provider/model migrations where all stored vectors must be rebuilt. + """ + target = get_memory_embedding_target(memory_id) + if target is None: + raise KeyError(str(memory_id)) + + if target["has_embedding"] and not force: + return { + "id": str(memory_id), + "status": "unchanged", + "reason": "embedding_exists", + } + + embedding = create_embedding(target["content"]) + updated = update_memory_embedding(memory_id, embedding, force=force) + if not updated: + # A concurrent writer may have filled a NULL embedding after our read. + return { + "id": str(memory_id), + "status": "unchanged", + "reason": "embedding_exists", + } + + return { + "id": str(memory_id), + "status": "regenerated", + "force": force, + } diff --git a/src/extractors/entities.py b/src/extractors/entities.py index 1fc7435..116756b 100644 --- a/src/extractors/entities.py +++ b/src/extractors/entities.py @@ -42,7 +42,7 @@ def _ensure_nltk_data(): # Network or disk failure is non-fatal; the chunker/tokenizer # will raise a clearer error at the call site if the data # is truly missing. - pass + continue _ensure_nltk_data() diff --git a/src/main.py b/src/main.py index ee38ad2..77f102e 100644 --- a/src/main.py +++ b/src/main.py @@ -20,6 +20,7 @@ get_memory_stats, ) from .embedder import create_embedding +from .embedding_regeneration import regenerate_memory_embedding from .extractors.entities import extract_entities from .extractors.tagger import auto_tag from .analytics.weekly_report import generate_weekly_report @@ -89,6 +90,21 @@ async def list_tools() -> List[Tool]: "required": ["content"], }, ), + Tool( + name="memory_regenerate_embedding", + description=( + "Regenerate a missing embedding. Set force=true to replace an " + "existing embedding after changing provider or model." + ), + inputSchema={ + "type": "object", + "properties": { + "memory_id": {"type": "string", "description": "UUID of the memory"}, + "force": {"type": "boolean", "description": "Replace an existing embedding", "default": False}, + }, + "required": ["memory_id"], + }, + ), Tool( name="memory_get_related", description="Get memories related to a specific memory by ID.", @@ -150,6 +166,8 @@ async def call_tool(name: str, arguments: Any) -> List[TextContent]: return await handle_memory_search(arguments) if name == "memory_store": return await handle_memory_store(arguments) + if name == "memory_regenerate_embedding": + return await handle_memory_regenerate_embedding(arguments) if name == "memory_get_related": return await handle_memory_get_related(arguments) if name == "memory_get_entity": @@ -236,6 +254,23 @@ async def handle_memory_store(args: Dict) -> List[TextContent]: )] +async def handle_memory_regenerate_embedding(args: Dict) -> List[TextContent]: + """Handle memory_regenerate_embedding tool.""" + import uuid + + memory_id = args["memory_id"] + try: + parsed_id = uuid.UUID(memory_id) + except ValueError: + return [TextContent(type="text", text=f"Invalid memory ID: {memory_id}")] + + try: + result = regenerate_memory_embedding(parsed_id, force=args.get("force", False)) + except KeyError: + return [TextContent(type="text", text=f"Memory not found: {memory_id}")] + return [TextContent(type="text", text=str(result))] + + async def handle_memory_get_related(args: Dict) -> List[TextContent]: """Handle memory_get_related tool.""" memory_id = args["memory_id"] diff --git a/tests/e2e/test_api.py b/tests/e2e/test_api.py index ee33d9c..ab1ef55 100644 --- a/tests/e2e/test_api.py +++ b/tests/e2e/test_api.py @@ -257,7 +257,7 @@ def test_continuity(): s, _ = req("POST", f"/v1/sessions/{S.session_id}/close", { "summary": "E2E test session — verified identity resolution, event recording, and continuity.", }, expect=(200, 201, 204)) - assert s in (200, 201, 204), f"session close failed" + assert s in (200, 201, 204), "session close failed" def test_context(): @@ -336,7 +336,7 @@ def _copy_markdown_to_container(container_md: str) -> None: ) result = subprocess.run( ["docker", "cp", str(host_md), f"openbrain-api:{container_md}"], - capture_output=True, text=True, timeout=30, + capture_output=True, text=True, timeout=30, check=False, ) if result.returncode != 0: raise FileNotFoundError( @@ -494,4 +494,4 @@ def _parse_argv() -> None: # print FAIL but don't always raise). The summary prints both. test_summary() failures += S.failed - sys.exit(0 if failures == 0 else 1) + sys.exit(0 if failures == 0 else 1) \ No newline at end of file diff --git a/tests/test_embedding_resilience.py b/tests/test_embedding_resilience.py new file mode 100644 index 0000000..f9aa491 --- /dev/null +++ b/tests/test_embedding_resilience.py @@ -0,0 +1,172 @@ +"""Regression tests for oversized and failed embedding recovery.""" +from contextlib import contextmanager +from types import SimpleNamespace +from unittest.mock import Mock, patch +import asyncio +import uuid + +import pytest +import requests + + +class _FakeEmbedder: + def __init__(self, *, max_chars=8, overlap=2, fail_on_call=None): + self.config = SimpleNamespace( + max_chars=max_chars, + chunk_overlap=overlap, + ) + self.calls = [] + self.fail_on_call = fail_on_call + + def embed(self, text): + self.calls.append(text) + if self.fail_on_call == len(self.calls): + raise requests.ConnectionError("provider down") + return [3.0, 4.0] + + +def _cursor(*, row=None): + cursor = Mock() + cursor.fetchone.return_value = row + + @contextmanager + def manager(): + yield cursor + + return cursor, manager + + +def test_long_embedding_is_chunked_and_pooled(): + from src import embedder + + fake = _FakeEmbedder(max_chars=8, overlap=2) + with patch.object(embedder, "_embedder", fake): + result = embedder.create_embedding("abcdefghijklmnopqr") + + assert len(fake.calls) == 3 + assert all(len(chunk) <= 8 for chunk in fake.calls) + assert fake.calls[0][-2:] == fake.calls[1][:2] + assert result == pytest.approx([0.6, 0.8]) + + +def test_chunk_provider_failure_propagates_instead_of_silently_degrading(): + from src import embedder + + fake = _FakeEmbedder(max_chars=8, overlap=2, fail_on_call=2) + with patch.object(embedder, "_embedder", fake): + with pytest.raises(requests.ConnectionError): + embedder.create_embedding("abcdefghijklmnopqr") + + assert len(fake.calls) == 2 + + +def test_regeneration_retries_memory_after_provider_failure(): + from src import embedding_regeneration as regeneration + + memory_id = uuid.uuid4() + with ( + patch.object( + regeneration, + "get_memory_embedding_target", + return_value={"content": "recover me", "has_embedding": False}, + ), + patch.object(regeneration, "create_embedding", return_value=[0.1, 0.2]) as create, + patch.object(regeneration, "update_memory_embedding", return_value=True) as update, + ): + result = regeneration.regenerate_memory_embedding(memory_id) + + create.assert_called_once_with("recover me") + update.assert_called_once_with(memory_id, [0.1, 0.2], force=False) + assert result["status"] == "regenerated" + + +def test_model_migration_requires_force_before_overwriting_embedding(): + from src import embedding_regeneration as regeneration + + memory_id = uuid.uuid4() + target = {"content": "migrate me", "has_embedding": True} + + with ( + patch.object(regeneration, "get_memory_embedding_target", return_value=target), + patch.object(regeneration, "create_embedding") as create, + patch.object(regeneration, "update_memory_embedding") as update, + ): + result = regeneration.regenerate_memory_embedding(memory_id) + + assert result == { + "id": str(memory_id), + "status": "unchanged", + "reason": "embedding_exists", + } + create.assert_not_called() + update.assert_not_called() + + with ( + patch.object(regeneration, "get_memory_embedding_target", return_value=target), + patch.object(regeneration, "create_embedding", return_value=[0.9, 0.8]) as create, + patch.object(regeneration, "update_memory_embedding", return_value=True) as update, + ): + result = regeneration.regenerate_memory_embedding(memory_id, force=True) + + create.assert_called_once_with("migrate me") + update.assert_called_once_with(memory_id, [0.9, 0.8], force=True) + assert result == { + "id": str(memory_id), + "status": "regenerated", + "force": True, + } + + +def test_provider_failure_does_not_replace_existing_embedding(): + from src import embedding_regeneration as regeneration + + memory_id = uuid.uuid4() + with ( + patch.object( + regeneration, + "get_memory_embedding_target", + return_value={"content": "retry later", "has_embedding": False}, + ), + patch.object( + regeneration, + "create_embedding", + side_effect=requests.ConnectionError("provider down"), + ), + patch.object(regeneration, "update_memory_embedding") as update, + ): + with pytest.raises(requests.ConnectionError): + regeneration.regenerate_memory_embedding(memory_id) + + update.assert_not_called() + + +def test_database_update_is_null_only_unless_forced(): + from src.db import attribution + + memory_id = uuid.uuid4() + cursor, manager = _cursor(row={"id": memory_id}) + with patch.object(attribution, "get_db_cursor", manager): + assert attribution.update_memory_embedding( + memory_id, [0.1, 0.2], force=False + ) + + statement, params = cursor.execute.call_args.args + assert "embedding IS NULL" in statement + assert params == ([0.1, 0.2], memory_id, False) + + +def test_rest_and_mcp_expose_regeneration_controls(): + from src.api.main import EmbeddingRegenerationRequest, app + from src import main + + assert EmbeddingRegenerationRequest(force=True).force is True + assert any( + route.path == "/memories/{memory_id}/regenerate-embedding" + for route in app.routes + ) + + tools = asyncio.run(main.list_tools()) + schemas = {tool.name: tool.inputSchema for tool in tools} + schema = schemas["memory_regenerate_embedding"] + assert schema["required"] == ["memory_id"] + assert schema["properties"]["force"]["default"] is False