nisse/nisse-core/tests/test_llm.py

import json
from pathlib import Path
from typing import Any

import httpx
import pytest

from nisse_core.canonical import (
    AssistantText,
    StoredEvent,
    ToolCall,
    UserMessage,
)
from nisse_core.envfile import load_env_file
from nisse_core.llm import (
    DEFAULT_MAX_OUTPUT_TOKENS,
    OUTAGE_BACKOFF_S,
    RATE_LIMIT_RETRY_S,
    AnthropicPassthroughClient,
    BudgetExhausted,
    LiteLLMClient,
    LLMError,
    RoleMap,
    TruncatedResponse,
    UpstreamUnavailable,
)

EVENTS = [StoredEvent(id=1, ts="t", event=UserMessage(text="hi", source="m:1"))]


def make_client(handler, sleep=lambda _s: None) -> LiteLLMClient:
    return LiteLLMClient(
        "http://litellm.test/v1",
        "sk-test",
        transport=httpx.MockTransport(handler),
        sleep=sleep,  # never wait in unit tests
    )


def completion(message: dict) -> httpx.Response:
    return httpx.Response(
        200, json={"model": "gpt-test-1", "choices": [{"message": message}]}
    )


def test_request_shape_and_auth() -> None:
    seen: dict = {}

    def handler(request: httpx.Request) -> httpx.Response:
        seen["url"] = str(request.url)
        seen["auth"] = request.headers.get("authorization")
        seen["body"] = json.loads(request.content)
        return completion({"content": "hello"})

    make_client(handler).complete(EVENTS, tools=None, model="local/qwen", system="s")
    assert seen["url"] == "http://litellm.test/v1/chat/completions"
    assert seen["auth"] == "Bearer sk-test"
    assert seen["body"]["model"] == "local/qwen"
    assert seen["body"]["messages"][0] == {"role": "system", "content": "s"}
    assert seen["body"]["messages"][1] == {"role": "user", "content": "hi"}


def test_assistant_text_parsed_with_producing_model() -> None:
    client = make_client(lambda request: completion({"content": "hello"}))
    [event] = client.complete(EVENTS, tools=None, model="local/qwen")
    assert event == AssistantText(text="hello", model="gpt-test-1")


def test_tool_calls_get_locally_minted_canonical_ids() -> None:
    message = {
        "content": None,
        "tool_calls": [
            {
                "id": "call_provider_abc",
                "type": "function",
                "function": {
                    "name": "run_tests",
                    "arguments": '{"path": "tests/"}',
                },
            }
        ],
    }
    client = make_client(lambda request: completion(message))
    [event] = client.complete(EVENTS, tools=None, model="local/qwen")
    assert isinstance(event, ToolCall)
    assert event.tool == "run_tests"
    assert event.arguments == {"path": "tests/"}
    assert event.call_id != "call_provider_abc"  # provider ID never a join key
    assert event.call_id.startswith("call-")


def test_truncated_tool_arguments_raise_llm_error() -> None:
    message = {
        "content": None,
        "tool_calls": [
            {
                "id": "c",
                "type": "function",
                "function": {"name": "submit_plan", "arguments": '{"content": "'},
            }
        ],
    }
    client = make_client(lambda request: completion(message))
    with pytest.raises(LLMError, match="non-JSON arguments"):
        client.complete(EVENTS, tools=None, model="local/qwen")


def test_non_object_tool_arguments_raise_llm_error() -> None:
    message = {
        "content": None,
        "tool_calls": [
            {
                "id": "c",
                "type": "function",
                "function": {"name": "submit_plan", "arguments": '"just a string"'},
            }
        ],
    }
    client = make_client(lambda request: completion(message))
    with pytest.raises(LLMError, match="expected object"):
        client.complete(EVENTS, tools=None, model="local/qwen")


def test_http_error_status_raises_llm_error() -> None:
    client = make_client(lambda request: httpx.Response(500, text="boom"))
    with pytest.raises(LLMError, match="500"):
        client.complete(EVENTS, tools=None, model="local/qwen")


def test_rate_limit_is_retried_once_then_surfaced() -> None:
    # A 429 is transient. Losing a long delegation to one costs far more
    # than waiting out the window, but retrying forever would hide a real
    # ceiling, so it is retried exactly once.
    attempts: list[int] = []
    waited: list[float] = []

    def handler(request: httpx.Request) -> httpx.Response:
        attempts.append(1)
        return httpx.Response(429, text="rate limited")

    client = make_client(handler, sleep=waited.append)
    with pytest.raises(LLMError, match="429"):
        client.complete(EVENTS, tools=None, model="local/qwen")
    assert len(attempts) == 2
    assert waited == [RATE_LIMIT_RETRY_S]
    # Reported even though the call failed, so a caller can leave the wait
    # out of how long the model took.
    assert client.last_rate_limit_wait_s == RATE_LIMIT_RETRY_S


def test_rate_limit_that_clears_succeeds() -> None:
    responses = [httpx.Response(429, text="slow down"), completion({"content": "hi"})]
    client = make_client(lambda request: responses.pop(0), sleep=lambda _s: None)
    [event] = client.complete(EVENTS, tools=None, model="local/qwen")
    assert isinstance(event, AssistantText)
    assert client.last_rate_limit_wait_s == RATE_LIMIT_RETRY_S


def test_the_rate_limit_wait_is_per_call() -> None:
    responses = [
        httpx.Response(429, text="slow down"),
        completion({"content": "one"}),
        completion({"content": "two"}),
    ]
    client = make_client(lambda request: responses.pop(0), sleep=lambda _s: None)
    client.complete(EVENTS, tools=None, model="local/qwen")
    client.complete(EVENTS, tools=None, model="local/qwen")
    assert client.last_rate_limit_wait_s == 0.0


def test_connection_failure_raises_llm_error() -> None:
    def handler(request: httpx.Request) -> httpx.Response:
        raise httpx.ConnectError("refused")

    with pytest.raises(LLMError, match="request failed"):
        make_client(handler).complete(EVENTS, tools=None, model="local/qwen")


def test_malformed_response_raises_llm_error() -> None:
    client = make_client(lambda request: httpx.Response(200, json={"oops": True}))
    with pytest.raises(LLMError, match="malformed"):
        client.complete(EVENTS, tools=None, model="local/qwen")


ROLES = RoleMap(
    driver="local/qwen",
    helper="anthropic/claude-sonnet-5",
    expert="anthropic/claude-opus-5",
)


def test_role_map_resolves_and_pins() -> None:
    roles = ROLES
    assert roles.resolve("driver") == "local/qwen"
    assert roles.pinned() == {
        "driver": "local/qwen",
        "helper": "anthropic/claude-sonnet-5",
        "expert": "anthropic/claude-opus-5",
        # Unset, the classifier falls back to the expert's model (L9).
        "classifier": "anthropic/claude-opus-5",
    }


def test_a_pinned_map_from_before_the_driver_reads_its_own_roles() -> None:
    # Pinned before M2.8: the old loop's roles are there, and are ignored.
    old = {"planner": "p", "executor": "e", "driver": "d", "helper": "h"}
    roles = RoleMap.from_pinned(old, ROLES)
    assert (roles.driver, roles.helper, roles.expert) == ("d", "h", ROLES.expert)


def test_unknown_role_raises() -> None:
    with pytest.raises(LLMError, match="unknown role"):
        ROLES.resolve("boss")


def test_load_env_file(tmp_path: Path) -> None:
    envfile = tmp_path / ".env"
    envfile.write_text(
        "# comment\n\nNISSE_LLM_TOKEN=sk-abc\nNISSE_LLM_BASE_URL = http://x \n"
    )
    env: dict[str, str] = {"NISSE_LLM_TOKEN": "sk-already-set"}
    load_env_file(envfile, env)
    assert env["NISSE_LLM_TOKEN"] == "sk-already-set"  # real env wins
    assert env["NISSE_LLM_BASE_URL"] == "http://x"


def test_load_env_file_missing_is_noop(tmp_path: Path) -> None:
    env: dict[str, str] = {}
    load_env_file(tmp_path / "nope.env", env)
    assert env == {}


# --- truncation (silent until measured against Opus) ---------------------


def test_openai_route_raises_on_truncation() -> None:
    # finish_reason "length" means the answer was cut off; parsing it as
    # complete yields a tool call with required arguments simply absent.
    truncated = httpx.Response(
        200,
        json={
            "model": "gpt-test-1",
            "choices": [
                {
                    "finish_reason": "length",
                    "message": {
                        "content": None,
                        "tool_calls": [
                            {
                                "id": "c",
                                "type": "function",
                                "function": {
                                    "name": "submit_plan",
                                    "arguments": '{"type": "arch"}',
                                },
                            }
                        ],
                    },
                }
            ],
        },
    )
    client = make_client(lambda request: truncated)
    with pytest.raises(TruncatedResponse, match="output limit"):
        client.complete(EVENTS, tools=None, model="gpt-6-astra")


def test_anthropic_route_raises_on_truncation() -> None:
    # Observed for real: Opus spent 4,064 output tokens against a 4,096
    # ceiling, and a longer plan came back as submit_plan with no content.
    def handler(request: httpx.Request) -> httpx.Response:
        return httpx.Response(
            200,
            json={
                "model": "claude-opus-5",
                "stop_reason": "max_tokens",
                "content": [
                    {
                        "type": "tool_use",
                        "id": "t",
                        "name": "submit_plan",
                        "input": {"type": "arch"},
                    }
                ],
            },
        )

    client = AnthropicPassthroughClient(
        "https://litellm.test",
        "tok",
        transport=httpx.MockTransport(handler),
        sleep=lambda _s: None,
    )
    with pytest.raises(TruncatedResponse, match="output limit"):
        client.complete(EVENTS, tools=None, model="claude-opus-5")


def test_a_complete_answer_is_not_mistaken_for_truncation() -> None:
    def handler(request: httpx.Request) -> httpx.Response:
        return httpx.Response(
            200,
            json={
                "model": "claude-opus-5",
                "stop_reason": "tool_use",
                "content": [
                    {
                        "type": "tool_use",
                        "id": "t",
                        "name": "submit_plan",
                        "input": {"type": "arch", "content": "a real plan"},
                    }
                ],
            },
        )

    client = AnthropicPassthroughClient(
        "https://litellm.test", "tok", transport=httpx.MockTransport(handler)
    )
    [call] = client.complete(EVENTS, tools=None, model="claude-opus-5")
    assert call.arguments["content"] == "a real plan"


def test_output_ceiling_is_generous_by_default() -> None:
    # Output is billed by what is produced, so a low ceiling buys nothing and
    # costs silent truncation. The configured value is tested in test_config.
    assert DEFAULT_MAX_OUTPUT_TOKENS >= 8192
    assert LiteLLMClient("http://x", "t")._max_tokens == DEFAULT_MAX_OUTPUT_TOKENS


def test_truncation_is_an_llm_error_so_callers_already_handle_it() -> None:
    # The planner retries, the executor and delegation gate. None of them
    # need to learn a new exception type.
    assert issubclass(TruncatedResponse, LLMError)


def test_a_5xx_is_upstream_trouble() -> None:
    client = make_client(lambda request: httpx.Response(503, text="bad gateway"))
    with pytest.raises(UpstreamUnavailable, match="503"):
        client.complete(EVENTS, tools=None, model="local-fast")


def test_a_connection_failure_is_upstream_trouble() -> None:
    def handler(request: httpx.Request) -> httpx.Response:
        raise httpx.ConnectError("refused")

    with pytest.raises(UpstreamUnavailable, match="request failed"):
        make_client(handler).complete(EVENTS, tools=None, model="local-fast")


def test_a_persistent_rate_limit_is_upstream_trouble() -> None:
    # It cleared for neither attempt, so the caller should keep its work
    # rather than discard it.
    client = make_client(lambda request: httpx.Response(429, text="slow down"))
    with pytest.raises(UpstreamUnavailable):
        client.complete(EVENTS, tools=None, model="local-fast")


def test_a_bad_request_is_not_upstream_trouble() -> None:
    client = make_client(lambda request: httpx.Response(400, text="bad schema"))
    with pytest.raises(LLMError) as excinfo:
        client.complete(EVENTS, tools=None, model="local-fast")
    assert not isinstance(excinfo.value, UpstreamUnavailable)


def test_budget_refusal_is_a_kind_of_upstream_trouble() -> None:
    # So one catch covers the family, but it can still be told apart where
    # the difference matters — it will not clear within the hour.
    assert issubclass(BudgetExhausted, UpstreamUnavailable)


# --- tool_choice (L9: constrained decoding) --------------------------------

TOOLS = [
    {
        "type": "function",
        "function": {
            "name": "finish",
            "description": "done",
            "parameters": {"type": "object", "properties": {}},
        },
    }
]


def capture_bodies() -> tuple[list[dict], Any]:
    bodies: list[dict] = []

    def handler(request: httpx.Request) -> httpx.Response:
        bodies.append(json.loads(request.content))
        return completion({"role": "assistant", "content": "ok"})

    return bodies, handler


def test_required_is_sent_only_when_asked_and_tools_are_offered() -> None:
    bodies, handler = capture_bodies()
    client = make_client(handler)
    client.complete(EVENTS, tools=TOOLS, model="m", require_tool_call=True)
    client.complete(EVENTS, tools=TOOLS, model="m")
    client.complete(EVENTS, tools=None, model="m", require_tool_call=True)
    assert bodies[0]["tool_choice"] == "required"
    assert "tool_choice" not in bodies[1]  # the server's default, auto
    assert "tool_choice" not in bodies[2]  # nothing to choose from


def test_the_passthrough_asks_in_anthropics_terms() -> None:
    bodies: list[dict] = []

    def handler(request: httpx.Request) -> httpx.Response:
        bodies.append(json.loads(request.content))
        return httpx.Response(
            200,
            json={
                "model": "claude-haiku-4-5",
                "stop_reason": "tool_use",
                "content": [
                    {"type": "tool_use", "id": "t", "name": "finish", "input": {}}
                ],
            },
        )

    client = AnthropicPassthroughClient(
        "https://litellm.test", "tok", transport=httpx.MockTransport(handler)
    )
    client.complete(
        EVENTS, tools=TOOLS, model="claude-haiku-4-5", require_tool_call=True
    )
    assert bodies[0]["tool_choice"] == {"type": "any"}


def test_the_answering_model_is_recorded() -> None:
    # A proxy fallback answers under another model's name; this is how a
    # caller can tell.
    client = make_client(
        lambda request: completion({"role": "assistant", "content": "hi"})
    )
    client.complete(EVENTS, tools=None, model="local-fast")
    assert client.last_model == "gpt-test-1"


# --- outages delay, never fail (L5) --------------------------------------------


def test_a_gateway_error_is_waited_out_with_backoff() -> None:
    responses = [
        httpx.Response(502, text="bad gateway"),
        httpx.Response(503, text="unavailable"),
        completion({"content": "back"}),
    ]
    waited: list[float] = []
    client = make_client(lambda request: responses.pop(0), sleep=waited.append)
    [event] = client.complete(EVENTS, tools=None, model="local/qwen")
    assert isinstance(event, AssistantText) and event.text == "back"
    assert waited == list(OUTAGE_BACKOFF_S[:2])
    # Waiting isn't the model's time.
    assert client.last_rate_limit_wait_s == sum(OUTAGE_BACKOFF_S[:2])


def test_an_outage_that_outlasts_the_backoff_is_surfaced() -> None:
    attempts: list[int] = []
    waited: list[float] = []

    def handler(request: httpx.Request) -> httpx.Response:
        attempts.append(1)
        return httpx.Response(502, text="bad gateway")

    client = make_client(handler, sleep=waited.append)
    with pytest.raises(UpstreamUnavailable, match="502"):
        client.complete(EVENTS, tools=None, model="local/qwen")
    assert len(attempts) == len(OUTAGE_BACKOFF_S) + 1
    assert waited == list(OUTAGE_BACKOFF_S)


def test_a_dropped_connection_is_retried() -> None:
    calls: list[int] = []

    def handler(request: httpx.Request) -> httpx.Response:
        calls.append(1)
        if len(calls) == 1:
            raise httpx.ConnectError("refused")
        return completion({"content": "ok"})

    [event] = make_client(handler).complete(EVENTS, tools=None, model="local/qwen")
    assert isinstance(event, AssistantText) and len(calls) == 2


def test_a_500_or_a_timeout_is_not_retried() -> None:
    # A 500 may be the request's fault; a timeout already had its full
    # request_timeout_s.
    for failure in (
        lambda request: httpx.Response(500, text="boom"),
        lambda request: (_ for _ in ()).throw(httpx.ReadTimeout("slow")),
    ):
        calls: list[int] = []

        def handler(
            request: httpx.Request, failure=failure, calls=calls
        ) -> httpx.Response:
            calls.append(1)
            return failure(request)

        with pytest.raises(LLMError):
            make_client(handler).complete(EVENTS, tools=None, model="local/qwen")
        assert len(calls) == 1


def test_a_budget_refusal_is_not_waited_on() -> None:
    calls: list[int] = []

    def handler(request: httpx.Request) -> httpx.Response:
        calls.append(1)
        return httpx.Response(503, text="Budget has been exceeded! Current cost: 5.1")

    with pytest.raises(BudgetExhausted):
        make_client(handler).complete(EVENTS, tools=None, model="local/qwen")
    assert len(calls) == 1