From 5823f0fae930f2da9f1b45d3a0f6a16b148e70bc Mon Sep 17 00:00:00 2001 From: Ken Weiner Date: Tue, 25 Aug 2026 06:45:01 -0700 Subject: [PATCH] fix: only classify genuine oversize input as a token-limit error (#791) Callers wrapped every ValueError from the embedding client in a "exceeds maximum token limit" message, so provider and configuration failures (dimension mismatch, empty response, upstream error) surfaced to users as though their input were too long. Add EmbeddingTokenLimitError, raised only by the pre-flight token checks in embed() and simple_batch_embed(), and narrow the remaps in search.py, agent_tools.py, document.py and representation.py to catch it. It subclasses ValueError so existing broad handlers keep working. Both simple_batch_embed() remap sites pass on_oversize="truncate" and so could never raise a token-limit error at all; their handlers only ever mislabelled provider failures. Fixes #568 Co-authored-by: Claude Opus 5 --- src/crud/document.py | 6 +-- src/crud/representation.py | 4 +- src/embedding_client.py | 19 +++++++-- src/utils/agent_tools.py | 8 +++- src/utils/search.py | 4 +- tests/llm/test_embedding_client.py | 68 ++++++++++++++++++++++++++++++ 6 files changed, 96 insertions(+), 13 deletions(-) diff --git a/src/crud/document.py b/src/crud/document.py index 0cec85c9..1b04bb0a 100644 --- a/src/crud/document.py +++ b/src/crud/document.py @@ -18,7 +18,7 @@ from src.crud.collection import get_or_create_collection from src.crud.peer import get_peer, reject_scope_observed from src.crud.session import get_session from src.dependencies import tracked_db -from src.embedding_client import embedding_client +from src.embedding_client import EmbeddingTokenLimitError, embedding_client from src.exceptions import ( ResourceNotFoundException, ValidationException, @@ -362,7 +362,7 @@ async def query_documents( if embedding is None: try: embedding = await embedding_client.embed(query) - except ValueError as e: + except EmbeddingTokenLimitError as e: raise ValidationException( "Query exceeds maximum token limit of " + f"{settings.EMBEDDING.MAX_INPUT_TOKENS}." @@ -987,7 +987,7 @@ async def create_observations( embeddings = await embedding_client.simple_batch_embed( contents, on_oversize="truncate" ) - except ValueError as e: + except EmbeddingTokenLimitError as e: raise ValidationException(str(e)) from e # Create document objects and track embeddings for vector store diff --git a/src/crud/representation.py b/src/crud/representation.py index 3b7070e8..6fafb842 100644 --- a/src/crud/representation.py +++ b/src/crud/representation.py @@ -13,7 +13,7 @@ from src import crud, exceptions, models, schemas from src.config import settings from src.dependencies import tracked_db from src.dreamer.dream_scheduler import check_and_schedule_dream -from src.embedding_client import embedding_client +from src.embedding_client import EmbeddingTokenLimitError, embedding_client from src.schemas import ResolvedConfiguration from src.telemetry.events import EmbeddingCallPurpose from src.telemetry.logging import accumulate_metric @@ -109,7 +109,7 @@ class RepresentationManager: embeddings = await embedding_client.simple_batch_embed( observation_texts, on_oversize="truncate" ) - except ValueError as e: + except EmbeddingTokenLimitError as e: raise exceptions.ValidationException( "Observation content exceeds maximum token limit of " + f"{settings.EMBEDDING.MAX_INPUT_TOKENS}." diff --git a/src/embedding_client.py b/src/embedding_client.py index e55e09fd..d3f4e369 100644 --- a/src/embedding_client.py +++ b/src/embedding_client.py @@ -159,6 +159,17 @@ def _publish_embedding_event( logger.debug("Failed to emit EmbeddingCallCompletedEvent", exc_info=True) +class EmbeddingTokenLimitError(ValueError): + """Raised when input text genuinely exceeds the model's token limit. + + Subclasses ``ValueError`` so existing broad handlers keep working, while + letting callers tell a real "content too long" condition apart from a + transient provider or configuration failure (dimension mismatch, empty + response, upstream error). Only the pre-flight token checks raise this; + provider failures keep raising plain ``ValueError``. + """ + + class BatchItem(NamedTuple): """A single item in a batch with its metadata.""" @@ -272,7 +283,7 @@ class _EmbeddingClient: token_count = len(self.encoding.encode(query)) if token_count > self.max_embedding_tokens: - raise ValueError( + raise EmbeddingTokenLimitError( f"Query exceeds maximum token limit of {self.max_embedding_tokens} tokens (got {token_count} tokens)" ) @@ -358,8 +369,8 @@ class _EmbeddingClient: List of embedding vectors, one per input text (in order) Raises: - ValueError: If any text exceeds token limits and `on_oversize` is - ``"raise"`` + EmbeddingTokenLimitError: If any text exceeds token limits and + `on_oversize` is ``"raise"`` """ if not texts: return [] @@ -380,7 +391,7 @@ class _EmbeddingClient: tokens, ) else: - raise ValueError( + raise EmbeddingTokenLimitError( f"Text at index {idx} exceeds maximum token limit of " + f"{self.max_embedding_tokens} tokens (got {len(token_ids)} tokens)" ) diff --git a/src/utils/agent_tools.py b/src/utils/agent_tools.py index 0dbdd304..5f87d455 100644 --- a/src/utils/agent_tools.py +++ b/src/utils/agent_tools.py @@ -13,7 +13,7 @@ from sqlalchemy.ext.asyncio import AsyncSession from src import crud, models, schemas from src.config import settings from src.dependencies import tracked_db -from src.embedding_client import embedding_client +from src.embedding_client import EmbeddingTokenLimitError, embedding_client from src.exceptions import ResourceNotFoundException from src.models import Document from src.schemas import ResolvedConfiguration @@ -1884,11 +1884,15 @@ async def _handle_search_memory( parent_category=ctx.parent_category, ): query_embedding = await embedding_client.embed(query) - except ValueError: + except EmbeddingTokenLimitError: return ( "ERROR: Query exceeds maximum token limit of " + f"{settings.EMBEDDING.MAX_INPUT_TOKENS}. Please use a shorter query." ) + except ValueError as e: + # Provider/config failure, not an oversized query. Keep returning a + # string so the tool loop can continue, but don't blame the query. + return f"ERROR: Embedding the query failed: {e}" # Base telemetry metadata; results_count gets filled in below. search_meta: dict[str, Any] = { diff --git a/src/utils/search.py b/src/utils/search.py index 761b63e0..329e082b 100644 --- a/src/utils/search.py +++ b/src/utils/search.py @@ -14,7 +14,7 @@ from sqlalchemy.ext.asyncio import AsyncSession from src import models from src.config import settings from src.dependencies import tracked_db -from src.embedding_client import embedding_client +from src.embedding_client import EmbeddingTokenLimitError, embedding_client from src.exceptions import ValidationException from src.models import session_peers_table from src.telemetry.events import EmbeddingCallPurpose @@ -388,7 +388,7 @@ async def search( parent_category="api", ): query_embedding = await embedding_client.embed(query) - except ValueError as e: + except EmbeddingTokenLimitError as e: raise ValidationException( f"Query exceeds maximum token limit of {settings.EMBEDDING.MAX_INPUT_TOKENS}." ) from e diff --git a/tests/llm/test_embedding_client.py b/tests/llm/test_embedding_client.py index fc2ff411..7fd9237d 100644 --- a/tests/llm/test_embedding_client.py +++ b/tests/llm/test_embedding_client.py @@ -14,6 +14,7 @@ from src.config import ( from src.embedding_client import ( BatchItem, EmbeddingClient, + EmbeddingTokenLimitError, _EmbeddingClient, # pyright: ignore[reportPrivateUsage] ) @@ -1173,3 +1174,70 @@ async def test_gemini_process_batch_wraps_contents_as_content_part( assert all(isinstance(c, genai_types.Content) for c in contents) assert contents[0].parts[0].text == "hello" assert contents[1].parts[0].text == "world" + + +# --- Token-limit classification (issue #568) ------------------------------- +# +# Only genuine "content too long" conditions may raise +# EmbeddingTokenLimitError. Provider/config failures must stay plain +# ValueError so callers don't rewrite them as token-limit errors. + + +def test_embedding_token_limit_error_is_value_error() -> None: + """Subclassing ValueError keeps pre-existing broad handlers working.""" + assert issubclass(EmbeddingTokenLimitError, ValueError) + + +@pytest.mark.asyncio +async def test_embed_raises_token_limit_error_before_calling_provider( + monkeypatch: pytest.MonkeyPatch, +) -> None: + client, fake_embeddings = _build_openai_client( + monkeypatch, + embedding=[0.1, 0.2], + model="text-embedding-3-small", + send_dimensions=False, + vector_dimensions=2, + ) + + with pytest.raises(EmbeddingTokenLimitError): + await client.embed("word " * 20_000) + + assert fake_embeddings.calls == [], "provider must not be called on oversize input" + + +@pytest.mark.asyncio +async def test_simple_batch_embed_raises_token_limit_error_before_calling_provider( + monkeypatch: pytest.MonkeyPatch, +) -> None: + client, fake_embeddings = _build_openai_client( + monkeypatch, + embedding=[0.1, 0.2], + model="text-embedding-3-small", + send_dimensions=False, + vector_dimensions=2, + ) + + with pytest.raises(EmbeddingTokenLimitError): + await client.simple_batch_embed(["fine", "word " * 20_000]) + + assert fake_embeddings.calls == [], "provider must not be called on oversize input" + + +@pytest.mark.asyncio +async def test_provider_dimension_mismatch_is_not_a_token_limit_error( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """A wrong-width vector is a provider/config fault, not an oversized input.""" + client, _ = _build_openai_client( + monkeypatch, + embedding=[0.1, 0.2, 0.3], # 3 wide, client expects 2 + model="text-embedding-3-small", + send_dimensions=False, + vector_dimensions=2, + ) + + with pytest.raises(ValueError) as excinfo: + await client.embed("short query") + + assert not isinstance(excinfo.value, EmbeddingTokenLimitError)