From eeb938fc65e31f62637daed63811e6971f655cbe Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Nicol=C3=B2=20Boschi?= Date: Fri, 13 Mar 2026 10:18:18 +0100 Subject: [PATCH] fix: truncate documents exceeding LiteLLM reranker context limit (#549) * fix: register embedded profiles in CLI metadata on daemon start When HindsightEmbedded(profile="myapp") starts a daemon, the profile was never written to metadata.json or given a .env file, making it invisible to `hindsight-embed profile list` and other CLI commands. Add _register_profile() to DaemonEmbedManager which saves HINDSIGHT_API_* config to ~/.hindsight/profiles/{name}.env and registers the port in metadata.json. Called after a successful new daemon start and when the daemon is already running, so orphaned profiles also get registered on next use. * fix: truncate documents exceeding LiteLLM reranker context limit Add HINDSIGHT_API_RERANKER_LITELLM_MAX_TOKENS_PER_DOC env var for both litellm and litellm-sdk reranker providers. When set, documents are truncated to the configured token limit using tiktoken (cl100k_base) before being sent to the reranker, preventing BadRequestError for models with small context windows (e.g. 1024-token limit). * refactor: use shared _tiktoken_encoder for doc truncation in LiteLLM reranker * refactor: use _get_tiktoken_encoding() consistently, remove eager module-level encoder instance * doc: add HINDSIGHT_API_RERANKER_LITELLM_MAX_TOKENS_PER_DOC to configuration reference --- hindsight-api/hindsight_api/config.py | 6 ++++ .../hindsight_api/engine/cross_encoder.py | 28 +++++++++++++++++++ .../hindsight_api/engine/memory_engine.py | 8 +----- hindsight-api/hindsight_api/main.py | 1 + .../docs/developer/configuration.md | 1 + 5 files changed, 37 insertions(+), 7 deletions(-) diff --git a/hindsight-api/hindsight_api/config.py b/hindsight-api/hindsight_api/config.py index 97aa407b..53d102b1 100644 --- a/hindsight-api/hindsight_api/config.py +++ b/hindsight-api/hindsight_api/config.py @@ -193,6 +193,7 @@ ENV_EMBEDDINGS_LITELLM_MODEL = "HINDSIGHT_API_EMBEDDINGS_LITELLM_MODEL" ENV_RERANKER_LITELLM_API_BASE = "HINDSIGHT_API_RERANKER_LITELLM_API_BASE" ENV_RERANKER_LITELLM_API_KEY = "HINDSIGHT_API_RERANKER_LITELLM_API_KEY" ENV_RERANKER_LITELLM_MODEL = "HINDSIGHT_API_RERANKER_LITELLM_MODEL" +ENV_RERANKER_LITELLM_MAX_TOKENS_PER_DOC = "HINDSIGHT_API_RERANKER_LITELLM_MAX_TOKENS_PER_DOC" # LiteLLM SDK configuration (direct API access, no proxy needed) ENV_EMBEDDINGS_LITELLM_SDK_API_KEY = "HINDSIGHT_API_EMBEDDINGS_LITELLM_SDK_API_KEY" @@ -409,6 +410,7 @@ DEFAULT_TEXT_SEARCH_EXTENSION = "native" # Options: "native", "vchord", "pg_tex DEFAULT_LITELLM_API_BASE = "http://localhost:4000" DEFAULT_EMBEDDINGS_LITELLM_MODEL = "text-embedding-3-small" DEFAULT_RERANKER_LITELLM_MODEL = "cohere/rerank-english-v3.0" +DEFAULT_RERANKER_LITELLM_MAX_TOKENS_PER_DOC: int | None = None # LiteLLM SDK defaults DEFAULT_EMBEDDINGS_LITELLM_SDK_MODEL = "cohere/embed-english-v3.0" @@ -676,6 +678,7 @@ class HindsightConfig: reranker_litellm_api_base: str reranker_litellm_api_key: str | None reranker_litellm_model: str + reranker_litellm_max_tokens_per_doc: int | None reranker_litellm_sdk_api_key: str | None reranker_litellm_sdk_model: str reranker_litellm_sdk_api_base: str | None @@ -1104,6 +1107,9 @@ class HindsightConfig: or os.getenv(ENV_LITELLM_API_BASE, DEFAULT_LITELLM_API_BASE), reranker_litellm_api_key=os.getenv(ENV_RERANKER_LITELLM_API_KEY) or os.getenv(ENV_LITELLM_API_KEY), reranker_litellm_model=os.getenv(ENV_RERANKER_LITELLM_MODEL, DEFAULT_RERANKER_LITELLM_MODEL), + reranker_litellm_max_tokens_per_doc=int(v) + if (v := os.getenv(ENV_RERANKER_LITELLM_MAX_TOKENS_PER_DOC)) + else DEFAULT_RERANKER_LITELLM_MAX_TOKENS_PER_DOC, # LiteLLM SDK reranker (direct API access) reranker_litellm_sdk_api_key=os.getenv(ENV_RERANKER_LITELLM_SDK_API_KEY), reranker_litellm_sdk_model=os.getenv(ENV_RERANKER_LITELLM_SDK_MODEL, DEFAULT_RERANKER_LITELLM_SDK_MODEL), diff --git a/hindsight-api/hindsight_api/engine/cross_encoder.py b/hindsight-api/hindsight_api/engine/cross_encoder.py index 043dbdbd..ed64c1fd 100644 --- a/hindsight-api/hindsight_api/engine/cross_encoder.py +++ b/hindsight-api/hindsight_api/engine/cross_encoder.py @@ -20,6 +20,7 @@ from ..config import ( DEFAULT_RERANKER_COHERE_MODEL, DEFAULT_RERANKER_FLASHRANK_CACHE_DIR, DEFAULT_RERANKER_FLASHRANK_MODEL, + DEFAULT_RERANKER_LITELLM_MAX_TOKENS_PER_DOC, DEFAULT_RERANKER_LITELLM_MODEL, DEFAULT_RERANKER_LITELLM_SDK_MODEL, DEFAULT_RERANKER_LOCAL_FORCE_CPU, @@ -820,6 +821,17 @@ class FlashRankCrossEncoder(CrossEncoderModel): return await loop.run_in_executor(FlashRankCrossEncoder._executor, self._predict_sync, pairs) +def _truncate_to_tokens(text: str, max_tokens: int) -> str: + """Truncate text to at most max_tokens using the shared tiktoken encoder.""" + from .memory_engine import _get_tiktoken_encoding + + enc = _get_tiktoken_encoding() + tokens = enc.encode(text) + if len(tokens) <= max_tokens: + return text + return enc.decode(tokens[:max_tokens]) + + class LiteLLMCrossEncoder(CrossEncoderModel): """ LiteLLM cross-encoder implementation using LiteLLM proxy's /rerank endpoint. @@ -843,6 +855,7 @@ class LiteLLMCrossEncoder(CrossEncoderModel): api_key: str | None = None, model: str = DEFAULT_RERANKER_LITELLM_MODEL, timeout: float = 60.0, + max_tokens_per_doc: int | None = DEFAULT_RERANKER_LITELLM_MAX_TOKENS_PER_DOC, ): """ Initialize LiteLLM cross-encoder client. @@ -853,11 +866,15 @@ class LiteLLMCrossEncoder(CrossEncoderModel): model: Reranking model name (default: cohere/rerank-english-v3.0) Use provider prefix (e.g., cohere/, together_ai/, voyage/) timeout: Request timeout in seconds (default: 60.0) + max_tokens_per_doc: If set, truncate each document to this many tokens before + sending to the reranker (uses tiktoken cl100k_base encoding). + Useful for models with small context windows (e.g. 1024 tokens). """ self.api_base = api_base.rstrip("/") self.api_key = api_key self.model = model self.timeout = timeout + self.max_tokens_per_doc = max_tokens_per_doc self._async_client: httpx.AsyncClient | None = None @property @@ -905,6 +922,8 @@ class LiteLLMCrossEncoder(CrossEncoderModel): for query, indexed_texts in query_groups.items(): texts = [text for _, text in indexed_texts] + if self.max_tokens_per_doc is not None: + texts = [_truncate_to_tokens(t, self.max_tokens_per_doc) for t in texts] indices = [idx for idx, _ in indexed_texts] # LiteLLM /rerank follows Cohere API format @@ -950,6 +969,7 @@ class LiteLLMSDKCrossEncoder(CrossEncoderModel): model: str = DEFAULT_RERANKER_LITELLM_SDK_MODEL, api_base: str | None = None, timeout: float = 60.0, + max_tokens_per_doc: int | None = DEFAULT_RERANKER_LITELLM_MAX_TOKENS_PER_DOC, ): """ Initialize LiteLLM SDK cross-encoder client. @@ -959,11 +979,15 @@ class LiteLLMSDKCrossEncoder(CrossEncoderModel): model: Model name with provider prefix (e.g., "deepinfra/Qwen3-reranker-8B") api_base: Custom base URL for API (optional) timeout: Request timeout in seconds (default: 60.0) + max_tokens_per_doc: If set, truncate each document to this many tokens before + sending to the reranker (uses tiktoken cl100k_base encoding). + Useful for models with small context windows (e.g. 1024 tokens). """ self.api_key = api_key self.model = model self.api_base = api_base self.timeout = timeout + self.max_tokens_per_doc = max_tokens_per_doc self._initialized = False self._litellm = None # Will be set during initialization @@ -1017,6 +1041,8 @@ class LiteLLMSDKCrossEncoder(CrossEncoderModel): for query, indexed_texts in query_groups.items(): texts = [text for _, text in indexed_texts] + if self.max_tokens_per_doc is not None: + texts = [_truncate_to_tokens(t, self.max_tokens_per_doc) for t in texts] indices = [idx for idx, _ in indexed_texts] # Build kwargs for rerank call @@ -1189,6 +1215,7 @@ def create_cross_encoder_from_env() -> CrossEncoderModel: api_base=config.reranker_litellm_api_base, api_key=config.reranker_litellm_api_key, model=config.reranker_litellm_model, + max_tokens_per_doc=config.reranker_litellm_max_tokens_per_doc, ) elif provider == "litellm-sdk": api_key = config.reranker_litellm_sdk_api_key @@ -1200,6 +1227,7 @@ def create_cross_encoder_from_env() -> CrossEncoderModel: api_key=api_key, model=config.reranker_litellm_sdk_model, api_base=config.reranker_litellm_sdk_api_base, + max_tokens_per_doc=config.reranker_litellm_max_tokens_per_doc, ) elif provider == "zeroentropy": api_key = config.reranker_zeroentropy_api_key diff --git a/hindsight-api/hindsight_api/engine/memory_engine.py b/hindsight-api/hindsight_api/engine/memory_engine.py index a5e0b778..142a8432 100644 --- a/hindsight-api/hindsight_api/engine/memory_engine.py +++ b/hindsight-api/hindsight_api/engine/memory_engine.py @@ -51,13 +51,9 @@ def get_current_schema() -> str: return schema -# Initialize tiktoken encoder once at module level for efficiency -_tiktoken_encoder = tiktoken.get_encoding("cl100k_base") # GPT-4/GPT-3.5-turbo encoding - - def count_tokens(text: str) -> int: """Count tokens in text using tiktoken (cl100k_base encoding for GPT-4/3.5).""" - return len(_tiktoken_encoder.encode(text)) + return len(_get_tiktoken_encoding().encode(text)) def fq_table(table_name: str) -> str: @@ -208,8 +204,6 @@ def utcnow(): # Logger for memory system logger = logging.getLogger(__name__) -import tiktoken - from .db_utils import acquire_with_retry # Cache tiktoken encoding for token budget filtering (module-level singleton) diff --git a/hindsight-api/hindsight_api/main.py b/hindsight-api/hindsight_api/main.py index 43f0b9f1..a48dac14 100644 --- a/hindsight-api/hindsight_api/main.py +++ b/hindsight-api/hindsight_api/main.py @@ -229,6 +229,7 @@ def main(): reranker_litellm_api_base=config.reranker_litellm_api_base, reranker_litellm_api_key=config.reranker_litellm_api_key, reranker_litellm_model=config.reranker_litellm_model, + reranker_litellm_max_tokens_per_doc=config.reranker_litellm_max_tokens_per_doc, reranker_litellm_sdk_api_key=config.reranker_litellm_sdk_api_key, reranker_litellm_sdk_model=config.reranker_litellm_sdk_model, reranker_litellm_sdk_api_base=config.reranker_litellm_sdk_api_base, diff --git a/hindsight-docs/docs/developer/configuration.md b/hindsight-docs/docs/developer/configuration.md index 8f6e8555..012746b7 100644 --- a/hindsight-docs/docs/developer/configuration.md +++ b/hindsight-docs/docs/developer/configuration.md @@ -426,6 +426,7 @@ Supported OpenAI embedding dimensions: | `HINDSIGHT_API_RERANKER_LITELLM_SDK_API_KEY` | LiteLLM **SDK** API key for direct reranking (no proxy needed) | - | | `HINDSIGHT_API_RERANKER_LITELLM_SDK_MODEL` | LiteLLM SDK rerank model (e.g., `deepinfra/Qwen3-reranker-8B`) | `cohere/rerank-english-v3.0` | | `HINDSIGHT_API_RERANKER_LITELLM_SDK_API_BASE` | Custom API base URL for LiteLLM SDK (optional) | - | +| `HINDSIGHT_API_RERANKER_LITELLM_MAX_TOKENS_PER_DOC` | Truncate documents to this many tokens before sending to the reranker (applies to both `litellm` and `litellm-sdk`). Use for models with small context windows (e.g. set to `900` for a 1024-token limit model). Unset by default (no truncation). | - | | `HINDSIGHT_API_RERANKER_ZEROENTROPY_API_KEY` | ZeroEntropy API key for reranking | - | | `HINDSIGHT_API_RERANKER_ZEROENTROPY_MODEL` | ZeroEntropy rerank model (`zerank-2`, `zerank-2-small`) | `zerank-2` | | `HINDSIGHT_API_RERANKER_FLASHRANK_MODEL` | FlashRank model for fast CPU-based reranking | `ms-marco-MiniLM-L-12-v2` |