Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
9 changes: 9 additions & 0 deletions hindsight-api-slim/hindsight_api/config.py
Original file line number Diff line number Diff line change
Expand Up @@ -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"
Expand Down Expand Up @@ -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"
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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),
Expand Down
53 changes: 53 additions & 0 deletions hindsight-api-slim/hindsight_api/engine/embeddings.py
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand Down Expand Up @@ -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.
Expand All @@ -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
Expand All @@ -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

Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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
Expand Down
59 changes: 59 additions & 0 deletions hindsight-api-slim/tests/test_litellm_sdk_embeddings.py
Original file line number Diff line number Diff line change
Expand Up @@ -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(
Expand Down
3 changes: 3 additions & 0 deletions hindsight-docs/docs/developer/configuration.md
Original file line number Diff line number Diff line change
Expand Up @@ -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` |
Expand Down Expand Up @@ -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)
Expand Down
3 changes: 3 additions & 0 deletions skills/hindsight-docs/references/developer/configuration.md
Original file line number Diff line number Diff line change
Expand Up @@ -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` |
Expand Down Expand Up @@ -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)
Expand Down
2 changes: 1 addition & 1 deletion skills/hindsight-docs/references/openapi.json
Original file line number Diff line number Diff line change
Expand Up @@ -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": {
Expand Down
Loading