Skip to content

context-graph: embed messages, entities and edges inside Memgraph for recall - #422

Merged
antejavor merged 1 commit into
mainfrom
feat/416-embeddings
Oct 5, 2026
Merged

antejavor merged 1 commit into
mainfrom
feat/416-embeddings

Conversation

@antejavor

Copy link
Copy Markdown
Contributor

Closes #416. Implements the decisions in #393; part of map #390 (#415 → #416 → #417 → #418).

What

sessions-graph

  • New embeddings.py:
    • embed_session embeds three units through MAGE's embeddings.text, 64 texts per call:
    • Model: BAAI/bge-small-en-v1.5 by default, which is what the benchmark measured. Each vector records embedding_model. A vector from another model counts as missing and is replaced, so vectors from two models never mix.
    • Errors: EmbeddingUnavailableError covers a missing module (plain Memgraph) or a model that can't load.
  • SessionsGraph:
    • embed_session() records embedding_status: completed with embedding_model, or failed with embedding_error.
    • get_pending_embedding_sessions() returns sessions whose embedding failed, never ran, or used another model.
    • Reconciliation, both single and batch, embeds what it wrote. An embedding failure never fails the reconciliation.
  • CLI: sessions-graph embed (--session ID | --pending) [--limit N] [--model M]. The model comes from the flag, then recall.embedding_model in the config, then the default.
  • Connector:
    • SESSION_END always spawns a detached embed --session, independent of auto_reconcile. It costs no LLM, and the hook never waits on the model.
    • The embed child gets the Memgraph config but not the LLM keys (_child_env(llm=False)).

agent-context-graph

  • Config: a recall.embedding_model key, preserved across bootstrap rewrites like auto_reconcile.
  • doctor --connector sessions-graph: adds an embeddings check that loads the model inside Memgraph and reports its dimension, or explains that MAGE (2 GiB or more) is needed.
  • Bootstrap and docs: bootstrap's hint and the plugin READMEs and skills suggest memgraph/memgraph-mage.

Tests

  • Suites:

    Suite Result
    sessions-graph 78 passed, 1 skipped (OpenAI)
    agent-context-graph 115 passed
    eval 280 passed, 5 skipped
  • New end-to-end tests on real MAGE:

    • all three units embedded and tool calls skipped;
    • a second call embeds nothing;
    • a vector from another model is replaced;
    • a model that can't load is recorded on the session;
    • the CLI's success and failure paths.

    The chunk, entity and edge are hand-made in extraction's shape. Reconciliation's own embed call runs only in the existing OpenAI-gated end-to-end test.

  • Checked by hand:

    • On plain memgraph/memgraph:3.10.0, embedding raises "There is no procedure named 'embeddings.text'" and the session records failed.
    • doctor's check reports BAAI/bge-small-en-v1.5 (384 dimensions) on MAGE, and fails clearly for a model that can't load.

… recall

Implements the #393 decisions. sessions-graph's new embeddings module
embeds a session's messages (Action.text), the entities mentioned in its
chunks, and its extracted edges (r.text) through MAGE's embeddings.text,
default BAAI/bge-small-en-v1.5. Each vector records embedding_model; one
from another model counts as missing and is replaced, so models never mix.

- SESSION_END always spawns a detached `sessions-graph embed --session`,
  independent of auto_reconcile; it gets the Memgraph config, not the LLM
  keys. Reconciliation embeds what it wrote; a failure there never fails
  the reconciliation.
- The Session records embedding_status (completed with the model, or
  failed with the error); `sessions-graph embed --pending` retries.
- agent-context-graph: `recall.embedding_model` config key, kept across
  bootstrap rewrites; doctor checks embeddings when sessions-graph is
  enabled; bootstrap and the setup docs suggest memgraph-mage.

Closes #416.
@antejavor
antejavor merged commit 37e27b4 into main Oct 5, 2026
25 checks passed
antejavor added a commit that referenced this pull request Oct 5, 2026
Resolve the claude_code/codex adapter conflicts with #427 by keeping this
branch's declarative adapters and carrying #427's reply rules into the
shared turn_end rule: a blank reply records nothing, and the reply's text
is not duplicated into its metadata. #427's tests now expect a turn end,
not a session end, after the reply.

main's sessions-graph spawns the no-LLM `embed` step on SessionEnd (#422),
which Codex recall (#426, #427) relied on via the per-turn Stop. Since a
Stop is now a TurnEnd, sessions-graph spawns `embed` on TurnEnd too, so
every runtime's turns are embedded as they end, including runtimes with no
session-end hook. LLM reconciliation stays on SessionEnd only.
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.

sessions-graph: embed messages, entities and edges in Memgraph

1 participant