diff --git a/src/crud/document.py b/src/crud/document.py index ba80712d..0cec85c9 100644 --- a/src/crud/document.py +++ b/src/crud/document.py @@ -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 diff --git a/src/crud/representation.py b/src/crud/representation.py index 57e6e8e1..3b7070e8 100644 --- a/src/crud/representation.py +++ b/src/crud/representation.py @@ -107,7 +107,7 @@ class RepresentationManager: parent_category="representation", ): embeddings = await embedding_client.simple_batch_embed( - observation_texts + observation_texts, on_oversize="truncate" ) except ValueError as e: raise exceptions.ValidationException( diff --git a/src/deriver/deriver.py b/src/deriver/deriver.py index 2ad7d14f..f76c4d52 100644 --- a/src/deriver/deriver.py +++ b/src/deriver/deriver.py @@ -7,6 +7,7 @@ from src import crud from src.config import ConfiguredModelSettings, settings from src.crud.representation import RepresentationManager from src.dependencies import tracked_db +from src.exceptions import RepresentationSaveError from src.llm import honcho_llm_call from src.llm.types import LLMTelemetryContext from src.models import Message @@ -196,6 +197,7 @@ async def process_representation_tasks_batch( agg_representation_result = crud.CreateDocumentsResult() successful_observer_count = 0 + save_errors: list[tuple[str, Exception]] = [] if observations.is_empty() or not message_ids: logger.warning( "Deriver generated zero observations for messages %s:%s in %s/%s!", @@ -236,10 +238,11 @@ async def process_representation_tasks_batch( representation_result.semantic_dup_replaced_count ) successful_observer_count += 1 - except Exception as e: - logger.error( - "Failed to save representation for observer %s: %s", observer, e + except Exception as e: # noqa: BLE001 + logger.exception( + "Failed to save representation for observer %s", observer ) + save_errors.append((observer, e)) # Log metrics overall_duration = (time.perf_counter() - overall_start) * 1000 @@ -337,5 +340,16 @@ async def process_representation_tasks_batch( exact_dup_in_batch_count=agg_representation_result.exact_dup_in_batch_count, semantic_dup_rejected_count=agg_representation_result.semantic_dup_rejected_count, semantic_dup_replaced_count=agg_representation_result.semantic_dup_replaced_count, + failed_observer_count=len(save_errors), ) ) + + if save_errors and successful_observer_count == 0: + details = "; ".join( + f"{observer}: {exc.__class__.__name__}: {exc}" + for observer, exc in save_errors + ) + raise RepresentationSaveError( + f"save_representation failed for all {len(save_errors)} observer(s): " + + details + ) from save_errors[0][1] diff --git a/src/embedding_client.py b/src/embedding_client.py index 3cee4484..e55e09fd 100644 --- a/src/embedding_client.py +++ b/src/embedding_client.py @@ -321,39 +321,78 @@ class _EmbeddingClient: fn=_call_openai, ) - async def simple_batch_embed(self, texts: list[str]) -> list[list[float]]: + def _truncate_to_token_limit(self, text: str) -> tuple[str, int]: + """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. """ - Batch-embed a list of text strings. Each input must already fit within - `max_embedding_tokens`; this method does not sub-chunk oversized inputs. + token_ids = self.encoding.encode(text) + keep = self.max_embedding_tokens + while len(token_ids) > self.max_embedding_tokens: + keep = min(keep, len(token_ids) - 1) + if keep < 1: + return "", 0 + text = self.encoding.decode(token_ids[:keep]) + token_ids = self.encoding.encode(text) + keep -= 1 + return text, len(token_ids) + + async def simple_batch_embed( + self, + texts: list[str], + *, + on_oversize: Literal["raise", "truncate"] = "raise", + ) -> list[list[float]]: + """ + Batch-embed a list of text strings. Does not sub-chunk oversized inputs. Internally goes through the same token-aware batching pipeline as `batch_embed()` so the per-request token cap is respected. Args: texts: List of text strings to embed + on_oversize: ``"raise"`` (default) errors; ``"truncate"`` embeds a + token-capped prefix. Returns: List of embedding vectors, one per input text (in order) Raises: - ValueError: If any text exceeds token limits + ValueError: If any text exceeds token limits and `on_oversize` is + ``"raise"`` """ if not texts: return [] - # Validate per-input token limit and collect token counts for batching + # Validate / cap per-input token limit and collect counts for batching + prepared_texts: list[str] = [] token_counts: list[int] = [] for idx, text in enumerate(texts): - tokens = len(self.encoding.encode(text)) - if tokens > self.max_embedding_tokens: - raise ValueError( - f"Text at index {idx} exceeds maximum token limit of {self.max_embedding_tokens} tokens (got {tokens} tokens)" - ) + token_ids = self.encoding.encode(text) + if len(token_ids) > self.max_embedding_tokens: + if on_oversize == "truncate": + original_count = len(token_ids) + text, tokens = self._truncate_to_token_limit(text) + logger.warning( + "truncated oversize embedding input at idx %d: %d->%d tokens", + idx, + original_count, + tokens, + ) + else: + raise ValueError( + f"Text at index {idx} exceeds maximum token limit of " + + f"{self.max_embedding_tokens} tokens (got {len(token_ids)} tokens)" + ) + else: + tokens = len(token_ids) + prepared_texts.append(text) token_counts.append(tokens) # Use positional indices as text_ids so we can reassemble in input order. text_chunks: dict[str, list[tuple[str, int]]] = { - str(i): [(text, token_counts[i])] for i, text in enumerate(texts) + str(i): [(prepared_texts[i], token_counts[i])] + for i in range(len(prepared_texts)) } batches = self._create_batches(text_chunks) @@ -695,9 +734,16 @@ class EmbeddingClient: """Embed a single query string.""" return await self._get_client().embed(query) - async def simple_batch_embed(self, texts: list[str]) -> list[list[float]]: + async def simple_batch_embed( + self, + texts: list[str], + *, + on_oversize: Literal["raise", "truncate"] = "raise", + ) -> list[list[float]]: """Batch embed a list of text strings (each must fit token limit).""" - return await self._get_client().simple_batch_embed(texts) + return await self._get_client().simple_batch_embed( + texts, on_oversize=on_oversize + ) def prepare_chunks(self, id_resource_dict: dict[str, str]) -> dict[str, list[str]]: """Chunk texts using the same rules as `batch_embed` (no network).""" diff --git a/src/exceptions.py b/src/exceptions.py index 129324b6..41d775fe 100644 --- a/src/exceptions.py +++ b/src/exceptions.py @@ -133,6 +133,14 @@ class VectorStoreError(HonchoException): detail = "Vector store operation failed" +@final +class RepresentationSaveError(HonchoException): + """Raised when every observer's representation save fails in a batch.""" + + status_code: int = 500 + detail: str = "Representation save failed for all observers" + + class LLMError(Exception): """Exception raised when an LLM call fails. diff --git a/src/reconciler/sync_vectors.py b/src/reconciler/sync_vectors.py index 1ff94542..1a8e99b5 100644 --- a/src/reconciler/sync_vectors.py +++ b/src/reconciler/sync_vectors.py @@ -302,7 +302,9 @@ async def _sync_documents( EmbeddingCallPurpose.VECTOR_SYNC.value, parent_category="reconciliation", ): - new_embeddings = await embedding_client.simple_batch_embed(contents) + new_embeddings = await embedding_client.simple_batch_embed( + contents, on_oversize="truncate" + ) if len(new_embeddings) != len(docs_needing_embed): logger.warning( diff --git a/src/telemetry/events/representation.py b/src/telemetry/events/representation.py index ad6dfca4..c4291092 100644 --- a/src/telemetry/events/representation.py +++ b/src/telemetry/events/representation.py @@ -154,6 +154,10 @@ class RepresentationCompletedEvent(BaseEvent): default=0, description="Number of observers this representation was saved against", ) + failed_observer_count: int = Field( + default=0, + description="Number of observers whose save_representation failed (partial or total)", + ) def get_resource_id(self) -> str: """Resource ID includes workspace, session, and latest message for uniqueness.""" diff --git a/src/utils/agent_tools.py b/src/utils/agent_tools.py index 7f094154..b8af9eee 100644 --- a/src/utils/agent_tools.py +++ b/src/utils/agent_tools.py @@ -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) ) diff --git a/tests/conftest.py b/tests/conftest.py index f99859b3..8a078652 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -604,7 +604,9 @@ def mock_openai_embeddings(request: pytest.FixtureRequest): mock_embed.side_effect = embed_side_effect - async def mock_simple_batch_embed_func(texts: list[str]) -> list[list[float]]: + async def mock_simple_batch_embed_func( + texts: list[str], **_kwargs: object + ) -> list[list[float]]: return [_content_to_embedding(text) for text in texts] mock_simple_batch_embed.side_effect = mock_simple_batch_embed_func diff --git a/tests/crud/test_document.py b/tests/crud/test_document.py index 5f03c80b..6686e688 100644 --- a/tests/crud/test_document.py +++ b/tests/crud/test_document.py @@ -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. diff --git a/tests/crud/test_representation_manager.py b/tests/crud/test_representation_manager.py index 6c5e094c..3c2f57e4 100644 --- a/tests/crud/test_representation_manager.py +++ b/tests/crud/test_representation_manager.py @@ -520,7 +520,9 @@ class TestRepresentationManagerSave: ) assert len(saved.created_documents) == 1 - mock_embed.assert_awaited_once_with(["useful observation"]) + mock_embed.assert_awaited_once_with( + ["useful observation"], on_oversize="truncate" + ) saved_observations = _saved_observations(mock_save) assert len(saved_observations) == 1 assert saved_observations[0].content == "useful observation" @@ -576,7 +578,9 @@ class TestRepresentationManagerSave: ) assert len(saved.created_documents) == 1 - mock_embed.assert_awaited_once_with(["inferred conclusion"]) + mock_embed.assert_awaited_once_with( + ["inferred conclusion"], on_oversize="truncate" + ) saved_observations = _saved_observations(mock_save) assert len(saved_observations) == 1 assert isinstance(saved_observations[0], DeductiveObservation) @@ -630,6 +634,61 @@ class TestRepresentationManagerSave: mock_embed.assert_not_awaited() mock_save.assert_not_awaited() + @pytest.mark.asyncio + async def test_save_representation_embeds_with_truncate_on_oversize(self): + """One oversize observation must not drop the rest of the batch.""" + manager = RepresentationManager( + "workspace", + observer="observer", + observed="observed", + ) + representation = Representation( + explicit=[ + ExplicitObservation( + content="short fact", + created_at=datetime.now(timezone.utc), + message_ids=[1], + session_name="session", + ) + ], + deductive=[ + DeductiveObservation( + conclusion="inferred fact", + premises=["premise"], + source_ids=["doc-a"], + created_at=datetime.now(timezone.utc), + message_ids=[1], + session_name="session", + ) + ], + ) + + with ( + patch("src.crud.representation.tracked_db", _fake_tracked_db), + patch( + "src.crud.representation.embedding_client.simple_batch_embed", + new=AsyncMock(return_value=[[0.1], [0.2]]), + ) as mock_embed, + patch.object( + manager, + "_save_representation_internal", + new=AsyncMock( + return_value=CreateDocumentsResult(created_documents=[MagicMock()]) + ), + ), + ): + await manager.save_representation( + representation, + message_ids=[1], + session_name="session", + message_created_at=datetime.now(timezone.utc), + message_level_configuration=_resolved_config(), + ) + + mock_embed.assert_awaited_once_with( + ["inferred fact", "short fact"], on_oversize="truncate" + ) + class TestVectorQueryTopKFloor: """Regression for HONCHO-19Q / HONCHO-4Q4. diff --git a/tests/deriver/test_deriver_processing.py b/tests/deriver/test_deriver_processing.py index 6785d159..29a983b5 100644 --- a/tests/deriver/test_deriver_processing.py +++ b/tests/deriver/test_deriver_processing.py @@ -7,7 +7,9 @@ import pytest from src import crud, models from src.config import settings +from src.crud.representation import RepresentationManager from src.deriver.deriver import process_representation_tasks_batch +from src.exceptions import RepresentationSaveError from src.llm import HonchoLLMCallResponse from src.utils.representation import ( ExplicitObservationBase, @@ -70,6 +72,116 @@ class TestDeriverProcessing: assert kwargs["model_config"].stop_sequences == expected_config.stop_sequences assert "llm_settings" not in kwargs + async def test_all_observer_saves_failing_surfaces_failure(self): + """When every observer's save_representation fails, the batch must raise.""" + message = Mock( + id=1, + public_id="msg_1", + session_name="session-1", + workspace_name="workspace-1", + peer_name="alice", + content="hello", + token_count=5, + created_at=datetime.now(timezone.utc), + ) + configuration = Mock() + configuration.reasoning.enabled = True + + mock_response = HonchoLLMCallResponse( + content=PromptRepresentation( + explicit=[ + ExplicitObservationBase(content="The user has a dog named Rover") + ] + ), + input_tokens=10, + output_tokens=5, + finish_reasons=["STOP"], + ) + + failing_save = AsyncMock(side_effect=RuntimeError("429 RESOURCE_EXHAUSTED")) + emitted: list[Any] = [] + with ( + patch( + "src.deriver.deriver.honcho_llm_call", + new_callable=AsyncMock, + return_value=mock_response, + ), + patch.object(RepresentationManager, "save_representation", failing_save), + patch("src.deriver.deriver.emit", side_effect=emitted.append), + pytest.raises(RepresentationSaveError, match="save_representation failed"), + ): + await process_representation_tasks_batch( + messages=[message], + message_level_configuration=configuration, + observers=["bob"], + observed="alice", + queue_item_message_ids=[1], + ) + + # Telemetry must fire *before* the raise so a total save failure is still + # visible to metrics. Guards against emit() being moved after the raise. + assert emitted, "expected telemetry to be emitted before the raised failure" + assert emitted[-1].observer_count == 0 + assert emitted[-1].failed_observer_count == 1 + + async def test_partial_observer_failure_is_processed_and_surfaced(self): + """When some observers save and one fails, the batch does NOT raise + (saved observers are kept) and the failure is visible via telemetry. + """ + message = Mock( + id=1, + public_id="msg_1", + session_name="session-1", + workspace_name="workspace-1", + peer_name="alice", + content="hello", + token_count=5, + created_at=datetime.now(timezone.utc), + ) + configuration = Mock() + configuration.reasoning.enabled = True + + mock_response = HonchoLLMCallResponse( + content=PromptRepresentation( + explicit=[ + ExplicitObservationBase(content="The user has a dog named Rover") + ] + ), + input_tokens=10, + output_tokens=5, + finish_reasons=["STOP"], + ) + + # bob succeeds, carol fails. + partial_save = AsyncMock( + side_effect=[ + crud.CreateDocumentsResult(), + RuntimeError("429 RESOURCE_EXHAUSTED"), + ] + ) + emitted: list[Any] = [] + with ( + patch( + "src.deriver.deriver.honcho_llm_call", + new_callable=AsyncMock, + return_value=mock_response, + ), + patch.object(RepresentationManager, "save_representation", partial_save), + patch("src.deriver.deriver.emit", side_effect=emitted.append), + ): + await process_representation_tasks_batch( + messages=[message], + message_level_configuration=configuration, + observers=["bob", "carol"], + observed="alice", + queue_item_message_ids=[1], + ) + + assert emitted, "expected a telemetry event to be emitted" + event = emitted[-1] + assert event.observer_count == 1 + assert event.failed_observer_count == 1 + async def test_process_representation_tasks_batch_passes_custom_instructions_into_prompt( self, ) -> None: diff --git a/tests/deriver/test_vector_reconciliation.py b/tests/deriver/test_vector_reconciliation.py index 5e2e4712..0ecd4d5c 100644 --- a/tests/deriver/test_vector_reconciliation.py +++ b/tests/deriver/test_vector_reconciliation.py @@ -551,9 +551,12 @@ class TestReEmbedding: # Mock embedding client to track batch calls batch_call_count = 0 - async def track_batch_embed(contents: list[str]) -> list[list[float]]: + async def track_batch_embed( + contents: list[str], *, on_oversize: str, **_kwargs: object + ) -> list[list[float]]: nonlocal batch_call_count batch_call_count += 1 + assert on_oversize == "truncate" return [[1.0] * 1536 for _ in contents] with patch("src.reconciler.sync_vectors.embedding_client") as mock_embed_client: diff --git a/tests/live_llm/README.md b/tests/live_llm/README.md index 442049bc..5cf39e4b 100644 --- a/tests/live_llm/README.md +++ b/tests/live_llm/README.md @@ -68,5 +68,5 @@ Coverage by provider: - OpenAI transport → OpenRouter non-reasoning models (e.g. `inception/mercury-2`): non-chat / diffusion architectures must stay on `max_tokens`, no `reasoning_effort`, tool-calling parameter-schema compatibility is the canary for exotic OR-served providers - Gemini 2.5/3.0 classes: structured outputs, cached-content reuse, thought signatures, multi-turn tool replay - Gemini 3.1 class: thinking and tool replay coverage by default; structured-output/caching coverage should only be added once Google documents support for that path -- Embeddings (`test_live_embeddings.py`): single embed, batched embed, batch-vs-single alignment, and chunk-to-id mapping for every configured embedding model. `gemini-embedding-2*` is the reason this exists — those models collapse a list of bare strings into one document (#745), and only a live call catches it. Also covers first-class `EmbeddingModelConfig.timeout` plumbing (one representative model per transport): configured timeout lands on the SDK client, and a near-zero timeout aborts before the provider answers +- Embeddings (`test_live_embeddings.py`): single embed, batched embed, batch-vs-single alignment, chunk-to-id mapping, and oversize-truncate survival (`on_oversize="truncate"`) for every configured embedding model. `gemini-embedding-2*` is the reason this exists — those models collapse a list of bare strings into one document (#745), and only a live call catches it. Also covers first-class `EmbeddingModelConfig.timeout` plumbing (one representative model per transport): configured timeout lands on the SDK client, and a near-zero timeout aborts before the provider answers - OpenAI-compatible embedding providers (e.g. OpenRouter's `google/gemini-embedding-001`): the #932 surface. Those providers reject a base64 embedding request outright (HTTP 400) or answer HTTP 200 with empty data, so the whole matrix fails without `encoding_format="float"`. Real OpenAI accepts base64 happily, so only a third-party provider catches it. Note that OpenRouter load-balances across upstreams, so the base64 failure is per-attempt rather than guaranteed: a retry can land on an endpoint that accepts it. `test_live_openai_float_encoding_matches_base64` covers the other side, that the float switch must not move vectors on real OpenAI diff --git a/tests/live_llm/test_live_embeddings.py b/tests/live_llm/test_live_embeddings.py index 6ef2b048..7c57a2e3 100644 --- a/tests/live_llm/test_live_embeddings.py +++ b/tests/live_llm/test_live_embeddings.py @@ -197,6 +197,26 @@ async def test_live_openai_float_encoding_matches_base64( ), f"{spec.id}: float encoding diverges from base64 (cosine={similarity:.8f})" +@pytest.mark.asyncio +@pytest.mark.parametrize("spec", ALL_SPECS, ids=lambda spec: spec.id) +async def test_live_batch_embed_truncates_oversize_instead_of_dropping_batch( + spec: LiveEmbeddingSpec, +) -> None: + """on_oversize='truncate' keeps one vector per input when an item exceeds the cap.""" + # Tiny cap so the oversize input stays cheap to tokenize and send. + client = make_embedding_client(spec, max_input_tokens=32) + oversize = " ".join(f"oversize-token-{index}" for index in range(200)) + assert len(client.encoding.encode(oversize)) > client.max_embedding_tokens + + texts = [BATCH_TEXTS[0], oversize, BATCH_TEXTS[1]] + embeddings = await client.simple_batch_embed(texts, on_oversize="truncate") + + assert len(embeddings) == len(texts) + assert all(len(embedding) == spec.dimensions for embedding in embeddings) + # A collapsed or dropped batch would reuse a vector or return fewer. + assert len({tuple(embedding) for embedding in embeddings}) == len(texts) + + @pytest.mark.asyncio @pytest.mark.parametrize("spec", GEMINI_SPECS, ids=lambda spec: spec.id) async def test_live_gemini_batch_embed_survives_batch_split( diff --git a/tests/llm/test_embedding_client.py b/tests/llm/test_embedding_client.py index 9df80be1..fc2ff411 100644 --- a/tests/llm/test_embedding_client.py +++ b/tests/llm/test_embedding_client.py @@ -13,6 +13,7 @@ from src.config import ( ) from src.embedding_client import ( BatchItem, + EmbeddingClient, _EmbeddingClient, # pyright: ignore[reportPrivateUsage] ) @@ -910,6 +911,131 @@ async def test_simple_batch_embed_rejects_oversized_input( await client.simple_batch_embed([too_long]) +@pytest.mark.asyncio +async def test_simple_batch_embed_truncates_oversize_when_requested( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """on_oversize='truncate' embeds a prefix instead of failing the batch.""" + fake_embeddings = FakeOpenAIEmbeddingsAPI([0.1] * 4) + + class FakeOpenAIClient: + def __init__(self, *, api_key: str | None, base_url: str | None) -> None: + self.embeddings: FakeOpenAIEmbeddingsAPI = fake_embeddings + + monkeypatch.setattr("openai.AsyncOpenAI", FakeOpenAIClient) + + client = _EmbeddingClient( + EmbeddingModelConfig( + transport="openai", + model="text-embedding-3-small", + api_key="test-key", + base_url=None, + ), + vector_dimensions=4, + max_input_tokens=10, + max_tokens_per_request=1000, + send_dimensions=False, + ) + + short = "hello" + too_long = ("word " * 50).strip() + assert len(client.encoding.encode(too_long)) > client.max_embedding_tokens + + out = await client.simple_batch_embed([short, too_long], on_oversize="truncate") + + assert len(out) == 2 + assert fake_embeddings.calls, "expected a provider call after truncation" + received = fake_embeddings.calls[0]["input"] + assert received[0] == short + truncated = received[1] + assert isinstance(truncated, str) + assert truncated != too_long + assert len(client.encoding.encode(truncated)) <= client.max_embedding_tokens + + +@pytest.mark.asyncio +async def test_simple_batch_embed_truncate_reencodes_until_under_cap( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """decode(ids[:n]) can re-encode past n; truncate must re-verify the count.""" + fake_embeddings = FakeOpenAIEmbeddingsAPI([0.1] * 4) + + class FakeOpenAIClient: + def __init__(self, *, api_key: str | None, base_url: str | None) -> None: + self.embeddings: FakeOpenAIEmbeddingsAPI = fake_embeddings + + monkeypatch.setattr("openai.AsyncOpenAI", FakeOpenAIClient) + + client = _EmbeddingClient( + EmbeddingModelConfig( + transport="openai", + model="text-embedding-3-small", + api_key="test-key", + base_url=None, + ), + vector_dimensions=4, + max_input_tokens=10, + max_tokens_per_request=1000, + send_dimensions=False, + ) + + encode_calls = {"n": 0} + + def encode(text: str) -> list[int]: + encode_calls["n"] += 1 + if text.startswith("LONG"): + # 1: original oversize; 2: still over after first slice; 3+: fits. + if encode_calls["n"] == 1: + return list(range(20)) + if encode_calls["n"] == 2: + return list(range(12)) + return list(range(8)) + return [1] + + def decode(ids: list[int]) -> str: + return "LONG" + "x" * len(ids) + + monkeypatch.setattr(client.encoding, "encode", encode) + monkeypatch.setattr(client.encoding, "decode", decode) + + out = await client.simple_batch_embed(["LONG-input"], on_oversize="truncate") + + assert len(out) == 1 + received = fake_embeddings.calls[0]["input"][0] + assert isinstance(received, str) + # The provider must see the post-loop text, which encodes to 8 (<= cap). + assert encode(received) == list(range(8)) + assert encode_calls["n"] >= 3 + + +@pytest.mark.asyncio +async def test_public_embedding_client_forwards_on_oversize( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """The singleton wrapper must forward on_oversize to the inner client.""" + captured: dict[str, object] = {} + + class FakeInner: + async def simple_batch_embed( + self, + texts: list[str], + *, + on_oversize: str = "raise", + ) -> list[list[float]]: + captured["texts"] = texts + captured["on_oversize"] = on_oversize + return [[0.1]] + + wrapper = EmbeddingClient() + monkeypatch.setattr(wrapper, "_get_client", lambda: FakeInner()) + + out = await wrapper.simple_batch_embed(["hi"], on_oversize="truncate") + + assert out == [[0.1]] + assert captured["texts"] == ["hi"] + assert captured["on_oversize"] == "truncate" + + def test_prepare_chunks_returns_ordered_chunks( monkeypatch: pytest.MonkeyPatch, ) -> None: diff --git a/tests/telemetry/test_representation_v2_fields.py b/tests/telemetry/test_representation_v2_fields.py index a41bed41..6f2251ea 100644 --- a/tests/telemetry/test_representation_v2_fields.py +++ b/tests/telemetry/test_representation_v2_fields.py @@ -51,6 +51,7 @@ class TestRepresentationV2AdditiveFields: assert event.exact_dup_existing_count == 0 assert event.semantic_dup_rejected_count == 0 assert event.semantic_dup_replaced_count == 0 + assert event.failed_observer_count == 0 def test_input_tokens_semantics_preserved(self): """The downstream metering key must remain 'queued-message tokens'. @@ -161,6 +162,7 @@ class TestRepresentationV2AdditiveFields: "exact_dup_existing_count", "semantic_dup_rejected_count", "semantic_dup_replaced_count", + "failed_observer_count", ): assert field in data, f"missing field: {field}" assert data["hit_batch_token_cap"] is True diff --git a/tests/utils/test_agent_tools.py b/tests/utils/test_agent_tools.py index e53598fa..ac45cabf 100644 --- a/tests/utils/test_agent_tools.py +++ b/tests/utils/test_agent_tools.py @@ -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,