Store resolved source_message_ids instead of full batch ID list

This commit is contained in:
Matt Landers 2026-07-29 15:37:04 -04:00
parent aac8f52c3d
commit e70b412191
7 changed files with 60 additions and 32 deletions

View File

@ -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", []),

View File

@ -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,
)

View File

@ -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",
)

View File

@ -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,
)
)

View File

@ -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,
]

View File

@ -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

View File

@ -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),
)