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)