diff --git a/src/crud/representation.py b/src/crud/representation.py index 5ca7d308..e9b199a9 100644 --- a/src/crud/representation.py +++ b/src/crud/representation.py @@ -179,7 +179,7 @@ class RepresentationManager: metadata: schemas.DocumentMetadata = schemas.DocumentMetadata( message_ids=message_ids, - batch_message_ids=obs.batch_message_ids, + source_message_ids=obs.source_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 96b514da..b02cddb8 100644 --- a/src/deriver/deriver.py +++ b/src/deriver/deriver.py @@ -38,15 +38,15 @@ def _get_deriver_model_config() -> ConfiguredModelSettings: 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] = [] + prompt_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) + prompt_message_ids.append(message.id) - return "\n".join(formatted_messages), batch_message_ids + return "\n".join(formatted_messages), prompt_message_ids @with_sentry_transaction("minimal_deriver_batch", op="deriver") @@ -119,7 +119,7 @@ async def process_representation_tasks_batch( # 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) + formatted_messages, prompt_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) @@ -202,7 +202,7 @@ async def process_representation_tasks_batch( observations = Representation.from_prompt_representation( response.content, message_ids, - batch_message_ids, + prompt_message_ids, latest_message.session_name, latest_message.created_at, ) diff --git a/src/schemas/internal.py b/src/schemas/internal.py index a75b44c1..7f44ec49 100644 --- a/src/schemas/internal.py +++ b/src/schemas/internal.py @@ -33,9 +33,9 @@ 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( + source_message_ids: list[int] = Field( default_factory=list, - description="The full ordered message ID list enumerated in the deriver prompt", + description="Canonical citation message IDs resolved from source_indices", ) 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." @@ -62,7 +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="Per-derivation debugging provenance only: 0-based positions in the deriver batch that lose meaning across deduplication merges; source_message_ids is the canonical citation", ) diff --git a/src/utils/representation.py b/src/utils/representation.py index 4c82de2c..efbc9381 100644 --- a/src/utils/representation.py +++ b/src/utils/representation.py @@ -85,14 +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( + source_message_ids: list[int] = Field( default_factory=list, - description="The full ordered message ID list enumerated in the deriver prompt", + description="Canonical citation message IDs resolved from source_indices", ) 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="Per-derivation debugging provenance only: 0-based positions in the deriver batch that lose meaning across deduplication merges; source_message_ids is the canonical citation", ) @@ -642,8 +642,8 @@ class Representation(BaseModel): message_ids=flatten_message_ids( doc.internal_metadata.get("message_ids", []) ), - batch_message_ids=doc.internal_metadata.get( - "batch_message_ids", [] + source_message_ids=doc.internal_metadata.get( + "source_message_ids", [] ), session_name=doc.session_name, source_indices=doc.internal_metadata.get("source_indices", []), @@ -713,7 +713,7 @@ class Representation(BaseModel): cls, prompt_representation: "PromptRepresentation", message_ids: list[int], - batch_message_ids: list[int], + prompt_message_ids: list[int], session_name: str, created_at: datetime, ) -> "Representation": @@ -721,10 +721,12 @@ class Representation(BaseModel): explicit_observations: list[ExplicitObservation] = [] for explicit in prompt_representation.explicit: valid_source_indices: list[int] = [] + source_message_ids: list[int] = [] invalid_source_indices: list[int] = [] for source_index in explicit.source_indices: - if 0 <= source_index < len(batch_message_ids): + if 0 <= source_index < len(prompt_message_ids): valid_source_indices.append(source_index) + source_message_ids.append(prompt_message_ids[source_index]) else: invalid_source_indices.append(source_index) @@ -733,7 +735,7 @@ class Representation(BaseModel): "Dropping out-of-range source_indices %s for observation %r; deriver batch contains %d messages", invalid_source_indices, explicit.content, - len(batch_message_ids), + len(prompt_message_ids), ) explicit_observations.append( @@ -742,7 +744,7 @@ class Representation(BaseModel): source_indices=valid_source_indices, created_at=created_at, message_ids=message_ids, - batch_message_ids=batch_message_ids, + source_message_ids=source_message_ids, session_name=session_name, ) ) diff --git a/tests/crud/test_representation_manager.py b/tests/crud/test_representation_manager.py index 204dff3c..bf134fcf 100644 --- a/tests/crud/test_representation_manager.py +++ b/tests/crud/test_representation_manager.py @@ -483,7 +483,7 @@ class TestRepresentationManagerSave: source_indices=[0, 1], created_at=datetime.now(timezone.utc), message_ids=[20], - batch_message_ids=[10, 20], + source_message_ids=[10, 20], session_name="session", ) @@ -511,9 +511,11 @@ class TestRepresentationManagerSave: 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_message_ids == [10, 20] assert document.metadata.source_indices == [0, 1] - assert document.metadata.model_dump(exclude_none=True)["batch_message_ids"] == [ + assert document.metadata.model_dump(exclude_none=True)[ + "source_message_ids" + ] == [ 10, 20, ] diff --git a/tests/deriver/test_representation_crud.py b/tests/deriver/test_representation_crud.py index 6ab92d83..2e791331 100644 --- a/tests/deriver/test_representation_crud.py +++ b/tests/deriver/test_representation_crud.py @@ -111,7 +111,7 @@ def test_prompt_representation_conversion(): rep = Representation.from_prompt_representation( pr, message_ids=[1], - batch_message_ids=[1], + prompt_message_ids=[1], session_name="s", created_at=timestamp, ) @@ -135,7 +135,7 @@ def test_mixed_peer_source_indices_resolve_against_prompt_order( for message in messages: message.created_at = created_at - formatted_messages, batch_message_ids = _format_messages_for_prompt(messages) + formatted_messages, prompt_message_ids = _format_messages_for_prompt(messages) prompt_representation = PromptRepresentation( explicit=[ ExplicitObservationBase( @@ -149,25 +149,22 @@ def test_mixed_peer_source_indices_resolve_against_prompt_order( representation = Representation.from_prompt_representation( prompt_representation, message_ids=[20], - batch_message_ids=batch_message_ids, + prompt_message_ids=prompt_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 observation.source_message_ids == [10, 20] assert [ (line[:3], message_id, line.split(": ", 1)[1]) for message_id, line in zip( - batch_message_ids, formatted_messages.splitlines(), strict=True + prompt_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 619c3804..392bfc87 100644 --- a/tests/integration/test_representation.py +++ b/tests/integration/test_representation.py @@ -321,6 +321,8 @@ class TestDocumentCreationWorkflow: level="explicit", internal_metadata={ "message_ids": [1], + "source_message_ids": [2, 1], + "source_indices": [1, 0], }, session_name="test_session", created_at=datetime(2025, 1, 1, 10, 0, 0, tzinfo=timezone.utc), @@ -350,6 +352,8 @@ class TestDocumentCreationWorkflow: explicit_obs = representation.explicit[0] assert explicit_obs.content == "User said they like programming" assert explicit_obs.message_ids == [1] + assert explicit_obs.source_message_ids == [2, 1] + assert explicit_obs.source_indices == [1, 0] assert explicit_obs.session_name == "test_session" deductive_obs = representation.deductive[0] @@ -358,6 +362,24 @@ class TestDocumentCreationWorkflow: assert deductive_obs.message_ids == [1] assert deductive_obs.session_name == "test_session" + async def test_representation_from_old_document_defaults_source_message_ids(self): + """Old documents without resolved citations remain readable.""" + explicit_doc = models.Document( + id="old_explicit_doc_id", + workspace_name="test_workspace", + observer="test_peer", + observed="test_peer", + content="An older observation", + level="explicit", + internal_metadata={"message_ids": [1]}, + session_name="test_session", + created_at=datetime(2025, 1, 1, 10, 0, 0, tzinfo=timezone.utc), + ) + + representation = Representation.from_documents([explicit_doc]) + + assert representation.explicit[0].source_message_ids == [] + async def create_test_workspace_and_peer( self, db_session: AsyncSession, workspace_name: str | None = None ) -> tuple[models.Workspace, models.Peer]: @@ -412,7 +434,10 @@ class TestPromptRepresentationConversion: """ prompt_rep = PromptRepresentation( explicit=[ - ExplicitObservationBase(content="User likes coffee"), + ExplicitObservationBase( + source_indices=[0], + content="User likes coffee", + ), ExplicitObservationBase(content="User works remotely"), ], ) @@ -422,7 +447,7 @@ class TestPromptRepresentationConversion: representation = Representation.from_prompt_representation( prompt_rep, message_ids=[123], - batch_message_ids=[123], + prompt_message_ids=[123], session_name="test_session", created_at=timestamp, ) @@ -435,8 +460,10 @@ class TestPromptRepresentationConversion: # Check explicit observations assert representation.explicit[0].content == "User likes coffee" assert representation.explicit[0].message_ids == [123] + assert representation.explicit[0].source_message_ids == [123] assert representation.explicit[0].session_name == "test_session" assert representation.explicit[1].content == "User works remotely" + assert representation.explicit[1].source_message_ids == [] assert representation.explicit[0].created_at == timestamp async def test_empty_prompt_representation_conversion(self): @@ -445,7 +472,7 @@ class TestPromptRepresentationConversion: representation = Representation.from_prompt_representation( empty_prompt_rep, message_ids=[1], - batch_message_ids=[1], + prompt_message_ids=[1], session_name="test", created_at=datetime(2025, 1, 1, 12, 0, 0, tzinfo=timezone.utc), )