mirror of
https://github.com/usestrix/strix.git
synced 2026-08-23 11:22:37 +02:00
Ahmed Allam
4eedfc64b4
- Clamp the summary request's max_tokens to the model's output limit so a large STRIX_CONTEXT_SUMMARY_TOKENS can't get the request rejected (which left the overflowing session uncompacted). Applied consistently to the input-budget reservation and the request itself. - Replace the tokenizer-unavailable fallback with the UTF-8 byte length, a guaranteed upper bound on tokens for byte-level BPE, so budget checks can never under-count dense history.
87 lines
3.1 KiB
Python
87 lines
3.1 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 the UTF-8 byte
|
|
length as a guaranteed upper bound: byte-level BPE tokenizers (used by every
|
|
major provider) emit at least one byte per token, so token count can never
|
|
exceed the byte count. Over-counting is safe here — it makes budget checks
|
|
conservative — whereas any under-count could 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.encode("utf-8"))
|