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

684 lines
30 KiB
Python

import json
from urllib.parse import parse_qs
import httpx
from yuqei_sdk import (
AiPlatformClient,
AiQualityLabelReviewUpdate,
AiQualityRuleConfigUpdate,
KnowledgeDocumentCreate,
KnowledgeImportBatchCreate,
KnowledgeReindexRequest,
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": "Use Article 19 and verify the contract term.",
"citations": [
{
"source_type": "law",
"title": "Labor Contract Law",
"reference": "Article 19 chunk-1",
"quote": "Probation period must stay within the statutory limit.",
"score": 0.82,
"locator": "chunk-1",
"document_id": "doc-1",
"chunk_id": "chunk-1",
"knowledge_base_id": "laws-cn",
"keyword_score": 0.78,
"vector_score": 0.87,
"hybrid_score": 0.82,
"retrieval_strategy": "hybrid",
"matched_terms": ["probation"],
}
],
"primary_sources": [
{
"source_type": "law",
"title": "Labor Contract Law",
"reference": "Article 19 chunk-1",
"quote": "Probation period must stay within the statutory limit.",
"score": 0.82,
"locator": "chunk-1",
"document_id": "doc-1",
"chunk_id": "chunk-1",
"knowledge_base_id": "laws-cn",
"keyword_score": 0.78,
"vector_score": 0.87,
"hybrid_score": 0.82,
"retrieval_strategy": "hybrid",
"matched_terms": ["probation"],
}
],
"confidence": 0.72,
},
)
client = AiPlatformClient(
"http://ai-platform.test",
transport=httpx.MockTransport(handler),
)
response = client.run_legal_qa(
LegalQaRequest(question="How should probation period be set?"),
context=RequestContext(trace_id="trace-1", tenant_id="tenant-1", user_id="user-1"),
)
assert response.citations[0].reference == "Article 19 chunk-1"
assert response.primary_sources[0].score == 0.82
assert response.primary_sources[0].hybrid_score == 0.82
assert response.primary_sources[0].retrieval_strategy == "hybrid"
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": "Search evidence first, then answer.",
"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": "Search evidence first, then answer.",
"enabled": True,
"updated_at": "2026-06-22T00:00:00Z",
}
],
)
if request.url.path == "/api/v1/quality-rules/default" and request.method == "PUT":
payload = json.loads(request.read().decode("utf-8"))
assert payload["min_source_count"] == 1
return httpx.Response(
200,
json={
"id": "default",
"min_citation_count": 1,
"min_source_count": 1,
"high_token_threshold": 9000,
"high_latency_ms": 12000,
"critical_latency_ms": 24000,
"warning_penalty": 7,
"error_penalty": 21,
"updated_at": "2026-06-22T00:00:00Z",
},
)
if request.url.path == "/api/v1/quality-rules/default":
return httpx.Response(
200,
json={
"id": "default",
"min_citation_count": 1,
"min_source_count": 2,
"high_token_threshold": 8000,
"high_latency_ms": 15000,
"critical_latency_ms": 30000,
"warning_penalty": 10,
"error_penalty": 30,
"updated_at": "2026-06-22T00:00:00Z",
},
)
if request.url.path == "/api/v1/knowledge-import-batches" and request.method == "POST":
return httpx.Response(
200,
json={
"id": "kb-batch-1",
"knowledge_base_id": "laws-cn",
"source_name": "civil-code-lease",
"source_type": "law_import",
"status": "indexed",
"documents_count": 0,
"chunks_count": 0,
"created_at": "2026-06-22T00:00:00Z",
"updated_at": "2026-06-22T00:00:00Z",
},
)
if request.url.path == "/api/v1/knowledge-import-batches":
return httpx.Response(
200,
json=[
{
"id": "kb-batch-1",
"knowledge_base_id": "laws-cn",
"source_name": "civil-code-lease",
"source_type": "law_import",
"status": "indexed",
"documents_count": 1,
"chunks_count": 1,
"created_at": "2026-06-22T00:00:00Z",
"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",
"import_batch_id": "kb-batch-1",
"knowledge_base_id": "laws-cn",
"title": "Civil Code",
"source_type": "law",
"reference": "Article 703",
"content": "A lease contract defines rent and use of leased property.",
"index_status": "indexed",
"created_at": "2026-06-22T00:00:00Z",
},
)
if request.url.path == "/api/v1/knowledge-documents/doc-1/chunks":
return httpx.Response(
200,
json=[
{
"id": "chunk-1",
"document_id": "doc-1",
"import_batch_id": "kb-batch-1",
"knowledge_base_id": "laws-cn",
"title": "Civil Code",
"source_type": "law",
"reference": "Article 703",
"content": "A lease contract defines rent and use of leased property.",
"chunk_index": 1,
"locator": "chunk-1",
"embedding_provider": "local-hash",
"embedding_model": "local-hash-v1",
"embedding_dimension": 64,
"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"] == ["lease"]
assert query["limit"] == ["3"]
return httpx.Response(
200,
json=[
{
"id": "doc-1",
"import_batch_id": "kb-batch-1",
"knowledge_base_id": "laws-cn",
"title": "Civil Code",
"source_type": "law",
"reference": "Article 703",
"content": "A lease contract defines rent and use of leased property.",
"index_status": "indexed",
"created_at": "2026-06-22T00:00:00Z",
}
],
)
if request.url.path == "/api/v1/knowledge-documents/reindex":
payload = json.loads(request.read().decode("utf-8"))
assert payload["knowledge_base_id"] == "laws-cn"
return httpx.Response(
200,
json={
"documents_count": 1,
"chunks_count": 1,
"embedding_provider": "openai-compatible",
"embedding_model": "text-embedding-prod",
"embedding_dimension": 1536,
"document_ids": ["doc-1"],
},
)
if request.url.path == "/api/v1/knowledge-chunks/search":
return httpx.Response(
200,
json=[
{
"chunk": {
"id": "chunk-1",
"document_id": "doc-1",
"import_batch_id": "kb-batch-1",
"knowledge_base_id": "laws-cn",
"title": "Civil Code",
"source_type": "law",
"reference": "Article 703",
"content": "A lease contract defines rent and use of leased property.",
"chunk_index": 1,
"locator": "chunk-1",
"embedding_provider": "local-hash",
"embedding_model": "local-hash-v1",
"embedding_dimension": 64,
"created_at": "2026-06-22T00:00:00Z",
},
"score": 0.82,
"keyword_score": 0.76,
"vector_score": 0.89,
"hybrid_score": 0.82,
"retrieval_strategy": "hybrid",
"matched_terms": ["lease"],
}
],
)
if request.url.path == "/api/v1/knowledge-chunks/search/explain":
query = parse_qs(request.url.query.decode())
assert query["keyword"] == ["lease"]
assert query["limit"] == ["3"]
assert query["candidate_limit"] == ["9"]
assert query["strategy"] == ["hybrid"]
return httpx.Response(
200,
json={
"keyword": "lease",
"normalized_terms": ["lease"],
"retrieval_strategy": "hybrid",
"knowledge_base_ids": [],
"limit": 3,
"candidate_limit": 9,
"candidate_count": 1,
"included_count": 1,
"retrieval_backend": "json-explain",
"embedding_provider": "local-hash",
"embedding_model": "local-hash-v1",
"query_embedding_dimension": 64,
"keyword_threshold": 0.0,
"vector_threshold": 0.35,
"hybrid_keyword_weight": 0.55,
"hybrid_vector_weight": 0.45,
"results": [
{
"chunk": {
"id": "chunk-1",
"document_id": "doc-1",
"import_batch_id": "kb-batch-1",
"knowledge_base_id": "laws-cn",
"title": "Civil Code",
"source_type": "law",
"reference": "Article 703",
"content": "A lease contract defines rent and use of leased property.",
"chunk_index": 1,
"locator": "chunk-1",
"embedding_provider": "local-hash",
"embedding_model": "local-hash-v1",
"embedding_dimension": 64,
"created_at": "2026-06-22T00:00:00Z",
},
"score": 0.82,
"keyword_score": 0.76,
"vector_score": 0.89,
"hybrid_score": 0.82,
"retrieval_strategy": "hybrid",
"matched_terms": ["lease"],
"included": True,
"rank": 1,
"filter_reason": None,
}
],
},
)
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",
"provider": "deepseek",
"model": "deepseek-chat",
"provider_status": "succeeded",
"provider_error": None,
"prompt_preview": "Search evidence first, then answer. Question: How should lease deposit refund be handled?",
"question": "How should lease deposit refund be handled?",
"answer": "The contract should clearly state deposit refund conditions.",
"citations_count": 1,
"retrieval_query": "How should lease deposit refund be handled?",
"retrieval_strategy": "hybrid",
"retrieved_chunk_ids": ["chunk-1"],
"retrieved_sources": [
{
"chunk_id": "chunk-1",
"document_id": "doc-1",
"knowledge_base_id": "laws-cn",
"title": "Civil Code",
"source_type": "law",
"reference": "Article 703",
"locator": "chunk-1",
"score": 0.82,
"keyword_score": 0.76,
"vector_score": 0.89,
"hybrid_score": 0.82,
"retrieval_strategy": "hybrid",
"matched_terms": ["lease"],
}
],
"prompt_tokens": 12,
"completion_tokens": 8,
"total_tokens": 20,
"estimated_cost": 0.00028,
"latency_ms": 12,
"quality_status": "passed",
"quality_score": 100,
"quality_labels": [
{
"code": "citation_present",
"severity": "info",
"title": "Citation present",
"detail": "The answer returned 1 citation.",
"review_status": "pending",
"review_note": None,
"reviewed_by": None,
"reviewed_at": None,
}
],
"quality_summary": "passed: 0 error(s), 0 warning(s), score 100.",
"status": "succeeded",
"created_at": "2026-06-22T00:00:00Z",
}
],
)
if request.url.path == "/api/v1/audit/quality-dashboard":
query = parse_qs(request.url.query.decode())
assert query["days"] == ["7"]
assert query["pending_limit"] == ["5"]
return httpx.Response(
200,
json={
"summary": {
"days": 7,
"total_runs": 1,
"passed_runs": 1,
"needs_review_runs": 0,
"failed_quality_runs": 0,
"fallback_runs": 0,
"source_issue_runs": 0,
"token_issue_runs": 0,
"latency_issue_runs": 0,
"pending_labels": 1,
"pass_rate": 1.0,
"fallback_rate": 0.0,
"source_issue_rate": 0.0,
"latency_issue_rate": 0.0,
},
"daily_metrics": [
{
"date": "2026-06-22",
"total_runs": 1,
"passed_runs": 1,
"needs_review_runs": 0,
"failed_quality_runs": 0,
"fallback_runs": 0,
"source_issue_runs": 0,
"token_issue_runs": 0,
"latency_issue_runs": 0,
"pending_labels": 1,
"pass_rate": 1.0,
"fallback_rate": 0.0,
"source_issue_rate": 0.0,
"latency_issue_rate": 0.0,
}
],
"pending_labels": [
{
"run_id": "run-1",
"label_code": "citation_present",
"severity": "info",
"title": "Citation present",
"detail": "The answer returned 1 citation.",
"review_status": "pending",
"created_at": "2026-06-22T00:00:00Z",
"question": "How should lease deposit refund be handled?",
"provider": "deepseek",
"model": "deepseek-chat",
"quality_status": "passed",
"quality_score": 100,
}
],
"attribution": {
"by_model": [
{
"dimension": "model",
"key": "deepseek / deepseek-chat",
"label": "deepseek / deepseek-chat",
"total_runs": 1,
"affected_runs": 1,
"warning_labels": 1,
"error_labels": 0,
"pending_labels": 1,
"fallback_runs": 0,
"average_quality_score": 90.0,
"top_label_codes": ["source_limited"],
"root_cause_hint": "该模型相关质量异常较集中,优先对比其他模型的回答质量和成本。",
}
],
"by_prompt_template": [],
"by_knowledge_base": [],
"by_label_code": [
{
"dimension": "label_code",
"key": "source_limited",
"label": "Source limited",
"total_runs": 1,
"affected_runs": 1,
"warning_labels": 1,
"error_labels": 0,
"pending_labels": 1,
"fallback_runs": 0,
"average_quality_score": 90.0,
"top_label_codes": ["source_limited"],
"root_cause_hint": "来源数量偏少,优先补充文档或放宽检索召回策略。",
}
],
},
},
)
if request.url.path == "/api/v1/audit/ai-runs/run-1/quality-labels/citation_present":
payload = json.loads(request.read().decode("utf-8"))
assert request.method == "PATCH"
assert payload["review_status"] == "resolved"
return httpx.Response(
200,
json={
"id": "run-1",
"operation": "legal_qa",
"provider_config_id": "provider-1",
"prompt_template_id": "prompt-1",
"provider": "deepseek",
"model": "deepseek-chat",
"provider_status": "succeeded",
"provider_error": None,
"prompt_preview": "Search evidence first, then answer. Question: How should lease deposit refund be handled?",
"question": "How should lease deposit refund be handled?",
"answer": "The contract should clearly state deposit refund conditions.",
"citations_count": 1,
"retrieval_query": "How should lease deposit refund be handled?",
"retrieval_strategy": "hybrid",
"retrieved_chunk_ids": ["chunk-1"],
"retrieved_sources": [],
"prompt_tokens": 12,
"completion_tokens": 8,
"total_tokens": 20,
"estimated_cost": 0.00028,
"latency_ms": 12,
"quality_status": "passed",
"quality_score": 100,
"quality_labels": [
{
"code": "citation_present",
"severity": "info",
"title": "Citation present",
"detail": "The answer returned 1 citation.",
"review_status": "resolved",
"review_note": "已确认依据有效。",
"reviewed_by": "admin",
"reviewed_at": "2026-06-22T00:01:00Z",
}
],
"quality_summary": "passed: 0 error(s), 0 warning(s), score 100.",
"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="Search evidence first.")
).version == "v2"
assert client.list_prompt_templates()[0].name == "legal_qa.default"
assert client.get_quality_rule_config().min_source_count == 2
assert client.upsert_quality_rule_config(
AiQualityRuleConfigUpdate(
min_citation_count=1,
min_source_count=1,
high_token_threshold=9000,
high_latency_ms=12000,
critical_latency_ms=24000,
warning_penalty=7,
error_penalty=21,
)
).warning_penalty == 7
assert client.create_knowledge_import_batch(
KnowledgeImportBatchCreate(
knowledge_base_id="laws-cn",
source_name="civil-code-lease",
source_type="law_import",
)
).id == "kb-batch-1"
assert client.list_knowledge_import_batches(knowledge_base_id="laws-cn")[0].chunks_count == 1
assert client.add_knowledge_document(
KnowledgeDocumentCreate(
import_batch_id="kb-batch-1",
knowledge_base_id="laws-cn",
title="Civil Code",
source_type="law",
reference="Article 703",
content="A lease contract defines rent and use of leased property.",
)
).reference == "Article 703"
assert client.list_knowledge_chunks("doc-1")[0].locator == "chunk-1"
assert client.list_knowledge_chunks("doc-1")[0].embedding_provider == "local-hash"
assert client.list_knowledge_chunks("doc-1")[0].embedding_model == "local-hash-v1"
assert client.search_knowledge_documents("lease", limit=3)[0].title == "Civil Code"
reindex_result = client.reindex_knowledge_documents(KnowledgeReindexRequest(knowledge_base_id="laws-cn"))
assert reindex_result.embedding_provider == "openai-compatible"
assert reindex_result.embedding_dimension == 1536
assert reindex_result.document_ids == ["doc-1"]
chunk_result = client.search_knowledge_chunks("lease", limit=3)
assert chunk_result[0].score == 0.82
assert chunk_result[0].hybrid_score == 0.82
explain = client.explain_knowledge_search("lease", limit=3, candidate_limit=9)
assert explain.embedding_provider == "local-hash"
assert explain.results[0].included is True
assert explain.results[0].rank == 1
audit = client.list_ai_run_audits(limit=5)[0]
assert audit.operation == "legal_qa"
assert audit.provider == "deepseek"
assert audit.model == "deepseek-chat"
assert audit.provider_status == "succeeded"
assert audit.provider_error is None
assert audit.prompt_preview is not None
assert "lease deposit" in audit.prompt_preview
assert audit.prompt_tokens == 12
assert audit.completion_tokens == 8
assert audit.total_tokens == 20
assert audit.estimated_cost == 0.00028
assert audit.retrieved_chunk_ids == ["chunk-1"]
assert audit.retrieved_sources[0].reference == "Article 703"
assert audit.quality_status == "passed"
assert audit.quality_score == 100
assert audit.quality_labels[0].code == "citation_present"
assert audit.quality_labels[0].review_status == "pending"
assert audit.quality_summary == "passed: 0 error(s), 0 warning(s), score 100."
dashboard = client.get_quality_dashboard(days=7, pending_limit=5)
assert dashboard.summary.pass_rate == 1.0
assert dashboard.daily_metrics[0].date == "2026-06-22"
assert dashboard.pending_labels[0].run_id == "run-1"
assert dashboard.attribution.by_model[0].key == "deepseek / deepseek-chat"
assert dashboard.attribution.by_label_code[0].top_label_codes == ["source_limited"]
assert dashboard.attribution.by_label_code[0].root_cause_hint is not None
updated_audit = client.update_ai_run_quality_label(
"run-1",
"citation_present",
AiQualityLabelReviewUpdate(
review_status="resolved",
review_note="已确认依据有效。",
reviewed_by="admin",
),
)
assert updated_audit.quality_labels[0].review_status == "resolved"
assert updated_audit.quality_labels[0].reviewed_at == "2026-06-22T00:01:00Z"
assert "/api/v1/audit/ai-runs" in seen_paths
assert "/api/v1/audit/quality-dashboard" in seen_paths
assert "/api/v1/quality-rules/default" in seen_paths