RomanNum3ral
/
honcho
mirror of https://github.com/plastic-labs/honcho.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

fix: persist deriver batch message IDs for source index resolution

This commit is contained in:
Matt Landers 2026-07-28 15:38:22 -04:00
parent b53c68a078
commit 5bce6b951a
7 changed files with 169 additions and 15 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

1
src/crud/representation.py
View File

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

23
src/deriver/deriver.py
View File

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

7
src/schemas/internal.py
View File

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

46
src/utils/representation.py
View File

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

47
tests/crud/test_representation_manager.py
View File

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

58
tests/deriver/test_representation_crud.py
View File

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

2
tests/integration/test_representation.py
View File

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