mirror of
https://github.com/headroomlabs-ai/headroom.git
synced 2026-08-27 14:17:10 -04:00
fa4763761b
## Description Fixes #2504. `CostTracker.estimate_cost` runs on the per-request cost path and logs a WARNING whenever LiteLLM can't price the model: ```python except Exception as e: logger.warning(f"Failed to get pricing for model {model}: {e}") return None ``` For a custom / OpenAI-compatible model LiteLLM can't resolve (e.g. `glm-5.2` via `--backend anyllm --anyllm-provider openai`), this fires on **every single request**, flooding `proxy.log` with hundreds of identical lines and burying genuinely useful warnings. The `LiteLLM not available` branch above it has the same per-request flooding shape. ## Fix Track already-warned models in a small module-level set and emit each pricing-failure warning (and the LiteLLM-unavailable warning) once per process. The set is bounded by the number of distinct model names seen. No new dependencies or config. The cost result itself is unchanged (`None` on failure); only the log volume changes. ## Type of Change - [x] Bug fix (non-breaking change that fixes an issue) - [ ] New feature (non-breaking change that adds functionality) - [ ] Breaking change (fix or feature that would cause existing functionality to change) - [ ] Documentation update - [ ] Performance improvement - [ ] Code refactoring (no functional changes) ## Changes Made - `headroom/proxy/cost.py`: add a module-level `_warned_pricing_models` set and `_warn_pricing_once` helper; route the pricing-failure and LiteLLM-unavailable warnings in `estimate_cost` through it. - `tests/test_cost_pricing_warning_dedup.py` (new): assert a repeated unresolvable model warns once, distinct models each warn once, and the LiteLLM-unavailable warning is deduped too. ## Testing - [x] Unit tests pass (`pytest`) - [x] Linting passes (`ruff check .`) - [x] Type checking passes (`mypy headroom`) - [x] New tests added for new functionality - [ ] Manual testing performed ### Test Output ```text $ python -m pytest tests/test_cost_pricing_warning_dedup.py -q 3 passed # with the fix reverted, the module-level set does not exist, so the # dedup tests error/fail (the pre-fix code warned once per request) $ uvx ruff@0.15.17 check headroom/proxy/cost.py tests/test_cost_pricing_warning_dedup.py All checks passed! $ uvx mypy@1.20.2 --ignore-missing-imports headroom/proxy/cost.py Success: no issues found in 1 source file ``` ## Real Behavior Proof - Environment: Windows 11, Python 3.12, project venv (`uv sync --extra proxy`), `uvx ruff@0.15.17` / `uvx mypy@1.20.2`, pytest in the venv. - Exact command / steps: monkeypatched `_get_litellm_module` to a stub whose `cost_per_token` raises (and, separately, to `None`), called `CostTracker.estimate_cost("glm-5.2", ...)` five times and two distinct unresolvable models twice each, capturing `headroom.proxy` WARNING records with `caplog`. - Observed result: with the fix each model produces exactly one `Failed to get pricing for model ...` warning (and one `LiteLLM not available ...`) regardless of call count; the pre-fix code logged one per call. `estimate_cost` still returns `None` on failure. Ran against the actual module. - Not tested: a live multi-request session against a real unpriced model end to end. ## Review Readiness - [x] I have performed a self-review - [x] This PR is ready for human review ## Checklist - [x] My code follows the project's style guidelines - [x] I have performed a self-review of my code - [x] I have commented my code, particularly in hard-to-understand areas - [ ] I have made corresponding changes to the documentation - [x] My changes generate no new warnings - [x] I have added tests that prove my fix is effective or that my feature works - [x] New and existing unit tests pass locally with my changes - [ ] I have updated the CHANGELOG.md if applicable
65 lines
2.3 KiB
Python
65 lines
2.3 KiB
Python
"""Pricing-lookup warnings for an unresolvable model must fire once, not per request.
|
|
|
|
#2504: a custom / OpenAI-compatible model LiteLLM can't price (e.g. glm-5.2)
|
|
logged an identical WARNING on every single request, flooding proxy.log.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import logging
|
|
|
|
import pytest
|
|
|
|
|
|
@pytest.fixture
|
|
def cost_tracker(monkeypatch: pytest.MonkeyPatch):
|
|
import headroom.proxy.cost as cost_mod
|
|
|
|
# Reset the per-process dedup set so tests are order-independent.
|
|
cost_mod._warned_pricing_models.clear()
|
|
|
|
class _FakeLiteLLM:
|
|
@staticmethod
|
|
def cost_per_token(**_kwargs):
|
|
raise RuntimeError("LLM Provider NOT provided.")
|
|
|
|
monkeypatch.setattr(cost_mod, "_get_litellm_module", lambda: _FakeLiteLLM())
|
|
return cost_mod.CostTracker()
|
|
|
|
|
|
def test_pricing_failure_warns_once_per_model(cost_tracker, caplog):
|
|
with caplog.at_level(logging.WARNING, logger="headroom.proxy"):
|
|
for _ in range(5):
|
|
assert cost_tracker.estimate_cost("glm-5.2", 100, 50) is None
|
|
|
|
warnings = [
|
|
r for r in caplog.records if "Failed to get pricing for model glm-5.2" in r.getMessage()
|
|
]
|
|
assert len(warnings) == 1
|
|
|
|
|
|
def test_distinct_models_each_warn_once(cost_tracker, caplog):
|
|
with caplog.at_level(logging.WARNING, logger="headroom.proxy"):
|
|
cost_tracker.estimate_cost("glm-5.2", 10, 5)
|
|
cost_tracker.estimate_cost("glm-5.2", 10, 5)
|
|
cost_tracker.estimate_cost("mystery-model", 10, 5)
|
|
cost_tracker.estimate_cost("mystery-model", 10, 5)
|
|
|
|
msgs = [r.getMessage() for r in caplog.records if "Failed to get pricing" in r.getMessage()]
|
|
assert sum("for model glm-5.2:" in m for m in msgs) == 1
|
|
assert sum("for model mystery-model:" in m for m in msgs) == 1
|
|
|
|
|
|
def test_litellm_unavailable_warns_once_per_model(monkeypatch, caplog):
|
|
import headroom.proxy.cost as cost_mod
|
|
|
|
cost_mod._warned_pricing_models.clear()
|
|
monkeypatch.setattr(cost_mod, "_get_litellm_module", lambda: None)
|
|
tracker = cost_mod.CostTracker()
|
|
|
|
with caplog.at_level(logging.WARNING, logger="headroom.proxy"):
|
|
for _ in range(3):
|
|
assert tracker.estimate_cost("glm-5.2", 10, 5) is None
|
|
|
|
unavailable = [r for r in caplog.records if "LiteLLM not available" in r.getMessage()]
|
|
assert len(unavailable) == 1
|