scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-23 00:41:40 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 81

refactor(router): Enhance documentation and add prompt truncation for tokenization in prompt prefix affinity check

Browse source
This commit is contained in:
DragonAssassin-one 2026-05-11 10:37:54 +08:00
parent d73f8ee8a0
commit 55b7dbe84d
2 changed files with 40 additions and 4 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

8
litellm/router.py
View file

@ -348,9 +348,13 @@ class Router:
provider_budget_config (ProviderBudgetConfig): Provider budget configuration. Use this to set llm_provider budget limits. example $100/day to OpenAI, $100/day to Azure, etc. Defaults to None.
deployment_affinity_ttl_seconds (int): TTL for user-key -> deployment affinity mapping. Defaults to 3600.
prompt_prefix_affinity_tokens (int): Number of canonical prompt-prefix
tokens used for deterministic prompt-prefix routing. Defaults to 2048.
tokens used for deterministic prompt-prefix routing. This is
measured against tokenization of the canonical JSON representation
(not the provider-specific "raw prompt token count"). Defaults to 2048.
prompt_prefix_affinity_min_tokens (int): Minimum canonical prompt token
count before prompt-prefix affinity applies. Defaults to
count before prompt-prefix affinity applies. This is measured
against tokenization of the canonical JSON representation (not the
provider-specific "raw prompt token count"). Defaults to
MINIMUM_PROMPT_CACHE_TOKEN_COUNT.
ignore_invalid_deployments (bool): Ignores invalid deployments, and continues with other deployments. Default is to raise an error.
Returns:

36
litellm/router_utils/pre_call_checks/prompt_prefix_affinity_check.py
View file

@ -11,6 +11,7 @@ mapping in Redis. All Router instances with the same config and deployment IDs
will make the same routing decision.
"""
import asyncio
import hashlib
import json
from typing import Any, Dict, List, Optional, Tuple, cast
@ -28,9 +29,16 @@ class PromptPrefixAffinityCheck(CustomLogger):
This is intended to improve prompt-cache hit rate for providers where prompt
caching is scoped to the account/key behind a deployment.
Note: `prefix_tokens` / `min_tokens` are measured against the tokenization of
the canonical JSON representation produced by `_build_canonical_prompt()`
(e.g. includes field names, separators, quotes), not the provider-specific
"raw prompt token count" after rendering the final prompt.
"""
CACHE_KEY_EXCLUDED_FIELDS = frozenset({"encrypted_content"})
_CANONICAL_PROMPT_CHARS_PER_TOKEN_ESTIMATE = 32
_CANONICAL_PROMPT_ABSOLUTE_MAX_CHARS = 1_000_000
def __init__(
self,
@ -95,6 +103,26 @@ class PromptPrefixAffinityCheck(CustomLogger):
ensure_ascii=False,
)
def _truncate_canonical_prompt_for_tokenization(self, canonical_prompt: str) -> str:
"""
Bound tokenization work for very large prompts.
We only use the first `prefix_tokens` token IDs (and gate on `min_tokens`),
so we can safely cap the number of characters we feed into tokenization to
avoid worst-case CPU time on extremely large canonical prompts.
"""
token_target = max(self.prefix_tokens, self.min_tokens)
if token_target <= 0:
return canonical_prompt
soft_limit = token_target * self._CANONICAL_PROMPT_CHARS_PER_TOKEN_ESTIMATE
max_chars = min(soft_limit, self._CANONICAL_PROMPT_ABSOLUTE_MAX_CHARS)
if len(canonical_prompt) <= max_chars:
return canonical_prompt
return canonical_prompt[:max_chars]
def _get_prefix_hash(
self,
model: str,
@ -182,9 +210,13 @@ class PromptPrefixAffinityCheck(CustomLogger):
if canonical_prompt is None:
return typed_healthy_deployments
prefix_hash = self._get_prefix_hash(
canonical_prompt_for_tokenization = self._truncate_canonical_prompt_for_tokenization(
canonical_prompt
)
prefix_hash = await asyncio.to_thread(
self._get_prefix_hash,
model=model,
canonical_prompt=canonical_prompt,
canonical_prompt=canonical_prompt_for_tokenization,
)
if prefix_hash is None:
return typed_healthy_deployments