text_to_solver_program/tests/test_llm.py

import threading

import pytest

from core import llm
from core.llm import is_fatal_error, is_transient_error, model_for_step, provider_for
from core.model_schema import LLMUsage


class FakeCompletions:
    def __init__(self): self.kwargs = None

    def create_with_completion(self, **kwargs):
        self.kwargs = kwargs
        return "ok", object()


def fake_session(model="deepseek/deepseek-v4.1-flash", step_models=None):
    """LLMSession ohne echten Client – der Konstruktor würde einen API-Key verlangen."""
    session = llm.LLMSession.__new__(llm.LLMSession)
    session.model, session.provider = model, provider_for(model)
    session.step_models = step_models or {}
    session.usage = LLMUsage(model=model, provider=session.provider)
    session._lock = threading.Lock()
    session.session_id = "lauf-1"
    completions = {}
    session._clients = {}
    for provider in {session.provider, *(provider_for(m) for m in session.step_models.values())}:
        completions[provider] = FakeCompletions()
        session._clients[provider] = type("C", (), {"chat": type("Chat", (), {"completions": completions[provider]})()})()
    session.client = session._clients[session.provider]
    return session, completions


@pytest.mark.parametrize("message", [
    "429 RESOURCE_EXHAUSTED. {'error': {'code': 429, 'message': 'You exceeded your current quota'}}",  # Gemini-Kontingent
    "Error code: 402 - {'error': {'message': 'Insufficient credits'}}",                               # OpenRouter-Guthaben
    "Error code: 401 - {'error': {'message': 'No auth credentials found'}}",
])
def test_fatal_errors_are_not_retried(message):
    assert is_fatal_error(message)
    assert not is_transient_error(message)


@pytest.mark.parametrize("message", [
    "503 UNAVAILABLE. {'error': {'code': 503, 'message': 'model is currently experiencing high demand'}}",
    "Error code: 429 - {'error': {'message': 'Rate limit exceeded'}}",
    "Error code: 502 - Bad Gateway",
    "Connection error.",
])
def test_transient_errors_are_retried(message):
    assert is_transient_error(message)


def test_validation_error_is_neither():
    message = "1 validation error for OptimizationModelSpecification"
    assert not is_fatal_error(message) and not is_transient_error(message)


def test_provider_detection():
    assert provider_for("deepseek/deepseek-v4-flash-0731") == "openrouter"
    assert provider_for("gemini-2.5-flash") == "gemini"


def test_reasoning_effort_is_sent_only_when_configured(monkeypatch):
    def make_session(effort):
        monkeypatch.setattr(llm, "REASONING_EFFORT", effort)
        session, completions = fake_session()
        session.create(str, [{"role": "user", "content": "x"}])
        return completions["openrouter"].kwargs["extra_body"]

    assert "reasoning" not in make_session("")
    assert make_session("low")["reasoning"] == {"effort": "low"}


def test_truncated_answer_is_retried():
    # beobachtet: ein Anbieter schnitt die Modellspezifikation ab (max_tokens-Vorgabe des Anbieters)
    message = "IncompleteOutputException: The output is incomplete due to a max_tokens length limit."
    assert is_transient_error(message) and not is_fatal_error(message)


def test_max_tokens_is_sent_to_openrouter():
    session, completions = fake_session()
    session.create(str, [{"role": "user", "content": "x"}])
    assert completions["openrouter"].kwargs["max_tokens"] == llm.LLM_MAX_TOKENS >= 16000


# --- Modell je Pipeline-Stufe (LLM_MODEL_EXTRACTION, LLM_MODEL_SOLVER, ...) ---

def test_step_model_falls_back_to_session_model(monkeypatch):
    monkeypatch.delenv("LLM_MODEL_SOLVER", raising=False)
    assert model_for_step("Solver-Code", "a/b") == "a/b"
    assert model_for_step("unbenannt", "a/b") == "a/b"  # Stufe ohne Umgebungsvariable


def test_step_model_from_environment(monkeypatch):
    monkeypatch.setenv("LLM_MODEL_SOLVER", "schnell/modell")
    assert model_for_step("Solver-Code", "a/b") == "schnell/modell"
    assert model_for_step("Solver-Reparatur", "a/b") == "schnell/modell"  # dieselbe Variable
    assert model_for_step("Extraktion", "a/b") == "a/b"


def test_create_uses_step_model_and_records_it():
    session, completions = fake_session(step_models={"Extraktion": "stark/modell", "Solver-Code": "gemini-2.5-flash"})
    session.create(str, [{"role": "user", "content": "x"}], step="Extraktion")
    assert completions["openrouter"].kwargs["model"] == "stark/modell"

    # anderer Anbieter je Stufe: Gemini bekommt keine OpenRouter-Parameter
    session.create(str, [{"role": "user", "content": "x"}], step="Solver-Code")
    assert completions["gemini"].kwargs["model"] == "gemini-2.5-flash"
    assert "extra_body" not in completions["gemini"].kwargs
    assert [r.model for r in session.usage.call_records] == ["stark/modell", "gemini-2.5-flash"]


def test_overrides_lists_only_deviating_steps():
    session, _ = fake_session(model="a/b", step_models={"Extraktion": "a/b", "Solver-Code": "schnell/modell"})
    assert session.overrides() == {"Solver-Code": "schnell/modell"}


# --- Prompt-Caching: Kennung je Lauf (Sticky Routing) und erfasste Cache-Treffer ---

def test_all_calls_of_a_run_share_one_openrouter_session_id():
    session, completions = fake_session(step_models={"Solver-Code": "gemini-2.5-flash"})
    session.create(str, [{"role": "user", "content": "x"}], step="Extraktion")
    assert completions["openrouter"].kwargs["extra_body"]["session_id"] == "lauf-1"
    session.create(str, [{"role": "user", "content": "x"}], step="Solver-Code")
    assert "extra_body" not in completions["gemini"].kwargs          # Gemini kennt den Parameter nicht


def test_new_session_gets_a_new_id(monkeypatch):
    monkeypatch.setattr(llm, "_make_client", lambda provider: object())
    assert llm.LLMSession("a/b").session_id != llm.LLMSession("a/b").session_id


def test_cached_tokens_are_recorded_per_call_and_step():
    from types import SimpleNamespace as NS
    session, _ = fake_session()

    def antwort(prompt, cached):
        return NS(provider="Modal", usage=NS(prompt_tokens=prompt, completion_tokens=100, cost=0.001,
                                             completion_tokens_details=NS(reasoning_tokens=40),
                                             prompt_tokens_details=NS(cached_tokens=cached)))
    session._record(antwort(5000, 0), step="Prüfung")
    session._record(antwort(5200, 4800), step="Prüfung")      # zweite Runde liest den gleichen Präfix
    session._record(NS(provider=None, usage=NS(prompt_tokens=900, completion_tokens=10, cost=None,
                                               completion_tokens_details=None, prompt_tokens_details=None)),
                    step="Verstehen")                          # Anbieter ohne Cache-Angabe
    assert session.usage.cached_tokens == 4800 and session.usage.prompt_tokens == 11100
    assert [(r.prompt_tokens, r.cached_tokens) for r in session.usage.call_records] == [(5000, 0), (5200, 4800), (900, 0)]
    stufen = session.usage.by_step()
    assert stufen["Prüfung"]["prompt_tokens"] == 10200 and stufen["Prüfung"]["cached_tokens"] == 4800


def test_old_results_without_cache_fields_still_load():
    from core.model_schema import LLMUsage
    alt = {"model": "a/b", "provider": "openrouter", "calls": 1, "prompt_tokens": 10, "completion_tokens": 5,
           "reasoning_tokens": 0, "call_records": [{"step": "Verstehen", "seconds": 1.0, "completion_tokens": 5}]}
    usage = LLMUsage.model_validate(alt)
    assert usage.cached_tokens == 0 and usage.call_records[0].cached_tokens == 0


def test_understanding_and_figure_reading_share_their_own_model(monkeypatch):
    monkeypatch.setenv("LLM_MODEL_UNDERSTANDING", "bild/modell")
    assert model_for_step("Verstehen", "a/b") == "bild/modell"
    assert model_for_step("Abbildungen", "a/b") == "bild/modell"