holetron-lab/fleet-memory
1
0
Fork 0
Code Releases Packages Activity Actions

fix: truncate documents exceeding LiteLLM reranker context limit (#549)

Browse source 
* 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
This commit is contained in:
Nicolò Boschi 2026-03-13 10:18:18 +01:00 committed by GitHub
parent 2344484f77
commit eeb938fc65
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
5 changed files with 37 additions and 7 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

6
hindsight-api/hindsight_api/config.py
View file

@ -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),

28
hindsight-api/hindsight_api/engine/cross_encoder.py
View file

@ -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

8
hindsight-api/hindsight_api/engine/memory_engine.py
View file

@ -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)

1
hindsight-api/hindsight_api/main.py
View file

@ -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,

1
hindsight-docs/docs/developer/configuration.md
View file

@ -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` |