453 lines
19 KiB
Python
453 lines
19 KiB
Python
import json
|
|
from urllib.parse import parse_qs
|
|
|
|
import httpx
|
|
|
|
from yuqei_sdk import (
|
|
AiPlatformClient,
|
|
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/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,
|
|
"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,
|
|
"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.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_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 "/api/v1/audit/ai-runs" in seen_paths
|