nisse/nisse-dev/tests/test_llm.py

import json
from pathlib import Path

import httpx
import pytest

from nisse_dev.canonical import (
    AssistantText,
    StoredEvent,
    ToolCall,
    UserMessage,
)
from nisse_dev.envfile import load_env_file
from nisse_dev.llm import RATE_LIMIT_RETRY_S, LiteLLMClient, LLMError, RoleMap

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]


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)


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


ROLE_ENV = {
    "NISSE_ROLE_EXECUTOR": "local/qwen",
    "NISSE_ROLE_PLANNER": "anthropic/claude-opus-5",
    "NISSE_ROLE_DELEGATE": "anthropic/claude-opus-5",
    "NISSE_ROLE_CONSULTANT": "anthropic/claude-opus-5",
}


def test_role_map_resolves_and_pins() -> None:
    roles = RoleMap.from_env(ROLE_ENV)
    assert roles.resolve("executor") == "local/qwen"
    assert roles.pinned() == {
        "executor": "local/qwen",
        "planner": "anthropic/claude-opus-5",
        "delegate": "anthropic/claude-opus-5",
        "consultant": "anthropic/claude-opus-5",
    }


def test_role_map_missing_env_raises() -> None:
    with pytest.raises(LLMError, match="NISSE_ROLE_PLANNER"):
        RoleMap.from_env({"NISSE_ROLE_EXECUTOR": "local/qwen"})


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


def test_client_from_env_requires_config() -> None:
    with pytest.raises(LLMError, match="NISSE_LLM_BASE_URL"):
        LiteLLMClient.from_env({})


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 == {}