Persist segment context

edwinyyyu · edwinyyyu · commit 5d451f64d372 · 2026-03-17T11:34:52.000-07:00
Signed-off-by: Edwin Yu &lt;edwinyyyu@gmail.com&gt;
diff --git a/packages/server/server_tests/memmachine_server/episodic_memory/extra_memory/segment_linker/test_sqlalchemy_segment_linker.py b/packages/server/server_tests/memmachine_server/episodic_memory/extra_memory/segment_linker/test_sqlalchemy_segment_linker.py
@@ -14,6 +14,8 @@
 
 from memmachine_server.common.filter.filter_parser import Comparison
 from memmachine_server.episodic_memory.extra_memory.data_types import (
+    CitationContext,
+    MessageContext,
     Segment,
     Text,
 )
@@ -44,6 +46,7 @@ def _seg(
     offset: int = 0,
     ts_offset_seconds: int = 0,
     text: str = "hello",
+    context: MessageContext | CitationContext | None = None,
     properties: dict | None = None,
 ) -> Segment:
     return Segment(
@@ -53,6 +56,7 @@ def _seg(
         offset=offset,
         timestamp=BASE_TIME + timedelta(seconds=ts_offset_seconds),
         block=Text(text=text),
+        context=context,
         properties=properties or {},
     )
 
@@ -150,6 +154,71 @@ async def test_register_with_properties(
     assert returned.properties == {"color": "red", "score": 42}
 
 
+@pytest.mark.asyncio
+async def test_register_with_message_context(
+    partition: SQLAlchemySegmentLinkerPartition,
+) -> None:
+    ctx = MessageContext(source="User")
+    seg = _seg(context=ctx)
+    deriv = uuid4()
+    await partition.register_segments({seg: [deriv]})
+
+    result = await partition.get_segments_by_derivatives([deriv])
+    returned = next(iter(result[deriv]))
+    assert returned.context == ctx
+
+
+@pytest.mark.asyncio
+async def test_register_with_citation_context(
+    partition: SQLAlchemySegmentLinkerPartition,
+) -> None:
+    ctx = CitationContext(source="docs.txt", source_type="file", location="/tmp")
+    seg = _seg(context=ctx)
+    deriv = uuid4()
+    await partition.register_segments({seg: [deriv]})
+
+    result = await partition.get_segments_by_derivatives([deriv])
+    returned = next(iter(result[deriv]))
+    assert returned.context == ctx
+
+
+@pytest.mark.asyncio
+async def test_register_with_no_context(
+    partition: SQLAlchemySegmentLinkerPartition,
+) -> None:
+    seg = _seg()
+    deriv = uuid4()
+    await partition.register_segments({seg: [deriv]})
+
+    result = await partition.get_segments_by_derivatives([deriv])
+    returned = next(iter(result[deriv]))
+    assert returned.context is None
+
+
+@pytest.mark.asyncio
+async def test_context_preserved_in_segment_contexts(
+    partition: SQLAlchemySegmentLinkerPartition,
+) -> None:
+    """Context is preserved when retrieving segment contexts (backward/forward)."""
+    ep = uuid4()
+    ctx_user = MessageContext(source="User")
+    ctx_assistant = MessageContext(source="Assistant")
+    s0 = _seg(episode_uuid=ep, offset=0, ts_offset_seconds=0, context=ctx_user)
+    s1 = _seg(episode_uuid=ep, offset=1, ts_offset_seconds=1, context=ctx_assistant)
+    s2 = _seg(episode_uuid=ep, offset=2, ts_offset_seconds=2, context=ctx_user)
+    deriv = uuid4()
+    await partition.register_segments({s0: [deriv], s1: [deriv], s2: [deriv]})
+
+    result = await partition.get_segment_contexts(
+        [s1.uuid], max_backward_segments=1, max_forward_segments=1
+    )
+    ctx = list(result[s1.uuid])
+    assert len(ctx) == 3
+    assert ctx[0].context == ctx_user
+    assert ctx[1].context == ctx_assistant
+    assert ctx[2].context == ctx_user
+
+
 @pytest.mark.asyncio
 async def test_register_active_validation(
     partition: SQLAlchemySegmentLinkerPartition,
@@ -1109,3 +1178,68 @@ async def registerer() -> None:
 
     await asyncio.gather(purger(), registerer())
     assert errors == []
+
+
+@pytest.mark.integration
+@pytest.mark.asyncio
+async def test_pg_context_preserved_via_lateral_join(
+    pg_linker: SQLAlchemySegmentLinker,
+) -> None:
+    """Context is preserved when retrieved via the LATERAL join path (multiple seeds)."""
+    partition = pg_linker.get_partition(PARTITION_KEY)
+    ep = uuid4()
+    ctx_user = MessageContext(source="User")
+    ctx_assistant = MessageContext(source="Assistant")
+    s0 = _seg(episode_uuid=ep, offset=0, ts_offset_seconds=0, context=ctx_user)
+    s1 = _seg(episode_uuid=ep, offset=1, ts_offset_seconds=1, context=ctx_assistant)
+    s2 = _seg(episode_uuid=ep, offset=2, ts_offset_seconds=2, context=ctx_user)
+    s3 = _seg(episode_uuid=ep, offset=3, ts_offset_seconds=3, context=ctx_assistant)
+    s4 = _seg(episode_uuid=ep, offset=4, ts_offset_seconds=4, context=ctx_user)
+    deriv = uuid4()
+    await partition.register_segments(
+        {s0: [deriv], s1: [deriv], s2: [deriv], s3: [deriv], s4: [deriv]}
+    )
+
+    # Two seeds exercises the LATERAL join code path.
+    result = await partition.get_segment_contexts(
+        [s1.uuid, s3.uuid], max_backward_segments=1, max_forward_segments=1
+    )
+
+    ctx_a = list(result[s1.uuid])
+    assert len(ctx_a) == 3
+    assert ctx_a[0].context == ctx_user
+    assert ctx_a[1].context == ctx_assistant
+    assert ctx_a[2].context == ctx_user
+
+    ctx_b = list(result[s3.uuid])
+    assert len(ctx_b) == 3
+    assert ctx_b[0].context == ctx_user
+    assert ctx_b[1].context == ctx_assistant
+    assert ctx_b[2].context == ctx_user
+
+
+@pytest.mark.integration
+@pytest.mark.asyncio
+async def test_pg_mixed_context_types(
+    pg_linker: SQLAlchemySegmentLinker,
+) -> None:
+    """Different context types (message, citation, None) round-trip correctly on PG."""
+    partition = pg_linker.get_partition(PARTITION_KEY)
+    ctx_msg = MessageContext(source="User")
+    ctx_cite = CitationContext(source="paper.pdf", source_type="file", location="p.3")
+
+    s_msg = _seg(ts_offset_seconds=0, context=ctx_msg)
+    s_cite = _seg(ts_offset_seconds=1, context=ctx_cite)
+    s_none = _seg(ts_offset_seconds=2)
+
+    d1, d2, d3 = uuid4(), uuid4(), uuid4()
+    await partition.register_segments({s_msg: [d1], s_cite: [d2], s_none: [d3]})
+
+    r1 = await partition.get_segments_by_derivatives([d1])
+    assert next(iter(r1[d1])).context == ctx_msg
+
+    r2 = await partition.get_segments_by_derivatives([d2])
+    assert next(iter(r2[d2])).context == ctx_cite
+
+    r3 = await partition.get_segments_by_derivatives([d3])
+    assert next(iter(r3[d3])).context is None
diff --git a/packages/server/src/memmachine_server/episodic_memory/extra_memory/segment_linker/sqlalchemy_segment_linker.py b/packages/server/src/memmachine_server/episodic_memory/extra_memory/segment_linker/sqlalchemy_segment_linker.py
@@ -52,6 +52,7 @@
 )
 from memmachine_server.episodic_memory.extra_memory.data_types import (
     Block,
+    Context,
     Segment,
 )
 from memmachine_server.episodic_memory.extra_memory.segment_linker.segment_linker import (
@@ -63,6 +64,7 @@
 logger = logging.getLogger(__name__)
 
 _JSON_AUTO = JSON().with_variant(JSONB, "postgresql")
+_ContextAdapter = TypeAdapter(Context | None)
 _BlockAdapter = TypeAdapter(Block)
 
 
@@ -94,6 +96,7 @@ class SegmentRow(BaseSegmentLinker):
     timestamp: MappedColumn[datetime] = mapped_column(
         DateTime(timezone=True), nullable=False
     )
+    context = mapped_column(_JSON_AUTO, nullable=True)
     block: MappedColumn[Block] = mapped_column(_JSON_AUTO, nullable=False)
 
     __table_args__ = (
@@ -275,6 +278,7 @@ async def _insert_links(
                 "index": segment.index,
                 "offset": segment.offset,
                 "timestamp": segment.timestamp,
+                "context": segment.context.model_dump() if segment.context else None,
                 "block": segment.block.model_dump(),
             }
             for segment in links
@@ -534,6 +538,7 @@ async def _lateral_query(
                 lateral_subquery.c.index,
                 lateral_subquery.c.offset,
                 lateral_subquery.c.timestamp,
+                lateral_subquery.c.context,
                 lateral_subquery.c.block,
             ).select_from(seeds_subquery.join(lateral_subquery, true()))
 
@@ -549,6 +554,7 @@ async def _lateral_query(
                         index=row.index,
                         offset=row.offset,
                         timestamp=row.timestamp,
+                        context=row.context,
                         block=row.block,
                     )
                 )
@@ -1046,13 +1052,15 @@ def _segment_from_segment_row(
         properties: dict[str, PropertyValue],
     ) -> Segment:
         """Convert a SegmentRow and its properties into a Segment."""
+        context = _ContextAdapter.validate_python(row.context)
         block = _BlockAdapter.validate_python(row.block)
         return Segment(
             uuid=row.uuid,
             episode_uuid=row.episode_uuid,
             index=row.index,
             offset=row.offset,
             timestamp=row.timestamp,
+            context=context,
             block=block,
             properties=properties,
         )
