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"
    session.ignored_providers = []
    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


def test_solver_step_uses_higher_max_tokens():
    # beobachtet am Portfoliomanagement: 4× je ~26k Reasoning-Tokens am 32k-Limit abgeschnitten.
    # DeepSeek V4.1 Flash unterstützt bis 393k Output-Tokens (OpenRouter) – die Solver-Stufe bekommt
    # ein eigenes, höheres Limit, die übrigen Stufen bleiben beim globalen.
    s_solver, c_solver = fake_session()
    s_solver.create(str, [{"role": "user", "content": "x"}], step="Solver-Code")
    assert c_solver["openrouter"].kwargs["max_tokens"] == llm.LLM_MAX_TOKENS_SOLVER > llm.LLM_MAX_TOKENS
    s_repair, c_repair = fake_session()
    s_repair.create(str, [{"role": "user", "content": "y"}], step="Solver-Reparatur")
    assert c_repair["openrouter"].kwargs["max_tokens"] == llm.LLM_MAX_TOKENS_SOLVER
    s_extra, c_extra = fake_session()
    s_extra.create(str, [{"role": "user", "content": "z"}], step="Extraktion")
    assert c_extra["openrouter"].kwargs["max_tokens"] == llm.LLM_MAX_TOKENS


def test_solver_max_tokens_zero_falls_back_to_global(monkeypatch):
    monkeypatch.setattr(llm, "LLM_MAX_TOKENS_SOLVER", 0)
    session, completions = fake_session()
    session.create(str, [{"role": "user", "content": "x"}], step="Solver-Code")
    assert completions["openrouter"].kwargs["max_tokens"] == llm.LLM_MAX_TOKENS


# --- 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"


# --- Tatsächlich antwortendes Modell (wichtig bei Aliasen wie ~deepseek/deepseek-flash-latest) ---

def test_responded_model_is_recorded_per_call_and_summed():
    from types import SimpleNamespace as NS
    session, _ = fake_session(model="~deepseek/deepseek-flash-latest")

    def antwort(modell):
        return NS(provider="Together", model=modell,
                  usage=NS(prompt_tokens=10, completion_tokens=5, cost=0.0001,
                           completion_tokens_details=None, prompt_tokens_details=None))
    session._record(antwort("deepseek/deepseek-v4.1-flash"), step="Verstehen", model="~deepseek/deepseek-flash-latest")
    session._record(antwort("deepseek/deepseek-v4.2-flash"), step="Extraktion", model="~deepseek/deepseek-flash-latest")
    assert [(r.model, r.responded_model) for r in session.usage.call_records] == [
        ("~deepseek/deepseek-flash-latest", "deepseek/deepseek-v4.1-flash"),
        ("~deepseek/deepseek-flash-latest", "deepseek/deepseek-v4.2-flash")]
    assert session.usage.responded_models == {"deepseek/deepseek-v4.1-flash": 1, "deepseek/deepseek-v4.2-flash": 1}


def test_missing_or_odd_model_field_is_ignored():
    from types import SimpleNamespace as NS
    session, _ = fake_session()
    session._record(NS(provider=None, model=object(), usage=NS(prompt_tokens=1, completion_tokens=1, cost=None,
                                                            completion_tokens_details=None, prompt_tokens_details=None)))
    assert session.usage.responded_models == {} and session.usage.call_records[0].responded_model is None


def _antwort_mit_verbrauch(completion_tokens):
    from types import SimpleNamespace as NS
    return NS(provider="AtlasCloud", model="deepseek/deepseek-v4.1-flash",
              usage=NS(prompt_tokens=8000, completion_tokens=completion_tokens, cost=0.02,
                       completion_tokens_details=NS(reasoning_tokens=completion_tokens - 500),
                       prompt_tokens_details=NS(cached_tokens=0)))


def test_truncated_answer_is_counted_before_the_retry(monkeypatch):
    # beobachtet an wagner_whitin_dijkstra: an LLM_MAX_TOKENS abgeschnitten, bezahlt, aber nirgends gezählt
    from instructor.core.exceptions import IncompleteOutputException
    monkeypatch.setattr(llm, "RETRY_WAITS", (0,))
    session, completions = fake_session()
    versuche = iter([IncompleteOutputException(last_completion=_antwort_mit_verbrauch(32000)),
                     ("ok", _antwort_mit_verbrauch(9000))])

    def create_with_completion(**kwargs):
        naechster = next(versuche)
        if isinstance(naechster, Exception):
            raise naechster
        return naechster

    completions["openrouter"].create_with_completion = create_with_completion
    assert session.create(str, [{"role": "user", "content": "x"}], step="Solver-Code") == "ok"
    assert session.usage.calls == 2 and session.usage.completion_tokens == 41000
    assert session.usage.cost_usd == pytest.approx(0.04)
    assert [r.step for r in session.usage.call_records] == ["Solver-Code (verworfen)", "Solver-Code"]


def test_rejected_attempts_are_counted_even_when_the_call_fails(monkeypatch):
    from types import SimpleNamespace as NS
    from instructor.core.exceptions import InstructorRetryException
    session, completions = fake_session()
    fehler = InstructorRetryException("Schema verletzt", last_completion=_antwort_mit_verbrauch(3000), n_attempts=1,
                                      total_usage=0, failed_attempts=[NS(completion=_antwort_mit_verbrauch(3000))])

    def create_with_completion(**kwargs):
        raise fehler

    completions["openrouter"].create_with_completion = create_with_completion
    with pytest.raises(InstructorRetryException):
        session.create(str, [{"role": "user", "content": "x"}], step="Prüfung")
    assert session.usage.calls == 2 and session.usage.completion_tokens == 6000   # zwei verschiedene Antworten


def test_ordinary_errors_have_no_discarded_answers():
    assert llm.discarded_completions(RuntimeError("Error code: 503")) == []


# --- Langsamen Anbieter nicht festhalten (OPENROUTER_SWITCH_BELOW) ---

def _aufruf(session, anbieter, tokens, sekunden):
    from types import SimpleNamespace as NS
    session._record(NS(provider=anbieter, usage=NS(prompt_tokens=100, completion_tokens=tokens, cost=None,
                                                   completion_tokens_details=None, prompt_tokens_details=None)),
                    step="Verstehen", seconds=sekunden)
    session._avoid_if_slow()


def test_slow_provider_is_avoided_for_the_rest_of_the_run(monkeypatch):
    # beobachtet: alle Aufrufe bei AtlasCloud mit ~65 Tokens/s, 1575 s statt ~150 s
    monkeypatch.setattr(llm, "OPENROUTER_SWITCH_BELOW", 75)
    session, completions = fake_session()
    _aufruf(session, "AtlasCloud", 15201, 232.0)
    assert session.ignored_providers == ["AtlasCloud"] and session.session_id != "lauf-1"
    session.create(str, [{"role": "user", "content": "x"}], step="Extraktion")
    body = completions["openrouter"].kwargs["extra_body"]
    assert body["provider"]["ignore"] == ["AtlasCloud"] and body["session_id"] == session.session_id


@pytest.mark.parametrize("anbieter, tokens, sekunden", [
    ("Together", 15000, 60.0),      # schnell
    ("Fireworks", 400, 20.0),       # zu kurz, um die Rate zu beurteilen (Wartezeit bis zum ersten Token)
])
def test_fast_or_short_calls_keep_the_provider(monkeypatch, anbieter, tokens, sekunden):
    monkeypatch.setattr(llm, "OPENROUTER_SWITCH_BELOW", 75)
    session, _ = fake_session()
    _aufruf(session, anbieter, tokens, sekunden)
    assert session.ignored_providers == [] and session.session_id == "lauf-1"


def test_switching_can_be_disabled_and_is_capped(monkeypatch):
    monkeypatch.setattr(llm, "OPENROUTER_SWITCH_BELOW", 0)
    session, _ = fake_session()
    _aufruf(session, "AtlasCloud", 15000, 230.0)
    assert session.ignored_providers == []
    monkeypatch.setattr(llm, "OPENROUTER_SWITCH_BELOW", 75)
    for anbieter in ("A", "B", "C", "D"):
        _aufruf(session, anbieter, 15000, 230.0)
    assert session.ignored_providers == ["A", "B", "C"]      # höchstens MAX_IGNORED_PROVIDERS


def test_no_endpoints_left_lifts_the_ignore_list():
    session, completions = fake_session()
    session.ignored_providers = ["AtlasCloud", "Parasail"]
    gesendet = []

    def create_with_completion(**kwargs):
        gesendet.append(kwargs["extra_body"]["provider"].get("ignore"))
        if len(gesendet) == 1:
            raise RuntimeError("Error code: 404 - {'error': {'message': 'No endpoints found for deepseek/x.'}}")
        return "ok", object()

    completions["openrouter"].create_with_completion = create_with_completion
    assert session.create(str, [{"role": "user", "content": "x"}], step="Prüfung") == "ok"
    assert gesendet == [["AtlasCloud", "Parasail"], None] and session.ignored_providers == []


# --- Abschneiden am Tokenlimit: anderer Anbieter statt Warten, am Ende LLMUnavailable statt Absturz ---

def _abgeschnitten(anbieter):
    from instructor.core.exceptions import IncompleteOutputException
    antwort = _antwort_mit_verbrauch(32000)
    antwort.provider = anbieter
    return IncompleteOutputException(last_completion=antwort)


def test_truncation_avoids_the_provider_and_retries_without_waiting(monkeypatch):
    # gemessen am 18.09.2026: Venice 4 von 28 Aufrufen abgeschnitten, alle übrigen Anbieter fast nie
    gewartet = []
    monkeypatch.setattr(llm.time, "sleep", gewartet.append)
    session, completions = fake_session()
    gesendet, versuche = [], iter([_abgeschnitten("Venice"), ("ok", _antwort_mit_verbrauch(9000))])

    def create_with_completion(**kwargs):
        gesendet.append((kwargs["extra_body"]["provider"].get("ignore"), kwargs["extra_body"]["session_id"]))
        naechster = next(versuche)
        if isinstance(naechster, Exception):
            raise naechster
        return naechster

    completions["openrouter"].create_with_completion = create_with_completion
    assert session.create(str, [{"role": "user", "content": "x"}], step="Solver-Code") == "ok"
    assert gewartet == [0]
    assert gesendet[0] == (None, "lauf-1") and gesendet[1][0] == ["Venice"] and gesendet[1][1] != "lauf-1"


def test_exhausted_retries_raise_llm_unavailable(monkeypatch):
    monkeypatch.setattr(llm.time, "sleep", lambda s: None)
    session, completions = fake_session()
    anbieter = iter(["Venice", "Makora", "Parasail", "Together"])

    def create_with_completion(**kwargs):
        raise _abgeschnitten(next(anbieter))

    completions["openrouter"].create_with_completion = create_with_completion
    with pytest.raises(llm.LLMUnavailable) as exc:
        session.create(str, [{"role": "user", "content": "x"}], step="Extraktion")
    assert exc.value.step == "Extraktion"
    assert session.usage.calls == 4                                       # alle vier bezahlten Antworten gezählt
    assert session.ignored_providers == ["Venice", "Makora", "Parasail"]  # höchstens MAX_IGNORED_PROVIDERS


def test_fatal_errors_stay_unwrapped():
    session, completions = fake_session()

    def create_with_completion(**kwargs):
        raise RuntimeError("Error code: 402 - {'error': {'message': 'Insufficient credits'}}")

    completions["openrouter"].create_with_completion = create_with_completion
    with pytest.raises(RuntimeError) as exc:
        session.create(str, [{"role": "user", "content": "x"}])
    assert not isinstance(exc.value, llm.LLMUnavailable)