From 5bce6b951a6cabcf4743aceaeb8bdbb2b809b19b Mon Sep 17 00:00:00 2001 From: Matt Landers Date: Tue, 28 Jul 2026 15:38:22 -0400 Subject: [PATCH] fix: persist deriver batch message IDs for source index resolution --- src/crud/representation.py | 1 + src/deriver/deriver.py | 23 +++++++-- src/schemas/internal.py | 7 ++- src/utils/representation.py | 46 ++++++++++++++---- tests/crud/test_representation_manager.py | 47 ++++++++++++++++++ tests/deriver/test_representation_crud.py | 58 +++++++++++++++++++++++ tests/integration/test_representation.py | 2 + 7 files changed, 169 insertions(+), 15 deletions(-) diff --git a/src/crud/representation.py b/src/crud/representation.py index ffe18085..5ca7d308 100644 --- a/src/crud/representation.py +++ b/src/crud/representation.py @@ -179,6 +179,7 @@ class RepresentationManager: metadata: schemas.DocumentMetadata = schemas.DocumentMetadata( message_ids=message_ids, + batch_message_ids=obs.batch_message_ids, premises=obs_premises, message_created_at=format_datetime_utc(message_created_at), source_indices=getattr(obs, "source_indices", []), diff --git a/src/deriver/deriver.py b/src/deriver/deriver.py index 05e89d1e..96b514da 100644 --- a/src/deriver/deriver.py +++ b/src/deriver/deriver.py @@ -35,6 +35,20 @@ def _get_deriver_model_config() -> ConfiguredModelSettings: return settings.DERIVER.MODEL_CONFIG +def _format_messages_for_prompt(messages: list[Message]) -> tuple[str, list[int]]: + """Format one ordered message batch and retain its index-to-ID mapping.""" + formatted_messages: list[str] = [] + batch_message_ids: list[int] = [] + for index, message in enumerate(messages): + formatted_message = format_new_turn_with_timestamp( + message.content, message.created_at, message.peer_name + ) + formatted_messages.append(f"[{index}] {formatted_message}") + batch_message_ids.append(message.id) + + return "\n".join(formatted_messages), batch_message_ids + + @with_sentry_transaction("minimal_deriver_batch", op="deriver") async def process_representation_tasks_batch( messages: list[Message], @@ -103,11 +117,9 @@ async def process_representation_tasks_batch( "id", ) - # Format messages with timestamps and 0-based indices for source citation - formatted_messages = "\n".join( - f"[{i}] {format_new_turn_with_timestamp(msg.content, msg.created_at, msg.peer_name)}" - for i, msg in enumerate(messages) - ) + # Build the prompt text and its index-to-ID mapping in one pass so they + # cannot disagree about ordering. + formatted_messages, batch_message_ids = _format_messages_for_prompt(messages) # Track token usage - count only tokens from messages being processed prompt_tokens = estimate_deriver_prompt_tokens(custom_instructions) @@ -190,6 +202,7 @@ async def process_representation_tasks_batch( observations = Representation.from_prompt_representation( response.content, message_ids, + batch_message_ids, latest_message.session_name, latest_message.created_at, ) diff --git a/src/schemas/internal.py b/src/schemas/internal.py index 16890b24..a75b44c1 100644 --- a/src/schemas/internal.py +++ b/src/schemas/internal.py @@ -33,6 +33,10 @@ class DocumentMetadata(BaseModel): message_ids: list[int] = Field( description="The ID range(s) of the messages that this document was derived from. Acts as a link to the primary source of the document. Note that as a document gets deduplicated, additional ranges will be added, because the same document could be derived from completely separate message ranges." ) + batch_message_ids: list[int] = Field( + default_factory=list, + description="The full ordered message ID list enumerated in the deriver prompt", + ) message_created_at: str = Field( description="The timestamp of the message that this document was derived from. Note that this is not the same as the created_at timestamp of the document. This timestamp is usually only saved with second-level precision." ) @@ -58,8 +62,7 @@ class DocumentMetadata(BaseModel): ) source_indices: list[int] = Field( default_factory=list, - description="0-based indices into the deriver batch's message list " - "indicating which messages directly support this observation", + description="0-based indices into the deriver batch's message list indicating which messages directly support this observation", ) diff --git a/src/utils/representation.py b/src/utils/representation.py index 97ad27c2..a3c17945 100644 --- a/src/utils/representation.py +++ b/src/utils/representation.py @@ -1,3 +1,4 @@ +import logging from collections.abc import Sequence from datetime import datetime from typing import Any @@ -7,6 +8,8 @@ from pydantic import BaseModel, Field, field_validator from src import models from src.utils.formatting import parse_datetime_iso +logger = logging.getLogger(__name__) + # Conclusion levels whose `session_name` stamp is trustworthy enough to scope on. # # Explicit conclusions come from the deriver over a single session's message @@ -82,11 +85,14 @@ class ObservationMetadata(BaseModel): id: str = Field(default="", description="Document ID for this observation") created_at: datetime message_ids: list[int] + batch_message_ids: list[int] = Field( + default_factory=list, + description="The full ordered message ID list enumerated in the deriver prompt", + ) session_name: str | None = None source_indices: list[int] = Field( default_factory=list, - description="0-based indices into the deriver batch's message list " - "indicating which messages directly support this observation", + description="0-based indices into the deriver batch's message list indicating which messages directly support this observation", ) @@ -636,6 +642,9 @@ class Representation(BaseModel): message_ids=flatten_message_ids( doc.internal_metadata.get("message_ids", []) ), + batch_message_ids=doc.internal_metadata.get( + "batch_message_ids", [] + ), session_name=doc.session_name, source_indices=doc.internal_metadata.get("source_indices", []), ) @@ -704,21 +713,42 @@ class Representation(BaseModel): cls, prompt_representation: "PromptRepresentation", message_ids: list[int], + batch_message_ids: list[int], session_name: str, created_at: datetime, ) -> "Representation": """Convert PromptRepresentation to Representation.""" - return cls( - explicit=[ + explicit_observations: list[ExplicitObservation] = [] + for explicit in prompt_representation.explicit: + valid_source_indices: list[int] = [] + invalid_source_indices: list[int] = [] + for source_index in explicit.source_indices: + if 0 <= source_index < len(batch_message_ids): + valid_source_indices.append(source_index) + else: + invalid_source_indices.append(source_index) + + if invalid_source_indices: + logger.warning( + "Dropping out-of-range source_indices %s for observation %r; deriver batch contains %d messages", + invalid_source_indices, + explicit.content, + len(batch_message_ids), + ) + + explicit_observations.append( ExplicitObservation( - content=e.content, - source_indices=e.source_indices, + content=explicit.content, + source_indices=valid_source_indices, created_at=created_at, message_ids=message_ids, + batch_message_ids=batch_message_ids, session_name=session_name, ) - for e in prompt_representation.explicit - ], + ) + + return cls( + explicit=explicit_observations, deductive=[], inductive=[], ) diff --git a/tests/crud/test_representation_manager.py b/tests/crud/test_representation_manager.py index f551de76..204dff3c 100644 --- a/tests/crud/test_representation_manager.py +++ b/tests/crud/test_representation_manager.py @@ -471,6 +471,53 @@ class TestRepresentationManagerSessionScoping: class TestRepresentationManagerSave: + @pytest.mark.asyncio + async def test_save_representation_threads_source_trace_metadata(self): + manager = RepresentationManager( + "workspace", + observer="observer", + observed="alice", + ) + observation = ExplicitObservation( + content="Alice chose the first option", + source_indices=[0, 1], + created_at=datetime.now(timezone.utc), + message_ids=[20], + batch_message_ids=[10, 20], + session_name="session", + ) + + with ( + patch( + "src.crud.representation.crud.get_or_create_collection", + new=AsyncMock(return_value=MagicMock()), + ), + patch( + "src.crud.representation.crud.create_documents", + new=AsyncMock(return_value=CreateDocumentsResult()), + ) as mock_create_documents, + ): + await manager._save_representation_internal( # pyright: ignore[reportPrivateUsage] + MagicMock(spec=AsyncSession), + [observation], + [[0.1]], + message_ids=[20], + session_name="session", + message_created_at=datetime.now(timezone.utc), + message_level_configuration=_resolved_config(), + ) + + create_call = mock_create_documents.await_args + assert create_call is not None + document = create_call.args[1][0] + assert document.metadata.message_ids == [20] + assert document.metadata.batch_message_ids == [10, 20] + assert document.metadata.source_indices == [0, 1] + assert document.metadata.model_dump(exclude_none=True)["batch_message_ids"] == [ + 10, + 20, + ] + @pytest.mark.asyncio async def test_save_representation_filters_blank_observations_before_embedding( self, diff --git a/tests/deriver/test_representation_crud.py b/tests/deriver/test_representation_crud.py index cb4f974d..2d9d21ce 100644 --- a/tests/deriver/test_representation_crud.py +++ b/tests/deriver/test_representation_crud.py @@ -1,5 +1,12 @@ import datetime +import logging +import pytest + +from src import models +from src.deriver.deriver import ( + _format_messages_for_prompt, # pyright: ignore[reportPrivateUsage] +) from src.utils.representation import ( DeductiveObservation, ExplicitObservation, @@ -97,6 +104,7 @@ def test_prompt_representation_conversion(): rep = Representation.from_prompt_representation( pr, message_ids=[1], + batch_message_ids=[1], session_name="s", created_at=timestamp, ) @@ -106,3 +114,53 @@ def test_prompt_representation_conversion(): # (they would be created directly by the Dreamer via the create_observations tool) assert len(rep.deductive) == 0 assert rep.explicit[0].created_at == timestamp + + +def test_mixed_peer_source_indices_resolve_against_prompt_order( + caplog: pytest.LogCaptureFixture, +) -> None: + created_at = datetime.datetime(2025, 1, 1, 12, 0, 0, tzinfo=datetime.timezone.utc) + messages = [ + models.Message(id=10, peer_name="bob", content="Which option?"), + models.Message(id=20, peer_name="alice", content="The first one"), + models.Message(id=30, peer_name="bob", content="Got it"), + ] + for message in messages: + message.created_at = created_at + + formatted_messages, batch_message_ids = _format_messages_for_prompt(messages) + prompt_representation = PromptRepresentation( + explicit=[ + ExplicitObservationBase( + content="Alice chose the first option", + source_indices=[0, 1, 3], + ) + ] + ) + + with caplog.at_level(logging.WARNING, logger="src.utils.representation"): + representation = Representation.from_prompt_representation( + prompt_representation, + message_ids=[20], + batch_message_ids=batch_message_ids, + session_name="s", + created_at=created_at, + ) + + observation = representation.explicit[0] + assert observation.source_indices == [0, 1] + assert observation.batch_message_ids == [10, 20, 30] + assert [ + (line[:3], message_id, line.split(": ", 1)[1]) + for message_id, line in zip( + batch_message_ids, formatted_messages.splitlines(), strict=True + ) + ] == [ + ("[0]", 10, "Which option?"), + ("[1]", 20, "The first one"), + ("[2]", 30, "Got it"), + ] + assert [ + observation.batch_message_ids[index] for index in observation.source_indices + ] == [10, 20] + assert "Dropping out-of-range source_indices [3]" in caplog.text diff --git a/tests/integration/test_representation.py b/tests/integration/test_representation.py index 60f80c5c..619c3804 100644 --- a/tests/integration/test_representation.py +++ b/tests/integration/test_representation.py @@ -422,6 +422,7 @@ class TestPromptRepresentationConversion: representation = Representation.from_prompt_representation( prompt_rep, message_ids=[123], + batch_message_ids=[123], session_name="test_session", created_at=timestamp, ) @@ -444,6 +445,7 @@ class TestPromptRepresentationConversion: representation = Representation.from_prompt_representation( empty_prompt_rep, message_ids=[1], + batch_message_ids=[1], session_name="test", created_at=datetime(2025, 1, 1, 12, 0, 0, tzinfo=timezone.utc), )