98 lines
3.8 KiB
Python
98 lines
3.8 KiB
Python
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()
|