lmstudiofixed/services/litellm-lmstudio-adapter/tests/test_adapter.py

import json
import importlib.util
from pathlib import Path
import unittest


MODULE_PATH = (
    Path(__file__).resolve().parent.parent / "litellm_lmstudio_adapter.py"
)
if MODULE_PATH.exists():
    spec = importlib.util.spec_from_file_location("litellm_lmstudio_adapter", MODULE_PATH)
    module = importlib.util.module_from_spec(spec)
    assert spec is not None and spec.loader is not None
    spec.loader.exec_module(module)
    adapter = module
else:  # pragma: no cover
    adapter = None


class LiteLLMLMStudioAdapterTests(unittest.TestCase):
    def setUp(self):
        if adapter is None:
            self.fail("litellm_lmstudio_adapter.app is not available")

    def test_build_chat_native_payload_maps_reasoning(self):
        payload = {
            "model": "qwen/qwen3.6-35b-a3b",
            "messages": [
                {"role": "system", "content": "be precise"},
                {"role": "user", "content": "hello"},
                {"role": "assistant", "content": "hi"},
                {"role": "user", "content": "continue"},
            ],
            "reasoning": "off",
            "max_tokens": 64,
            "temperature": 0.2,
        }

        native_payload = adapter.build_chat_native_payload(payload)

        self.assertEqual(native_payload["model"], "qwen/qwen3.6-35b-a3b")
        self.assertEqual(
            native_payload["input"],
            "User: hello\n\nAssistant: hi\n\nUser: continue",
        )
        self.assertEqual(native_payload["system_prompt"], "be precise")
        self.assertEqual(native_payload["reasoning"], "off")
        self.assertEqual(native_payload["max_output_tokens"], 64)

    def test_translate_chat_stream_chunk_emits_reasoning_delta(self):
        event = {
            "type": "reasoning.delta",
            "content": "step 1",
            "model": "qwen/qwen3.6-35b-a3b",
        }

        chunk = adapter.translate_chat_stream_event(
            event=event,
            model="qwen/qwen3.6-35b-a3b",
            chunk_id="chatcmpl-test",
        )

        self.assertIsNotNone(chunk)
        parsed = json.loads(chunk.removeprefix("data: ").strip())
        self.assertEqual(parsed["object"], "chat.completion.chunk")
        self.assertEqual(parsed["choices"][0]["delta"]["reasoning_content"], "step 1")

    def test_translate_chat_stream_chunk_emits_content_delta(self):
        event = {
            "type": "message.delta",
            "content": "9193",
            "model": "qwen/qwen3.6-35b-a3b",
        }

        chunk = adapter.translate_chat_stream_event(
            event=event,
            model="qwen/qwen3.6-35b-a3b",
            chunk_id="chatcmpl-test",
        )

        parsed = json.loads(chunk.removeprefix("data: ").strip())
        self.assertEqual(parsed["choices"][0]["delta"]["content"], "9193")

    def test_translate_chat_stream_chunk_ignores_non_content_events(self):
        event = {
            "type": "prompt_processing.start",
        }

        chunk = adapter.translate_chat_stream_event(
            event=event,
            model="qwen/qwen3.6-35b-a3b",
            chunk_id="chatcmpl-test",
        )

        self.assertIsNone(chunk)

    def test_translate_responses_stream_event_emits_output_text_delta(self):
        event = {
            "type": "message.delta",
            "content": "9193",
            "model": "qwen/qwen3.6-35b-a3b",
        }

        frames = adapter.translate_responses_stream_event(
            event=event,
            model="qwen/qwen3.6-35b-a3b",
            response_id="resp_test",
        )

        self.assertTrue(frames)
        parsed = json.loads(frames[0].removeprefix("data: ").strip())
        self.assertEqual(parsed["type"], "response.output_text.delta")
        self.assertEqual(parsed["delta"], "9193")

    def test_parse_sse_event_blocks_extracts_event_and_data(self):
        raw = (
            "event: content_block.delta\n"
            "data: {\"type\":\"message\",\"content\":\"9193\"}\n\n"
            "event: done\n"
            "data: [DONE]\n\n"
        )

        blocks = list(adapter.parse_sse_event_blocks(raw))

        self.assertEqual(len(blocks), 2)
        self.assertEqual(blocks[0]["event"], "content_block.delta")
        self.assertEqual(blocks[0]["data"]["content"], "9193")
        self.assertEqual(blocks[1]["event"], "done")
        self.assertEqual(blocks[1]["data_raw"], "[DONE]")

    def test_build_stub_error_uses_openai_shape(self):
        response = adapter.build_stub_error("audio.transcriptions")

        self.assertEqual(response.status_code, 501)
        payload = json.loads(response.body.decode("utf-8"))
        self.assertEqual(payload["error"]["type"], "not_implemented")
        self.assertIn("audio.transcriptions", payload["error"]["message"])


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