PMFlow/backend/tests/test_llm_structured_extraction.py

import unittest
from unittest.mock import AsyncMock

from app.services.llm import ChatCompletionResult, LLMService, StructuredExtractionResult


class LLMStructuredExtractionTest(unittest.TestCase):
    def setUp(self) -> None:
        self.service = LLMService()

    def _extract(self, response: dict, *, function_name: str = "openclaw_command_dispatch") -> StructuredExtractionResult:
        return self.service._extract_structured_object_from_response(
            function_name=function_name,
            response=response,
            fallback_text="",
            provider="glm",
            model="GLM-4.7",
        )

    def test_extracts_object_from_plain_content(self) -> None:
        result = self._extract(
            {
                "choices": [
                    {
                        "message": {
                            "content": '{"summary":"ok","confidence":0.9,"operations":[]}',
                        }
                    }
                ]
            }
        )

        self.assertEqual(result.validation, "ok")
        self.assertEqual(result.payload["summary"], "ok")

    def test_extracts_object_from_typed_chunk_content(self) -> None:
        result = self._extract(
            {
                "choices": [
                    {
                        "message": {
                            "content": [
                                {"type": "output_text", "text": '{"summary":"chunk","confidence":0.8,"operations":[]}'}
                            ]
                        }
                    }
                ]
            }
        )

        self.assertEqual(result.validation, "ok")
        self.assertEqual(result.payload["summary"], "chunk")

    def test_extracts_object_from_tool_call_arguments(self) -> None:
        result = self._extract(
            {
                "choices": [
                    {
                        "message": {
                            "tool_calls": [
                                {
                                    "function": {
                                        "arguments": '{"summary":"args","confidence":0.7,"operations":[]}'
                                    }
                                }
                            ],
                            "content": "",
                        }
                    }
                ]
            }
        )

        self.assertEqual(result.source_field, "tool_calls.function.arguments")
        self.assertEqual(result.payload["summary"], "args")

    def test_extracts_json_from_reasoning_content_with_prefix_text(self) -> None:
        result = self._extract(
            {
                "choices": [
                    {
                        "message": {
                            "content": "",
                            "reasoning_content": "先分析上下文，再给结果：\n{\"summary\":\"from_reasoning\",\"confidence\":0.66,\"operations\":[]}",
                        }
                    }
                ]
            }
        )

        self.assertEqual(result.source_field, "message.reasoning_content")
        self.assertEqual(result.payload["summary"], "from_reasoning")

    def test_returns_json_not_found_when_reasoning_has_only_natural_language(self) -> None:
        result = self._extract(
            {
                "choices": [
                    {
                        "message": {
                            "content": "",
                            "reasoning_content": "我认为还需要更多上下文。",
                        }
                    }
                ]
            }
        )

        self.assertEqual(result.error_code, "structured_json_not_found")
        self.assertEqual(result.source_field, "message.reasoning_content")

    def test_unwraps_nested_json_field_once(self) -> None:
        result = self._extract(
            {
                "choices": [
                    {
                        "message": {
                            "content": '{"json":"{\\"summary\\":\\"nested\\",\\"confidence\\":0.55,\\"operations\\":[]}"}',
                        }
                    }
                ]
            }
        )

        self.assertEqual(result.validation, "ok")
        self.assertEqual(result.payload["summary"], "nested")

    def test_rejects_top_level_array_for_object_only_function(self) -> None:
        result = self._extract(
            {
                "choices": [
                    {
                        "message": {
                            "content": '[{"summary":"bad"}]',
                        }
                    }
                ]
            }
        )

        self.assertEqual(result.error_code, "structured_schema_invalid")


class LLMStructuredCacheTest(unittest.IsolatedAsyncioTestCase):
    async def test_generate_structured_object_uses_cache_after_first_hit(self) -> None:
        service = LLMService()
        calls: list[str] = []
        cache: dict[str, str] = {}

        class _FakeRedis:
            async def get(self, key: str):
                return cache.get(key)

            async def setex(self, key: str, _ttl: int, value: str):
                cache[key] = value

        service._get_redis = AsyncMock(return_value=_FakeRedis())
        service._resolve_engine_for_call = AsyncMock(
            return_value={
                "provider": "glm",
                "model": "GLM-4.7",
                "api_key": "fake",
                "base_url": "http://unit.test",
            }
        )
        service._resolve_active_prompt_profile = AsyncMock(return_value={"role_name": "", "system_prompt": ""})
        service._log_ai_call = AsyncMock()

        async def _fake_chat_completion_raw(**_: object) -> ChatCompletionResult:
            calls.append("hit")
            return ChatCompletionResult(
                text='{"summary":"cached","confidence":0.9,"operations":[]}',
                raw_response={
                    "choices": [
                        {
                            "message": {
                                "content": '{"summary":"cached","confidence":0.9,"operations":[]}',
                            }
                        }
                    ]
                },
            )

        service._chat_completion_raw = AsyncMock(side_effect=_fake_chat_completion_raw)

        first = await service.generate_structured_object(
            "planner prompt",
            db=object(),
            tenant_id=1,
            function_name="openclaw_command_dispatch",
            response_mode="fast",
        )
        second = await service.generate_structured_object(
            "planner prompt",
            db=object(),
            tenant_id=1,
            function_name="openclaw_command_dispatch",
            response_mode="fast",
        )

        self.assertEqual(first.payload["summary"], "cached")
        self.assertEqual(second.payload["summary"], "cached")
        self.assertTrue(second.from_cache)
        self.assertEqual(len(calls), 1)


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