refactor: centralize llm provider transport
This commit is contained in:
leefer
@@ -0,0 +1,187 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import io
|
||||
import json
|
||||
import unittest
|
||||
import urllib.error
|
||||
from pathlib import Path
|
||||
from unittest.mock import patch
|
||||
|
||||
from assistant_agent import ReviewAssistantError, stream_review_assistant
|
||||
from backend.llm import transport
|
||||
from heaven_agent import HeavenAgentError, interpret_heaven
|
||||
from llm_strategy import LLMCompilerError, test_llm_connection
|
||||
from mentor_agent import MentorAgentError, MentorSkill, stream_with_mentor
|
||||
|
||||
|
||||
ROOT = Path(__file__).resolve().parents[1]
|
||||
|
||||
|
||||
class FakeResponse:
|
||||
def __init__(self, *, payload: bytes = b"", lines: list[bytes] | None = None) -> None:
|
||||
self.payload = payload
|
||||
self.lines = lines or []
|
||||
|
||||
def __enter__(self):
|
||||
return self
|
||||
|
||||
def __exit__(self, exc_type, exc_value, traceback):
|
||||
return False
|
||||
|
||||
def read(self) -> bytes:
|
||||
return self.payload
|
||||
|
||||
def __iter__(self):
|
||||
return iter(self.lines)
|
||||
|
||||
|
||||
class OpenAITransportTests(unittest.TestCase):
|
||||
def test_chat_completion_builds_one_openai_compatible_request(self) -> None:
|
||||
captured = {}
|
||||
response = FakeResponse(
|
||||
payload=json.dumps(
|
||||
{"choices": [{"message": {"content": "OK"}}]}
|
||||
).encode("utf-8")
|
||||
)
|
||||
|
||||
def open_request(request, timeout):
|
||||
captured["url"] = request.full_url
|
||||
captured["headers"] = request.headers
|
||||
captured["payload"] = json.loads(request.data.decode("utf-8"))
|
||||
captured["timeout"] = timeout
|
||||
return response
|
||||
|
||||
with patch("backend.llm.transport.urllib.request.urlopen", side_effect=open_request):
|
||||
result = transport.chat_completion(
|
||||
api_key="secret",
|
||||
base_url="https://example.test/v1/",
|
||||
model="model",
|
||||
messages=[{"role": "user", "content": "ping"}],
|
||||
timeout=17,
|
||||
user_agent="XiaobaiReviewWeb/test",
|
||||
)
|
||||
|
||||
self.assertEqual(result.content, "OK")
|
||||
self.assertGreaterEqual(result.latency_ms, 0)
|
||||
self.assertEqual(captured["url"], "https://example.test/v1/chat/completions")
|
||||
self.assertEqual(captured["payload"]["stream"], False)
|
||||
self.assertEqual(captured["headers"]["Authorization"], "Bearer secret")
|
||||
self.assertEqual(captured["timeout"], 17)
|
||||
|
||||
def test_stream_completion_parses_deltas_and_ignores_final_snapshot(self) -> None:
|
||||
response = FakeResponse(
|
||||
lines=[
|
||||
b'data: {"choices":[{"delta":{"content":"first"}}]}\n',
|
||||
b'data: {"choices":[{"delta":{"content":" second"}}]}\n',
|
||||
b'data: {"choices":[{"message":{"content":"first second"}}]}\n',
|
||||
b"data: [DONE]\n",
|
||||
]
|
||||
)
|
||||
with patch("backend.llm.transport.urllib.request.urlopen", return_value=response):
|
||||
chunks = list(
|
||||
transport.stream_chat_completion(
|
||||
api_key="secret",
|
||||
base_url="https://example.test/v1",
|
||||
model="model",
|
||||
messages=[],
|
||||
timeout=17,
|
||||
user_agent="XiaobaiReviewWeb/test",
|
||||
)
|
||||
)
|
||||
self.assertEqual(chunks, ["first", " second"])
|
||||
|
||||
def test_empty_stream_has_a_stable_transport_error(self) -> None:
|
||||
response = FakeResponse(lines=[b"data: [DONE]\n"])
|
||||
with patch("backend.llm.transport.urllib.request.urlopen", return_value=response):
|
||||
with self.assertRaises(transport.OpenAIEmptyResponseError):
|
||||
list(
|
||||
transport.stream_chat_completion(
|
||||
api_key="secret",
|
||||
base_url="https://example.test/v1",
|
||||
model="model",
|
||||
messages=[],
|
||||
timeout=17,
|
||||
user_agent="XiaobaiReviewWeb/test",
|
||||
)
|
||||
)
|
||||
|
||||
def test_http_error_keeps_code_and_sanitized_provider_detail(self) -> None:
|
||||
error = urllib.error.HTTPError(
|
||||
"https://example.test/v1/chat/completions",
|
||||
429,
|
||||
"rate limited",
|
||||
{},
|
||||
io.BytesIO(b'{"error":{"message":"capacity"}}'),
|
||||
)
|
||||
self.addCleanup(error.close)
|
||||
with patch("backend.llm.transport.urllib.request.urlopen", side_effect=error):
|
||||
with self.assertRaises(transport.OpenAIHTTPError) as caught:
|
||||
transport.chat_completion(
|
||||
api_key="secret",
|
||||
base_url="https://example.test/v1",
|
||||
model="model",
|
||||
messages=[],
|
||||
timeout=17,
|
||||
user_agent="XiaobaiReviewWeb/test",
|
||||
)
|
||||
self.assertEqual(caught.exception.code, 429)
|
||||
self.assertEqual(
|
||||
caught.exception.describe("模型调用失败"),
|
||||
"模型调用失败（HTTP 429）：capacity",
|
||||
)
|
||||
|
||||
def test_feature_agents_have_no_direct_provider_transport(self) -> None:
|
||||
paths = (
|
||||
"backend/features/mentor/agent.py",
|
||||
"backend/features/heaven/agent.py",
|
||||
"backend/features/review/agent.py",
|
||||
"backend/features/screener/compiler.py",
|
||||
)
|
||||
for relative in paths:
|
||||
source = (ROOT / relative).read_text(encoding="utf-8")
|
||||
with self.subTest(path=relative):
|
||||
self.assertNotIn("urllib.request", source)
|
||||
self.assertNotIn("/chat/completions", source)
|
||||
self.assertIn("llm_transport.", source)
|
||||
|
||||
|
||||
class FeatureErrorMappingTests(unittest.TestCase):
|
||||
def test_feature_specific_http_messages_are_preserved(self) -> None:
|
||||
error = transport.OpenAIHTTPError(429, "capacity")
|
||||
skill = MentorSkill(
|
||||
skill_id="test",
|
||||
name="测试老师",
|
||||
description="",
|
||||
tagline="",
|
||||
focus=(),
|
||||
content="",
|
||||
path=Path("SKILL.md"),
|
||||
)
|
||||
with patch(
|
||||
"mentor_agent.llm_transport.stream_chat_completion", side_effect=error
|
||||
):
|
||||
with self.assertRaisesRegex(
|
||||
MentorAgentError, "问师模型调用失败（HTTP 429）：capacity"
|
||||
):
|
||||
list(stream_with_mentor(skill, {}, "问题", [], "key", "https://x", "m"))
|
||||
with patch("heaven_agent.llm_transport.chat_completion", side_effect=error):
|
||||
with self.assertRaisesRegex(
|
||||
HeavenAgentError, "问天模型调用失败（HTTP 429）：capacity"
|
||||
):
|
||||
interpret_heaven("heart", {}, "key", "https://x", "m")
|
||||
with patch(
|
||||
"assistant_agent.llm_transport.stream_chat_completion", side_effect=error
|
||||
):
|
||||
with self.assertRaisesRegex(
|
||||
ReviewAssistantError, "智能解读服务暂不可用（429）"
|
||||
):
|
||||
list(stream_review_assistant({}, "问题", [], "key", "https://x", "m"))
|
||||
with patch("llm_strategy.llm_transport.chat_completion", side_effect=error):
|
||||
with self.assertRaisesRegex(
|
||||
LLMCompilerError, "模型连接测试失败（HTTP 429）：capacity"
|
||||
):
|
||||
test_llm_connection("key", "https://x", "m")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
Reference in New Issue
Block a user