Store resolved source_message_ids instead of full batch ID list
This commit is contained in:
parent
aac8f52c3d
commit
e70b412191
|
|
@ -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", []),
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
)
|
||||
|
|
|
|||
|
|
@ -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",
|
||||
)
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
)
|
||||
)
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
]
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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),
|
||||
)
|
||||
|
|
|
|||
Loading…
Reference in New Issue