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/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." 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/quality-rules/default" in seen_paths