77 lines
3.1 KiB
Python
77 lines
3.1 KiB
Python
|
import json
|
||
|
|
|
||
|
|
import httpx
|
||
|
|
import pytest
|
||
|
|
from fastapi import FastAPI
|
||
|
|
|
||
|
|
from src.agent.collect import validated_reply
|
||
|
|
from src.agent.llm import LLMAPIError, LLMAuthError, LLMRateLimitError
|
||
|
|
from src.api.endpoints import router
|
||
|
|
from src.backends.langgraph import LangGraphBackend
|
||
|
|
from src.core.config import Config
|
||
|
|
from src.core.fastgpt_client import get_chat_backend
|
||
|
|
|
||
|
|
|
||
|
|
class FailingLLM:
|
||
|
|
def __init__(self, error_type, partial):
|
||
|
|
self.error_type = error_type
|
||
|
|
self.partial = partial
|
||
|
|
|
||
|
|
async def complete(self, *args, **kwargs):
|
||
|
|
raise self.error_type('private provider detail')
|
||
|
|
|
||
|
|
async def stream_tokens(self, *args, **kwargs):
|
||
|
|
if self.partial:
|
||
|
|
yield '<state>1002</state>'
|
||
|
|
yield 'partial reply'
|
||
|
|
raise self.error_type('private provider detail')
|
||
|
|
|
||
|
|
|
||
|
|
class Graph:
|
||
|
|
async def ainvoke(self, data, *, config, context):
|
||
|
|
await validated_reply(context.llm, 'system', 'user', '1002')
|
||
|
|
raise AssertionError('failed turn must not complete')
|
||
|
|
|
||
|
|
|
||
|
|
def sse_events(body):
|
||
|
|
events = []
|
||
|
|
for block in body.strip().split('\n\n'):
|
||
|
|
lines = block.splitlines()
|
||
|
|
kind = next(line[7:] for line in lines if line.startswith('event: '))
|
||
|
|
data = json.loads(next(line[6:] for line in lines if line.startswith('data: ')))
|
||
|
|
events.append((kind, data))
|
||
|
|
return events
|
||
|
|
|
||
|
|
|
||
|
|
@pytest.mark.asyncio
|
||
|
|
@pytest.mark.parametrize('partial', [False, True], ids=['before-text', 'after-text'])
|
||
|
|
@pytest.mark.parametrize('error_type,code,message', [
|
||
|
|
(LLMAuthError, '401', '认证失败'),
|
||
|
|
(LLMRateLimitError, '429', '请求过于频繁，请稍后重试'),
|
||
|
|
(LLMAPIError, '500', '大模型服务器无响应'),
|
||
|
|
])
|
||
|
|
async def test_http_sse_errors_match_non_stream_codes_and_terminate(
|
||
|
|
monkeypatch, partial, error_type, code, message,
|
||
|
|
):
|
||
|
|
# 独立协议测试 app：实际 LangGraphBackend/LLM 错误映射，不启动真实 provider。
|
||
|
|
monkeypatch.setattr(Config, 'AGENT_BACKEND', 'fastgpt')
|
||
|
|
app = FastAPI()
|
||
|
|
app.include_router(router)
|
||
|
|
backend = LangGraphBackend(graph=Graph(), llm=FailingLLM(error_type, partial))
|
||
|
|
app.dependency_overrides[get_chat_backend] = lambda: backend
|
||
|
|
payload = {'sessionId': 'error-audit', 'timeStamp': '20261022120000', 'text': 'hello'}
|
||
|
|
async with httpx.AsyncClient(transport=httpx.ASGITransport(app=app), base_url='http://test') as client:
|
||
|
|
response = await client.post('/chat?stream=true', json=payload)
|
||
|
|
assert response.status_code == 200 # 错误分类在 SSE data.code，不改已发送的 HTTP 头
|
||
|
|
assert response.headers['content-type'].startswith('text/event-stream')
|
||
|
|
events = sse_events(response.text)
|
||
|
|
assert events[-1] == ('error', {'code': code, 'msg': message})
|
||
|
|
assert sum(kind == 'error' for kind, _ in events) == 1
|
||
|
|
assert all(kind != 'done' for kind, _ in events)
|
||
|
|
assert any(kind == 'text_delta' for kind, _ in events) is partial
|
||
|
|
assert 'private provider detail' not in response.text
|
||
|
|
|
||
|
|
result = await client.post('/chat', json=payload)
|
||
|
|
assert result.json()['code'] == code
|
||
|
|
assert result.json()['msg'] == message
|