Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
10 changes: 9 additions & 1 deletion context-graph/sessions-graph/README.md
Original file line number Diff line number Diff line change
Expand Up @@ -137,6 +137,14 @@ back to the session that produced them — see
[`CONTEXT.md`](./CONTEXT.md#language) for the **Session Reconciliation** /
**Reconcilable Content** / **Reconciliation Status** terminology.

The same pass also writes the session's **episodic memory**: an
`(:Episode {summary, summarized_at})` node linked via
`(:Session)-[:HAS_EPISODE]->(:Episode)` (at most one per session — re-running
reconciliation updates it rather than adding another), produced by a second,
dedicated LLM call over the same deduped session text — a "what happened in
this session" gist, not the structured entity graph. This is what a "what did
we do last time?" recall query actually reads.

This requires the `sessions-graph[reconciliation]` extra and an LLM API key
(`OPENAI_API_KEY` or `ANTHROPIC_API_KEY`) for LightRAG — see the
[lightrag-memgraph README](../../integrations/lightrag-memgraph/README.md).
Expand Down Expand Up @@ -191,7 +199,7 @@ summary = await graph.reconcile_session(
lightrag_wrapper=lightrag_wrapper,
enforce_ontology=True, # match the CLI: promote entity_type to real labels
)
print(summary.status, summary.texts_considered, summary.texts_deduped)
print(summary.status, summary.texts_considered, summary.texts_deduped, summary.summary_written)
```

Label promotion is opt-in and mirrors unstructured2graph's flags: the default
Expand Down
6 changes: 5 additions & 1 deletion context-graph/sessions-graph/src/sessions_graph/cli.py
Original file line number Diff line number Diff line change
Expand Up @@ -120,7 +120,11 @@ async def _run_reconcile(parsed: argparse.Namespace) -> int:
session_id, lightrag_wrapper=lightrag_wrapper, enforce_ontology=True
)
if summary.status == "completed":
print(f"OK {session_id}: {summary.texts_deduped}/{summary.texts_considered} unique texts reconciled")
summarized = " (summary written)" if summary.summary_written else ""
print(
f"OK {session_id}: {summary.texts_deduped}/{summary.texts_considered} "
f"unique texts reconciled{summarized}"
)
else:
print(f"FAILED {session_id}: {summary.error}", file=sys.stderr)
exit_code = 1
Expand Down
32 changes: 30 additions & 2 deletions context-graph/sessions-graph/src/sessions_graph/core.py
Original file line number Diff line number Diff line change
Expand Up @@ -6,10 +6,12 @@
(:User {user_id})
(:Memory {memory_id, user_id, content, created_at, session_id?})
(:Session {session_id})
(:Episode {summary, summarized_at}) — written by reconcile_session(), see below

Relationships:
(:User)-[:HAS_MEMORY]->(:Memory)
(:Session)-[:PRODUCED_MEMORY]->(:Memory) — only when session_id is provided
(:Session)-[:HAS_EPISODE]->(:Episode) — at most one per session
"""

from __future__ import annotations
Expand All @@ -27,6 +29,7 @@
ReconciliationSummary,
build_reconciliation_sources,
content_hash,
summarize_session_texts,
)

if TYPE_CHECKING:
Expand Down Expand Up @@ -277,7 +280,7 @@ async def reconcile_session(
enforce_ontology: bool = False,
ontology_path: str | Path | None = None,
) -> ReconciliationSummary:
"""Batch-extract entities from a session's Action + Memory content.
"""Batch-extract entities and a narrative summary from a session's content.

Pulls all reconcilable Message/ToolCall/ToolResult text recorded for
*session_id* in Actions Graph, plus this session's Memories, dedupes
Expand All @@ -286,6 +289,15 @@ async def reconcile_session(
linked back to their source Action/Memory node via ``HAS_CHUNK`` so
entities trace back to the session that produced them.

This same pass also produces the session's episodic memory: an
``(:Episode {summary, summarized_at})`` node linked from the session
via ``HAS_EPISODE``, via a second, dedicated LLM call (entity
extraction and narrative summarization are different task shapes, so
this doesn't piggyback on LightRAG's own extraction prompt) -- but
it's still one trigger, one fetch/dedupe of session text, no separate
schedule. Re-reconciling a session updates its one Episode rather
than creating another.

This is deliberately not wired to run automatically inside the
``SESSION_END`` hook — LightRAG extraction is LLM-backed and slow, and
hook subprocesses run under a runtime timeout. Call this from a
Expand Down Expand Up @@ -346,6 +358,7 @@ async def reconcile_session(
unique_texts.setdefault(content_hash(source.text), source.text)

try:
summary_text: str | None = None
if unique_texts:
grouped_chunks = await from_texts(
list(unique_texts.values()),
Expand All @@ -358,19 +371,34 @@ async def reconcile_session(
)
chunks_by_text_hash = dict(zip(unique_texts.keys(), grouped_chunks, strict=True))
self._link_chunks_to_sources(sources, chunks_by_text_hash)
summary_text = await summarize_session_texts(lightrag_wrapper, list(unique_texts.values()))

reconciled_at = datetime.now(timezone.utc).isoformat()
self._db.query(
"""
MATCH (s:Session {session_id: $session_id})
SET s.reconciliation_status = 'completed', s.reconciled_at = $reconciled_at
""",
params={"session_id": session_id, "reconciled_at": datetime.now(timezone.utc).isoformat()},
params={"session_id": session_id, "reconciled_at": reconciled_at},
)
if summary_text:
# MERGE on the (Session)-[:HAS_EPISODE]->(Episode) pattern (not just CREATE)
# so re-reconciling a session updates its one Episode instead of accumulating
# duplicates -- Episode has no natural external id of its own to dedupe on.
self._db.query(
"""
MATCH (s:Session {session_id: $session_id})
MERGE (s)-[:HAS_EPISODE]->(e:Episode)
SET e.summary = $summary, e.summarized_at = $summarized_at
""",
params={"session_id": session_id, "summary": summary_text, "summarized_at": reconciled_at},
)
return ReconciliationSummary(
session_id=session_id,
status="completed",
texts_considered=len(sources),
texts_deduped=len(unique_texts),
summary_written=summary_text is not None,
)
except Exception as e:
self._db.query(
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -115,6 +115,35 @@ class ReconciliationSummary:
texts_considered: int
texts_deduped: int
error: str | None = None
summary_written: bool = False


_SUMMARY_PROMPT_TEMPLATE = (
"Summarize what happened in this agent session in 2-3 sentences, for future recall "
'(e.g. answering "what did we do last time?"). Be concrete: name entities, decisions, and '
"outcomes; skip meta-commentary about the summarization task itself.\n\n"
"Session content:\n{content}"
)


def build_session_summary_prompt(texts: list[str]) -> str:
"""Build the prompt for distilling a session's reconcilable text into a narrative gist."""
return _SUMMARY_PROMPT_TEMPLATE.format(content="\n\n".join(texts))


async def summarize_session_texts(lightrag_wrapper: Any, texts: list[str]) -> str:
"""Produce a narrative session summary via the same LLM the LightRAG wrapper is configured with.

This is a second, dedicated LLM call alongside the entity-extraction pass ``from_texts``
already makes -- narrative summarization and structured entity/relationship extraction are
different task shapes, so this doesn't piggyback on LightRAG's own extraction prompt. It's
still one ``reconcile_session`` pass: one trigger, one fetch/dedupe of session text, no
separate schedule.
"""
rag = lightrag_wrapper.get_lightrag()
prompt = build_session_summary_prompt(texts)
result = await rag.llm_model_func(prompt)
return result.strip()


def build_reconciliation_sources(
Expand Down
21 changes: 21 additions & 0 deletions context-graph/sessions-graph/tests/test_e2e_reconciliation.py
Original file line number Diff line number Diff line change
Expand Up @@ -87,13 +87,23 @@ async def test_reconcile_session_extracts_real_entities_from_session_content(
assert summary.status == "completed"
assert summary.texts_considered == 3
assert summary.texts_deduped == 3
assert summary.summary_written is True

rows = memgraph.query(
"MATCH (s:Session {session_id: $session_id}) RETURN s.reconciliation_status AS status",
params={"session_id": session_id},
)
assert rows[0]["status"] == "completed"

episode_rows = memgraph.query(
"MATCH (:Session {session_id: $session_id})-[:HAS_EPISODE]->(e:Episode) "
"RETURN e.summary AS summary, e.summarized_at AS summarized_at",
params={"session_id": session_id},
)
assert len(episode_rows) == 1
assert episode_rows[0]["summary"]
assert episode_rows[0]["summarized_at"]

has_chunk_rows = memgraph.query(
"""
MATCH (:Session {session_id: $session_id})-[:HAS_ACTION]->(:Action)-[:HAS_CHUNK]->(:Chunk)
Expand All @@ -106,3 +116,14 @@ async def test_reconcile_session_extracts_real_entities_from_session_content(
workspace = lightrag_wrapper.get_lightrag().chunk_entity_relation_graph.workspace
entity_rows = memgraph.query(f"MATCH (:{workspace})-[:MENTIONED_IN]->(:Chunk) RETURN count(*) AS count")
assert entity_rows[0]["count"] > 0

# Re-reconciling the same session must update the one Episode, not accumulate another.
second_summary = await graph.reconcile_session(
session_id, lightrag_wrapper=lightrag_wrapper, actions_graph=actions_graph
)
assert second_summary.summary_written is True
episode_rows_after_rerun = memgraph.query(
"MATCH (:Session {session_id: $session_id})-[:HAS_EPISODE]->(e:Episode) RETURN count(e) AS count",
params={"session_id": session_id},
)
assert episode_rows_after_rerun[0]["count"] == 1
59 changes: 56 additions & 3 deletions context-graph/sessions-graph/tests/test_reconciliation.py
Original file line number Diff line number Diff line change
Expand Up @@ -19,8 +19,10 @@
MAX_RECONCILABLE_CHARS,
ReconciliationSource,
build_reconciliation_sources,
build_session_summary_prompt,
content_hash,
extract_reconcilable_text,
summarize_session_texts,
)

from actions_graph.models import ErrorEvent, Message, MessageRole, ToolCall, ToolResult
Expand Down Expand Up @@ -103,6 +105,30 @@ def test_skips_actions_with_no_reconcilable_text(self):
assert sources == []


# ---------------------------------------------------------------------------
# build_session_summary_prompt / summarize_session_texts
# ---------------------------------------------------------------------------


class TestSummarizeSessionTexts:
def test_prompt_includes_all_texts(self):
prompt = build_session_summary_prompt(["Alice asked about the graph engine.", "Bob replied with a plan."])
assert "Alice asked about the graph engine." in prompt
assert "Bob replied with a plan." in prompt

@pytest.mark.asyncio
async def test_calls_lightrag_wrappers_llm_model_func_and_strips_result(self):
lightrag_wrapper = MagicMock()
lightrag_wrapper.get_lightrag.return_value.llm_model_func = AsyncMock(return_value=" A tidy summary. ")

result = await summarize_session_texts(lightrag_wrapper, ["Some session text."])

assert result == "A tidy summary."
lightrag_wrapper.get_lightrag.return_value.llm_model_func.assert_awaited_once()
prompt_arg = lightrag_wrapper.get_lightrag.return_value.llm_model_func.call_args.args[0]
assert "Some session text." in prompt_arg


# ---------------------------------------------------------------------------
# SessionsGraph.reconcile_session (stubbed Memgraph + mocked ActionsGraph/LightRAG)
# ---------------------------------------------------------------------------
Expand All @@ -128,13 +154,19 @@ def _fake_actions_graph(actions):
return ag


def _fake_lightrag_wrapper(summary_text: str = "A narrative summary of the session."):
wrapper = MagicMock()
wrapper.get_lightrag.return_value.llm_model_func = AsyncMock(return_value=summary_text)
return wrapper


@pytest.mark.asyncio
async def test_reconcile_session_success_marks_completed_and_links_chunks():
db = _stub_db()
g = _graph(db)
actions = [Message(session_id="s-1", role=MessageRole.ASSISTANT, content="Alice works on the graph engine.")]
actions_graph = _fake_actions_graph(actions)
lightrag_wrapper = MagicMock()
lightrag_wrapper = _fake_lightrag_wrapper()

fake_chunk = Chunk(text="Alice works on the graph engine.", hash=content_hash("Alice works on the graph engine."))
with patch("unstructured2graph.from_texts", new=AsyncMock(return_value=[[fake_chunk]])) as mock_from_texts:
Expand All @@ -146,13 +178,33 @@ async def test_reconcile_session_success_marks_completed_and_links_chunks():
mock_from_texts.assert_awaited_once()


@pytest.mark.asyncio
async def test_reconcile_session_writes_episode_from_dedicated_llm_call():
db = _stub_db()
g = _graph(db)
actions = [Message(session_id="s-1", role=MessageRole.ASSISTANT, content="Alice works on the graph engine.")]
actions_graph = _fake_actions_graph(actions)
lightrag_wrapper = _fake_lightrag_wrapper("Alice was discussed working on the graph engine.")

fake_chunk = Chunk(text="Alice works on the graph engine.", hash=content_hash("Alice works on the graph engine."))
with patch("unstructured2graph.from_texts", new=AsyncMock(return_value=[[fake_chunk]])):
summary = await g.reconcile_session("s-1", lightrag_wrapper=lightrag_wrapper, actions_graph=actions_graph)

assert summary.summary_written is True
lightrag_wrapper.get_lightrag.return_value.llm_model_func.assert_awaited_once()
episode_calls = [call for call in db.query.call_args_list if "HAS_EPISODE" in call.args[0]]
assert len(episode_calls) == 1
assert episode_calls[0].kwargs["params"]["summary"] == "Alice was discussed working on the graph engine."
assert episode_calls[0].kwargs["params"]["session_id"] == "s-1"


@pytest.mark.asyncio
async def test_reconcile_session_passes_promotion_and_ontology_kwargs_through_to_from_texts():
db = _stub_db()
g = _graph(db)
actions = [Message(session_id="s-1", role=MessageRole.ASSISTANT, content="Alice works on the graph engine.")]
actions_graph = _fake_actions_graph(actions)
lightrag_wrapper = MagicMock()
lightrag_wrapper = _fake_lightrag_wrapper()

fake_chunk = Chunk(text="Alice works on the graph engine.", hash=content_hash("Alice works on the graph engine."))
with patch("unstructured2graph.from_texts", new=AsyncMock(return_value=[[fake_chunk]])) as mock_from_texts:
Expand Down Expand Up @@ -181,7 +233,7 @@ async def test_reconcile_session_dedupes_identical_text_before_calling_lightrag(
Message(session_id="s-1", role=MessageRole.ASSISTANT, content="Same question"),
]
actions_graph = _fake_actions_graph(actions)
lightrag_wrapper = MagicMock()
lightrag_wrapper = _fake_lightrag_wrapper()

fake_chunk = Chunk(text="Same question", hash=content_hash("Same question"))
with patch("unstructured2graph.from_texts", new=AsyncMock(return_value=[[fake_chunk]])) as mock_from_texts:
Expand All @@ -206,6 +258,7 @@ async def test_reconcile_session_no_reconcilable_content_skips_lightrag_but_stil

assert summary.status == "completed"
assert summary.texts_considered == 0
assert summary.summary_written is False
mock_from_texts.assert_not_awaited()


Expand Down
Loading