import importlib
import sys
import types
from dataclasses import dataclass
from importlib.machinery import ModuleSpec
from json import loads
from pathlib import Path
from types import SimpleNamespace
from unittest import IsolatedAsyncioTestCase, TestCase
from unittest.mock import AsyncMock, MagicMock, patch
from uuid import UUID

from avalan.entities import (
    GenerationSettings,
    Message,
    MessageContentFile,
    MessageContentImage,
    MessageContentText,
    MessageRole,
    PromptCacheRetention,
    ReasoningEffort,
    ReasoningSettings,
    ReasoningToken,
    Token,
    ToolCall,
    ToolCallDiagnostic,
    ToolCallDiagnosticCode,
    ToolCallDiagnosticStage,
    ToolCallError,
    ToolCallResult,
    ToolCallToken,
    TransformerEngineSettings,
)
from avalan.task.usage import (
    usage_observation_from_response,
    usage_totals_from_response,
)
from avalan.tool.parser import ToolCallParser

FIXTURES = Path(__file__).resolve().parents[2] / "fixtures" / "tool_parsing"


class AsyncIter:
    def __init__(self, items):
        self._iter = iter(items)

    def __aiter__(self):
        return self

    async def __anext__(self):
        try:
            return next(self._iter)
        except StopIteration as exc:
            raise StopAsyncIteration from exc


def patch_openai_imports():
    class Omit:
        def __bool__(self):
            return False

    openai_stub = types.ModuleType("openai")
    openai_stub.__spec__ = ModuleSpec("openai", loader=None)
    openai_stub.AsyncOpenAI = MagicMock()
    openai_stub.AsyncStream = MagicMock()
    openai_stub.Omit = Omit
    openai_stub.AsyncOpenAI.return_value.responses = MagicMock()

    transformers_stub = types.ModuleType("transformers")
    transformers_stub.__spec__ = ModuleSpec("transformers", loader=None)
    transformers_stub.PreTrainedModel = MagicMock()
    transformers_stub.PreTrainedTokenizer = MagicMock()
    transformers_stub.PreTrainedTokenizerFast = MagicMock()
    transformers_stub.__getattr__ = lambda name: MagicMock()

    transformers_utils_stub = types.ModuleType("transformers.utils")
    transformers_utils_stub.get_json_schema = MagicMock()
    transformers_logging_stub = types.ModuleType("transformers.utils.logging")
    transformers_logging_stub.disable_progress_bar = MagicMock()
    transformers_logging_stub.enable_progress_bar = MagicMock()
    transformers_utils_stub.logging = transformers_logging_stub
    transformers_tokenization_stub = types.ModuleType(
        "transformers.tokenization_utils_base"
    )
    transformers_tokenization_stub.BatchEncoding = MagicMock()
    transformers_stub.tokenization_utils_base = transformers_tokenization_stub
    transformers_generation_stub = types.ModuleType("transformers.generation")
    transformers_generation_stub.__spec__ = ModuleSpec(
        "transformers.generation", loader=None, is_package=True
    )
    transformers_generation_stub.__path__ = []
    transformers_generation_stub.StoppingCriteria = MagicMock()
    transformers_generation_stub.StoppingCriteriaList = MagicMock()
    transformers_stopping_criteria_stub = types.ModuleType(
        "transformers.generation.stopping_criteria"
    )
    transformers_stopping_criteria_stub.__spec__ = ModuleSpec(
        "transformers.generation.stopping_criteria", loader=None
    )
    transformers_stopping_criteria_stub.StoppingCriteria = MagicMock()
    transformers_stopping_criteria_stub.StoppingCriteriaList = MagicMock()
    transformers_generation_stub.stopping_criteria = (
        transformers_stopping_criteria_stub
    )
    transformers_stub.generation = transformers_generation_stub
    transformers_stub.utils = transformers_utils_stub

    diffusers_stub = types.ModuleType("diffusers")
    diffusers_stub.__spec__ = ModuleSpec("diffusers", loader=None)
    diffusers_stub.DiffusionPipeline = MagicMock()

    patcher = patch.dict(
        sys.modules,
        {
            "openai": openai_stub,
            "transformers": transformers_stub,
            "transformers.utils": transformers_utils_stub,
            "transformers.utils.logging": transformers_logging_stub,
            "transformers.tokenization_utils_base": (
                transformers_tokenization_stub
            ),
            "transformers.generation": transformers_generation_stub,
            "transformers.generation.stopping_criteria": (
                transformers_stopping_criteria_stub
            ),
            "diffusers": diffusers_stub,
        },
    )
    patcher.start()
    return openai_stub, patcher


class OpenAITestCase(IsolatedAsyncioTestCase):
    def setUp(self):
        self.openai_stub, self.patch = patch_openai_imports()
        importlib.reload(
            importlib.import_module("avalan.model.nlp.text.vendor.openai")
        )
        self.mod = importlib.import_module(
            "avalan.model.nlp.text.vendor.openai"
        )

    def tearDown(self):
        self.patch.stop()

    async def test_stream_client_and_model(self):
        chunks = [
            SimpleNamespace(type="response.output_text.delta", delta="x"),
            SimpleNamespace(type="response.output_text.delta", delta="y"),
        ]
        stream = self.mod.OpenAIStream(AsyncIter(chunks))
        t1 = await stream.__anext__()
        self.assertIsInstance(t1, Token)
        self.assertEqual(t1.token, "x")
        t2 = await stream.__anext__()
        self.assertIsInstance(t2, Token)
        self.assertEqual(t2.token, "y")
        with self.assertRaises(StopAsyncIteration):
            await stream.__anext__()

        stream_instance = AsyncIter([])
        self.openai_stub.AsyncOpenAI.return_value.responses.create = AsyncMock(
            return_value=stream_instance
        )
        client = self.mod.OpenAIClient(api_key="k", base_url="b")
        client._template_messages = MagicMock(return_value=[{"c": 1}])
        with patch.object(self.mod, "OpenAIStream") as StreamMock:
            result = await client("m", [])
        self.openai_stub.AsyncOpenAI.assert_called_once_with(
            base_url="b", api_key="k"
        )
        client._client.responses.create.assert_awaited_once_with(
            extra_headers={
                "X-Title": "Avalan",
                "HTTP-Referer": "https://github.com/avalan-ai/avalan",
            },
            model="m",
            input=[{"c": 1}],
            store=False,
            stream=True,
            timeout=None,
        )
        StreamMock.assert_called_once_with(
            stream=stream_instance,
            provider_family="openai",
        )
        self.assertIs(result, StreamMock.return_value)

    async def test_client_sends_top_level_instructions(self):
        response = SimpleNamespace(
            output=[SimpleNamespace(content=[SimpleNamespace(text="x")])]
        )
        self.openai_stub.AsyncOpenAI.return_value.responses.create = AsyncMock(
            return_value=response
        )
        client = self.mod.OpenAIClient(api_key="k", base_url="b")
        message = Message(role=MessageRole.USER, content="hi")

        await client(
            "m",
            [message],
            instructions="top-level",
            use_async_generator=False,
        )

        kwargs = client._client.responses.create.await_args.kwargs
        self.assertEqual(kwargs["instructions"], "top-level")
        self.assertEqual(kwargs["input"], [{"role": "user", "content": "hi"}])
        self.assertNotIn("top-level", str(kwargs["input"]))

    async def test_responses_payload_combines_instructions_cache_and_options(
        self,
    ):
        response = SimpleNamespace(
            output=[SimpleNamespace(content=[SimpleNamespace(text="ok")])]
        )
        create_mock = AsyncMock(return_value=response)
        self.openai_stub.AsyncOpenAI.return_value.responses.create = (
            create_mock
        )
        client = self.mod.OpenAIClient(api_key="k", base_url="b")
        tool = MagicMock()
        tool.json_schemas.return_value = [
            {"type": "function", "function": {"name": "pkg.lookup"}}
        ]
        messages = [
            Message(role=MessageRole.SYSTEM, content="system path"),
            Message(role=MessageRole.DEVELOPER, content="developer path"),
            Message(
                role=MessageRole.USER,
                content=[
                    MessageContentText(type="text", text="Summarize"),
                    MessageContentFile(
                        type="file",
                        file={
                            "file_data": "YWJj",
                            "filename": "report.pdf",
                        },
                    ),
                    MessageContentImage(
                        type="image_url",
                        image_url={
                            "data": "aW1n",
                            "mime_type": "image/png",
                            "detail": "high",
                        },
                    ),
                ],
            ),
        ]
        settings = GenerationSettings(
            max_new_tokens=10,
            prompt_cache_retention=PromptCacheRetention.EXTENDED_24H,
            reasoning=ReasoningSettings(effort=ReasoningEffort.HIGH),
            response_format={"type": "json_object"},
            stop_strings=["END"],
        )

        await client(
            "gpt-5",
            messages,
            settings=settings,
            instructions="top-level policy",
            tool=tool,
            use_async_generator=False,
        )

        create_mock.assert_awaited_once_with(
            extra_headers={
                "X-Title": "Avalan",
                "HTTP-Referer": "https://github.com/avalan-ai/avalan",
            },
            model="gpt-5",
            input=[
                {"role": "system", "content": "system path"},
                {"role": "developer", "content": "developer path"},
                {
                    "role": "user",
                    "content": [
                        {"type": "input_text", "text": "Summarize"},
                        {
                            "type": "input_file",
                            "file_data": "data:application/pdf;base64,YWJj",
                            "filename": "report.pdf",
                        },
                        {
                            "type": "input_image",
                            "image_url": "data:image/png;base64,aW1n",
                            "detail": "high",
                        },
                    ],
                },
            ],
            store=False,
            stream=False,
            timeout=None,
            instructions="top-level policy",
            max_output_tokens=10,
            text={"format": {"type": "json_object"}, "stop": ["END"]},
            reasoning={"effort": "high"},
            prompt_cache_retention="24h",
            tools=[{"type": "function", "name": "avl_cGtnLmxvb2t1cA"}],
        )
        self.assertNotIn(
            "top-level policy", str(create_mock.await_args.kwargs["input"])
        )

    async def test_rejects_invalid_instructions_before_provider_call(self):
        create_mock = AsyncMock()
        self.openai_stub.AsyncOpenAI.return_value.responses.create = (
            create_mock
        )
        client = self.mod.OpenAIClient(api_key="k", base_url="b")

        with self.assertRaisesRegex(AssertionError, "instructions"):
            await client(
                "m",
                [],
                instructions={"raw": "prompt"},  # type: ignore[arg-type]
            )

        create_mock.assert_not_awaited()

    async def test_client_consumes_tokens(self):
        chunks = [
            SimpleNamespace(type="response.output_text.delta", delta="a"),
            SimpleNamespace(type="response.output_text.delta", delta="b"),
        ]
        stream_instance = AsyncIter(chunks)
        self.openai_stub.AsyncOpenAI.return_value.responses.create = AsyncMock(
            return_value=stream_instance
        )
        client = self.mod.OpenAIClient(api_key="k", base_url="b")
        client._template_messages = MagicMock(return_value=[{"c": 1}])
        result = await client("m", [])
        t1 = await result.__anext__()
        self.assertIsInstance(t1, Token)
        self.assertEqual(t1.token, "a")
        t2 = await result.__anext__()
        self.assertIsInstance(t2, Token)
        self.assertEqual(t2.token, "b")
        with self.assertRaises(StopAsyncIteration):
            await result.__anext__()

        with patch.object(self.mod, "OpenAIClient") as ClientMock:
            settings = TransformerEngineSettings(
                auto_load_model=False,
                auto_load_tokenizer=False,
                access_token="t",
                base_url="u",
            )
            model = self.mod.OpenAIModel("m", settings)
            loaded = model._load_model()
        ClientMock.assert_called_once_with(base_url="u", api_key="t")
        self.assertIs(loaded, ClientMock.return_value)

    async def test_stream_records_terminal_usage_after_completion(self):
        usage = SimpleNamespace(
            input_tokens=3,
            input_tokens_details=SimpleNamespace(cached_tokens=1),
            cache_creation_input_tokens=2,
            output_tokens=4,
            output_tokens_details=SimpleNamespace(reasoning_tokens=2),
            total_tokens=9,
        )
        chunks = [
            SimpleNamespace(type="response.output_text.delta", delta="a"),
            SimpleNamespace(
                type="response.completed",
                response=SimpleNamespace(usage=usage),
            ),
        ]
        stream = self.mod.OpenAIStream(AsyncIter(chunks))

        token = await stream.__anext__()
        self.assertIsInstance(token, Token)
        self.assertIsNone(stream.usage)
        with self.assertRaises(StopAsyncIteration):
            await stream.__anext__()

        self.assertIs(stream.usage, usage)
        self.assertEqual(stream.provider_family, "openai")
        observation = usage_observation_from_response(stream)
        self.assertIsNotNone(observation)
        assert observation is not None
        self.assertEqual(observation.metadata, {"provider_family": "openai"})
        totals = usage_totals_from_response(stream)
        self.assertIsNotNone(totals)
        assert totals is not None
        self.assertEqual(totals.input_tokens, 3)
        self.assertEqual(totals.cached_input_tokens, 1)
        self.assertEqual(totals.cache_creation_input_tokens, 2)
        self.assertEqual(totals.output_tokens, 4)
        self.assertEqual(totals.reasoning_tokens, 2)
        self.assertEqual(totals.total_tokens, 9)

    async def test_stream_records_dict_terminal_usage_after_exhaustion(self):
        stream = self.mod.OpenAIStream(
            AsyncIter(
                [
                    {
                        "type": "response.completed",
                        "response": {
                            "usage": {
                                "input_tokens": 0,
                                "prompt_tokens_details": {"cached_tokens": 0},
                                "completion_tokens": 0,
                                "completion_tokens_details": {
                                    "reasoning_tokens": 0
                                },
                                "total_tokens": 0,
                            }
                        },
                    }
                ]
            )
        )

        with self.assertRaises(StopAsyncIteration):
            await stream.__anext__()
        observation = usage_observation_from_response(stream)
        totals = usage_totals_from_response(stream)

        self.assertIsNotNone(observation)
        assert observation is not None
        self.assertEqual(observation.metadata, {"provider_family": "openai"})
        self.assertIsNotNone(totals)
        assert totals is not None
        self.assertEqual(totals.input_tokens, 0)
        self.assertEqual(totals.cached_input_tokens, 0)
        self.assertIsNone(totals.cache_creation_input_tokens)
        self.assertEqual(totals.output_tokens, 0)
        self.assertEqual(totals.reasoning_tokens, 0)
        self.assertEqual(totals.total_tokens, 0)

    async def test_stream_keeps_null_or_interrupted_usage_unavailable(self):
        null_usage = self.mod.OpenAIStream(
            AsyncIter(
                [
                    SimpleNamespace(
                        type="response.completed",
                        response=SimpleNamespace(usage=None),
                    )
                ]
            )
        )
        with self.assertRaises(StopAsyncIteration):
            await null_usage.__anext__()
        self.assertIsNone(null_usage.usage)

        class FailingIter:
            def __aiter__(self):
                return self

            async def __anext__(self):
                raise RuntimeError("provider failure")

        interrupted = self.mod.OpenAIStream(FailingIter())
        with self.assertRaises(RuntimeError):
            await interrupted.__anext__()
        self.assertIsNone(interrupted.usage)

        class FailingAfterUsageIter:
            def __init__(self):
                self._count = 0

            def __aiter__(self):
                return self

            async def __anext__(self):
                self._count += 1
                if self._count == 1:
                    return {
                        "type": "response.completed",
                        "response": {"usage": {"input_tokens": 1}},
                    }
                raise RuntimeError("provider failure")

        interrupted_after_usage = self.mod.OpenAIStream(
            FailingAfterUsageIter()
        )
        with self.assertRaises(RuntimeError):
            await interrupted_after_usage.__anext__()
        self.assertIsNone(interrupted_after_usage.usage)

    async def test_client_omits_auth_header_without_api_key(self):
        client = self.mod.OpenAIClient(
            api_key=None, base_url="http://localhost:9001/v1"
        )

        self.assertIsInstance(client, self.mod.OpenAIClient)
        self.openai_stub.AsyncOpenAI.assert_called_once()
        kwargs = self.openai_stub.AsyncOpenAI.call_args.kwargs
        self.assertEqual(kwargs["base_url"], "http://localhost:9001/v1")
        self.assertEqual(kwargs["api_key"], "")
        self.assertIsInstance(
            kwargs["default_headers"]["Authorization"], self.mod.Omit
        )

    async def test_client_uses_default_model_id_when_missing(self):
        response = SimpleNamespace(output=[])
        self.openai_stub.AsyncOpenAI.return_value.responses.create = AsyncMock(
            return_value=response
        )
        client = self.mod.OpenAIClient(
            api_key="key", base_url="http://localhost:9001/v1"
        )
        client._template_messages = MagicMock(return_value=[{"c": 1}])

        await client("", [], use_async_generator=False)

        self.openai_stub.AsyncOpenAI.return_value.responses.create.assert_awaited_once_with(
            extra_headers={
                "X-Title": "Avalan",
                "HTTP-Referer": "https://github.com/avalan-ai/avalan",
            },
            model="default",
            input=[{"c": 1}],
            store=False,
            stream=False,
            timeout=None,
        )

    async def test_model_loads_with_base_url_without_access_token(self):
        with patch.object(self.mod, "OpenAIClient") as ClientMock:
            settings = TransformerEngineSettings(
                auto_load_model=False,
                auto_load_tokenizer=False,
                base_url="http://localhost:9001/v1",
            )
            model = self.mod.OpenAIModel("m", settings)
            loaded = model._load_model()

        ClientMock.assert_called_once_with(
            base_url="http://localhost:9001/v1", api_key=None
        )
        self.assertIs(loaded, ClientMock.return_value)

    async def test_model_loads_with_azure_api_version(self):
        with patch.object(self.mod, "OpenAIClient") as ClientMock:
            settings = TransformerEngineSettings(
                auto_load_model=False,
                auto_load_tokenizer=False,
                access_token="token",
                base_url=(
                    "https://tenant.openai.azure.com/openai/deployments/d"
                ),
                provider_options={
                    "azure_api_version": "2025-04-01-preview",
                },
            )
            model = self.mod.OpenAIModel("deployment", settings)
            loaded = model._load_model()

        ClientMock.assert_called_once_with(
            base_url="https://tenant.openai.azure.com/openai/deployments/d",
            api_key="token",
            azure_api_version="2025-04-01-preview",
        )
        self.assertIs(loaded, ClientMock.return_value)

    async def test_model_rejects_invalid_provider_api_version(self):
        settings = TransformerEngineSettings(
            auto_load_model=False,
            auto_load_tokenizer=False,
            access_token="token",
            base_url="https://tenant.openai.azure.com/openai/deployments/d",
            provider_options={"azure_api_version": 20250401},
        )
        model = self.mod.OpenAIModel("deployment", settings)

        with self.assertRaises(AssertionError):
            model._load_model()

    async def test_stream_event_types(self):
        events = [
            SimpleNamespace(type="response.output_item.added"),
            SimpleNamespace(type="response.content_part.added"),
            SimpleNamespace(type="response.reasoning_text.delta", delta="r1"),
            SimpleNamespace(type="response.reasoning_text.delta", delta="r2"),
            SimpleNamespace(type="response.output_item.done"),
            SimpleNamespace(
                type="response.output_item.added",
                item=SimpleNamespace(
                    custom_tool_call=SimpleNamespace(id=object())
                ),
            ),
            SimpleNamespace(
                type="response.output_item.added",
                item=SimpleNamespace(
                    id="c1",
                    custom_tool_call=SimpleNamespace(
                        id="c1", name="pkg__func"
                    ),
                ),
            ),
            SimpleNamespace(
                type="response.function_call_arguments.delta",
                id="c1",
                delta="{",
            ),
            SimpleNamespace(
                type="response.function_call_arguments.delta",
                id="c1",
                delta="}",
            ),
            SimpleNamespace(
                type="response.output_item.done", item=SimpleNamespace(id="c1")
            ),
            SimpleNamespace(type="response.output_item.added"),
            SimpleNamespace(type="response.content_part.added"),
            SimpleNamespace(type="response.output_text.delta", delta="hi"),
            SimpleNamespace(type="response.output_item.done"),
        ]
        stream = self.mod.OpenAIStream(AsyncIter(events))
        t1 = await stream.__anext__()
        self.assertIsInstance(t1, ReasoningToken)
        self.assertEqual(t1.token, "r1")
        t2 = await stream.__anext__()
        self.assertIsInstance(t2, ReasoningToken)
        self.assertEqual(t2.token, "r2")
        t3 = await stream.__anext__()
        self.assertIsInstance(t3, ToolCallToken)
        self.assertEqual(t3.token, "{")
        t4 = await stream.__anext__()
        self.assertIsInstance(t4, ToolCallToken)
        self.assertEqual(t4.token, "}")
        t5 = await stream.__anext__()
        self.assertIsInstance(t5, ToolCallToken)
        self.assertEqual(t5.call.id, "c1")
        self.assertEqual(t5.call.name, "pkg__func")
        self.assertEqual(t5.call.arguments, {})
        self.assertEqual(
            t5.token,
            '<tool_call>{"name": "pkg__func", "arguments": {}, "id":'
            ' "c1"}</tool_call>',
        )
        t6 = await stream.__anext__()
        self.assertIsInstance(t6, Token)
        self.assertEqual(t6.token, "hi")
        with self.assertRaises(StopAsyncIteration):
            await stream.__anext__()

    async def test_function_call_events(self):
        events = [
            SimpleNamespace(
                type="response.function_call_arguments.delta",
                id="c2",
                delta="{",
            ),
            SimpleNamespace(
                type="response.function_call_arguments.delta",
                id="c2",
                delta="}",
            ),
            SimpleNamespace(
                type="response.output_item.done", item=SimpleNamespace(id="c2")
            ),
            SimpleNamespace(
                type="response.output_item.done",
                item=SimpleNamespace(
                    type="function_call",
                    id="c3",
                    name="pkg__f",
                    arguments='{"p": 1}',
                ),
            ),
        ]
        stream = self.mod.OpenAIStream(AsyncIter(events))
        await stream.__anext__()
        await stream.__anext__()
        t3 = await stream.__anext__()
        self.assertEqual(t3.call.id, "c2")
        self.assertEqual(t3.call.arguments, {})
        t4 = await stream.__anext__()
        self.assertEqual(t4.call.id, "c3")
        self.assertEqual(t4.call.name, "pkg__f")
        self.assertEqual(t4.call.arguments, {"p": 1})
        with self.assertRaises(StopAsyncIteration):
            await stream.__anext__()

    async def test_provider_argument_deltas_match_serialized_call(self):
        fixture = loads(
            (FIXTURES / "provider_openai_argument_deltas.json").read_text(
                encoding="utf-8"
            )
        )
        call_id = fixture["call_id"]
        events = [
            SimpleNamespace(
                type="response.output_item.added",
                item=SimpleNamespace(
                    id=call_id,
                    custom_tool_call=SimpleNamespace(
                        id=call_id,
                        name=fixture["provider_name"],
                    ),
                ),
            ),
            *[
                SimpleNamespace(
                    type="response.function_call_arguments.delta",
                    id=call_id,
                    delta=delta,
                )
                for delta in fixture["argument_deltas"]
            ],
            SimpleNamespace(
                type="response.output_item.done",
                item=SimpleNamespace(id=call_id),
            ),
            SimpleNamespace(
                type="response.output_text.delta",
                delta=fixture["assistant_after"],
            ),
        ]

        stream = self.mod.OpenAIStream(AsyncIter(events))
        outputs = []
        async for output in stream:
            outputs.append(output)

        delta_outputs = outputs[: len(fixture["argument_deltas"])]
        self.assertTrue(
            all(isinstance(output, ToolCallToken) for output in delta_outputs)
        )
        self.assertEqual(
            [output.token for output in delta_outputs],
            fixture["argument_deltas"],
        )

        final_token = outputs[len(fixture["argument_deltas"])]
        self.assertIsInstance(final_token, ToolCallToken)
        self.assertIsNotNone(final_token.call)
        self.assertEqual(final_token.call.id, call_id)
        self.assertEqual(final_token.call.name, fixture["provider_name"])
        self.assertEqual(
            final_token.call.arguments,
            {"city": "Paris", "unit": "c"},
        )

        parsed = ToolCallParser().parse(final_token.token)
        self.assertEqual(len(parsed.calls), 1)
        self.assertEqual(parsed.diagnostics, [])
        self.assertEqual(parsed.calls[0].id, final_token.call.id)
        self.assertEqual(parsed.calls[0].name, final_token.call.name)
        self.assertEqual(
            parsed.calls[0].arguments,
            final_token.call.arguments,
        )
        self.assertEqual(outputs[-1].token, fixture["assistant_after"])

    async def test_generation_settings_and_tools(self):
        stream_instance = AsyncIter([])
        self.openai_stub.AsyncOpenAI.return_value.responses.create = AsyncMock(
            return_value=stream_instance
        )
        client = self.mod.OpenAIClient(api_key="k", base_url="b")
        client._template_messages = MagicMock(return_value=[{"c": 1}])
        tool = MagicMock()
        tool.json_schemas.return_value = [
            {"type": "function", "function": {"name": "pkg.func"}}
        ]
        settings = GenerationSettings(
            temperature=0.5,
            top_p=0.8,
            max_new_tokens=10,
            stop_strings=["stop"],
            response_format={
                "type": "json_schema",
                "json_schema": {
                    "name": "result",
                    "schema": {"type": "object"},
                    "strict": True,
                },
            },
        )
        await client(
            "m",
            [],
            settings=settings,
            tool=tool,
        )
        self.openai_stub.AsyncOpenAI.return_value.responses.create.assert_awaited_once_with(
            extra_headers={
                "X-Title": "Avalan",
                "HTTP-Referer": "https://github.com/avalan-ai/avalan",
            },
            model="m",
            input=[{"c": 1}],
            store=False,
            stream=True,
            timeout=None,
            max_output_tokens=10,
            temperature=0.5,
            top_p=0.8,
            text={
                "format": {
                    "type": "json_schema",
                    "name": "result",
                    "schema": {"type": "object"},
                    "strict": True,
                },
                "stop": ["stop"],
            },
            tools=[{"type": "function", "name": "avl_cGtnLmZ1bmM"}],
        )

    async def test_azure_responses_payload_uses_text_format(self):
        response = SimpleNamespace(
            output=[SimpleNamespace(content=[SimpleNamespace(text="ok")])]
        )
        create_mock = AsyncMock(return_value=response)
        self.openai_stub.AsyncOpenAI.return_value.responses.create = (
            create_mock
        )
        client = self.mod.OpenAIClient(
            api_key="key",
            base_url="https://tenant.openai.azure.com/openai/v1/",
        )
        client._template_messages = MagicMock(return_value=[{"c": 1}])
        settings = GenerationSettings(
            max_new_tokens=20,
            temperature=0.3,
            top_p=0.4,
            stop_strings="END",
            response_format={"type": "json_object"},
            reasoning=ReasoningSettings(effort=ReasoningEffort.HIGH),
        )

        await client(
            "claims-extractor-prod",
            [],
            settings=settings,
            use_async_generator=False,
        )

        create_mock.assert_awaited_once_with(
            extra_headers={
                "X-Title": "Avalan",
                "HTTP-Referer": "https://github.com/avalan-ai/avalan",
            },
            model="claims-extractor-prod",
            input=[{"c": 1}],
            store=False,
            stream=False,
            timeout=None,
            max_output_tokens=20,
            text={"format": {"type": "json_object"}, "stop": "END"},
            reasoning={"effort": "high"},
        )

    async def test_azure_legacy_api_version_uses_extra_query(self):
        response = SimpleNamespace(output=[])
        create_mock = AsyncMock(return_value=response)
        self.openai_stub.AsyncOpenAI.return_value.responses.create = (
            create_mock
        )
        client = self.mod.OpenAIClient(
            api_key="key",
            base_url="https://tenant.openai.azure.com/openai/deployments/dep",
            azure_api_version="2025-04-01-preview",
        )
        client._template_messages = MagicMock(return_value=[])

        await client("dep", [], use_async_generator=False)

        self.assertEqual(
            create_mock.await_args.kwargs["extra_query"],
            {"api-version": "2025-04-01-preview"},
        )

    async def test_rejects_invalid_responses_options_before_provider_call(
        self,
    ):
        create_mock = AsyncMock()
        self.openai_stub.AsyncOpenAI.return_value.responses.create = (
            create_mock
        )
        client = self.mod.OpenAIClient(api_key="key", base_url="url")
        invalid_settings = GenerationSettings(
            response_format={
                "type": "json_schema",
                "json_schema": {"schema": {"type": "object"}},
                "schema": {"type": "object"},
            }
        )

        with self.assertRaisesRegex(AssertionError, "ambiguous"):
            await client("model", [], settings=invalid_settings)

        create_mock.assert_not_awaited()

    async def test_rejects_invalid_prompt_cache_retention_before_provider_call(
        self,
    ):
        create_mock = AsyncMock()
        self.openai_stub.AsyncOpenAI.return_value.responses.create = (
            create_mock
        )
        client = self.mod.OpenAIClient(api_key="key", base_url="url")

        with self.assertRaisesRegex(AssertionError, "not supported"):
            await client(
                "model",
                [],
                settings=GenerationSettings(
                    prompt_cache_retention="forever",
                ),
            )
        with self.assertRaisesRegex(AssertionError, "must be a string"):
            await client(
                "model",
                [],
                settings=GenerationSettings(
                    prompt_cache_retention=object(),  # type: ignore[arg-type]
                ),
            )

        create_mock.assert_not_awaited()

    async def test_rejects_invalid_reasoning_effort(self):
        create_mock = AsyncMock()
        self.openai_stub.AsyncOpenAI.return_value.responses.create = (
            create_mock
        )
        client = self.mod.OpenAIClient(
            api_key="key",
            base_url="https://tenant.openai.azure.com/openai/v1/",
        )
        settings = GenerationSettings(
            reasoning=ReasoningSettings(
                effort="ultra",  # type: ignore[arg-type]
            )
        )

        with self.assertRaisesRegex(AssertionError, "reasoning effort"):
            await client("deployment", [], settings=settings)

        create_mock.assert_not_awaited()

    async def test_gpt_55_reasoning_xhigh_effort_is_forwarded(self):
        response = SimpleNamespace(
            output=[SimpleNamespace(content=[SimpleNamespace(text="x")])]
        )
        create_mock = AsyncMock(return_value=response)
        self.openai_stub.AsyncOpenAI.return_value.responses.create = (
            create_mock
        )
        client = self.mod.OpenAIClient(api_key="key", base_url="url")
        settings = GenerationSettings(
            reasoning=ReasoningSettings(effort=ReasoningEffort.XHIGH)
        )

        await client(
            "gpt-5.5",
            [Message(role=MessageRole.USER, content="hi")],
            settings,
            use_async_generator=False,
        )

        kwargs = create_mock.await_args.kwargs
        self.assertEqual(kwargs["reasoning"], {"effort": "xhigh"})

    def test_response_text_format_normalizes_supported_shapes(self):
        self.assertEqual(
            self.mod.OpenAIClient._prompt_cache_retention_config(
                GenerationSettings(prompt_cache_retention="in_memory")
            ),
            "in_memory",
        )
        self.assertEqual(
            self.mod.OpenAIClient._response_text_format({"type": "text"}),
            {"type": "text"},
        )
        self.assertEqual(
            self.mod.OpenAIClient._response_text_format(
                {
                    "type": "json_schema",
                    "name": "document",
                    "schema": {"type": "object"},
                    "strict": False,
                }
            ),
            {
                "type": "json_schema",
                "name": "document",
                "schema": {"type": "object"},
                "strict": False,
            },
        )
        with self.assertRaisesRegex(AssertionError, "not supported"):
            self.mod.OpenAIClient._response_text_format({"type": "xml"})
        with self.assertRaises(AssertionError):
            self.mod.OpenAIClient._response_text_format(
                {
                    "type": "json_schema",
                    "schema": {"type": "object"},
                }
            )
        with self.assertRaises(AssertionError):
            self.mod.OpenAIClient._response_text_format(
                {
                    "type": "json_schema",
                    "json_schema": {"name": "bad"},
                }
            )

    def test_azure_configuration_validation(self):
        self.assertFalse(self.mod.OpenAIClient._is_azure_base_url(None))

        with self.assertRaisesRegex(AssertionError, "/openai/v1/"):
            self.mod.OpenAIClient(
                api_key="key",
                base_url=(
                    "https://tenant.openai.azure.com/openai/deployments/dep"
                ),
            )

        try:
            self.mod.OpenAIClient(
                api_key="key",
                base_url=(
                    "https://tenant.openai.azure.com/openai/v1/"
                    "?api-version=super-secret"
                ),
            )
        except AssertionError as exc:
            self.assertIn("query parameters", str(exc))
            self.assertNotIn("super-secret", str(exc))

        with self.assertRaisesRegex(AssertionError, "api-key"):
            self.mod.OpenAIClient(
                api_key=None,
                base_url="https://tenant.openai.azure.com/openai/v1/",
            )

        with self.assertRaisesRegex(AssertionError, "only supported"):
            self.mod.OpenAIClient(
                api_key="key",
                base_url="https://api.openai.com/v1",
                azure_api_version="2025-04-01-preview",
            )


class VendorClientsTestCase(TestCase):
    def setUp(self):
        self.openai_stub, self.patch = patch_openai_imports()
        importlib.reload(
            importlib.import_module("avalan.model.nlp.text.vendor.openai")
        )

    def tearDown(self):
        self.patch.stop()

    def test_openrouter_client_and_model(self):
        mod = importlib.import_module(
            "avalan.model.nlp.text.vendor.openrouter"
        )
        importlib.reload(mod)
        self.openai_stub.AsyncOpenAI.reset_mock()
        client = mod.OpenRouterClient(api_key="k", base_url=None)
        self.openai_stub.AsyncOpenAI.assert_called_once_with(
            base_url="https://openrouter.ai/api/v1", api_key="k"
        )
        self.assertIsInstance(client, mod.OpenRouterClient)
        with patch.object(mod, "OpenRouterClient") as ClientMock:
            settings = TransformerEngineSettings(
                auto_load_model=False,
                auto_load_tokenizer=False,
                access_token="t",
                base_url="b",
            )
            model = mod.OpenRouterModel("m", settings)
            loaded = model._load_model()
        ClientMock.assert_called_once_with(base_url="b", api_key="t")
        self.assertIs(loaded, ClientMock.return_value)

    def test_together_client_and_model(self):
        mod = importlib.import_module("avalan.model.nlp.text.vendor.together")
        importlib.reload(mod)
        self.openai_stub.AsyncOpenAI.reset_mock()
        mod.TogetherClient(api_key="k", base_url=None)
        self.openai_stub.AsyncOpenAI.assert_called_once_with(
            base_url="https://api.together.xyz/v1", api_key="k"
        )
        with patch.object(mod, "TogetherClient") as ClientMock:
            settings = TransformerEngineSettings(
                auto_load_model=False,
                auto_load_tokenizer=False,
                access_token="t",
                base_url="b",
            )
            model = mod.TogetherModel("m", settings)
            loaded = model._load_model()
        ClientMock.assert_called_once_with(base_url="b", api_key="t")
        self.assertIs(loaded, ClientMock.return_value)


class NonStreamingResponseTestCase(IsolatedAsyncioTestCase):
    def setUp(self):
        self.openai_stub, self.patch = patch_openai_imports()
        importlib.reload(
            importlib.import_module("avalan.model.nlp.text.vendor.openai")
        )
        self.mod = importlib.import_module(
            "avalan.model.nlp.text.vendor.openai"
        )

    def tearDown(self):
        self.patch.stop()

    async def test_response_single_stream(self):
        resp = SimpleNamespace(
            output=[SimpleNamespace(content=[SimpleNamespace(text="ok")])],
            usage=SimpleNamespace(input_tokens=1),
        )
        self.openai_stub.AsyncOpenAI.return_value.responses.create = AsyncMock(
            return_value=resp
        )
        settings = TransformerEngineSettings(
            auto_load_model=False,
            auto_load_tokenizer=False,
            access_token="tok",
            base_url="url",
        )
        model = self.mod.OpenAIModel("m", settings)
        model._model = model._load_model()
        gen = GenerationSettings(use_async_generator=False)
        response = await model("hi", settings=gen)
        self.openai_stub.AsyncOpenAI.assert_called_once_with(
            base_url="url", api_key="tok"
        )
        self.openai_stub.AsyncOpenAI.return_value.responses.create.assert_awaited_once_with(
            extra_headers={
                "X-Title": "Avalan",
                "HTTP-Referer": "https://github.com/avalan-ai/avalan",
            },
            model="m",
            input=[{"role": "user", "content": "hi"}],
            store=False,
            stream=False,
            timeout=None,
            temperature=1.0,
            top_p=1.0,
        )
        from avalan.model.stream import TextGenerationSingleStream

        self.assertIsInstance(response._output_fn, TextGenerationSingleStream)
        self.assertFalse(response._use_async_generator)
        self.assertEqual(response.usage.input_tokens, 1)
        self.assertEqual(response.provider_family, "openai")
        observation = usage_observation_from_response(response)
        self.assertIsNotNone(observation)
        assert observation is not None
        self.assertEqual(observation.metadata, {"provider_family": "openai"})
        self.assertEqual(await response.to_str(), "ok")

    async def test_azure_response_usage_preserves_safe_metadata(self):
        usage = SimpleNamespace(
            input_tokens=11,
            input_tokens_details=SimpleNamespace(cached_tokens=4),
            output_tokens=7,
            output_tokens_details=SimpleNamespace(reasoning_tokens=3),
            total_tokens=18,
            model="private-deployment-name",
            response_id="private-response-id",
        )
        resp = SimpleNamespace(
            output=[
                SimpleNamespace(content=[SimpleNamespace(text="azure ok")])
            ],
            usage=usage,
        )
        self.openai_stub.AsyncOpenAI.return_value.responses.create = AsyncMock(
            return_value=resp
        )
        client = self.mod.OpenAIClient(
            api_key="key",
            base_url="https://acct.openai.azure.com/openai/v1/",
        )
        response = await client(
            "private-deployment-name",
            [Message(role=MessageRole.USER, content="hi")],
            use_async_generator=False,
        )

        observation = usage_observation_from_response(response)

        self.assertEqual(response.provider_family, "azure_openai")
        self.assertIsNotNone(observation)
        assert observation is not None
        self.assertEqual(
            observation.metadata, {"provider_family": "azure_openai"}
        )
        self.assertEqual(observation.totals.input_tokens, 11)
        self.assertEqual(observation.totals.cached_input_tokens, 4)
        self.assertIsNone(observation.totals.cache_creation_input_tokens)
        self.assertEqual(observation.totals.output_tokens, 7)
        self.assertEqual(observation.totals.reasoning_tokens, 3)
        self.assertEqual(observation.totals.total_tokens, 18)
        self.assertNotIn("private-deployment-name", str(observation))
        self.assertNotIn("private-response-id", str(observation))


class TemplateMessagesFormatTestCase(IsolatedAsyncioTestCase):
    def setUp(self):
        self.openai_stub, self.patch = patch_openai_imports()
        importlib.reload(
            importlib.import_module("avalan.model.nlp.text.vendor.openai")
        )
        self.mod = importlib.import_module(
            "avalan.model.nlp.text.vendor.openai"
        )

    def tearDown(self):
        self.patch.stop()

    async def _assert_messages(self, content, expected_content):
        resp = SimpleNamespace(
            output=[SimpleNamespace(content=[SimpleNamespace(text="x")])]
        )
        self.openai_stub.AsyncOpenAI.return_value.responses.create = AsyncMock(
            return_value=resp
        )
        client = self.mod.OpenAIClient(api_key="key", base_url="url")
        message = Message(role=MessageRole.USER, content=content)
        await client("model", [message], use_async_generator=False)
        create_mock = (
            self.openai_stub.AsyncOpenAI.return_value.responses.create
        )
        create_mock.assert_awaited_once()
        kwargs = create_mock.await_args.kwargs
        self.assertEqual(
            kwargs["input"], [{"role": "user", "content": expected_content}]
        )

    async def test_string_message(self):
        await self._assert_messages("hi", "hi")

    async def test_text_message_content(self):
        content = MessageContentText(type="text", text="hi")
        await self._assert_messages(content, "hi")

    async def test_image_message_content(self):
        content = MessageContentImage(type="image_url", image_url={"url": "u"})
        await self._assert_messages(
            content, [{"type": "input_image", "image_url": "u"}]
        )

    async def test_image_message_content_from_file_id(self):
        content = MessageContentImage(
            type="image_url", image_url={"file_id": "file-img"}
        )
        await self._assert_messages(
            content, [{"type": "input_image", "file_id": "file-img"}]
        )

    async def test_image_message_content_from_base64(self):
        content = MessageContentImage(
            type="image_url",
            image_url={
                "data": "YWJj",
                "detail": "high",
                "mime_type": "image/jpeg",
            },
        )
        await self._assert_messages(
            content,
            [
                {
                    "type": "input_image",
                    "image_url": "data:image/jpeg;base64,YWJj",
                    "detail": "high",
                }
            ],
        )

    async def test_image_message_content_preserves_data_url(self):
        content = MessageContentImage(
            type="image_url",
            image_url={
                "data": "data:image/png;base64,YWJj",
                "mime_type": "image/png",
            },
        )
        await self._assert_messages(
            content,
            [
                {
                    "type": "input_image",
                    "image_url": "data:image/png;base64,YWJj",
                }
            ],
        )

    async def test_image_message_content_rejects_non_image_mime(self):
        resp = SimpleNamespace(
            output=[SimpleNamespace(content=[SimpleNamespace(text="x")])]
        )
        create_mock = AsyncMock(return_value=resp)
        self.openai_stub.AsyncOpenAI.return_value.responses.create = (
            create_mock
        )
        client = self.mod.OpenAIClient(api_key="key", base_url="url")
        message = Message(
            role=MessageRole.USER,
            content=MessageContentImage(
                type="image_url",
                image_url={"data": "YWJj", "mime_type": "application/pdf"},
            ),
        )

        with self.assertRaisesRegex(AssertionError, "image MIME type"):
            await client("model", [message], use_async_generator=False)

        create_mock.assert_not_awaited()

    async def test_image_message_content_rejects_invalid_mime_type(self):
        resp = SimpleNamespace(
            output=[SimpleNamespace(content=[SimpleNamespace(text="x")])]
        )
        create_mock = AsyncMock(return_value=resp)
        self.openai_stub.AsyncOpenAI.return_value.responses.create = (
            create_mock
        )
        client = self.mod.OpenAIClient(api_key="key", base_url="url")
        message = Message(
            role=MessageRole.USER,
            content=MessageContentImage(
                type="image_url",
                image_url={"data": "YWJj", "mime_type": object()},
            ),
        )

        with self.assertRaisesRegex(AssertionError, "image MIME type"):
            await client("model", [message], use_async_generator=False)

        create_mock.assert_not_awaited()

    async def test_mixed_message_content(self):
        content = [
            MessageContentText(type="text", text="hi"),
            MessageContentImage(type="image_url", image_url={"url": "u"}),
        ]
        await self._assert_messages(
            content,
            [
                {"type": "input_text", "text": "hi"},
                {"type": "input_image", "image_url": "u"},
            ],
        )

    async def test_file_message_content(self):
        content = MessageContentFile(
            type="file", file={"file_url": "https://example.com/a.pdf"}
        )
        await self._assert_messages(
            content,
            [
                {
                    "type": "input_file",
                    "file_url": "https://example.com/a.pdf",
                }
            ],
        )

    async def test_mixed_message_content_with_file_data(self):
        content = [
            MessageContentText(type="text", text="hi"),
            MessageContentFile(
                type="file",
                file={"file_data": "YWJj", "filename": "report.pdf"},
            ),
        ]
        await self._assert_messages(
            content,
            [
                {"type": "input_text", "text": "hi"},
                {
                    "type": "input_file",
                    "file_data": "data:application/pdf;base64,YWJj",
                    "filename": "report.pdf",
                },
            ],
        )

    async def test_mixed_message_content_with_pdf_file_data_uses_data_url(
        self,
    ):
        content = [
            MessageContentText(type="text", text="hi"),
            MessageContentFile(
                type="file",
                file={
                    "file_data": "YWJj",
                    "filename": "report.pdf",
                    "mime_type": "application/pdf",
                },
            ),
        ]
        await self._assert_messages(
            content,
            [
                {"type": "input_text", "text": "hi"},
                {
                    "type": "input_file",
                    "file_data": "data:application/pdf;base64,YWJj",
                    "filename": "report.pdf",
                },
            ],
        )

    async def test_pdf_file_data_uses_data_url_when_inferred_from_filename(
        self,
    ):
        content = MessageContentFile(
            type="file",
            file={"file_data": "YWJj", "filename": "report.pdf"},
        )
        await self._assert_messages(
            content,
            [
                {
                    "type": "input_file",
                    "file_data": "data:application/pdf;base64,YWJj",
                    "filename": "report.pdf",
                }
            ],
        )

    async def test_reasoning_effort_is_forwarded(self):
        response = SimpleNamespace(
            output=[SimpleNamespace(content=[SimpleNamespace(text="x")])]
        )
        create_mock = AsyncMock(return_value=response)
        self.openai_stub.AsyncOpenAI.return_value.responses.create = (
            create_mock
        )
        client = self.mod.OpenAIClient(api_key="key", base_url="url")
        settings = GenerationSettings(
            reasoning=ReasoningSettings(effort=ReasoningEffort.XHIGH)
        )

        await client(
            "model",
            [Message(role=MessageRole.USER, content="hi")],
            settings,
            use_async_generator=False,
        )

        kwargs = create_mock.await_args.kwargs
        self.assertEqual(kwargs["reasoning"], {"effort": "xhigh"})

    def test_helper_methods_cover_additional_reasoning_and_file_paths(self):
        self.assertEqual(
            self.mod.OpenAIClient._reasoning_config(
                GenerationSettings(
                    reasoning=ReasoningSettings(effort=ReasoningEffort.MAX)
                )
            ),
            {"effort": "xhigh"},
        )
        self.assertEqual(
            self.mod.OpenAIClient._content_block(
                {
                    "type": "unknown",
                    "value": 1,
                }
            ),
            {"type": "unknown", "value": 1},
        )
        self.assertEqual(
            self.mod.OpenAIClient._file_block(
                {
                    "file_id": "file-1",
                    "filename": "report.pdf",
                }
            ),
            {
                "type": "input_file",
                "file_id": "file-1",
                "filename": "report.pdf",
            },
        )
        with self.assertRaises(AssertionError):
            self.mod.OpenAIClient._file_block({})
        with self.assertRaises(AssertionError):
            self.mod.OpenAIClient._image_block({})
        self.assertEqual(
            self.mod.OpenAIClient._non_stream_response_content(
                {"output": None}
            ),
            "",
        )

    def test_template_messages_skips_entries_without_content_key(self):
        client = self.mod.OpenAIClient(api_key="key", base_url="url")

        with patch.object(
            self.mod.TextGenerationVendor,
            "_template_messages",
            return_value=[{"role": "assistant"}],
        ):
            templated = client._template_messages([])

        self.assertEqual(templated, [{"role": "assistant"}])

    async def test_non_stream_tool_call_output(self):
        response = SimpleNamespace(
            output=[
                SimpleNamespace(
                    type="output_text",
                    content=[SimpleNamespace(text="hello ")],
                ),
                SimpleNamespace(
                    type="tool_call",
                    call=SimpleNamespace(
                        id="call1",
                        function=SimpleNamespace(
                            name="pkg__tool", arguments='{"a":1}'
                        ),
                    ),
                ),
            ]
        )

        create_mock = AsyncMock(return_value=response)
        self.openai_stub.AsyncOpenAI.return_value.responses.create = (
            create_mock
        )

        with patch.object(
            self.mod.TextGenerationVendor,
            "build_tool_call_token",
            return_value=ToolCallToken(token="<tool_call />"),
        ) as build_token:
            client = self.mod.OpenAIClient(api_key="key", base_url="url")
            message = Message(role=MessageRole.USER, content="hi")
            stream = await client(
                "model", [message], use_async_generator=False
            )

        from avalan.model.stream import TextGenerationSingleStream

        self.assertIsInstance(stream, TextGenerationSingleStream)
        self.assertEqual(stream.content, "hello <tool_call />")
        build_token.assert_called_once_with("call1", "pkg__tool", '{"a":1}')


class TemplateAndToolSchemaTestCase(TestCase):
    def setUp(self):
        self.openai_stub, self.patch = patch_openai_imports()
        importlib.reload(
            importlib.import_module("avalan.model.nlp.text.vendor.openai")
        )
        self.mod = importlib.import_module(
            "avalan.model.nlp.text.vendor.openai"
        )

    def tearDown(self):
        self.patch.stop()

    def test_tool_schemas_none(self):
        tool = MagicMock()
        tool.json_schemas.return_value = None
        self.assertIsNone(self.mod.OpenAIClient._tool_schemas(tool))
        tool.json_schemas.return_value = [{"type": "x"}]
        self.assertEqual(self.mod.OpenAIClient._tool_schemas(tool), [])

    def test_template_messages_tool_results(self):
        client = self.mod.OpenAIClient(api_key="k", base_url="b")

        @dataclass
        class R:
            v: int

        call1 = ToolCall(id="c1", name="pkg.func", arguments={"a": 1})
        result1 = ToolCallResult(
            id="c1",
            name="pkg.func",
            arguments={"a": 1},
            call=call1,
            result=R(v=2),
        )
        msg1 = Message(role=MessageRole.TOOL, tool_call_result=result1)

        call2 = ToolCall(id="c2", name="pkg.func2")
        result2 = ToolCallResult(
            id="c2", name="pkg.func2", call=call2, result={"x": 3}
        )
        msg2 = Message(role=MessageRole.TOOL, tool_call_result=result2)

        templated = client._template_messages([msg1, msg2])
        self.assertEqual(
            templated,
            [
                {
                    "type": "function_call",
                    "name": "avl_cGtnLmZ1bmM",
                    "call_id": "c1",
                    "arguments": '{"a": 1}',
                },
                {
                    "type": "function_call_output",
                    "call_id": "c1",
                    "output": '{"v": 2}',
                },
                {
                    "type": "function_call",
                    "name": "avl_cGtnLmZ1bmMy",
                    "call_id": "c2",
                    "arguments": "null",
                },
                {
                    "type": "function_call_output",
                    "call_id": "c2",
                    "output": '{"x": 3}',
                },
            ],
        )

    def test_template_messages_tool_results_stringify_uuid_call_id(self):
        client = self.mod.OpenAIClient(api_key="k", base_url="b")
        call_id = UUID("11111111-1111-1111-1111-111111111111")
        call = ToolCall(
            id=call_id,
            name="pkg.func",
            arguments={"entity_id": call_id},
        )
        result = ToolCallResult(
            id="c1",
            name="pkg.func",
            call=call,
            result={"entity_id": call_id},
        )
        message = Message(role=MessageRole.TOOL, tool_call_result=result)

        templated = client._template_messages([message])

        self.assertEqual(
            templated,
            [
                {
                    "type": "function_call",
                    "name": "avl_cGtnLmZ1bmM",
                    "call_id": str(call_id),
                    "arguments": (
                        '{"entity_id": "11111111-1111-1111-1111-111111111111"}'
                    ),
                },
                {
                    "type": "function_call_output",
                    "call_id": str(call_id),
                    "output": (
                        '{"entity_id": "11111111-1111-1111-1111-111111111111"}'
                    ),
                },
            ],
        )

    def test_template_messages_tool_error(self):
        client = self.mod.OpenAIClient(api_key="k", base_url="b")
        call = ToolCall(id="c1", name="pkg.func", arguments={"a": 1})
        error = ToolCallError(
            id="c1",
            name="pkg.func",
            call=call,
            error=ValueError("boom"),
            message="boom",
        )
        msg = Message(role=MessageRole.TOOL, tool_call_error=error)
        templated = client._template_messages([msg])
        self.assertEqual(
            templated,
            [
                {
                    "type": "function_call",
                    "name": "avl_cGtnLmZ1bmM",
                    "call_id": "c1",
                    "arguments": '{"a": 1}',
                },
                {
                    "type": "function_call_output",
                    "call_id": "c1",
                    "output": '{"error": "boom"}',
                },
            ],
        )

    def test_template_messages_tool_diagnostic(self):
        client = self.mod.OpenAIClient(api_key="k", base_url="b")
        diagnostic = ToolCallDiagnostic(
            id="d1",
            call_id="c1",
            requested_name="missing",
            code=ToolCallDiagnosticCode.UNKNOWN_TOOL,
            stage=ToolCallDiagnosticStage.RESOLVE,
            message="Tool is unknown.",
        )
        msg = Message(
            role=MessageRole.TOOL,
            name="missing",
            arguments={"a": 1},
            tool_call_diagnostic=diagnostic,
        )

        templated = client._template_messages([msg])

        self.assertEqual(templated[0]["call_id"], "c1")
        self.assertEqual(templated[0]["name"], "missing")
        self.assertEqual(templated[0]["arguments"], '{"a": 1}')
        output = loads(templated[1]["output"])
        self.assertEqual(output["code"], "tool.unknown")
        self.assertEqual(output["requested_name"], "missing")

    def test_template_messages_unanchored_tool_diagnostic(self):
        client = self.mod.OpenAIClient(api_key="k", base_url="b")
        diagnostic = ToolCallDiagnostic(
            id="d1",
            code=ToolCallDiagnosticCode.MALFORMED_CALL,
            stage=ToolCallDiagnosticStage.PARSE,
            message="Tool call could not be parsed.",
        )
        msg = Message(
            role=MessageRole.TOOL,
            tool_call_diagnostic=diagnostic,
        )

        templated = client._template_messages([msg])

        self.assertEqual(templated[0]["role"], "assistant")
        output = loads(templated[0]["content"])
        self.assertEqual(output["code"], "tool_call.malformed")


class OpenAIAdditionalCoverageTestCase(TestCase):
    def setUp(self):
        self.openai_stub, self.patch = patch_openai_imports()
        importlib.reload(
            importlib.import_module("avalan.model.nlp.text.vendor.openai")
        )
        self.mod = importlib.import_module(
            "avalan.model.nlp.text.vendor.openai"
        )

    def tearDown(self):
        self.patch.stop()

    def test_non_stream_response_content_handles_dict(self):
        response = {
            "output": [
                {
                    "type": "message",
                    "content": [{"text": "hello"}],
                },
                {
                    "type": "tool_call",
                    "call": {
                        "id": "call-id",
                        "function": {
                            "name": "pkg.tool",
                            "arguments": '{"a":1}',
                        },
                    },
                },
            ]
        }
        token = SimpleNamespace(token="<tool>")
        with patch.object(
            self.mod.TextGenerationVendor,
            "build_tool_call_token",
            return_value=token,
        ) as build:
            text = self.mod.OpenAIClient._non_stream_response_content(response)

        self.assertEqual(text, "hello<tool>")
        build.assert_called_once_with("call-id", "pkg.tool", '{"a":1}')

    def test_non_streaming_response_str_variants(self):
        settings = GenerationSettings()
        response = self.mod.OpenAINonStreamingResponse(
            lambda **_: "value",
            logger=MagicMock(),
            generation_settings=settings,
            settings=settings,
            use_async_generator=False,
            static_response_text="cached",
        )
        self.assertEqual(str(response), "cached")

        buffered = self.mod.OpenAINonStreamingResponse(
            lambda **_: "value",
            logger=MagicMock(),
            generation_settings=settings,
            settings=settings,
            use_async_generator=False,
        )
        buffered._buffer.write("buffered")
        self.assertEqual(str(buffered), "buffered")

        fallback = self.mod.OpenAINonStreamingResponse(
            lambda **_: "value",
            logger=MagicMock(),
            generation_settings=settings,
            settings=settings,
            use_async_generator=False,
        )
        fallback._buffer = SimpleNamespace(getvalue=lambda: None)
        self.assertIn("OpenAINonStreamingResponse", str(fallback))


class OpenAIModelStreamingFlagTestCase(IsolatedAsyncioTestCase):
    def setUp(self):
        self.openai_stub, self.patch = patch_openai_imports()
        importlib.reload(
            importlib.import_module("avalan.model.nlp.text.vendor.openai")
        )
        self.mod = importlib.import_module(
            "avalan.model.nlp.text.vendor.openai"
        )

    def tearDown(self):
        self.patch.stop()

    async def test_call_returns_streaming_response(self):
        settings = TransformerEngineSettings(
            access_token="tok",
            auto_load_model=False,
            auto_load_tokenizer=False,
        )
        model = self.mod.OpenAIModel("model-id", settings)
        model._model = AsyncMock(
            return_value=lambda *_args, **_kwargs: AsyncIter([])
        )

        response = await model(
            "prompt",
            system_prompt="sys",
            developer_prompt="dev",
            settings=GenerationSettings(),
        )

        self.assertIsInstance(response, self.mod.TextGenerationResponse)


if __name__ == "__main__":
    from unittest import main

    main()
