mirror of
https://github.com/usestrix/strix.git
synced 2026-08-21 18:52:47 +02:00
Ahmed Allam
2163a66b78
The chars/4 fallback under-counts dense text (code, base64, CJK), which could let a summary request be packed past the real context window and get rejected. Use a conservative ~3-chars/token estimate instead so budget checks never under-count.
86 lines
3.0 KiB
Python
86 lines
3.0 KiB
Python
"""Model-aware token budgets.
|
|
|
|
The context window and output cap vary widely by model (128k for gpt-4o, 272k
|
|
for gpt-5, 1M for claude-sonnet-4, 131k for deepseek). We resolve them from
|
|
LiteLLM's model metadata so compaction triggers at the right point for the
|
|
selected model instead of a fixed guess, falling back to a large configurable
|
|
default for models LiteLLM doesn't map.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import logging
|
|
from functools import lru_cache
|
|
from typing import Any
|
|
|
|
import litellm
|
|
|
|
from strix.config import load_settings
|
|
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
# LiteLLM keys models without the routing prefix users type (``openai/``,
|
|
# ``litellm/``, ``ollama/`` ...). Strip a leading provider segment on lookup.
|
|
_STRIPPABLE_PREFIXES = ("openai/", "litellm/", "any-llm/", "ollama/", "ollama_chat/")
|
|
|
|
_DEFAULT_OUTPUT_TOKENS = 8_192
|
|
|
|
|
|
def _lookup_key(model: str) -> str:
|
|
for prefix in _STRIPPABLE_PREFIXES:
|
|
if model.startswith(prefix):
|
|
return model[len(prefix) :]
|
|
return model
|
|
|
|
|
|
def _safe_get_model_info(model: str) -> dict[str, Any] | None:
|
|
try:
|
|
return dict(litellm.get_model_info(model))
|
|
except Exception: # noqa: BLE001 - unmapped models raise; caller falls back.
|
|
return None
|
|
|
|
|
|
@lru_cache(maxsize=128)
|
|
def _model_info(model: str) -> dict[str, int]:
|
|
for candidate in (model, _lookup_key(model)):
|
|
info = _safe_get_model_info(candidate)
|
|
if info is not None:
|
|
return {
|
|
"max_input_tokens": int(
|
|
info.get("max_input_tokens") or info.get("max_tokens") or 0
|
|
),
|
|
"max_output_tokens": int(info.get("max_output_tokens") or 0),
|
|
}
|
|
logger.debug("No LiteLLM model info for %r; using configured fallbacks", model)
|
|
return {"max_input_tokens": 0, "max_output_tokens": 0}
|
|
|
|
|
|
def context_window(model: str) -> int:
|
|
"""Input token capacity for ``model`` (configured fallback when unmapped)."""
|
|
resolved = _model_info(model)["max_input_tokens"]
|
|
return resolved or load_settings().context.fallback_context_tokens
|
|
|
|
|
|
def output_limit(model: str) -> int:
|
|
"""Max output tokens for ``model`` (a conservative default when unmapped)."""
|
|
return _model_info(model)["max_output_tokens"] or _DEFAULT_OUTPUT_TOKENS
|
|
|
|
|
|
def count_tokens(model: str, text: str) -> int:
|
|
"""Token count for ``text`` under ``model``.
|
|
|
|
LiteLLM's counter handles known tokenizers (and defaults to a tiktoken
|
|
encoding otherwise). If it still can't count, fall back to a *conservative*
|
|
estimate: token density varies, and dense text (code, base64, CJK) can run
|
|
well under 4 chars/token, so we assume ~3 to over-estimate rather than
|
|
under-estimate — an under-estimate would let a summary request be packed
|
|
past the real context window and get rejected.
|
|
"""
|
|
if not text:
|
|
return 0
|
|
try:
|
|
return int(litellm.token_counter(model=_lookup_key(model), text=text))
|
|
except Exception: # noqa: BLE001 - tokenizer may be unavailable for some models.
|
|
return -(-len(text) // 3)
|