tests/test_llm.py

from __future__ import annotations

import json
import sys

from tiny_trigger.automation import AutomationDocument
from tiny_trigger.llm import (
    DEFAULT_ANTHROPIC_MODEL,
    DEFAULT_OPENAI_MODEL,
    SYSTEM_PROMPT,
    _anthropic_message,
    _build_user_prompt,
    _chat_payload,
    _clean_api_key,
    _openai_chat_completion,
    compile_automation_with_anthropic,
    compile_automation_with_openai,
    compile_automation_with_replicate,
    _validate_compile_result,
    extract_json_object,
)


def test_extract_json_object_from_fenced_response() -> None:
    raw = """```json
{"rules": [{"name": "notify", "when": {"all": [{"present": {"label": "package"}}]}, "then": [{"type": "simulate", "name": "notify"}]}]}
```"""
    data = json.loads(extract_json_object(raw))
    document = AutomationDocument.model_validate(data)

    assert document.rules[0].name == "notify"


def test_openai_payload_uses_json_object() -> None:
    payload = _chat_payload(model="qwen", user_prompt="compile this", response_format="json_object")

    assert payload["response_format"] == {"type": "json_object"}
    assert payload["max_tokens"] == 512


def test_prompt_includes_validation_schema() -> None:
    prompt = _build_user_prompt(instruction="if package at door notify me", class_names=["package", "door"])

    assert "Available detection labels: package, door" in prompt
    assert "Full validation schema" in prompt
    assert "AutomationDocument" in prompt


def test_prompt_teaches_near_relations() -> None:
    prompt = _build_user_prompt(
        instruction="If person near steering wheel then turn on pc.",
        class_names=["person", "steering wheel"],
    )

    assert '"near": {"a": "person", "b": "steering wheel", "max_gap_percent": 16}' in prompt
    assert "Do not replace a near relation with two present conditions." in SYSTEM_PROMPT
    assert "Use max_gap_percent for near/far box-edge distance." in SYSTEM_PROMPT
    assert "Use state gates in gate: enabled, cooldown." in SYSTEM_PROMPT
    assert 'Use trigger.on="enter" for state assertions' in SYSTEM_PROMPT
    assert 'trigger.on="change"' in SYSTEM_PROMPT
    assert '"trigger": {"on": "enter"}' in prompt
    assert '"gate": {"enabled": true}' in prompt
    assert '"gate": {"enabled": true, "cooldown": {"key": "package-at-door", "minutes": 15}}' in prompt


def test_api_key_rejects_pasted_traceback() -> None:
    try:
        _clean_api_key('   File "/tmp/example.py", line 1\nrequests.exceptions.HTTPError', "Anthropic")
    except ValueError as exc:
        assert "not an API key" in str(exc) or "single-line token" in str(exc)
    else:
        raise AssertionError("Tracebacks should not be accepted as API keys")


def test_invalid_empty_object_fails_validation() -> None:
    try:
        _validate_compile_result("{}")
    except Exception as exc:
        assert "rules" in str(exc)
    else:
        raise AssertionError("{} should not validate as an automation document")


def test_replicate_compile_uses_stream_api(monkeypatch) -> None:
    class Replicate:
        class Client:
            def __init__(self, api_token):
                assert api_token == "test-token"

            def stream(self, model, input):
                assert model == "openai/gpt-5.2"
                assert input["messages"] == []
                assert input["verbosity"] == "medium"
                assert input["reasoning_effort"] == "low"
                assert "Return only the JSON object." in input["prompt"]
                yield '{"rules":[{"name":"notify","when":{"all":[{"present":{"label":"person"}}]},'
                yield '"then":[{"type":"simulate","name":"notify"}]}]}'

    monkeypatch.setitem(sys.modules, "replicate", Replicate)

    result = compile_automation_with_replicate(
        instruction="if person present notify",
        class_names=["person"],
        api_token="test-token",
        model="openai/gpt-5.2",
        reasoning_effort="low",
    )

    assert result.document.rules[0].name == "notify"


def test_replicate_invalid_json_does_not_auto_repair(monkeypatch) -> None:
    class Replicate:
        class Client:
            calls = 0

            def __init__(self, api_token):
                return None

            def stream(self, model, input):
                self.__class__.calls += 1
                yield "not json"

    monkeypatch.setitem(sys.modules, "replicate", Replicate)

    try:
        compile_automation_with_replicate(
            instruction="if person present notify",
            class_names=["person"],
            api_token="test-token",
            model="openai/gpt-5.2",
        )
    except ValueError as exc:
        assert "Raw response: not json" in str(exc)
    else:
        raise AssertionError("Invalid provider output should fail validation")

    assert Replicate.Client.calls == 1


def test_replicate_rate_limit_error_is_human_readable(monkeypatch) -> None:
    class RateLimitError(Exception):
        status = 429
        detail = "Request was throttled. Your rate limit"

    class Replicate:
        class Client:
            def __init__(self, api_token):
                return None

            def stream(self, model, input):
                raise RateLimitError()
                yield ""

    monkeypatch.setitem(sys.modules, "replicate", Replicate)

    try:
        compile_automation_with_replicate(
            instruction="if person present notify",
            class_names=["person"],
            api_token="test-token",
            model="openai/gpt-5.2",
        )
    except RuntimeError as exc:
        message = str(exc)
        assert "Replicate API request failed" in message
        assert "status 429" in message
        assert "provider rate limit" in message
    else:
        raise AssertionError("Rate limits should surface as RuntimeError")


def test_openai_compile_uses_chat_completions(monkeypatch) -> None:
    monkeypatch.setitem(sys.modules, "openai", None)

    class Response:
        def raise_for_status(self) -> None:
            return None

        def json(self):
            return {
                "choices": [
                    {
                        "message": {
                            "content": (
                                '{"rules":[{"name":"notify","when":{"all":[{"present":{"label":"person"}}]},'
                                '"then":[{"type":"simulate","name":"notify"}]}]}'
                            )
                        }
                    }
                ]
            }

    class Requests:
        @staticmethod
        def post(endpoint, headers, json, timeout):
            assert endpoint == "https://api.openai.com/v1/chat/completions"
            assert headers["Authorization"] == "Bearer test-key"
            assert json["model"] == "gpt-test"
            assert json["response_format"] == {"type": "json_object"}
            assert timeout == 120.0
            return Response()

    monkeypatch.setitem(sys.modules, "requests", Requests)

    result = compile_automation_with_openai(
        instruction="if person present notify",
        class_names=["person"],
        api_key="test-key",
        model="gpt-test",
    )

    assert result.document.rules[0].name == "notify"


def test_openai_sdk_completion() -> None:
    class Message:
        content = '{"rules":[{"name":"notify","when":{"all":[{"present":{"label":"person"}}]},"then":[{"type":"simulate","name":"notify"}]}]}'

    class Choice:
        message = Message()

    class Response:
        choices = [Choice()]

    class Completions:
        @staticmethod
        def create(**payload):
            assert payload["model"] == DEFAULT_OPENAI_MODEL
            assert payload["response_format"] == {"type": "json_object"}
            return Response()

    class Chat:
        completions = Completions()

    class OpenAIClient:
        chat = Chat()

        def __init__(self, api_key, timeout):
            assert api_key == "test-key"
            assert timeout == 120.0

    class OpenAIModule:
        OpenAI = OpenAIClient

    raw = _openai_chat_completion(
        api_key="test-key",
        model=DEFAULT_OPENAI_MODEL,
        user_prompt="compile",
        timeout=120.0,
        openai_module=OpenAIModule,
    )

    assert json.loads(raw)["rules"][0]["name"] == "notify"


def test_anthropic_compile_uses_messages_api(monkeypatch) -> None:
    monkeypatch.setitem(sys.modules, "anthropic", None)

    class Response:
        def raise_for_status(self) -> None:
            return None

        def json(self):
            return {
                "content": [
                    {
                        "type": "text",
                        "text": (
                            '{"rules":[{"name":"notify","when":{"all":[{"present":{"label":"person"}}]},'
                            '"then":[{"type":"simulate","name":"notify"}]}]}'
                        ),
                    }
                ]
            }

    class Requests:
        @staticmethod
        def post(endpoint, headers, json, timeout):
            assert endpoint == "https://api.anthropic.com/v1/messages"
            assert headers["x-api-key"] == "test-key"
            assert headers["anthropic-version"] == "2023-06-01"
            assert json["model"] == "claude-test"
            assert json["system"] == SYSTEM_PROMPT
            assert timeout == 120.0
            return Response()

    monkeypatch.setitem(sys.modules, "requests", Requests)

    result = compile_automation_with_anthropic(
        instruction="if person present notify",
        class_names=["person"],
        api_key="test-key",
        model="claude-test",
    )

    assert result.document.rules[0].name == "notify"


def test_anthropic_sdk_message() -> None:
    class TextBlock:
        text = '{"rules":[{"name":"notify","when":{"all":[{"present":{"label":"person"}}]},"then":[{"type":"simulate","name":"notify"}]}]}'

    class Response:
        content = [TextBlock()]

    class Messages:
        @staticmethod
        def create(**payload):
            assert payload["model"] == DEFAULT_ANTHROPIC_MODEL
            assert payload["system"] == SYSTEM_PROMPT
            return Response()

    class AnthropicClient:
        messages = Messages()

        def __init__(self, api_key, timeout):
            assert api_key == "test-key"
            assert timeout == 120.0

    class AnthropicModule:
        Anthropic = AnthropicClient

    raw = _anthropic_message(
        api_key="test-key",
        model=DEFAULT_ANTHROPIC_MODEL,
        user_prompt="compile",
        timeout=120.0,
        anthropic_module=AnthropicModule,
    )

    assert json.loads(raw)["rules"][0]["name"] == "notify"