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()
    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"}