fix: bound embedder input and make message writes resilient to embedding failures

src/neo4j_agent_memory/embeddings/openai.py

@@ -2,6 +2,7 @@
 
 from __future__ import annotations
 
+import logging
 from typing import TYPE_CHECKING, Any
 
 from neo4j_agent_memory.core.exceptions import EmbeddingError
@@ -18,6 +19,63 @@
     "text-embedding-ada-002": 1536,
 }
 
+logger = logging.getLogger(__name__)
+
+# Max input tokens per model. The API hard-fails (400) above this; we truncate to stay under it.
+# Target slightly below the true ceiling (8192) for safety headroom.
+MODEL_MAX_INPUT_TOKENS = {
+    "text-embedding-3-small": 8192,
+    "text-embedding-3-large": 8192,
+    "text-embedding-ada-002": 8192,
+}
+_DEFAULT_MAX_INPUT_TOKENS = 8192
+_TOKEN_HEADROOM = 192
+
+
+def _max_tokens_for(model: str) -> int:
+    """Return the safe token budget for ``model`` (ceiling minus headroom)."""
+    return MODEL_MAX_INPUT_TOKENS.get(model, _DEFAULT_MAX_INPUT_TOKENS) - _TOKEN_HEADROOM
+
+
+def _truncate_to_tokens(text: str, model: str) -> str:
+    """Truncate ``text`` so it fits within the token budget of ``model``.
+
+    Uses ``tiktoken`` when available for an accurate token count. Falls back to
+    a character-based estimate (4 chars per token) when ``tiktoken`` is not
+    installed or the encoding lookup fails, since ``tiktoken`` is an optional
+    dependency of this package.
+    """
+    budget = _max_tokens_for(model)
+    try:
+        import tiktoken  # tiktoken is an optional dependency, not declared in pyproject
+
+        try:
+            enc = tiktoken.encoding_for_model(model)
+        except KeyError:
+            enc = tiktoken.get_encoding("cl100k_base")
+        toks: list[int] = enc.encode(text)
+        if len(toks) <= budget:
+            return text
+        logger.warning(
+            "embedding input exceeds token budget (%d tokens > %d budget) for model %s; truncating",
+            len(toks),
+            budget,
+            model,
+        )
+        return str(enc.decode(toks[:budget]))
+    except Exception:
+        char_budget = budget * 4
+        if len(text) <= char_budget:
+            return text
+        logger.warning(
+            "embedding input exceeds char-estimate budget (%d chars > %d) for model %s "
+            "(tiktoken unavailable); truncating",
+            len(text),
+            char_budget,
+            model,
+        )
+        return text[:char_budget]
+
 
 class OpenAIEmbedder(BaseEmbedder):
     """OpenAI embedding provider."""
@@ -71,6 +129,7 @@ def dimensions(self) -> int:
     async def embed(self, text: str) -> list[float]:
         """Generate embedding for a single text."""
         client = self._ensure_client()
+        text = _truncate_to_tokens(text, self._model)
 
         try:
             kwargs: dict[str, Any] = {"input": text, "model": self._model}
@@ -87,6 +146,7 @@ async def embed_batch(self, texts: list[str]) -> list[list[float]]:
         if not texts:
             return []
 
+        texts = [_truncate_to_tokens(t, self._model) for t in texts]
         client = self._ensure_client()
         all_embeddings: list[list[float]] = []
 

src/neo4j_agent_memory/memory/short_term.py

@@ -3,6 +3,7 @@
 from __future__ import annotations
 
 import json
+import logging
 from collections.abc import Awaitable, Callable
 from datetime import datetime, timezone
 from enum import Enum
@@ -17,6 +18,8 @@
 from neo4j_agent_memory.graph import queries
 from neo4j_agent_memory.graph.query_builder import build_create_entity_query
 
+logger = logging.getLogger(__name__)
+
 
 def _llm_summarizer(
     provider: LLMProvider,
@@ -456,12 +459,23 @@ async def add_messages_batch(
                     )
                 )
 
-            # Generate embeddings in batch if enabled
+            # Generate embeddings in batch if enabled. A batch embed failure must never abort the
+            # whole batch insert: I fall back to null vectors for every message in this batch and
+            # continue, matching the single-message add_message() resilience behavior above.
             if generate_embeddings and self._embedder is not None:
-                embeddings = await self._embedder.embed_batch(contents_for_embedding)
-                for j, emb in enumerate(embeddings):
-                    batch_data[j]["embedding"] = emb
-                    batch_messages[j].embedding = emb
+                try:
+                    embeddings = await self._embedder.embed_batch(contents_for_embedding)
+                except Exception as e:  # noqa: BLE001
+                    logger.warning(
+                        "batch embedding failed for %d message(s); storing this batch without "
+                        "vectors and continuing: %s",
+                        len(contents_for_embedding),
+                        str(e)[:200],
+                    )
+                else:
+                    for j, emb in enumerate(embeddings):
+                        batch_data[j]["embedding"] = emb
+                        batch_messages[j].embedding = emb
 
             # Insert batch into database
             await self._client.execute_write(
@@ -713,10 +727,21 @@ async def add_message(
             session_id, conversation_id, user_identifier=user_identifier
         )
 
-        # Generate embedding if enabled
-        embedding = None
+        # Generate embedding if enabled. An embedding failure must never drop the message or its
+        # extracted entities: I degrade to a null message-vector and continue. Entities extracted
+        # from this message carry their own embeddings, so the memory remains fully recallable even
+        # without the message-level vector.
+        embedding: list[float] | None = None
         if generate_embedding and self._embedder is not None:
-            embedding = await self._embedder.embed(content)
+            try:
+                embedding = await self._embedder.embed(content)
+            except Exception as e:  # noqa: BLE001
+                logger.warning(
+                    "message embedding failed; storing message without a vector and continuing "
+                    "with extraction: %s",
+                    str(e)[:200],
+                )
+                embedding = None
 
         # Create message
         message = Message(

tests/unit/embeddings/test_openai.py

@@ -0,0 +1,239 @@
+"""Unit tests for OpenAI embedder token-budget truncation."""
+
+from __future__ import annotations
+
+import logging
+import sys
+from unittest.mock import AsyncMock, MagicMock, patch
+
+import pytest
+
+
+@pytest.fixture
+def force_char_fallback():
+    """Force the char-estimate fallback path by making ``import tiktoken`` fail.
+
+    Setting ``sys.modules["tiktoken"] = None`` causes a subsequent
+    ``import tiktoken`` to raise ``ImportError`` inside ``_truncate_to_tokens``,
+    which is the path we want to exercise since tiktoken is optional.
+    """
+    with patch.dict(sys.modules, {"tiktoken": None}):
+        yield
+
+
+class TestTruncateToTokens:
+    """Tests for the module-level ``_truncate_to_tokens`` helper."""
+
+    def test_under_budget_returns_input_unchanged(self, force_char_fallback):
+        """Char-estimate path: short input passes through untouched."""
+        from neo4j_agent_memory.embeddings.openai import _truncate_to_tokens
+
+        text = "hello world"
+        assert _truncate_to_tokens(text, "text-embedding-3-small") == text
+
+    def test_over_budget_truncates_and_warns(self, force_char_fallback, caplog):
+        """Char-estimate path: oversize input is truncated to char_budget and a warning is logged."""
+        from neo4j_agent_memory.embeddings.openai import (
+            _DEFAULT_MAX_INPUT_TOKENS,
+            _TOKEN_HEADROOM,
+            _truncate_to_tokens,
+        )
+
+        char_budget = (_DEFAULT_MAX_INPUT_TOKENS - _TOKEN_HEADROOM) * 4
+        oversize = "x" * (char_budget + 5000)
+
+        with caplog.at_level(logging.WARNING, logger="neo4j_agent_memory.embeddings.openai"):
+            result = _truncate_to_tokens(oversize, "text-embedding-3-small")
+
+        assert len(result) == char_budget
+        assert result == "x" * char_budget
+        assert any(
+            "exceeds char-estimate budget" in rec.message and "truncating" in rec.message
+            for rec in caplog.records
+        )
+
+    def test_unknown_model_uses_default_budget(self, force_char_fallback):
+        """Unknown models fall back to the default token budget."""
+        from neo4j_agent_memory.embeddings.openai import _truncate_to_tokens
+
+        text = "some short text"
+        assert _truncate_to_tokens(text, "made-up-model-name") == text
+
+
+class TestOpenAIEmbedderTruncation:
+    """Tests that ``OpenAIEmbedder`` truncates before hitting the API."""
+
+    @pytest.mark.asyncio
+    async def test_embed_truncates_oversize_input(self, force_char_fallback):
+        """``embed()`` must truncate before calling the OpenAI client."""
+        from neo4j_agent_memory.embeddings.openai import (
+            _DEFAULT_MAX_INPUT_TOKENS,
+            _TOKEN_HEADROOM,
+            OpenAIEmbedder,
+        )
+
+        char_budget = (_DEFAULT_MAX_INPUT_TOKENS - _TOKEN_HEADROOM) * 4
+        oversize = "y" * (char_budget + 10_000)
+
+        embedder = OpenAIEmbedder(model="text-embedding-3-small", api_key="test-key")
+
+        mock_response = MagicMock()
+        mock_response.data = [MagicMock(embedding=[0.1, 0.2, 0.3])]
+        mock_client = MagicMock()
+        mock_client.embeddings = MagicMock()
+        mock_client.embeddings.create = AsyncMock(return_value=mock_response)
+        embedder._client = mock_client
+
+        result = await embedder.embed(oversize)
+
+        assert result == [0.1, 0.2, 0.3]
+        mock_client.embeddings.create.assert_awaited_once()
+        sent_input = mock_client.embeddings.create.await_args.kwargs["input"]
+        assert len(sent_input) == char_budget
+
+    @pytest.mark.asyncio
+    async def test_embed_batch_truncates_every_item(self, force_char_fallback):
+        """``embed_batch()`` must truncate every text before calling the API."""
+        from neo4j_agent_memory.embeddings.openai import (
+            _DEFAULT_MAX_INPUT_TOKENS,
+            _TOKEN_HEADROOM,
+            OpenAIEmbedder,
+        )
+
+        char_budget = (_DEFAULT_MAX_INPUT_TOKENS - _TOKEN_HEADROOM) * 4
+        oversize_a = "a" * (char_budget + 500)
+        oversize_b = "b" * (char_budget + 1500)
+        small = "c" * 10
+        texts = [oversize_a, small, oversize_b]
+
+        embedder = OpenAIEmbedder(model="text-embedding-3-small", api_key="test-key")
+
+        mock_response = MagicMock()
+        mock_response.data = [
+            MagicMock(index=0, embedding=[0.1]),
+            MagicMock(index=1, embedding=[0.2]),
+            MagicMock(index=2, embedding=[0.3]),
+        ]
+        mock_client = MagicMock()
+        mock_client.embeddings = MagicMock()
+        mock_client.embeddings.create = AsyncMock(return_value=mock_response)
+        embedder._client = mock_client
+
+        result = await embedder.embed_batch(texts)
+
+        assert result == [[0.1], [0.2], [0.3]]
+        mock_client.embeddings.create.assert_awaited_once()
+        sent_batch = mock_client.embeddings.create.await_args.kwargs["input"]
+        assert len(sent_batch) == 3
+        assert len(sent_batch[0]) == char_budget
+        assert sent_batch[1] == small
+        assert len(sent_batch[2]) == char_budget
+
+    @pytest.mark.asyncio
+    async def test_embed_batch_empty_list_short_circuits(self):
+        """Empty list must not touch the client or the truncator."""
+        from neo4j_agent_memory.embeddings.openai import OpenAIEmbedder
+
+        embedder = OpenAIEmbedder(model="text-embedding-3-small", api_key="test-key")
+        mock_client = MagicMock()
+        mock_client.embeddings = MagicMock()
+        mock_client.embeddings.create = AsyncMock()
+        embedder._client = mock_client
+
+        result = await embedder.embed_batch([])
+
+        assert result == []
+        mock_client.embeddings.create.assert_not_called()
+
+
+class TestTruncateToTokensWithRealTiktoken:
+    """Tests that exercise the real ``tiktoken`` path (no patching).
+
+    ``tiktoken`` is installed in the dev environment, so calling
+    ``_truncate_to_tokens`` without any fixture that hides it exercises the
+    primary code path: ``tiktoken.encoding_for_model`` / ``enc.encode`` /
+    ``enc.decode``, including the ``KeyError`` fallback to ``cl100k_base``.
+    """
+
+    def test_real_tiktoken_under_budget_passthrough(self):
+        """Real tiktoken path: short input returns unchanged."""
+        import tiktoken
+
+        from neo4j_agent_memory.embeddings.openai import _truncate_to_tokens
+
+        text = "hello world"
+        enc = tiktoken.encoding_for_model("text-embedding-3-small")
+        assert len(enc.encode(text)) < 8000
+
+        assert _truncate_to_tokens(text, "text-embedding-3-small") is text
+
+    def test_real_tiktoken_over_budget_truncates_to_exact_token_count(self, caplog):
+        """Real tiktoken path: oversize input is truncated to at-most-budget tokens."""
+        import tiktoken
+
+        from neo4j_agent_memory.embeddings.openai import (
+            _DEFAULT_MAX_INPUT_TOKENS,
+            _TOKEN_HEADROOM,
+            _truncate_to_tokens,
+        )
+
+        budget = _DEFAULT_MAX_INPUT_TOKENS - _TOKEN_HEADROOM
+        enc = tiktoken.encoding_for_model("text-embedding-3-small")
+        oversize = "hello world " * 5000
+        assert len(enc.encode(oversize)) > budget, (
+            "test fixture is not oversized under the real encoder; adjust the multiplier"
+        )
+
+        with caplog.at_level(logging.WARNING, logger="neo4j_agent_memory.embeddings.openai"):
+            result = _truncate_to_tokens(oversize, "text-embedding-3-small")
+
+        re_encoded = enc.encode(result)
+        assert len(re_encoded) <= budget
+        assert any(
+            "exceeds token budget" in rec.message and "truncating" in rec.message
+            for rec in caplog.records
+        )
+
+    def test_real_tiktoken_unknown_model_falls_back_to_cl100k_base_encoding(self, caplog):
+        """Unknown model: ``KeyError`` from ``encoding_for_model`` is caught, falls back to cl100k_base."""
+        import tiktoken
+
+        from neo4j_agent_memory.embeddings.openai import (
+            _DEFAULT_MAX_INPUT_TOKENS,
+            _TOKEN_HEADROOM,
+            _truncate_to_tokens,
+        )
+
+        budget = _DEFAULT_MAX_INPUT_TOKENS - _TOKEN_HEADROOM
+        cl100k = tiktoken.get_encoding("cl100k_base")
+
+        short_text = "hello world"
+        result_short = _truncate_to_tokens(short_text, "some-made-up-nonexistent-model-xyz")
+        assert result_short == short_text
+
+        oversize = "hello world " * 5000
+        assert len(cl100k.encode(oversize)) > budget, (
+            "test fixture is not oversized under cl100k_base; adjust the multiplier"
+        )
+
+        with caplog.at_level(logging.WARNING, logger="neo4j_agent_memory.embeddings.openai"):
+            result_long = _truncate_to_tokens(oversize, "some-made-up-nonexistent-model-xyz")
+
+        re_encoded = cl100k.encode(result_long)
+        assert len(re_encoded) <= budget
+
+    def test_real_tiktoken_truncated_output_is_never_longer_in_tokens_than_original(self):
+        """Sanity invariant: truncation strictly shrinks token count for oversize input."""
+        import tiktoken
+
+        from neo4j_agent_memory.embeddings.openai import _truncate_to_tokens
+
+        enc = tiktoken.encoding_for_model("text-embedding-3-small")
+        oversize = "hello world " * 5000
+        original_token_count = len(enc.encode(oversize))
+        assert original_token_count > 8000
+
+        result = _truncate_to_tokens(oversize, "text-embedding-3-small")
+        truncated_token_count = len(enc.encode(result))
+
+        assert truncated_token_count < original_token_count