mirror of
https://github.com/usestrix/strix.git
synced 2026-08-16 17:27:26 +02:00
yoni
21afdaea9e
LiteLLM maps xAI models only provider-qualified, so neither "grok/grok-4" nor bare "grok-4" resolves: subscription runs fell back to the generic 200k context window and an 8k output cap instead of Grok's 256k/256k. subscription.litellm_model_name() now owns the routing-prefix -> LiteLLM name mapping (grok/ -> xai/, chatgpt/ -> bare) and context_budget uses it.
77 lines
2.6 KiB
Python
77 lines
2.6 KiB
Python
"""Tests for model-aware token budgets."""
|
|
|
|
from __future__ import annotations
|
|
|
|
from typing import TYPE_CHECKING
|
|
|
|
from strix.config import load_settings
|
|
from strix.llm import context_budget
|
|
|
|
|
|
if TYPE_CHECKING:
|
|
import pytest
|
|
|
|
|
|
def test_context_window_known_model() -> None:
|
|
# gpt-4o is mapped by LiteLLM at 128k input tokens.
|
|
assert context_budget.context_window("gpt-4o") == 128_000
|
|
|
|
|
|
def test_context_window_strips_provider_prefix() -> None:
|
|
assert context_budget.context_window("openai/gpt-4o") == 128_000
|
|
|
|
|
|
def test_context_window_chatgpt_prefix_skips_provider_auth(
|
|
monkeypatch: pytest.MonkeyPatch,
|
|
) -> None:
|
|
context_budget._model_info.cache_clear()
|
|
calls: list[str] = []
|
|
|
|
def _model_info(model: str) -> dict[str, int]:
|
|
calls.append(model)
|
|
return {"max_input_tokens": 1_050_000, "max_output_tokens": 128_000}
|
|
|
|
monkeypatch.setattr("strix.llm.context_budget.litellm.get_model_info", _model_info)
|
|
try:
|
|
assert context_budget.context_window("chatgpt/gpt-5.6-luna") == 1_050_000
|
|
assert calls == ["gpt-5.6-luna"]
|
|
finally:
|
|
context_budget._model_info.cache_clear()
|
|
|
|
|
|
def test_context_window_grok_prefix_resolves_to_xai() -> None:
|
|
# LiteLLM maps xAI models only provider-qualified: neither "grok/grok-4" nor
|
|
# bare "grok-4" resolves, so the subscription prefix becomes "xai/".
|
|
context_budget._model_info.cache_clear()
|
|
try:
|
|
assert context_budget.context_window("grok/grok-4") == 256_000
|
|
assert context_budget.output_limit("grok/grok-4") == 256_000
|
|
finally:
|
|
context_budget._model_info.cache_clear()
|
|
|
|
|
|
def test_context_window_unmapped_uses_fallback(monkeypatch: pytest.MonkeyPatch) -> None:
|
|
context_budget._model_info.cache_clear()
|
|
|
|
def _raise(_model: str) -> dict[str, int]:
|
|
raise ValueError("This model isn't mapped yet.")
|
|
|
|
monkeypatch.setattr("strix.llm.context_budget.litellm.get_model_info", _raise)
|
|
expected = load_settings().context.fallback_context_tokens
|
|
assert context_budget.context_window("totally-made-up-model") == expected
|
|
context_budget._model_info.cache_clear()
|
|
|
|
|
|
def test_count_tokens_fallback_on_error(monkeypatch: pytest.MonkeyPatch) -> None:
|
|
def _raise(**_kwargs: object) -> int:
|
|
raise RuntimeError("no tokenizer")
|
|
|
|
monkeypatch.setattr("strix.llm.context_budget.litellm.token_counter", _raise)
|
|
# Falls back to UTF-8 byte length (upper bound on tokens).
|
|
assert context_budget.count_tokens("weird-model", "x" * 400) == 400
|
|
assert context_budget.count_tokens("weird-model", "😀" * 10) == 40
|
|
|
|
|
|
def test_count_tokens_empty_is_zero() -> None:
|
|
assert context_budget.count_tokens("gpt-4o", "") == 0
|