fix(embedding): truncate on conclusion create paths and document BPE loop
Storage callers in create_observations (API + agent tools) now pass on_oversize="truncate" so a single oversize item cannot drop the batch. Docstring on _truncate_to_token_limit notes why decode/re-encode is load-bearing.
This commit is contained in:
parent
0cf316e6ac
commit
0342900ff1
|
|
@ -984,7 +984,9 @@ async def create_observations(
|
|||
# Generate embeddings in batch
|
||||
contents = [obs.content for obs in observations]
|
||||
try:
|
||||
embeddings = await embedding_client.simple_batch_embed(contents)
|
||||
embeddings = await embedding_client.simple_batch_embed(
|
||||
contents, on_oversize="truncate"
|
||||
)
|
||||
except ValueError as e:
|
||||
raise ValidationException(str(e)) from e
|
||||
|
||||
|
|
|
|||
|
|
@ -318,7 +318,10 @@ class _EmbeddingClient:
|
|||
)
|
||||
|
||||
def _truncate_to_token_limit(self, text: str) -> tuple[str, int]:
|
||||
"""Return a prefix of `text` whose re-encoded token count fits the cap."""
|
||||
"""Return a prefix of `text` whose re-encoded token count fits the cap.
|
||||
|
||||
Decode/re-encode after slicing: BPE boundaries can re-expand past the cap.
|
||||
"""
|
||||
token_ids = self.encoding.encode(text)
|
||||
keep = self.max_embedding_tokens
|
||||
while len(token_ids) > self.max_embedding_tokens:
|
||||
|
|
|
|||
|
|
@ -930,7 +930,9 @@ async def create_observations(
|
|||
run_id=run_id,
|
||||
parent_category=parent_category,
|
||||
):
|
||||
embeddings = await embedding_client.simple_batch_embed(contents)
|
||||
embeddings = await embedding_client.simple_batch_embed(
|
||||
contents, on_oversize="truncate"
|
||||
)
|
||||
embeddings_by_index = dict(
|
||||
zip(range(len(normalized_observations)), embeddings, strict=True)
|
||||
)
|
||||
|
|
|
|||
|
|
@ -1006,6 +1006,46 @@ class TestDocumentCRUD:
|
|||
assert documents[0].content in ["Observation 1", "Observation 2"]
|
||||
assert documents[1].content in ["Observation 1", "Observation 2"]
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_create_observations_embeds_with_truncate_on_oversize(
|
||||
self,
|
||||
db_session: AsyncSession,
|
||||
sample_data: tuple[models.Workspace, models.Peer],
|
||||
):
|
||||
"""API conclusion creates must opt into truncation on oversize content."""
|
||||
test_workspace, test_peer = sample_data
|
||||
test_peer2, test_session, _ = await self._setup_test_data(
|
||||
db_session, test_workspace, test_peer
|
||||
)
|
||||
|
||||
with patch(
|
||||
"src.crud.document.embedding_client.simple_batch_embed",
|
||||
new=AsyncMock(return_value=[[0.1] * 1536, [0.2] * 1536]),
|
||||
) as mock_embed:
|
||||
created = await crud.create_observations(
|
||||
db_session,
|
||||
observations=[
|
||||
schemas.ConclusionCreate(
|
||||
content="short conclusion",
|
||||
observer_id=test_peer.name,
|
||||
observed_id=test_peer2.name,
|
||||
session_id=test_session.name,
|
||||
),
|
||||
schemas.ConclusionCreate(
|
||||
content="another conclusion",
|
||||
observer_id=test_peer.name,
|
||||
observed_id=test_peer2.name,
|
||||
session_id=test_session.name,
|
||||
),
|
||||
],
|
||||
workspace_name=test_workspace.name,
|
||||
)
|
||||
|
||||
assert len(created) == 2
|
||||
mock_embed.assert_awaited_once_with(
|
||||
["short conclusion", "another conclusion"], on_oversize="truncate"
|
||||
)
|
||||
|
||||
|
||||
class TestSessionPurityInvariant:
|
||||
"""Regression tests for the explicit-document session-purity invariant.
|
||||
|
|
|
|||
|
|
@ -334,7 +334,9 @@ class TestCreateObservations:
|
|||
"""If batch embedding fails but individual embeds succeed, all observations are created."""
|
||||
workspace, peer1, peer2, session, _, _ = tool_test_data
|
||||
|
||||
async def fail_batch_embed(_texts: list[str]) -> list[list[float]]:
|
||||
async def fail_batch_embed(
|
||||
_texts: list[str], **_kwargs: object
|
||||
) -> list[list[float]]:
|
||||
raise RuntimeError("embedding provider timeout")
|
||||
|
||||
async def succeed_single_embed(_content: str) -> list[float]:
|
||||
|
|
@ -393,7 +395,9 @@ class TestCreateObservations:
|
|||
"""If batch embedding fails and some individual embeds also fail, only successful ones are created."""
|
||||
workspace, peer1, peer2, session, _, _ = tool_test_data
|
||||
|
||||
async def fail_batch_embed(_texts: list[str]) -> list[list[float]]:
|
||||
async def fail_batch_embed(
|
||||
_texts: list[str], **_kwargs: object
|
||||
) -> list[list[float]]:
|
||||
raise RuntimeError("embedding provider timeout")
|
||||
|
||||
async def embed_per_observation(content: str) -> list[float]:
|
||||
|
|
@ -458,7 +462,9 @@ class TestCreateObservations:
|
|||
workspace, peer1, peer2, session, _, _ = tool_test_data
|
||||
created_documents: list[Any] = []
|
||||
|
||||
async def fake_batch_embed(texts: list[str]) -> list[list[float]]:
|
||||
async def fake_batch_embed(
|
||||
texts: list[str], **_kwargs: object
|
||||
) -> list[list[float]]:
|
||||
assert texts == ["trimmed observation"]
|
||||
return [[0.4, 0.5, 0.6]]
|
||||
|
||||
|
|
@ -504,6 +510,60 @@ class TestCreateObservations:
|
|||
assert len(created_documents) == 1
|
||||
assert created_documents[0].content == "trimmed observation"
|
||||
|
||||
async def test_create_observations_embeds_with_truncate_on_oversize(
|
||||
self,
|
||||
tool_test_data: Any,
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
):
|
||||
"""Storage path must opt into truncation so one long obs cannot drop the batch."""
|
||||
workspace, peer1, peer2, session, _, _ = tool_test_data
|
||||
captured: dict[str, object] = {}
|
||||
|
||||
async def fake_batch_embed(
|
||||
texts: list[str], *, on_oversize: str = "raise", **_kwargs: object
|
||||
) -> list[list[float]]:
|
||||
captured["texts"] = texts
|
||||
captured["on_oversize"] = on_oversize
|
||||
return [[0.1] for _ in texts]
|
||||
|
||||
async def fake_create_documents(
|
||||
_db: AsyncSession,
|
||||
documents: list[Any],
|
||||
workspace_name: str,
|
||||
*,
|
||||
observer: str,
|
||||
observed: str,
|
||||
deduplicate: bool = False,
|
||||
) -> crud.CreateDocumentsResult:
|
||||
_ = (workspace_name, observer, observed, deduplicate)
|
||||
return crud.CreateDocumentsResult(created_documents=documents)
|
||||
|
||||
monkeypatch.setattr(
|
||||
"src.utils.agent_tools.embedding_client.simple_batch_embed",
|
||||
fake_batch_embed,
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
"src.utils.agent_tools.crud.create_documents", fake_create_documents
|
||||
)
|
||||
|
||||
result = await create_observations(
|
||||
observations=[
|
||||
schemas.ObservationInput(content="short fact", level="explicit"),
|
||||
schemas.ObservationInput(content="long fact", level="explicit"),
|
||||
],
|
||||
observer=peer1.name,
|
||||
observed=peer2.name,
|
||||
session_name=session.name,
|
||||
workspace_name=workspace.name,
|
||||
message_ids=[],
|
||||
message_created_at=str(datetime.now(timezone.utc)),
|
||||
)
|
||||
|
||||
assert isinstance(result, ObservationsCreatedResult)
|
||||
assert result.created_count == 2
|
||||
assert captured["on_oversize"] == "truncate"
|
||||
assert captured["texts"] == ["short fact", "long fact"]
|
||||
|
||||
async def test_create_observations_skips_all_blank_content(
|
||||
self,
|
||||
tool_test_data: Any,
|
||||
|
|
|
|||
Loading…
Reference in New Issue