44 lines
1.6 KiB
Python
44 lines
1.6 KiB
Python
import httpx
|
|
|
|
from yuqei_sdk import AiPlatformClient, LegalQaRequest, RequestContext
|
|
|
|
|
|
def test_ai_platform_client_posts_legal_qa_with_context_headers() -> None:
|
|
captured_request: httpx.Request | None = None
|
|
|
|
def handler(request: httpx.Request) -> httpx.Response:
|
|
nonlocal captured_request
|
|
captured_request = request
|
|
return httpx.Response(
|
|
200,
|
|
json={
|
|
"answer": "依据劳动合同法第十九条，需要核对合同期限。",
|
|
"citations": [
|
|
{
|
|
"source_type": "law",
|
|
"title": "中华人民共和国劳动合同法",
|
|
"reference": "第十九条",
|
|
"quote": "劳动合同期限三个月以上不满一年的，试用期不得超过一个月。",
|
|
}
|
|
],
|
|
"confidence": 0.72,
|
|
},
|
|
)
|
|
|
|
client = AiPlatformClient(
|
|
"http://ai-platform.test",
|
|
transport=httpx.MockTransport(handler),
|
|
)
|
|
|
|
response = client.run_legal_qa(
|
|
LegalQaRequest(question="试用期如何约定？"),
|
|
context=RequestContext(trace_id="trace-1", tenant_id="tenant-1", user_id="user-1"),
|
|
)
|
|
|
|
assert response.citations[0].reference == "第十九条"
|
|
assert captured_request is not None
|
|
assert captured_request.url.path == "/api/v1/legal/qa"
|
|
assert captured_request.headers["X-Trace-Id"] == "trace-1"
|
|
assert captured_request.headers["X-Tenant-Id"] == "tenant-1"
|
|
assert captured_request.headers["X-User-Id"] == "user-1"
|