Skip to content

fix(bolt): embed entities during extraction + add generate_entity_embeddings_batch() - #154

Open
steveonjava wants to merge 1 commit into
neo4j-labs:mainfrom
steveonjava:feat/agent-memory-001
Open

fix(bolt): embed entities during extraction + add generate_entity_embeddings_batch()#154
steveonjava wants to merge 1 commit into
neo4j-labs:mainfrom
steveonjava:feat/agent-memory-001

Conversation

@steveonjava

Copy link
Copy Markdown

What & why

On the Bolt backend, the LLM extraction pipeline wrote every :Entity with embedding = NULL, while
search_entities() is vector-only (db.index.vector.queryNodes('entity_embedding_idx', …)). Net
effect: a knowledge graph built via add_message(extract_entities=True) /
extract_entities_from_session() populates the graph but is invisible to semantic recall — every
query returns empty, with no error. The NAMS backend already embeds extracted entities server-side, so
this fixes a Bolt-vs-NAMS asymmetry and aligns Bolt with the documented prerequisite
("embedding model configured" for automatic extraction). Fixes #NNNNN.

Changes

  1. Embed entities inline during extraction (memory/short_term.py):
    • _extract_and_link_entities() and extract_entities_from_session() now embed each extracted
      entity's name (matching add_entity()) when an embedder is configured, via a new
      _embed_names() helper that batches with embedder.embed_batch() and tolerates individual
      embed failures (leaves NULL for later backfill rather than aborting the message).
    • New generate_embeddings: bool = True control flag, threaded from add_message(...) and exposed
      on extract_entities_from_session(...). Default True = embed (expected behavior); pass False
      for fast structural-only import + deferred batch embedding.
  2. New public batch helper (memory/long_term.py):
    • LongTermMemory.generate_entity_embeddings_batch(batch_size=128, on_progress=None) — the entity
      analogue of the existing ShortTermMemory.generate_embeddings_batch(). Uses the already-defined
      GET_ENTITIES_WITHOUT_EMBEDDINGS / UPDATE_ENTITY_EMBEDDING / COUNT_ENTITIES_WITHOUT_EMBEDDINGS
      queries (previously dead code). Idempotent; repairs existing graphs without re-extraction.

No new dependencies, no new imports, no schema/index changes (the entity_embedding_idx vector index
and the three queries already exist).

Behavior change / compat notes

  • Extraction now embeds by default, adding embedding round-trips during extraction. This increases
    extraction latency but makes recall actually work. Opt out with generate_embeddings=False (then
    call generate_entity_embeddings_batch() later) for the previous fast-import behavior. Flagged in
    CHANGELOG.
  • Embeds name only, identical to add_entity, so manually-added and extracted entities share one
    vector space.
  • Failures on individual entity-name embeds are logged and skipped (entity persists with NULL
    embedding), so extraction never hard-fails on a single bad embed.

Testing

  • Unit: extraction with a stub embedder writes non-null embedding on every extracted entity;
    with generate_embeddings=False, embeddings are NULL.
  • Unit: generate_entity_embeddings_batch() embeds only NULL-embedding entities, is idempotent on
    a second run (embedded=0, remaining=0), and reports failed on embedder errors without raising.
  • Integration (real Neo4j + embedder): add_message(extract_entities=True)search_entities()
    returns the extracted entities by semantic query (the issue's repro now passes).
  • Integration: bulk extract_entities_from_session(generate_embeddings=False) then
    generate_entity_embeddings_batch() → recall works.
  • Regression: NAMS backend unaffected (no code touched there).

Notes for reviewers

  • Diffs in the linked issue/handoff were anchored to v0.5.0 line numbers; I've re-anchored to main
    here.
  • Open question for maintainers: should add_message reuse its existing generate_embeddings flag
    (currently message-only) to also gate entity embedding, or get a dedicated
    generate_entity_embeddings flag? This PR reuses the existing flag for a single intuitive knob;
    happy to split if you prefer.

Related Issue

Fixes #NNNNN

…eddings_batch()

On the Bolt backend, the LLM extraction pipeline wrote every :Entity with
embedding=NULL, hardcoded, ignoring the configured embedder. Recall
(long_term.search_entities) is vector-only with no keyword/fulltext fallback,
so extracted entities are mathematically unreachable by semantic recall.

Fixes:
1. Embed entity names inline during extraction (_extract_and_link_entities and
   extract_entities_from_session), gated on generate_embeddings=True (default).
   Uses embed_batch for efficiency, falls back to per-name embed on failure.
2. Add LongTermMemory.generate_entity_embeddings_batch(), the entity analogue
   of ShortTermMemory.generate_embeddings_batch(). Uses the already-defined
   GET_ENTITIES_WITHOUT_EMBEDDINGS / UPDATE_ENTITY_EMBEDDING queries.

Embeds name only (matching add_entity). Tolerates individual embed failures.
Backward compatible: pass generate_embeddings=False for fast structural-only
import + deferred batch embedding.
@vercel

vercel Bot commented Jun 27, 2026

Copy link
Copy Markdown

@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
steveonjava marked this pull request as ready for review August 10, 2026 05:27
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant