shaguabijia-app-server/tests/test_llm_cost.py

"""LLM 调用成本计算 compute_llm_cost:按模型分桶累加 token × 单价;error/无 usage 跳过。"""
from __future__ import annotations

from app.services.llm_cost import compute_llm_cost

_PRICE = {
    "per_model": {"qwen3.5-flash": {"input_per_1m": 0.8, "output_per_1m": 2.0}},
    "default": {"input_per_1m": 3.0, "output_per_1m": 15.0},
}


def test_sums_per_model_single_model():
    # 真实样本:4 次 qwen3.5-flash;Σprompt=6888、Σcompletion=337
    calls = [
        {"model": "qwen3.5-flash", "error": None, "usage": {"prompt_tokens": 1512, "completion_tokens": 22}},
        {"model": "qwen3.5-flash", "error": None, "usage": {"prompt_tokens": 2111, "completion_tokens": 160}},
        {"model": "qwen3.5-flash", "error": None, "usage": {"prompt_tokens": 1940, "completion_tokens": 142}},
        {"model": "qwen3.5-flash", "error": None, "usage": {"prompt_tokens": 1325, "completion_tokens": 13}},
    ]
    cost, snapshot = compute_llm_cost(calls, _PRICE)
    # 6888/1e6*0.8 + 337/1e6*2.0 = 0.0055104 + 0.000674 = 0.0061844 → round(6)
    assert cost == 0.006184
    assert snapshot == {
        "mode": "per_model",
        "prices": {
            "qwen3.5-flash": {"input_per_1m": 0.8, "output_per_1m": 2.0, "_source": "per_model"},
        },
    }


def test_multi_model_prices_each_bucket_separately():
    calls = [
        {"model": "qwen3.5-flash", "error": None, "usage": {"prompt_tokens": 1_000_000, "completion_tokens": 0}},
        {"model": "gpt-x", "error": None, "usage": {"prompt_tokens": 0, "completion_tokens": 1_000_000}},
    ]
    price = {
        "per_model": {
            "qwen3.5-flash": {"input_per_1m": 0.8, "output_per_1m": 2.0},
            "gpt-x": {"input_per_1m": 10.0, "output_per_1m": 30.0},
        },
        "default": {"input_per_1m": 3.0, "output_per_1m": 15.0},
    }
    cost, snap = compute_llm_cost(calls, price)
    assert cost == 30.8  # qwen 1M入×0.8=0.8 + gpt-x 1M出×30=30.0
    assert set(snap["prices"]) == {"qwen3.5-flash", "gpt-x"}


def test_unknown_model_falls_back_to_default():
    calls = [{"model": "mystery", "error": None, "usage": {"prompt_tokens": 1_000_000, "completion_tokens": 0}}]
    price = {"per_model": {}, "default": {"input_per_1m": 3.0, "output_per_1m": 15.0}}
    cost, snap = compute_llm_cost(calls, price)
    assert cost == 3.0
    assert snap["prices"]["mystery"]["_source"] == "default"


def test_unpriced_model_marked_and_zero_cost():
    calls = [{"model": "mystery", "error": None, "usage": {"prompt_tokens": 1_000_000, "completion_tokens": 999}}]
    cost, snap = compute_llm_cost(calls, {"per_model": {}})  # 无 default
    assert cost == 0.0
    assert snap["prices"]["mystery"]["unpriced"] is True


def test_error_and_missing_usage_calls_skipped():
    calls = [
        {"model": "qwen3.5-flash", "error": "boom", "usage": {"prompt_tokens": 9_999_999, "completion_tokens": 9_999_999}},
        {"model": "qwen3.5-flash", "error": None, "usage": None},  # 无 usage
        {"model": "qwen3.5-flash", "error": None, "usage": {"prompt_tokens": 1_000_000, "completion_tokens": 0}},
    ]
    cost, _ = compute_llm_cost(calls, _PRICE)
    assert cost == 0.8  # 只有第 3 条计入


def test_empty_or_all_error_returns_none():
    assert compute_llm_cost([], _PRICE) == (None, None)
    assert compute_llm_cost(None, _PRICE) == (None, None)
    all_error = [{"model": "x", "error": "boom", "usage": {"prompt_tokens": 100, "completion_tokens": 100}}]
    assert compute_llm_cost(all_error, _PRICE) == (None, None)


def test_malformed_price_entry_is_treated_as_unpriced_not_raised():
    # 手改配置页可能存出残缺/非法单价(缺 output_per_1m、非 dict);不能抛异常连累 token 回填。
    calls = [
        {"model": "bad-a", "error": None, "usage": {"prompt_tokens": 1_000_000, "completion_tokens": 5}},
        {"model": "bad-b", "error": None, "usage": {"prompt_tokens": 1_000_000, "completion_tokens": 5}},
        {"model": "ok", "error": None, "usage": {"prompt_tokens": 1_000_000, "completion_tokens": 0}},
    ]
    price = {
        "per_model": {
            "bad-a": {"input_per_1m": 0.8},   # 缺 output_per_1m
            "bad-b": 5,                        # 非 dict
            "ok": {"input_per_1m": 3.0, "output_per_1m": 15.0},
        },
    }
    cost, snap = compute_llm_cost(calls, price)  # 不得抛异常
    assert cost == 3.0  # 只有 ok(1M 入 × 3.0)计入;两个残缺项按 unpriced
    assert snap["prices"]["bad-a"].get("unpriced") is True
    assert snap["prices"]["bad-b"].get("unpriced") is True


def test_get_llm_prices_falls_back_to_default_then_uses_override():
    from app.db.session import SessionLocal
    from app.models.app_config import AppConfig
    from app.repositories import app_config
    from app.services.llm_cost import get_llm_prices

    db = SessionLocal()
    try:
        # 无 override → CONFIG_DEFS 默认(含 per_model / default)
        prices = get_llm_prices(db)
        assert "per_model" in prices and "default" in prices
        # 有 override → 用 DB 值
        app_config.set_value(
            db, "llm_token_price",
            {"per_model": {"m": {"input_per_1m": 1.0, "output_per_1m": 2.0}},
             "default": {"input_per_1m": 0.0, "output_per_1m": 0.0}},
            admin_id=1,
        )
        assert get_llm_prices(db)["per_model"]["m"]["input_per_1m"] == 1.0
    finally:
        row = db.get(AppConfig, "llm_token_price")
        if row is not None:
            db.delete(row)
            db.commit()
        db.close()


def test_backfill_llm_calls_stores_cost_and_snapshot(monkeypatch):
    from datetime import UTC, datetime

    from app.api.v1 import compare_record
    from app.db.session import SessionLocal
    from app.models.app_config import AppConfig
    from app.models.comparison import ComparisonRecord
    from app.repositories import app_config

    sample = [
        {"model": "qwen3.5-flash", "error": None, "usage": {"prompt_tokens": 1512, "completion_tokens": 22}},
        {"model": "qwen3.5-flash", "error": None, "usage": {"prompt_tokens": 2111, "completion_tokens": 160}},
        {"model": "qwen3.5-flash", "error": None, "usage": {"prompt_tokens": 1940, "completion_tokens": 142}},
        {"model": "qwen3.5-flash", "error": None, "usage": {"prompt_tokens": 1325, "completion_tokens": 13}},
    ]
    monkeypatch.setattr(compare_record, "fetch_llm_calls", lambda trace_id: sample)

    db = SessionLocal()
    try:
        app_config.set_value(
            db, "llm_token_price",
            {"per_model": {"qwen3.5-flash": {"input_per_1m": 0.8, "output_per_1m": 2.0}},
             "default": {"input_per_1m": 3.0, "output_per_1m": 15.0}},
            admin_id=1,
        )
        rec = ComparisonRecord(
            trace_id="llmcost-bf-1", status="success",
            created_at=datetime.now(UTC).replace(tzinfo=None),
        )
        db.add(rec)
        db.commit()
        rid = rec.id
    finally:
        db.close()

    compare_record._backfill_llm_calls(rid, "llmcost-bf-1")  # 独立 session 内回填

    db = SessionLocal()
    try:
        rec = db.get(ComparisonRecord, rid)
        assert rec.llm_cost_yuan == 0.006184
        assert rec.llm_price_snapshot["prices"]["qwen3.5-flash"]["input_per_1m"] == 0.8
        assert rec.input_tokens == 6888  # 现有 token 派生仍在
    finally:
        db.delete(db.get(ComparisonRecord, rid))
        row = db.get(AppConfig, "llm_token_price")
        if row is not None:
            db.delete(row)
        db.commit()
        db.close()


def test_admin_detail_schema_exposes_llm_cost_fields():
    from app.admin.schemas.comparison import AdminComparisonDetail

    fields = AdminComparisonDetail.model_fields
    assert "llm_cost_yuan" in fields
    assert "llm_price_snapshot" in fields