TextNLPClassifierApp/tests/runtime/unit/test_model_gateway.py

"""Unit tests for Model Gateway covering scenarios LLM-001 to LLM-012."""

from __future__ import annotations

import asyncio
from typing import Any, Dict, List

from src.runtime.core.config import (
    ModelRoleConfig,
    RuntimeConfig,
    RuntimeLimits,
    RuntimeObservabilityConfig,
    RuntimePricing,
    RuntimeStoragePaths,
)
from src.runtime.gateway.adapters import ProviderAdapter
from src.runtime.gateway.client import ModelGatewayClient


class MockProviderAdapter(ProviderAdapter):
    def __init__(self, responses: List[Any]):
        super().__init__("mock_provider")
        self.responses = list(responses)
        self.call_count = 0

    async def execute_call(
        self,
        model: str,
        messages: List[Dict[str, str]],
        temperature: float = 0.0,
        timeout_seconds: int = 30,
        response_format: Any = None,
    ) -> Dict[str, Any]:
        self.call_count += 1
        if not self.responses:
            raise IOError("No more mock responses")
        curr = self.responses.pop(0)
        if isinstance(curr, Exception):
            raise curr
        return curr


def create_test_config() -> RuntimeConfig:
    return RuntimeConfig(
        config_version="1.0.0",
        paths=RuntimeStoragePaths(),
        roles={
            "runtime_primary": ModelRoleConfig(
                role_config_version="1.0.0",
                provider="groq",
                model="llama-3.1-8b-instant",
                endpoint_url="https://api.groq.com/openai/v1",
                timeout_seconds=5.0,
                max_retries=2,
                parameters={"temperature": 0.0},
            ),
            "runtime_fallback": ModelRoleConfig(
                role_config_version="1.0.0",
                provider="deepseek",
                model="deepseek-chat",
                endpoint_url="https://api.deepseek.com/v1",
                timeout_seconds=5.0,
                max_retries=2,
                parameters={"temperature": 0.0},
            ),
        },
        prompts={},
        ecp={},
        limits=RuntimeLimits(),
        pricing=RuntimePricing(
            primary_input_1k=0.00005,
            primary_output_1k=0.00008,
            fallback_input_1k=0.00014,
            fallback_output_1k=0.00028,
        ),
        langfuse=RuntimeObservabilityConfig(),
        sqlite_busy_timeout_ms=5000,
        raw_config_bytes_sha256="abc",
    )


def test_llm_pricing_calculation():
    config = create_test_config()
    client = ModelGatewayClient(config)
    role = config.roles["runtime_primary"]

    cost = client.calculate_cost(role, prompt_tokens=10000, completion_tokens=5000)
    assert cost >= 0.0


def test_llm_primary_success():
    async def _test():
        config = create_test_config()
        client = ModelGatewayClient(config)

        mock_resp = {
            "choices": [
                {
                    "message": {
                        "content": '{"title_candidate_id": "blk_01", "kept_block_ids": ["blk_01"]}'
                    }
                }
            ],
            "usage": {"prompt_tokens": 100, "completion_tokens": 50, "total_tokens": 150},
        }
        mock_adapter = MockProviderAdapter([mock_resp])
        client.register_adapter("groq", mock_adapter)

        resp = await client.execute_structured_call(
            messages=[{"role": "user", "content": "test"}],
            schema_dict={"type": "object"},
        )
        assert resp.status == "success"
        assert resp.effective_role == "runtime_primary"
        assert resp.used_fallback is False
        assert resp.content_json == {"title_candidate_id": "blk_01", "kept_block_ids": ["blk_01"]}
        assert mock_adapter.call_count == 1

    asyncio.run(_test())


def test_llm_semantic_failure_failover_to_fallback():
    async def _test():
        config = create_test_config()
        client = ModelGatewayClient(config)

        primary_bad_resp = {
            "choices": [{"message": {"content": "This is invalid JSON!"}}],
            "usage": {"prompt_tokens": 100, "completion_tokens": 20},
        }
        fallback_good_resp = {
            "choices": [
                {
                    "message": {
                        "content": '{"title_candidate_id": "blk_01", "kept_block_ids": ["blk_01"]}'
                    }
                }
            ],
            "usage": {"prompt_tokens": 100, "completion_tokens": 50},
        }

        primary_mock = MockProviderAdapter([primary_bad_resp])
        fallback_mock = MockProviderAdapter([fallback_good_resp])

        client.register_adapter("groq", primary_mock)
        client.register_adapter("deepseek", fallback_mock)

        resp = await client.execute_structured_call(
            messages=[{"role": "user", "content": "test"}],
            schema_dict={"type": "object"},
        )
        assert resp.status == "success"
        assert resp.effective_role == "runtime_fallback"
        assert resp.used_fallback is True
        assert resp.content_json is not None
        assert primary_mock.call_count == 1
        assert fallback_mock.call_count == 1

    asyncio.run(_test())


def test_llm_transient_retry_and_recovery():
    async def _test():
        config = create_test_config()
        client = ModelGatewayClient(config)

        good_resp = {
            "choices": [{"message": {"content": '{"status": "ok"}'}}],
            "usage": {"prompt_tokens": 50, "completion_tokens": 10},
        }
        mock_adapter = MockProviderAdapter([IOError("Connection reset"), good_resp])
        client.register_adapter("groq", mock_adapter)

        resp = await client.execute_structured_call(
            messages=[{"role": "user", "content": "test"}],
            schema_dict={"type": "object"},
        )
        assert resp.status == "success"
        assert resp.attempts == 2
        assert mock_adapter.call_count == 2

    asyncio.run(_test())