fix: bound embedder input and make message writes resilient to embedding failures - #175
Open
steveonjava wants to merge 3 commits into
Open
fix: bound embedder input and make message writes resilient to embedding failures#175steveonjava wants to merge 3 commits into
steveonjava wants to merge 3 commits into
Conversation
I hit this on real oversize messages: the OpenAI API rejects any input over the model's token ceiling with a raw 400, and that error was propagating straight up to the caller. I add a token-budget truncation step in embed() and embed_batch() so oversize text gets truncated instead of rejected. The truncation uses tiktoken when it is available for an accurate count, and falls back to a 4-chars-per-token estimate when it is not, since tiktoken is an optional dependency of this package. Related to neo4j-labs#104.
add_message embedded content before creating the message node. When the embed call raised, for example on oversize input against the model's context limit, the exception aborted the whole write and the message plus its extracted entities were silently dropped. I wrap the embed call in try/except and fall back to a null vector on failure so the node write and entity extraction always proceed. I apply the same fallback to the batch path in add_messages_batch, so a single bad message in a batch does not sink the rest of the batch. Related to neo4j-labs#104.
|
@steveonjava is attempting to deploy a commit to the lyonwj's projects Team on Vercel. A member of the Team first needs to authorize it. |
steveonjava
marked this pull request as ready for review
August 7, 2026 03:21
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Oversize input to the OpenAI embedder caused a silent data loss bug: if the input exceeded the token budget, the API returned a 400 error, the message write aborted, and the extracted entities were lost. This PR closes the gap with a three-layer fix.
Root cause
OpenAIEmbedder.embed()andembed_batch()passed raw text to the OpenAI API with no token budget check. When the input was too long, the API returned a 400 error, which bubbled up throughadd_message()andadd_messages_batch(), aborting the whole write. Same for the batch path: one failing batch meant all messages in it were dropped.Fix
Three layers, each addressing a different failure mode:
Layer 1: Input-bounded embedder.
_truncate_to_tokens()truncates text to the model token budget before calling the API. Usestiktokenwhen available for an accurate count, falls back to a 4-char-per-token estimate otherwise. Called in bothembed()andembed_batch().Layer 2: Resilient
add_message(). Wraps the single-message embed in a try/except. On failure, stores the message with a null vector and continues to entity extraction. The entities carry their own embeddings, so the memory stays fully recallable.Layer 3: Batch-level degradation in
add_messages_batch(). Same pattern at batch level: if the whole batch embed fails, all messages in the batch get null vectors and insertion continues.Verification
uv buildproduces a wheel that installs and imports cleanly.Files changed
src/neo4j_agent_memory/embeddings/openai.py(added_truncate_to_tokens(),MODEL_MAX_INPUT_TOKENS,_max_tokens_for(),_DEFAULT_MAX_INPUT_TOKENS,_TOKEN_HEADROOM; wired truncation intoembed()andembed_batch())src/neo4j_agent_memory/memory/short_term.py(try/except wrappers aroundembed()andembed_batch()inadd_message()andadd_messages_batch(), with warning logs on failure)tests/unit/embeddings/test_openai.py(new file: 93 lines covering the char-fallback path, the real tiktoken path, embedder truncation inembed()andembed_batch(), and empty-list short-circuit)tests/unit/test_relation_storage.py(new tests for embedding resilience:add_messagesurvives embed failure with entity extraction intact,add_messages_batchsurvives embed_batch failure, happy-path vector flow confirmed)Fixes #104