wehub-resource-sync
413 lines
17 KiB
Python
413 lines
17 KiB
Python
"""
|
|
Tests for model pricing and cache-aware LLM cost computation.
|
|
|
|
Covers :class:`ModelPricing`, :func:`compute_llm_cost` (the cache-aware
|
|
cost formula), and :func:`fetch_model_pricing`'s parsing of cache-read /
|
|
cache-write rates from a catalog entry.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
from typing import Any
|
|
|
|
import pytest
|
|
|
|
from omnigent.llms import context_window
|
|
from omnigent.llms.context_window import (
|
|
ModelPricing,
|
|
_registry_context_window,
|
|
compute_llm_cost,
|
|
fetch_model_pricing,
|
|
get_model_context_window,
|
|
resolve_effective_context_window,
|
|
)
|
|
|
|
|
|
def test_resolve_effective_context_window_prefers_declared_window(
|
|
monkeypatch: pytest.MonkeyPatch,
|
|
) -> None:
|
|
"""
|
|
A spec-declared ``executor.context_window`` wins over the catalog lookup.
|
|
|
|
Regression for the runner over-compaction bug: an agent that declares a
|
|
1M window (e.g. Polly) must be budgeted against 1M, not the 128K catalog
|
|
default. If the resolver fell back to the catalog here, the compaction
|
|
budget would be ~8x too small and fire constantly.
|
|
"""
|
|
|
|
def _boom(_model: str) -> int:
|
|
raise AssertionError("catalog lookup must not run when a window is declared")
|
|
|
|
monkeypatch.setattr(context_window, "get_model_context_window", _boom)
|
|
assert resolve_effective_context_window(1_000_000, "claude-opus-4-8") == 1_000_000
|
|
# Declared window applies even when the spec pins no model.
|
|
assert resolve_effective_context_window(1_000_000, None) == 1_000_000
|
|
|
|
|
|
def test_resolve_effective_context_window_falls_back_to_catalog(
|
|
monkeypatch: pytest.MonkeyPatch,
|
|
) -> None:
|
|
"""With no declared window, resolve via the model catalog lookup."""
|
|
monkeypatch.setattr(context_window, "get_model_context_window", lambda model: 200_000)
|
|
assert resolve_effective_context_window(None, "claude-opus-4-8") == 200_000
|
|
|
|
|
|
def test_resolve_effective_context_window_none_when_no_window_and_no_model() -> None:
|
|
"""No declared window and no model → ``None`` (caller skips budgeting)."""
|
|
assert resolve_effective_context_window(None, None) is None
|
|
|
|
|
|
def test_resolve_effective_context_window_override_bypasses_declared_window(
|
|
monkeypatch: pytest.MonkeyPatch,
|
|
) -> None:
|
|
"""
|
|
An active model override sizes against the override model's catalog window,
|
|
NOT the spec-declared window.
|
|
|
|
Matches the server ring: ``executor.context_window`` describes only the
|
|
spec model, so overriding a 1M-window agent down to a 200K model must
|
|
budget against 200K — otherwise the runner under-compacts past the real
|
|
model's limit.
|
|
"""
|
|
seen: list[str] = []
|
|
|
|
def _catalog(model: str) -> int:
|
|
seen.append(model)
|
|
return 200_000
|
|
|
|
monkeypatch.setattr(context_window, "get_model_context_window", _catalog)
|
|
result = resolve_effective_context_window(
|
|
1_000_000, "claude-opus-4-8", model_override="small-200k-model"
|
|
)
|
|
assert result == 200_000
|
|
# The override model — not the spec model — drives the catalog lookup.
|
|
assert seen == ["small-200k-model"]
|
|
|
|
|
|
def test_resolve_effective_context_window_declared_window_wins_without_override(
|
|
monkeypatch: pytest.MonkeyPatch,
|
|
) -> None:
|
|
"""An explicit ``model_override=None`` keeps the declared-window fast path."""
|
|
|
|
def _boom(_model: str) -> int:
|
|
raise AssertionError("catalog lookup must not run when no override is active")
|
|
|
|
monkeypatch.setattr(context_window, "get_model_context_window", _boom)
|
|
assert (
|
|
resolve_effective_context_window(1_000_000, "claude-opus-4-8", model_override=None)
|
|
== 1_000_000
|
|
)
|
|
|
|
|
|
def test_compute_llm_cost_prices_cache_tokens_at_their_own_rates() -> None:
|
|
"""
|
|
Cache reads/writes are billed at their own rates, not the input rate.
|
|
|
|
Anthropic reports ``input_tokens`` as the non-cached portion and
|
|
breaks out ``cache_read_input_tokens`` (cheap) / cache creation
|
|
(pricey). A correct cost sums all four priced parts. If the formula
|
|
reverted to ``input*price + output*price`` it would drop the 8000
|
|
cache-read + 2000 cache-write tokens entirely (0.0136 -> 0.007).
|
|
"""
|
|
pricing = ModelPricing(
|
|
input_per_token=2e-6,
|
|
output_per_token=1e-5,
|
|
cache_read_per_token=2e-7, # 0.1x input
|
|
cache_write_per_token=2.5e-6, # 1.25x input
|
|
)
|
|
usage: dict[str, Any] = {
|
|
"input_tokens": 1000,
|
|
"output_tokens": 500,
|
|
"cache_read_input_tokens": 8000,
|
|
"cache_creation_input_tokens": 2000,
|
|
}
|
|
# 1000*2e-6 + 500*1e-5 + 8000*2e-7 + 2000*2.5e-6
|
|
# = 0.002 + 0.005 + 0.0016 + 0.005 = 0.0136
|
|
assert compute_llm_cost(usage, pricing) == pytest.approx(0.0136)
|
|
|
|
|
|
def test_compute_llm_cost_derives_cache_rates_from_input_when_unpublished() -> None:
|
|
"""
|
|
With no published cache rates, derive them from the input rate via the
|
|
standard ratios: cache read at 0.10x input, cache write at 1.25x input.
|
|
|
|
``databricks-*`` catalog entries omit cache pricing, so this fallback is
|
|
what every relay/native session on the gateway is billed by. Pricing cache
|
|
reads at the full input rate (the old fallback) over-charged cache-heavy
|
|
sessions ~10x — the bug this fixes.
|
|
"""
|
|
pricing = ModelPricing(
|
|
input_per_token=2e-6,
|
|
output_per_token=1e-5,
|
|
cache_read_per_token=None,
|
|
cache_write_per_token=None,
|
|
)
|
|
usage: dict[str, Any] = {
|
|
"input_tokens": 1000,
|
|
"output_tokens": 500,
|
|
"cache_read_input_tokens": 8000,
|
|
"cache_creation_input_tokens": 2000,
|
|
}
|
|
# cache read at 0.10x input (2e-7), cache write at 1.25x input (2.5e-6):
|
|
# 1000*2e-6 + 500*1e-5 + 8000*2e-7 + 2000*2.5e-6
|
|
# = 0.002 + 0.005 + 0.0016 + 0.005 = 0.0136
|
|
# The old full-input fallback would give 0.027 (cache read at 1.6e-2),
|
|
# so a value of 0.027 here means the ratio fallback regressed.
|
|
assert compute_llm_cost(usage, pricing) == pytest.approx(0.0136)
|
|
|
|
|
|
def test_compute_llm_cost_without_cache_tokens_is_the_flat_formula() -> None:
|
|
"""
|
|
No cache-token keys -> reduces to ``input*price + output*price``.
|
|
|
|
Regression guard for the common / OpenAI case (no cache breakdown):
|
|
the cache-aware formula must not change the number when there are no
|
|
cache tokens.
|
|
"""
|
|
pricing = ModelPricing(
|
|
input_per_token=2e-6,
|
|
output_per_token=1e-5,
|
|
cache_read_per_token=2e-7,
|
|
cache_write_per_token=2.5e-6,
|
|
)
|
|
usage: dict[str, Any] = {"input_tokens": 1000, "output_tokens": 500}
|
|
# 1000*2e-6 + 500*1e-5 = 0.002 + 0.005 = 0.007 (cache terms are 0)
|
|
assert compute_llm_cost(usage, pricing) == pytest.approx(0.007)
|
|
|
|
|
|
def test_fetch_model_pricing_parses_cache_rates(monkeypatch: pytest.MonkeyPatch) -> None:
|
|
"""
|
|
``fetch_model_pricing`` surfaces catalog cache-read/write rates.
|
|
|
|
The MLflow catalog publishes ``cache_read_per_million_tokens`` /
|
|
``cache_write_per_million_tokens`` for Anthropic models; this pins
|
|
that they reach :class:`ModelPricing` (per-token), so cost can be
|
|
cache-accurate. A failure means the cache rates were dropped and
|
|
cost would fall back to the derived input-ratio default.
|
|
"""
|
|
# Catalog lookup is disabled globally in tests (conftest); re-enable
|
|
# for this one and stub the network fetch with a cache-priced entry.
|
|
monkeypatch.delenv("OMNIGENT_DISABLE_CATALOG_LOOKUP", raising=False)
|
|
monkeypatch.setattr(
|
|
context_window,
|
|
"_fetch_mlflow_provider_catalog",
|
|
lambda provider: {
|
|
"claude-x": {
|
|
"pricing": {
|
|
"input_per_million_tokens": 2.5,
|
|
"output_per_million_tokens": 10.0,
|
|
"cache_read_per_million_tokens": 0.25,
|
|
"cache_write_per_million_tokens": 3.125,
|
|
}
|
|
}
|
|
},
|
|
)
|
|
pricing = fetch_model_pricing("anthropic/claude-x")
|
|
assert pricing is not None
|
|
assert pricing.input_per_token == pytest.approx(2.5e-6)
|
|
assert pricing.output_per_token == pytest.approx(1e-5)
|
|
assert pricing.cache_read_per_token == pytest.approx(0.25e-6)
|
|
assert pricing.cache_write_per_token == pytest.approx(3.125e-6)
|
|
|
|
|
|
def test_fetch_model_pricing_omits_cache_rates_when_absent(
|
|
monkeypatch: pytest.MonkeyPatch,
|
|
) -> None:
|
|
"""
|
|
A catalog entry with no cache fields yields ``None`` cache rates.
|
|
|
|
OpenAI entries in the catalog carry only input/output rates;
|
|
``compute_llm_cost`` then derives cache rates from the input rate via
|
|
the standard ratios. If these came back as ``0.0`` instead of ``None``,
|
|
cache tokens would be billed free.
|
|
"""
|
|
monkeypatch.delenv("OMNIGENT_DISABLE_CATALOG_LOOKUP", raising=False)
|
|
monkeypatch.setattr(
|
|
context_window,
|
|
"_fetch_mlflow_provider_catalog",
|
|
lambda provider: {
|
|
"gpt-x": {
|
|
"pricing": {
|
|
"input_per_million_tokens": 1.25,
|
|
"output_per_million_tokens": 10.0,
|
|
}
|
|
}
|
|
},
|
|
)
|
|
pricing = fetch_model_pricing("openai/gpt-x")
|
|
assert pricing is not None
|
|
assert pricing.input_per_token == pytest.approx(1.25e-6)
|
|
assert pricing.cache_read_per_token is None
|
|
assert pricing.cache_write_per_token is None
|
|
|
|
|
|
def test_fetch_model_pricing_databricks_alias_falls_back_to_base_model(
|
|
monkeypatch: pytest.MonkeyPatch,
|
|
) -> None:
|
|
"""A ``databricks-<base>`` alias absent from the Databricks catalog is
|
|
priced from the base model's underlying-provider catalog.
|
|
|
|
Models served through the Databricks gateway are reported as
|
|
``databricks-claude-opus-4-8``, which the Databricks catalog may not
|
|
list even though anthropic's ``claude-opus-4-8`` is priced. Without the
|
|
de-prefix fallback, every unpinned claude-sdk agent on the Databricks
|
|
gateway (which defaults to ``databricks-claude-opus-4-8``) would show
|
|
"unpriced" — the exact gap reported for the debbie/debby supervisors.
|
|
"""
|
|
monkeypatch.delenv("OMNIGENT_DISABLE_CATALOG_LOOKUP", raising=False)
|
|
|
|
def _catalog(provider: str) -> dict[str, Any] | None:
|
|
"""Databricks catalog lacks opus; the base (anthropic) catalog prices it."""
|
|
if provider == "databricks":
|
|
# Has some databricks models, but NOT the opus alias under test.
|
|
return {
|
|
"databricks-claude-sonnet-4-6": {
|
|
"pricing": {
|
|
"input_per_million_tokens": 3.0,
|
|
"output_per_million_tokens": 15.0,
|
|
}
|
|
}
|
|
}
|
|
# The underlying provider (anthropic) prices the de-prefixed base.
|
|
return {
|
|
"claude-opus-4-8": {
|
|
"pricing": {
|
|
"input_per_million_tokens": 15.0,
|
|
"output_per_million_tokens": 75.0,
|
|
}
|
|
}
|
|
}
|
|
|
|
monkeypatch.setattr(context_window, "_fetch_mlflow_provider_catalog", _catalog)
|
|
|
|
pricing = fetch_model_pricing("databricks-claude-opus-4-8")
|
|
assert pricing is not None, (
|
|
"databricks-claude-opus-4-8 was not priced — the databricks→base "
|
|
"fallback did not reach anthropic's claude-opus-4-8."
|
|
)
|
|
# Priced from the base model's rates (15 / 75 per million), not the
|
|
# databricks sonnet entry (3 / 15).
|
|
assert pricing.input_per_token == pytest.approx(15e-6)
|
|
assert pricing.output_per_token == pytest.approx(75e-6)
|
|
|
|
|
|
def test_provider_catalog_is_cached_across_calls(
|
|
monkeypatch: pytest.MonkeyPatch,
|
|
) -> None:
|
|
"""
|
|
The per-provider catalog is downloaded once, then served from cache.
|
|
|
|
This pins the perf fix: the response builder calls
|
|
``get_model_context_window`` on every ``GET /v1/sessions/{id}``
|
|
snapshot, and each call used to re-issue a ~490ms GitHub fetch.
|
|
With the TTL cache, repeated lookups for the same provider must hit
|
|
the network exactly once. A regression (cache removed) would show as
|
|
a download count > 1. Asserting the resolved window also proves the
|
|
cached payload still flows through the resolver unchanged.
|
|
"""
|
|
monkeypatch.delenv("OMNIGENT_DISABLE_CATALOG_LOOKUP", raising=False)
|
|
# Clear any residue from earlier tests so the count starts clean.
|
|
context_window._catalog_cache.clear()
|
|
calls: list[str] = []
|
|
|
|
def _fake_download(provider: str) -> dict[str, Any]:
|
|
"""Record each network hit and return a one-model catalog."""
|
|
calls.append(provider)
|
|
return {"claude-z": {"context_window": {"max_input": 200_000, "max_output": 8_192}}}
|
|
|
|
monkeypatch.setattr(context_window, "_download_mlflow_provider_catalog", _fake_download)
|
|
|
|
# litellm resolves many real names; force the catalog path by using a
|
|
# name it won't know, so the fetch is exercised deterministically.
|
|
first = context_window.get_model_context_window("claude-z")
|
|
second = context_window.get_model_context_window("claude-z")
|
|
assert first == 208_192 # max_input + max_output from the stub
|
|
assert second == 208_192
|
|
# Exactly one network download despite two resolver calls.
|
|
assert calls == ["anthropic"]
|
|
|
|
|
|
def test_provider_catalog_caches_fetch_failure(
|
|
monkeypatch: pytest.MonkeyPatch,
|
|
) -> None:
|
|
"""
|
|
A failed download (``None``) is cached too, not retried every call.
|
|
|
|
A transient GitHub outage returns ``None``; without caching that
|
|
result, every subsequent snapshot would re-pay the 5s timeout for an
|
|
hour. Pinning that ``None`` is cached keeps a single failure from
|
|
amplifying into per-request latency. The caller still falls back to
|
|
the 128K default, which this also checks.
|
|
"""
|
|
monkeypatch.delenv("OMNIGENT_DISABLE_CATALOG_LOOKUP", raising=False)
|
|
context_window._catalog_cache.clear()
|
|
calls: list[str] = []
|
|
|
|
def _fail(provider: str) -> None:
|
|
"""Record the hit and simulate a network/parse failure (returns None)."""
|
|
calls.append(provider)
|
|
|
|
monkeypatch.setattr(context_window, "_download_mlflow_provider_catalog", _fail)
|
|
first = context_window.get_model_context_window("claude-z")
|
|
second = context_window.get_model_context_window("claude-z")
|
|
assert first == 128_000 # _DEFAULT_CONTEXT_WINDOW fallback
|
|
assert second == 128_000
|
|
assert calls == ["anthropic"]
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Omnigent's authoritative context-window registry (supersedes litellm/catalog)
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def test_registry_context_window_normalizes_id() -> None:
|
|
"""The registry strips provider prefixes and ``:tag`` suffixes before matching."""
|
|
assert _registry_context_window("qwen3-coder-plus") == 1_048_576
|
|
assert _registry_context_window("qwen/qwen3-coder") == 262_144
|
|
assert _registry_context_window("qwen3-coder:free") == 262_144
|
|
assert _registry_context_window("openrouter/qwen/qwen3-coder:free") == 262_144
|
|
assert _registry_context_window("QWEN3-CODER-PLUS") == 1_048_576 # case-insensitive
|
|
# A model the registry doesn't own → None (caller falls back to litellm).
|
|
assert _registry_context_window("qwen-nonexistent-xyz") is None
|
|
assert _registry_context_window("gpt-5.4") is None
|
|
|
|
|
|
def test_registry_resolves_anthropic_1m_beta_suffix() -> None:
|
|
"""The Anthropic ``[1m]`` beta marker resolves to a 1M window via the registry.
|
|
|
|
The suffix *is* the window — we read it, not strip it — so any
|
|
``<model>[1m]`` resolves to 1,000,000 while the bare base defers to the
|
|
upstream backends (which may size it differently).
|
|
"""
|
|
assert _registry_context_window("claude-opus-4-8[1m]") == 1_000_000
|
|
assert _registry_context_window("anthropic/claude-opus-4-8[1m]") == 1_000_000
|
|
assert _registry_context_window("claude-sonnet-4-6[1m]") == 1_000_000
|
|
assert _registry_context_window("CLAUDE-OPUS-4-8[1M]") == 1_000_000 # case-insensitive
|
|
# Databricks-hosted Claude (contains "claude") also resolves.
|
|
assert _registry_context_window("databricks-claude-opus-4-8[1m]") == 1_000_000
|
|
# Without the suffix the registry defers (None → caller uses litellm/catalog).
|
|
assert _registry_context_window("claude-opus-4-8") is None
|
|
# The rule is Claude-scoped: a non-Claude id ending in [1m] is NOT forced to
|
|
# 1M (it defers to litellm/catalog), so custom/self-hosted ids are safe.
|
|
assert _registry_context_window("my-local-model[1m]") is None
|
|
assert _registry_context_window("gpt-5.4[1m]") is None
|
|
|
|
|
|
def test_get_model_context_window_uses_registry_first(monkeypatch: pytest.MonkeyPatch) -> None:
|
|
"""Registry-curated ids resolve to their window with NO network.
|
|
|
|
Catalog lookup is disabled to prove hermeticity: the registry is consulted
|
|
before litellm and the catalog, so qwen models and the Anthropic ``[1m]``
|
|
beta resolve correctly even offline (the meter / overflow-threshold
|
|
bug was that these collapsed to the 128K default).
|
|
"""
|
|
monkeypatch.setenv("OMNIGENT_DISABLE_CATALOG_LOOKUP", "1")
|
|
monkeypatch.delenv("AP_CONTEXT_WINDOW_OVERRIDE", raising=False)
|
|
# Anthropic 1M beta: resolves via the registry, not the 128K default.
|
|
assert get_model_context_window("claude-opus-4-8[1m]") == 1_000_000
|
|
assert get_model_context_window("anthropic/claude-opus-4-8[1m]") == 1_000_000
|
|
# Qwen: curated window, not the default.
|
|
assert get_model_context_window("qwen3-coder-plus") == 1_048_576
|
|
# A model the registry doesn't own still falls back to the conservative default.
|
|
assert get_model_context_window("qwen-nonexistent-xyz") == 128_000
|