YG-Rules/tests/test_llm.py

import unittest
from pathlib import Path
from tempfile import TemporaryDirectory
from unittest.mock import MagicMock, patch

import requests

from app.utils.llm import LLMClient, LLMError, strip_thinking


class LLMRetryTest(unittest.TestCase):
    @patch("app.utils.llm.time.sleep")
    @patch("app.utils.llm.requests.post")
    def test_retries_http_529_then_succeeds(self, mock_post, mock_sleep):
        retry_response = MagicMock()
        retry_response.raise_for_status.side_effect = requests.exceptions.HTTPError(response=MagicMock(status_code=529))

        success_response = MagicMock()
        success_response.raise_for_status.return_value = None
        success_response.json.return_value = {
            "choices": [{"message": {"content": "最终答案"}}]
        }

        mock_post.side_effect = [retry_response, retry_response, success_response]

        client = LLMClient(api_key="test-key", retry_attempts=3, retry_delay_seconds=5)
        result = client.chat(messages=[{"role": "user", "content": "hello"}])

        self.assertEqual(result, "最终答案")
        self.assertEqual(mock_post.call_count, 3)
        self.assertEqual(mock_sleep.call_count, 2)
        mock_sleep.assert_called_with(5)

    @patch("app.utils.llm.time.sleep")
    @patch("app.utils.llm.requests.post")
    def test_raises_after_exhausting_http_529_retries(self, mock_post, mock_sleep):
        retry_response = MagicMock()
        retry_response.raise_for_status.side_effect = requests.exceptions.HTTPError(response=MagicMock(status_code=529))
        mock_post.side_effect = [retry_response, retry_response, retry_response]

        client = LLMClient(api_key="test-key", retry_attempts=3, retry_delay_seconds=5)

        with self.assertRaises(LLMError) as context:
            client.chat(messages=[{"role": "user", "content": "hello"}])

        self.assertIn("LLM 请求失败", str(context.exception))
        self.assertEqual(mock_post.call_count, 3)
        self.assertEqual(mock_sleep.call_count, 2)

    def test_reads_model_config_from_project_env_file(self):
        with TemporaryDirectory() as temp_dir:
            env_file = Path(temp_dir) / ".env"
            env_file.write_text(
                "\n".join([
                    "OPENAI_API_KEY=env-file-key",
                    "OPENAI_BASE_URL=https://api.minimaxi.com/v1",
                    "OPENAI_MODEL=MiniMax-M3",
                ]),
                encoding="utf-8",
            )
            with patch.dict("os.environ", {}, clear=True):
                with patch("app.utils.llm.os.getcwd", return_value=temp_dir):
                    client = LLMClient()

        self.assertEqual(client.api_key, "env-file-key")
        self.assertEqual(client.base_url, "https://api.minimaxi.com/v1")
        self.assertEqual(client.model, "MiniMax-M3")

    def test_strip_thinking_preserves_json_after_unclosed_think_tag(self):
        content = "<think>\nreasoning...\n{\"ok\": true}"

        self.assertEqual(strip_thinking(content), "{\"ok\": true}")

    def test_strip_thinking_removes_common_think_tag_variants(self):
        content = (
            "<THINK>第一段推理</THINK>"
            "正式答案"
            "<thinking>第二段推理</thinking>"
            "\n```json\n{\"ok\": true}\n```"
        )

        cleaned = strip_thinking(content)

        self.assertNotIn("THINK", cleaned.upper())
        self.assertNotIn("第一段推理", cleaned)
        self.assertNotIn("第二段推理", cleaned)
        self.assertIn("正式答案", cleaned)
        self.assertIn("{\"ok\": true}", cleaned)

    def test_strip_thinking_keeps_json_after_unclosed_thinking_tag(self):
        content = "<thinking>\n推理内容很长\n```json\n{\"ok\": true}\n```"

        self.assertEqual(strip_thinking(content), "```json\n{\"ok\": true}\n```")


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