diff --git a/hindsight-api-slim/hindsight_api/config.py b/hindsight-api-slim/hindsight_api/config.py index 44d2e9e81e..11243f0d34 100644 --- a/hindsight-api-slim/hindsight_api/config.py +++ b/hindsight-api-slim/hindsight_api/config.py @@ -373,6 +373,7 @@ def _resolve_operation_temperature(operation_env: str, default: float) -> float ENV_EMBEDDINGS_LITELLM_SDK_API_BASE = "HINDSIGHT_API_EMBEDDINGS_LITELLM_SDK_API_BASE" ENV_EMBEDDINGS_LITELLM_SDK_OUTPUT_DIMENSIONS = "HINDSIGHT_API_EMBEDDINGS_LITELLM_SDK_OUTPUT_DIMENSIONS" ENV_EMBEDDINGS_LITELLM_SDK_ENCODING_FORMAT = "HINDSIGHT_API_EMBEDDINGS_LITELLM_SDK_ENCODING_FORMAT" +ENV_EMBEDDINGS_LITELLM_SDK_MAX_INPUT_TOKENS = "HINDSIGHT_API_EMBEDDINGS_LITELLM_SDK_MAX_INPUT_TOKENS" ENV_RERANKER_LITELLM_SDK_API_KEY = "HINDSIGHT_API_RERANKER_LITELLM_SDK_API_KEY" ENV_RERANKER_LITELLM_SDK_MODEL = "HINDSIGHT_API_RERANKER_LITELLM_SDK_MODEL" ENV_RERANKER_LITELLM_SDK_API_BASE = "HINDSIGHT_API_RERANKER_LITELLM_SDK_API_BASE" @@ -909,6 +910,10 @@ def _parse_strategy_boosts(raw: str | None) -> dict[str, str]: # LiteLLM SDK defaults DEFAULT_EMBEDDINGS_LITELLM_SDK_MODEL = "cohere/embed-english-v3.0" DEFAULT_EMBEDDINGS_LITELLM_SDK_ENCODING_FORMAT = "float" +# Opt-in per-text input truncation (tiktoken cl100k_base tokens). Off by default; +# set to the embedding model's real input limit (e.g. 8192 for Bedrock Titan V2) +# to keep oversized content from permanently failing the embed call. See #2501. +DEFAULT_EMBEDDINGS_LITELLM_SDK_MAX_INPUT_TOKENS: int | None = None DEFAULT_RERANKER_LITELLM_SDK_MODEL = "cohere/rerank-english-v3.0" DEFAULT_HOST = "0.0.0.0" @@ -1685,6 +1690,7 @@ class HindsightConfig: embeddings_litellm_sdk_api_base: str | None embeddings_litellm_sdk_output_dimensions: int | None embeddings_litellm_sdk_encoding_format: str | None + embeddings_litellm_sdk_max_input_tokens: int | None # Gemini/Vertex AI embeddings embeddings_gemini_api_key: str | None embeddings_gemini_model: str @@ -2555,6 +2561,9 @@ def from_env(cls) -> "HindsightConfig": embeddings_litellm_sdk_encoding_format=os.getenv( ENV_EMBEDDINGS_LITELLM_SDK_ENCODING_FORMAT, DEFAULT_EMBEDDINGS_LITELLM_SDK_ENCODING_FORMAT ), + embeddings_litellm_sdk_max_input_tokens=int(v) + if (v := os.getenv(ENV_EMBEDDINGS_LITELLM_SDK_MAX_INPUT_TOKENS)) + else DEFAULT_EMBEDDINGS_LITELLM_SDK_MAX_INPUT_TOKENS, # Gemini/Vertex AI embeddings (with fallback to LLM keys) embeddings_gemini_api_key=os.getenv(ENV_EMBEDDINGS_GEMINI_API_KEY) or os.getenv(ENV_LLM_API_KEY), embeddings_gemini_model=os.getenv(ENV_EMBEDDINGS_GEMINI_MODEL, DEFAULT_EMBEDDINGS_GEMINI_MODEL), diff --git a/hindsight-api-slim/hindsight_api/engine/embeddings.py b/hindsight-api-slim/hindsight_api/engine/embeddings.py index ba9657910f..97c387a585 100644 --- a/hindsight-api-slim/hindsight_api/engine/embeddings.py +++ b/hindsight-api-slim/hindsight_api/engine/embeddings.py @@ -76,6 +76,25 @@ class _ZeroEntropyEmbedResponse(BaseModel): results: list[_ZeroEntropyEmbedResult] +def _truncate_to_tokens(text: str, max_tokens: int) -> tuple[str, int]: + """Truncate ``text`` to at most ``max_tokens`` cl100k_base tokens. + + tiktoken is an approximation of any given provider's tokenizer, so set + ``max_tokens`` with a little headroom below the model's real limit. + + Returns the (possibly truncated) text and the original token count (so the + caller can report how much was dropped); the count equals ``len(tokens)`` + whether or not truncation occurred. + """ + from .token_encoding import get_token_encoding + + enc = get_token_encoding() + tokens = enc.encode(text) + if len(tokens) <= max_tokens: + return text, len(tokens) + return enc.decode(tokens[:max_tokens]), len(tokens) + + class Embeddings(ABC): """ Abstract base class for embedding generation. @@ -1202,6 +1221,7 @@ def __init__( batch_size: int = 100, timeout: float = 60.0, encoding_format: str | None = "float", + max_input_tokens: int | None = None, ): """ Initialize LiteLLM SDK embeddings client. @@ -1216,6 +1236,10 @@ def __init__( timeout: Request timeout in seconds (default: 60.0) encoding_format: Encoding format for embeddings (default: "float"). Set to None or empty string to omit (needed for Voyage AI, Gemini). + max_input_tokens: If set, truncate each input text to this many tokens + (tiktoken cl100k_base) before embedding. Needed for models with a + fixed input-token limit (e.g. Bedrock Titan V2's hard 8192 cap), + where an oversized text would otherwise fail permanently (#2501). """ self.api_key = api_key self.model = model @@ -1224,6 +1248,7 @@ def __init__( self.batch_size = batch_size self.timeout = timeout self.encoding_format = encoding_format or None + self.max_input_tokens = max_input_tokens self._litellm = None # Will be set during initialization self._dimension: int | None = None @@ -1300,6 +1325,33 @@ def encode(self, texts: list[str]) -> list[list[float]]: if not texts: return [] + # Truncate oversized inputs before hitting the provider. Models with a + # fixed input-token limit (e.g. Bedrock Titan V2, 8192) reject an + # oversized text with a permanent error rather than truncating it + # server-side, which strands the caller (e.g. a delta mental model whose + # content grew past the cap) with no recovery path. See #2501. + if self.max_input_tokens is not None: + truncated_texts = [] + original_token_counts = [] + for t in texts: + new_text, original_tokens = _truncate_to_tokens(t, self.max_input_tokens) + truncated_texts.append(new_text) + if original_tokens > self.max_input_tokens: + original_token_counts.append(original_tokens) + texts = truncated_texts + if original_token_counts: + logger.warning( + "Embeddings: truncated %d of %d input(s) to %d tokens for model %s " + "(largest was ~%d tokens); embedded content is incomplete. " + "This usually means a mental model's content has grown past the model's " + "input limit — see issue #2501.", + len(original_token_counts), + len(texts), + self.max_input_tokens, + self.model, + max(original_token_counts), + ) + all_embeddings = [] # Process in batches @@ -1691,6 +1743,7 @@ def create_embeddings_from_env() -> Embeddings: api_base=config.embeddings_litellm_sdk_api_base, output_dimensions=config.embeddings_litellm_sdk_output_dimensions, encoding_format=config.embeddings_litellm_sdk_encoding_format, + max_input_tokens=config.embeddings_litellm_sdk_max_input_tokens, ) elif provider == "google": vertexai_project_id = config.embeddings_vertexai_project_id diff --git a/hindsight-api-slim/tests/test_litellm_sdk_embeddings.py b/hindsight-api-slim/tests/test_litellm_sdk_embeddings.py index 16c9752454..5a2ddba7b1 100644 --- a/hindsight-api-slim/tests/test_litellm_sdk_embeddings.py +++ b/hindsight-api-slim/tests/test_litellm_sdk_embeddings.py @@ -244,6 +244,65 @@ async def test_encode_empty_list(self, embeddings): assert isinstance(result, list) assert len(result) == 0 + async def test_encode_truncates_oversized_input_when_max_tokens_set(self, mock_litellm, caplog): + """Oversized inputs are truncated to max_input_tokens before embedding, with a warning (#2501).""" + import logging + + from hindsight_api.engine.token_encoding import get_token_encoding + + emb = LiteLLMSDKEmbeddings( + api_key="test_key", + model="bedrock/amazon.titan-embed-text-v2:0", + api_base=None, + batch_size=100, + timeout=60.0, + max_input_tokens=50, + ) + emb._litellm = mock_litellm + emb._dimension = 768 + mock_litellm.embedding.return_value.data = [{"embedding": [0.1] * 768, "index": 0}] + + long_text = "word " * 500 # far more than 50 tokens + with caplog.at_level(logging.WARNING): + emb.encode([long_text]) + + sent = mock_litellm.embedding.call_args.kwargs["input"][0] + enc = get_token_encoding() + assert len(enc.encode(sent)) <= 50 + assert sent != long_text # actually truncated + assert any("truncated" in r.message and r.levelno == logging.WARNING for r in caplog.records) + + async def test_encode_no_warning_when_input_within_limit(self, mock_litellm, caplog): + """No truncation warning when every input already fits under max_input_tokens.""" + import logging + + emb = LiteLLMSDKEmbeddings( + api_key="test_key", + model="bedrock/amazon.titan-embed-text-v2:0", + api_base=None, + batch_size=100, + timeout=60.0, + max_input_tokens=50, + ) + emb._litellm = mock_litellm + emb._dimension = 768 + mock_litellm.embedding.return_value.data = [{"embedding": [0.1] * 768, "index": 0}] + + with caplog.at_level(logging.WARNING): + emb.encode(["short text well under the limit"]) + + assert not any("truncated" in r.message for r in caplog.records) + + async def test_encode_does_not_truncate_when_max_tokens_unset(self, embeddings, mock_litellm): + """Without max_input_tokens the text is passed through verbatim (default behavior).""" + assert embeddings.max_input_tokens is None + mock_litellm.embedding.return_value.data = [{"embedding": [0.1] * 768, "index": 0}] + + long_text = "word " * 500 + embeddings.encode([long_text]) + + assert mock_litellm.embedding.call_args.kwargs["input"] == [long_text] + async def test_encode_before_initialization(self, mock_litellm): """Test that encode raises error if not initialized.""" emb = LiteLLMSDKEmbeddings( diff --git a/hindsight-docs/docs/developer/configuration.md b/hindsight-docs/docs/developer/configuration.md index 52cb8e4bbc..d8deae742e 100644 --- a/hindsight-docs/docs/developer/configuration.md +++ b/hindsight-docs/docs/developer/configuration.md @@ -580,6 +580,7 @@ server-level only (not overridable per tenant/bank) and a change requires a rest | `HINDSIGHT_API_EMBEDDINGS_LITELLM_SDK_API_BASE` | Custom base URL for LiteLLM SDK embeddings (optional) | - | | `HINDSIGHT_API_EMBEDDINGS_LITELLM_SDK_OUTPUT_DIMENSIONS` | Optional output embedding dimensions (provider-dependent, e.g., `768` for Gemini embedding models) | - | | `HINDSIGHT_API_EMBEDDINGS_LITELLM_SDK_ENCODING_FORMAT` | Encoding format for embedding responses. Set to empty string to omit the parameter (needed for Voyage AI, Gemini). | `float` | +| `HINDSIGHT_API_EMBEDDINGS_LITELLM_SDK_MAX_INPUT_TOKENS` | If set, truncate each text to this many tokens (tiktoken `cl100k_base`, approximate) before embedding. Set it to the model's real input limit (e.g. `8192` for Bedrock Titan V2, with a little headroom) so oversized content is truncated instead of failing the embed call permanently. Off by default. | - | | `HINDSIGHT_API_EMBEDDINGS_GEMINI_API_KEY` | Gemini API key for embeddings (falls back to `HINDSIGHT_API_LLM_API_KEY`) | - | | `HINDSIGHT_API_EMBEDDINGS_GEMINI_MODEL` | Gemini embedding model. The `gemini-embedding-2` family (e.g. `gemini-embedding-2-preview`) is supported on both the Gemini API and Vertex AI — because these multimodal models aggregate a multi-input request into one embedding, Hindsight automatically embeds one input per call to keep per-fact vectors. | `gemini-embedding-001` | | `HINDSIGHT_API_EMBEDDINGS_GEMINI_OUTPUT_DIMENSIONALITY` | Output embedding dimensions (Gemini supports configurable dimensionality) | `768` | @@ -796,6 +797,8 @@ export HINDSIGHT_API_EMBEDDINGS_LITELLM_SDK_API_KEY=your-provider-api-key export HINDSIGHT_API_EMBEDDINGS_LITELLM_SDK_MODEL=cohere/embed-english-v3.0 # Optional: request a specific output dimension when the provider supports it # export HINDSIGHT_API_EMBEDDINGS_LITELLM_SDK_OUTPUT_DIMENSIONS=768 +# Optional: truncate oversized inputs to the model's input-token limit (e.g. Bedrock Titan V2) +# export HINDSIGHT_API_EMBEDDINGS_LITELLM_SDK_MAX_INPUT_TOKENS=8192 # Supported LiteLLM SDK embedding providers: # - cohere/embed-english-v3.0 (1024 dimensions) diff --git a/skills/hindsight-docs/references/developer/configuration.md b/skills/hindsight-docs/references/developer/configuration.md index 42240315a0..459c604f6f 100644 --- a/skills/hindsight-docs/references/developer/configuration.md +++ b/skills/hindsight-docs/references/developer/configuration.md @@ -580,6 +580,7 @@ server-level only (not overridable per tenant/bank) and a change requires a rest | `HINDSIGHT_API_EMBEDDINGS_LITELLM_SDK_API_BASE` | Custom base URL for LiteLLM SDK embeddings (optional) | - | | `HINDSIGHT_API_EMBEDDINGS_LITELLM_SDK_OUTPUT_DIMENSIONS` | Optional output embedding dimensions (provider-dependent, e.g., `768` for Gemini embedding models) | - | | `HINDSIGHT_API_EMBEDDINGS_LITELLM_SDK_ENCODING_FORMAT` | Encoding format for embedding responses. Set to empty string to omit the parameter (needed for Voyage AI, Gemini). | `float` | +| `HINDSIGHT_API_EMBEDDINGS_LITELLM_SDK_MAX_INPUT_TOKENS` | If set, truncate each text to this many tokens (tiktoken `cl100k_base`, approximate) before embedding. Set it to the model's real input limit (e.g. `8192` for Bedrock Titan V2, with a little headroom) so oversized content is truncated instead of failing the embed call permanently. Off by default. | - | | `HINDSIGHT_API_EMBEDDINGS_GEMINI_API_KEY` | Gemini API key for embeddings (falls back to `HINDSIGHT_API_LLM_API_KEY`) | - | | `HINDSIGHT_API_EMBEDDINGS_GEMINI_MODEL` | Gemini embedding model. The `gemini-embedding-2` family (e.g. `gemini-embedding-2-preview`) is supported on both the Gemini API and Vertex AI — because these multimodal models aggregate a multi-input request into one embedding, Hindsight automatically embeds one input per call to keep per-fact vectors. | `gemini-embedding-001` | | `HINDSIGHT_API_EMBEDDINGS_GEMINI_OUTPUT_DIMENSIONALITY` | Output embedding dimensions (Gemini supports configurable dimensionality) | `768` | @@ -796,6 +797,8 @@ export HINDSIGHT_API_EMBEDDINGS_LITELLM_SDK_API_KEY=your-provider-api-key export HINDSIGHT_API_EMBEDDINGS_LITELLM_SDK_MODEL=cohere/embed-english-v3.0 # Optional: request a specific output dimension when the provider supports it # export HINDSIGHT_API_EMBEDDINGS_LITELLM_SDK_OUTPUT_DIMENSIONS=768 +# Optional: truncate oversized inputs to the model's input-token limit (e.g. Bedrock Titan V2) +# export HINDSIGHT_API_EMBEDDINGS_LITELLM_SDK_MAX_INPUT_TOKENS=8192 # Supported LiteLLM SDK embedding providers: # - cohere/embed-english-v3.0 (1024 dimensions) diff --git a/skills/hindsight-docs/references/openapi.json b/skills/hindsight-docs/references/openapi.json index 4fce846c52..36a689d0c8 100644 --- a/skills/hindsight-docs/references/openapi.json +++ b/skills/hindsight-docs/references/openapi.json @@ -10,7 +10,7 @@ "name": "Apache 2.0", "url": "https://www.apache.org/licenses/LICENSE-2.0.html" }, - "version": "0.8.3" + "version": "0.8.4" }, "paths": { "/health": {