test_app_contract.py

"""HTTP/OpenAI contract tests for app.py without loading the real GPU model."""

from __future__ import annotations

import asyncio
import contextvars
import importlib
import json
import sys
import types
import unittest
from unittest.mock import patch


class _FakeEncoding(dict):
    def to(self, _device: str):
        return self


class _FakeTokenizer:
    eos_token_id = 1
    pad_token_id = 0

    def apply_chat_template(self, messages, *, tokenize=False, add_generation_prompt=True, tools=None):
        payload = {"messages": messages, "tools": tools or []}
        return json.dumps(payload, ensure_ascii=False, sort_keys=True)

    def __call__(self, text, **_kwargs):
        # Deterministic byte-level surrogate that is reversible for context
        # compaction tests. Contract tests care about control flow, not BPE ids.
        return {"input_ids": list(str(text).encode("utf-8")) or [0]}

    def decode(self, generated, **_kwargs):
        raw = bytes(int(item) for item in generated if 0 <= int(item) <= 255)
        return raw.decode("utf-8", errors="ignore")


class _FakeAutoTokenizer:
    @classmethod
    def from_pretrained(cls, _model):
        return _FakeTokenizer()


class _FakeAutoModel:
    @classmethod
    def from_pretrained(cls, *_args, **_kwargs):
        raise AssertionError("real model loading must never happen in contract tests")


class _FakeStoppingCriteria:
    pass


class _FakeStoppingCriteriaList(list):
    pass


class _FakeInterface:
    def __init__(self, *args, **kwargs):
        pass

    def queue(self, *args, **kwargs):
        return self

    def launch(self, *args, **kwargs):
        return self


class _FakeComponent:
    def __init__(self, *args, **kwargs):
        pass


class _FakeApp:
    @staticmethod
    def create_app(*args, **kwargs):
        return types.SimpleNamespace(add_middleware=lambda *_a, **_k: None)


def _install_import_stubs() -> None:
    spaces = types.ModuleType("spaces")

    def gpu(*_args, **_kwargs):
        def decorator(fn):
            return fn
        return decorator

    spaces.GPU = gpu
    sys.modules["spaces"] = spaces

    transformers = types.ModuleType("transformers")
    transformers.AutoModelForCausalLM = _FakeAutoModel
    transformers.AutoTokenizer = _FakeAutoTokenizer
    transformers.StoppingCriteria = _FakeStoppingCriteria
    transformers.StoppingCriteriaList = _FakeStoppingCriteriaList
    sys.modules["transformers"] = transformers

    gradio = types.ModuleType("gradio")
    gradio.Interface = _FakeInterface
    gradio.Textbox = _FakeComponent
    gradio.Number = _FakeComponent
    gradio.Checkbox = _FakeComponent
    routes = types.ModuleType("gradio.routes")
    routes.App = _FakeApp
    gradio.routes = routes
    context = types.ModuleType("gradio.context")

    class _FakeLocalContext:
        request = contextvars.ContextVar("gradio_request", default=None)

    context.LocalContext = _FakeLocalContext
    gradio.context = context
    sys.modules["gradio"] = gradio
    sys.modules["gradio.routes"] = routes
    sys.modules["gradio.context"] = context


_install_import_stubs()
app = importlib.import_module("app")


READ = {
    "type": "function",
    "function": {
        "name": "Read",
        "description": "Read a file",
        "parameters": {
            "type": "object",
            "properties": {"file_path": {"type": "string"}},
            "required": ["file_path"],
        },
    },
}
GLOB = {
    "type": "function",
    "function": {
        "name": "Glob",
        "description": "Find files",
        "parameters": {
            "type": "object",
            "properties": {"pattern": {"type": "string"}},
            "required": ["pattern"],
        },
    },
}


class AppContractTests(unittest.TestCase):
    def test_default_long_context_uses_official_qwen_yarn_factor(self) -> None:
        self.assertEqual(app.NATIVE_CONTEXT_TOKENS, 32768)
        self.assertEqual(app.MAX_CONTEXT_TOKENS, 131072)
        self.assertTrue(app.YARN_ENABLED)
        self.assertEqual(app.YARN_FACTOR, 4.0)
        self.assertEqual(app.ZERO_GPU_SIZE, "xlarge")

        class LoadedModel:
            def eval(self):
                return self

            def parameters(self):
                return iter([types.SimpleNamespace(device="cuda")])

        previous_model = app.model
        app.model = None
        try:
            with patch.object(
                app.AutoModelForCausalLM,
                "from_pretrained",
                return_value=LoadedModel(),
            ) as loader:
                app._ensure_model_loaded()
            kwargs = loader.call_args.kwargs
            self.assertEqual(
                kwargs["rope_parameters"],
                {
                    "rope_type": "yarn",
                    "factor": 4.0,
                    "original_max_position_embeddings": 32768,
                    "rope_theta": 1_000_000.0,
                },
            )
            self.assertNotIn("rope_scaling", kwargs)
            self.assertEqual(kwargs["max_position_embeddings"], 131072)
        finally:
            app.model = previous_model

    def test_health_and_models_report_real_context_limit(self) -> None:
        health = app.health()
        self.assertEqual(health["context_length"], 131072)
        self.assertTrue(health["yarn_enabled"])
        self.assertEqual(health["yarn_factor"], 4.0)
        self.assertEqual(health["zero_gpu_size"], "xlarge")
        for item in app.models()["data"]:
            self.assertEqual(item["context_length"], 131072)
            self.assertEqual(item["max_input_tokens"], 131072)

    def test_default_temperature_is_greedy(self) -> None:
        request = app.ChatCompletionRequest(messages=[{"role": "user", "content": "oi"}])
        self.assertEqual(request.temperature, 0.0)

    def test_required_tool_uses_temperature_zero_and_structured_finish(self) -> None:
        request = app.ChatCompletionRequest(
            messages=[{"role": "user", "content": "Leia README.md"}],
            tools=[READ],
            tool_choice="required",
        )
        qwen = '<tool_call>{"name":"Read","arguments":{"file_path":"README.md"}}</tool_call>'
        with patch.object(app, "gerar", return_value=qwen) as gerar_mock:
            completion = app._completion_payload(request)

        self.assertEqual(gerar_mock.call_args.args[1], 0.0)
        choice = completion["choices"][0]
        self.assertEqual(choice["finish_reason"], "tool_calls")
        call = choice["message"]["tool_calls"][0]
        self.assertEqual(call["function"]["name"], "Read")
        self.assertEqual(json.loads(call["function"]["arguments"]), {"file_path": "README.md"})

    def test_required_tool_never_succeeds_as_plain_text(self) -> None:
        request = app.ChatCompletionRequest(
            messages=[{"role": "user", "content": "Use Read para README.md"}],
            tools=[READ],
            tool_choice="required",
        )
        with patch.object(app, "gerar", return_value="README content would be here"):
            with self.assertRaises(app.HTTPException) as raised:
                app._completion_payload(request)
        self.assertEqual(raised.exception.status_code, 502)
        self.assertIn("required tool call", str(raised.exception.detail))

    def test_required_remains_required_after_previous_read_result(self) -> None:
        history = [
            {"role": "user", "content": "Compare README.md and app.py"},
            {
                "role": "assistant",
                "content": None,
                "tool_calls": [{
                    "id": "call_read_1",
                    "type": "function",
                    "function": {"name": "Read", "arguments": '{"file_path":"README.md"}'},
                }],
            },
            {
                "role": "tool",
                "tool_call_id": "call_read_1",
                "name": "Read",
                "content": "README content",
            },
        ]
        request = app.ChatCompletionRequest(
            messages=history,
            tools=[READ, GLOB],
            tool_choice="required",
        )
        qwen = '<tool_call>{"name":"Read","arguments":{"file_path":"app.py"}}</tool_call>'
        with patch.object(app, "gerar", return_value=qwen) as gerar_mock:
            completion = app._completion_payload(request)

        self.assertEqual(completion["choices"][0]["finish_reason"], "tool_calls")
        # Both advertised tools remain available under `required`; the server
        # must not downgrade the request to none after one successful Read.
        passed_tools = json.loads(gerar_mock.call_args.args[3])
        self.assertEqual({t["function"]["name"] for t in passed_tools}, {"Read", "Glob"})
        self.assertEqual(gerar_mock.call_args.args[1], 0.0)

    def test_auto_rejects_complete_unadvertised_tool_instead_of_leaking_xml(self) -> None:
        request = app.ChatCompletionRequest(
            messages=[{"role": "user", "content": "Inspect the project if useful"}],
            tools=[READ],
            tool_choice="auto",
        )
        qwen = '<tool_call>{"name":"DeleteEverything","arguments":{}}</tool_call>'
        with patch.object(app, "gerar", return_value=qwen):
            with self.assertRaises(app.HTTPException) as raised:
                app._completion_payload(request)
        self.assertEqual(raised.exception.status_code, 502)
        self.assertIn("unadvertised tool call", str(raised.exception.detail))

    def test_tool_context_compaction_preserves_tool_catalog(self) -> None:
        messages = [
            {"role": "system", "content": "SYSTEM " + ("x" * 1800)},
            {"role": "user", "content": "Compare files " + ("y" * 900)},
        ]
        with patch.object(app, "MAX_CONTEXT_TOKENS", 1800):
            fitted = app._fit_messages_to_context(messages, [READ], 100)
            prompt = app._render_prompt(fitted, [READ])
            token_count = len(app.tokenizer(prompt, add_special_tokens=False)["input_ids"])
        self.assertLessEqual(token_count, 1700)
        self.assertIn('"name": "Read"', prompt)
        self.assertIn(app.CONTEXT_TRUNCATION_MARKER.strip(), prompt)

    def test_tool_context_overflow_fails_instead_of_slicing_schema(self) -> None:
        huge_tool = {
            "type": "function",
            "function": {
                "name": "Huge",
                "description": "x",
                "parameters": {
                    "type": "object",
                    "properties": {
                        "value": {"type": "string", "enum": ["z" * 3000]},
                    },
                },
            },
        }
        with patch.object(app, "MAX_CONTEXT_TOKENS", 500):
            with self.assertRaises(ValueError) as raised:
                app._fit_messages_to_context(
                    [{"role": "user", "content": "do it"}],
                    [huge_tool],
                    100,
                )
        self.assertIn("refusing to slice", str(raised.exception))


    def test_custom_openai_route_propagates_request_context_to_thread(self) -> None:
        http_request = types.SimpleNamespace(headers={"x-ip-token": "opaque-hf-token"})
        parsed = app.ChatCompletionRequest(messages=[{"role": "user", "content": "oi"}])

        def observe_context(_request):
            return app.LocalContext.request.get(None)

        with patch.object(app, "chat_completions", side_effect=observe_context):
            observed = asyncio.run(
                app._chat_completions_with_request_context(http_request, parsed)
            )

        self.assertIs(observed, http_request)
        self.assertIsNone(app.LocalContext.request.get(None))

    def test_zerogpu_limit_is_reported_as_429(self) -> None:
        response = app._zerogpu_limit_response(
            RuntimeError(
                "Space app has reached its GPU limit. Try re-running outside of examples"
            )
        )
        self.assertIsNotNone(response)
        self.assertEqual(response.status_code, 429)
        self.assertIn(b"Hugging Face token", response.body)
        self.assertIn(b"2x", response.body)

    def test_streaming_tool_delta_and_usage_match_openai_contract(self) -> None:
        request = app.ChatCompletionRequest(
            messages=[{"role": "user", "content": "Leia README.md"}],
            tools=[READ],
            tool_choice="required",
            stream=True,
            stream_options={"include_usage": True},
        )
        qwen = '<tool_call>{"name":"Read","arguments":{"file_path":"README.md"}}</tool_call>'
        with patch.object(app, "gerar", return_value=qwen):
            response = app.chat_completions(request)

        async def collect() -> str:
            pieces = []
            async for piece in response.body_iterator:
                if isinstance(piece, bytes):
                    piece = piece.decode("utf-8")
                pieces.append(piece)
            return "".join(pieces)

        stream = asyncio.run(collect())
        frames = [line[6:] for line in stream.splitlines() if line.startswith("data: ")]
        self.assertEqual(frames[-1], "[DONE]")
        payloads = [json.loads(frame) for frame in frames[:-1]]
        tool_chunks = [
            chunk for chunk in payloads
            if chunk.get("choices")
            and chunk["choices"][0].get("delta", {}).get("tool_calls")
        ]
        self.assertEqual(len(tool_chunks), 1)
        streamed_call = tool_chunks[0]["choices"][0]["delta"]["tool_calls"][0]
        self.assertEqual(streamed_call["index"], 0)
        self.assertTrue(streamed_call["id"].startswith("call_"))
        self.assertEqual(streamed_call["function"]["name"], "Read")
        self.assertTrue(any(chunk.get("choices") == [] and "usage" in chunk for chunk in payloads))
        self.assertTrue(any(
            chunk.get("choices")
            and chunk["choices"][0].get("finish_reason") == "tool_calls"
            for chunk in payloads
        ))


if __name__ == "__main__":
    unittest.main()