128 lines
4.0 KiB
Python
128 lines
4.0 KiB
Python
"""LLM 模块测试。"""
|
|
|
|
import sys
|
|
from pathlib import Path
|
|
from unittest.mock import MagicMock, patch
|
|
|
|
sys.path.insert(0, str(Path(__file__).parent.parent / "src"))
|
|
|
|
|
|
class TestCreateLLM:
|
|
"""测试 LLM 实例创建。"""
|
|
|
|
@patch("config.get_llm_config")
|
|
@patch("llama_index.llms.openai_like.OpenAILike")
|
|
def test_create_llm_uses_config(self, mock_openai_like, mock_get_config):
|
|
"""测试 _create_llm 使用配置文件中的参数。"""
|
|
mock_get_config.return_value = {
|
|
"model": "qwen/qwen3.6-27b",
|
|
"api_base": "http://100.123.83.113:1234/v1",
|
|
"api_key": "123456",
|
|
}
|
|
|
|
from llm import _create_llm
|
|
|
|
_create_llm()
|
|
|
|
mock_openai_like.assert_called_once_with(
|
|
model="qwen/qwen3.6-27b",
|
|
api_base="http://100.123.83.113:1234/v1",
|
|
api_key="123456",
|
|
temperature=0.1,
|
|
max_tokens=128000,
|
|
request_timeout=300.0,
|
|
is_chat_model=True,
|
|
)
|
|
|
|
|
|
class TestQueryAnalysis:
|
|
"""测试 query_analysis 端到端流程。"""
|
|
|
|
@patch("llm._create_llm")
|
|
@patch("llm.build_analysis_system_prompt")
|
|
def test_query_analysis_calls_llm(self, mock_prompt, mock_create_llm):
|
|
"""测试 query_analysis 正确调用 prompt 构建和 LLM stream_chat。"""
|
|
mock_prompt.return_value = "test prompt"
|
|
mock_llm = MagicMock()
|
|
mock_resp = MagicMock()
|
|
mock_resp.delta = "analysis result"
|
|
mock_llm.stream_chat.return_value = [mock_resp]
|
|
mock_create_llm.return_value = mock_llm
|
|
|
|
from llm import query_analysis
|
|
|
|
result = query_analysis(purpose="test purpose", index="test index")
|
|
|
|
mock_prompt.assert_called_once_with(
|
|
"test purpose",
|
|
"test index",
|
|
source_content="",
|
|
configured_language=None,
|
|
)
|
|
assert mock_llm.stream_chat.called
|
|
assert result == "analysis result"
|
|
|
|
@patch("llm._create_llm")
|
|
@patch("llm.build_analysis_system_prompt")
|
|
def test_query_analysis_passes_source_content(self, mock_prompt, mock_create_llm):
|
|
"""测试 source_content 正确传递。"""
|
|
mock_prompt.return_value = "test prompt"
|
|
mock_llm = MagicMock()
|
|
mock_resp = MagicMock()
|
|
mock_resp.delta = "result"
|
|
mock_llm.stream_chat.return_value = [mock_resp]
|
|
mock_create_llm.return_value = mock_llm
|
|
|
|
from llm import query_analysis
|
|
|
|
query_analysis(
|
|
purpose="",
|
|
index="",
|
|
source_content="some source",
|
|
configured_language="English",
|
|
)
|
|
|
|
mock_prompt.assert_called_once_with(
|
|
"",
|
|
"",
|
|
source_content="some source",
|
|
configured_language="English",
|
|
)
|
|
|
|
@patch("llm._create_llm")
|
|
@patch("llm.build_analysis_system_prompt")
|
|
def test_query_analysis_returns_text(self, mock_prompt, mock_create_llm):
|
|
"""测试返回值是纯文本字符串。"""
|
|
mock_prompt.return_value = "prompt"
|
|
mock_llm = MagicMock()
|
|
mock_resp1 = MagicMock()
|
|
mock_resp1.delta = "## 关键实体\n"
|
|
mock_resp2 = MagicMock()
|
|
mock_resp2.delta = "- 测试实体"
|
|
mock_llm.stream_chat.return_value = [mock_resp1, mock_resp2]
|
|
mock_create_llm.return_value = mock_llm
|
|
|
|
from llm import query_analysis
|
|
|
|
result = query_analysis(purpose="", index="")
|
|
|
|
assert isinstance(result, str)
|
|
assert "关键实体" in result
|
|
|
|
@patch("llm._create_llm")
|
|
@patch("llm.build_analysis_system_prompt")
|
|
def test_query_analysis_empty_response(self, mock_prompt, mock_create_llm):
|
|
"""测试 LLM 返回空字符串时不会崩溃。"""
|
|
mock_prompt.return_value = "prompt"
|
|
mock_llm = MagicMock()
|
|
mock_resp = MagicMock()
|
|
mock_resp.delta = ""
|
|
mock_llm.stream_chat.return_value = [mock_resp]
|
|
mock_create_llm.return_value = mock_llm
|
|
|
|
from llm import query_analysis
|
|
|
|
result = query_analysis(purpose="", index="")
|
|
|
|
assert result == ""
|