Files
ai-contract-platform/yuqei-ai-sdk-python/tests/test_ai_platform.py
T

199 lines
7.7 KiB
Python

from urllib.parse import parse_qs
import httpx
from yuqei_sdk import (
AiPlatformClient,
KnowledgeDocumentCreate,
LegalQaRequest,
PromptTemplateCreate,
ProviderConfigCreate,
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"
def test_ai_platform_client_covers_admin_data_endpoints() -> None:
seen_paths: list[str] = []
def handler(request: httpx.Request) -> httpx.Response:
seen_paths.append(request.url.path)
if request.url.path == "/api/v1/model-configs" and request.method == "POST":
return httpx.Response(
200,
json={
"id": "provider-1",
"provider": "deepseek",
"model": "deepseek-chat",
"display_name": "DeepSeek Chat",
"base_url": None,
"temperature": 0.1,
"max_tokens": 4096,
"enabled": True,
"is_default": True,
"updated_at": "2026-06-22T00:00:00Z",
},
)
if request.url.path == "/api/v1/model-configs":
return httpx.Response(
200,
json=[
{
"id": "provider-1",
"provider": "deepseek",
"model": "deepseek-chat",
"display_name": "DeepSeek Chat",
"base_url": None,
"temperature": 0.1,
"max_tokens": 4096,
"enabled": True,
"is_default": True,
"updated_at": "2026-06-22T00:00:00Z",
}
],
)
if request.url.path == "/api/v1/prompt-templates" and request.method == "POST":
return httpx.Response(
200,
json={
"id": "prompt-1",
"name": "legal_qa.default",
"version": "v2",
"content": "先检索依据,再回答。",
"enabled": True,
"updated_at": "2026-06-22T00:00:00Z",
},
)
if request.url.path == "/api/v1/prompt-templates":
return httpx.Response(
200,
json=[
{
"id": "prompt-1",
"name": "legal_qa.default",
"version": "v2",
"content": "先检索依据,再回答。",
"enabled": True,
"updated_at": "2026-06-22T00:00:00Z",
}
],
)
if request.url.path == "/api/v1/knowledge-documents" and request.method == "POST":
return httpx.Response(
200,
json={
"id": "doc-1",
"knowledge_base_id": "laws-cn",
"title": "中华人民共和国民法典",
"source_type": "law",
"reference": "第七百零三条",
"content": "租赁合同是出租人将租赁物交付承租人使用、收益,承租人支付租金的合同。",
"created_at": "2026-06-22T00:00:00Z",
},
)
if request.url.path == "/api/v1/knowledge-documents/search":
query = parse_qs(request.url.query.decode())
assert query["keyword"] == ["租赁"]
assert query["limit"] == ["3"]
return httpx.Response(
200,
json=[
{
"id": "doc-1",
"knowledge_base_id": "laws-cn",
"title": "中华人民共和国民法典",
"source_type": "law",
"reference": "第七百零三条",
"content": "租赁合同是出租人将租赁物交付承租人使用、收益,承租人支付租金的合同。",
"created_at": "2026-06-22T00:00:00Z",
}
],
)
if request.url.path == "/api/v1/audit/ai-runs":
return httpx.Response(
200,
json=[
{
"id": "run-1",
"operation": "legal_qa",
"provider_config_id": "provider-1",
"prompt_template_id": "prompt-1",
"question": "租赁合同如何约定押金?",
"answer": "应明确押金返还条件。",
"citations_count": 1,
"latency_ms": 12,
"status": "succeeded",
"created_at": "2026-06-22T00:00:00Z",
}
],
)
return httpx.Response(404)
client = AiPlatformClient("http://ai-platform.test", transport=httpx.MockTransport(handler))
assert client.upsert_model_config(
ProviderConfigCreate(
provider="deepseek",
model="deepseek-chat",
display_name="DeepSeek Chat",
temperature=0.1,
is_default=True,
)
).model == "deepseek-chat"
assert client.list_model_configs()[0].provider == "deepseek"
assert client.upsert_prompt_template(
PromptTemplateCreate(name="legal_qa.default", version="v2", content="先检索依据,再回答。")
).version == "v2"
assert client.list_prompt_templates()[0].name == "legal_qa.default"
assert client.add_knowledge_document(
KnowledgeDocumentCreate(
knowledge_base_id="laws-cn",
title="中华人民共和国民法典",
source_type="law",
reference="第七百零三条",
content="租赁合同是出租人将租赁物交付承租人使用、收益,承租人支付租金的合同。",
)
).reference == "第七百零三条"
assert client.search_knowledge_documents("租赁", limit=3)[0].title == "中华人民共和国民法典"
assert client.list_ai_run_audits(limit=5)[0].operation == "legal_qa"
assert "/api/v1/audit/ai-runs" in seen_paths