Repair embedding resilience for oversized and failed requests

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:

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."""

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,

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 $$;

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

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,
+    }

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

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"]

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)