proposa-ai/scripts/tests/test_llm_empty_response.py

import json
import os
import tempfile
import unittest
from pathlib import Path
from types import SimpleNamespace
from unittest.mock import patch

from llm_client.cache import build_cache_key
from llm_client.client import LLMClient


class _FakeCompletions:
    def __init__(self, contents):
        self.contents = iter(contents)
        self.calls = 0
        self.kwargs = []

    def create(self, **_kwargs):
        self.calls += 1
        self.kwargs.append(_kwargs)
        content = next(self.contents)
        return SimpleNamespace(
            choices=[SimpleNamespace(message=SimpleNamespace(content=content))]
        )


class LlmEmptyResponseTests(unittest.TestCase):
    def _client(self, contents):
        completions = _FakeCompletions(contents)
        client = object.__new__(LLMClient)
        client.model = "test-model"
        client.cfg = SimpleNamespace(generate_max_tokens=65536)
        client.client = SimpleNamespace(
            chat=SimpleNamespace(completions=completions)
        )
        return client, completions

    @staticmethod
    def _config():
        return SimpleNamespace(
            llm=SimpleNamespace(
                generate_max_tokens=100,
                generate_temperature=0.4,
            ),
            max_retries=2,
        )

    def test_blank_api_response_is_retried_and_not_cached(self):
        client, completions = self._client(["   ", "有效正文"])
        with tempfile.TemporaryDirectory() as temp_dir, patch.dict(
            os.environ,
            {"BID_LLM_CACHE": "1", "BID_LLM_CACHE_DIR": temp_dir},
        ), patch("llm_client.client.get_config", return_value=self._config()), patch(
            "llm_client.client.time.sleep"
        ):
            result = client.generate("system", "user", max_tokens=100)

            self.assertEqual(result, "有效正文")
            self.assertEqual(completions.calls, 2)
            cache_files = list(Path(temp_dir).glob("*.json"))
            self.assertEqual(len(cache_files), 1)
            cached = json.loads(cache_files[0].read_text(encoding="utf-8"))
            self.assertEqual(cached["content"], "有效正文")

    def test_existing_blank_cache_is_removed_before_request(self):
        client, completions = self._client(["重新生成的正文"])
        with tempfile.TemporaryDirectory() as temp_dir, patch.dict(
            os.environ,
            {"BID_LLM_CACHE": "1", "BID_LLM_CACHE_DIR": temp_dir},
        ), patch("llm_client.client.get_config", return_value=self._config()):
            key = build_cache_key(
                model="test-model",
                system_prompt="system",
                user_prompt="user",
                max_tokens=100,
                temperature=0.4,
                response_format=None,
            )
            cache_path = Path(temp_dir) / f"{key}.json"
            cache_path.write_text(
                json.dumps({"model": "test-model", "content": ""}),
                encoding="utf-8",
            )

            result = client.generate("system", "user", max_tokens=100)

            self.assertEqual(result, "重新生成的正文")
            self.assertEqual(completions.calls, 1)
            cached = json.loads(cache_path.read_text(encoding="utf-8"))
            self.assertEqual(cached["content"], "重新生成的正文")

    def test_blank_response_retries_double_max_tokens_four_times(self):
        client, completions = self._client(
            ["   ", "   ", "   ", "   ", "有效正文"]
        )
        with tempfile.TemporaryDirectory() as temp_dir, patch.dict(
            os.environ,
            {"BID_LLM_CACHE": "1", "BID_LLM_CACHE_DIR": temp_dir},
        ), patch(
            "llm_client.client.get_config",
            return_value=SimpleNamespace(
                llm=SimpleNamespace(
                    generate_max_tokens=65536,
                    generate_temperature=0.4,
                ),
                max_retries=4,
            ),
        ), patch("llm_client.client.time.sleep"):
            result = client.generate("system", "user", max_tokens=8192)

            self.assertEqual(result, "有效正文")
            self.assertEqual(completions.calls, 5)
            self.assertEqual(
                [kwargs["max_tokens"] for kwargs in completions.kwargs],
                [8192, 16384, 32768, 65536, 65536],
            )

    def test_json_parse_failure_retries_double_max_tokens_four_times(self):
        client, completions = self._client([
            "{not json",
            "{not json",
            "{not json",
            "{not json",
            "{not json",
        ])
        with tempfile.TemporaryDirectory() as temp_dir, patch.dict(
            os.environ,
            {"BID_LLM_CACHE": "1", "BID_LLM_CACHE_DIR": temp_dir},
        ), patch(
            "llm_client.client.get_config",
            return_value=SimpleNamespace(
                llm=SimpleNamespace(
                    generate_max_tokens=65536,
                    generate_temperature=0.4,
                ),
                max_retries=4,
            ),
        ), patch("llm_client.client.time.sleep"):
            result = client.extract_json(
                system_prompt="system",
                user_prompt="user",
                max_tokens=8192,
            )

            self.assertEqual(result, {})
            self.assertEqual(completions.calls, 5)
            self.assertEqual(
                [kwargs["max_tokens"] for kwargs in completions.kwargs],
                [8192, 16384, 32768, 65536, 65536],
            )


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