import unittest import httpx from app.clients.shared_ai_platform import RemoteSharedAIPlatformClient from shared_ai_platform_sdk.client import SharedAIPlatformClient as VendoredSharedAiPlatformClient from shared_ai_platform_sdk.types import RequestContext, SSEEvent class VendoredSharedAiPlatformSdkTestCase(unittest.TestCase): def test_stream_conversation_uses_configured_api_prefix(self): seen_paths = [] def handler(request: httpx.Request) -> httpx.Response: seen_paths.append(request.url.path) return httpx.Response( 200, headers={"content-type": "text/event-stream"}, content=b'data: {"type":"end","conversation_id":"conv-1"}\n\n', ) transport = httpx.MockTransport(handler) http_client = httpx.Client(base_url="http://testserver", transport=transport, timeout=30.0) sdk = VendoredSharedAiPlatformClient( base_url="http://testserver", api_key="test-key", app_id="test-app", api_prefix="/api/v0.1", client=http_client, ) try: events = list( sdk.stream_conversation( RequestContext(tenant_id="tenant-1", user_id="user-1", user_role="admin"), message="hello", ) ) finally: http_client.close() self.assertEqual(seen_paths, ["/api/v0.1/conversations/stream"]) self.assertEqual(events[0].payload["conversation_id"], "conv-1") def test_stream_conversation_supports_split_data_and_json_lines(self): sse_payload = ( 'data:\n' '{"type":"delta","content":"hello"}\n\n' 'data:\n' '{"type":"end","conversation_id":"conv-1"}\n\n' ).encode("utf-8") def handler(request: httpx.Request) -> httpx.Response: return httpx.Response( 200, headers={"content-type": "text/event-stream"}, content=sse_payload, ) transport = httpx.MockTransport(handler) http_client = httpx.Client(base_url="http://testserver", transport=transport, timeout=30.0) sdk = VendoredSharedAiPlatformClient( base_url="http://testserver", api_key="test-key", app_id="test-app", client=http_client, ) try: events = list( sdk.stream_conversation( RequestContext(tenant_id="tenant-1", user_id="user-1", user_role="admin"), message="hello", ) ) finally: http_client.close() self.assertEqual([event.event_type for event in events], ["delta", "end"]) self.assertEqual(events[0].payload["content"], "hello") self.assertEqual(events[1].payload["conversation_id"], "conv-1") class RemoteSharedAiPlatformClientStreamTestCase(unittest.TestCase): def test_remote_stream_ignores_start_and_usage_and_maps_delta(self): calls = [] class StubSdk: def stream_conversation(self, context, *, message, chat_context=None, provider=None, template_key=None, conversation_id=None): calls.append( { "message": message, "chat_context": chat_context, "template_key": template_key, "conversation_id": conversation_id, } ) yield SSEEvent(event_type="start", payload={"type": "start", "conversation_id": "conv-1"}) yield SSEEvent(event_type="delta", payload={"type": "delta", "content": "hello"}) yield SSEEvent(event_type="usage", payload={"type": "usage", "usage": {"provider_mode": "stream"}}) yield SSEEvent(event_type="sources", payload={"type": "sources", "sources": [{"title": "source", "content": "body"}]}) yield SSEEvent(event_type="end", payload={"type": "end", "conversation_id": "conv-1"}) def run_legal_review(self, context, *, title, text, contract_type=None, review_as="party_a", provider=None): return { "basic_info": { "contract_type": contract_type or "采购合同", "contract_title": title, "party_a": "甲方", "party_b": "乙方", "amount": "100万元", "term": "自2026-01-01至2026-12-31", }, "risk_summary": {"high": 1, "medium": 0, "low": 0, "overall": "high"}, "clauses": [{"title": "付款条款", "content": "签署后5日内支付50%", "clause_type": "payment", "clause_label": "付款条款", "summary": "签署后5日内支付50%"}], "risk_items": [{"title": "存在自动续约安排", "level": "high", "description": "需核查自动续约", "suggestion": "补充退出机制"}], "review_opinion": "建议人工复核", "recommendation": "manual_review", "confidence_summary": {"provider_mode": "provider-enhanced"}, "contract_structure": {"title": title, "payment_milestones": [], "clause_relationships": [], "missing_sections": [], "clause_type_counts": {"payment": 1}}, } client = RemoteSharedAIPlatformClient( base_url="http://testserver", app_id="test-app", api_key="test-key", timeout_seconds=30.0, ) client.sdk = StubSdk() current_user = {"id": "user-1", "tenant_id": "tenant-1", "role": "admin"} chunks = list(client.stream_answer(current_user, None, "hello", {"metadata": {"reviewTitle": "demo"}}, template_key="legal_chat_default")) self.assertEqual( chunks, [ 'data: {"type": "text", "content": "hello"}\n\n', 'data: {"type": "sources", "sources": [{"title": "source", "content": "body"}]}\n\n', 'data: {"type": "end", "conversation_id": "conv-1"}\n\n', ], ) self.assertEqual(calls[0]["template_key"], "legal_chat_default") self.assertEqual(calls[0]["chat_context"]["metadata"]["reviewTitle"], "demo") def test_remote_document_payloads_are_normalized_for_source_metadata(self): class StubSdk: def list_documents(self, context, knowledge_base_id, *, page=1, page_size=20): return { "items": [ { "id": "doc-1", "knowledge_base_id": knowledge_base_id, "file_name": "civil-code.md", "mime_type": "text/markdown", "file_size": 128, "status": "indexed", "metadata": { "source_type": "statute", "issuer": "全国人民代表大会", "effective_date": "2021-01-01", "jurisdiction": "CN", "tags": ["民法典", "合同"], }, } ], "page": page, "page_size": page_size, "total": 1, } def upload_document(self, context, knowledge_base_id, *, file_name, mime_type, content): return { "id": "doc-2", "knowledge_base_id": knowledge_base_id, "file_name": file_name, "mime_type": mime_type, "file_size": len(content), "status": "indexed", "tags_json": ["制度"], "metadata_json": { "source_type": "policy", "authority": "企业法务部", "effective_date": "2026-05-01", "region": "CN-SH", }, } def retry_document(self, context, knowledge_base_id, document_id, *, file_name=None, mime_type=None, content=None): return { "id": document_id, "knowledge_base_id": knowledge_base_id, "file_name": file_name or "retry.txt", "mime_type": mime_type or "text/plain", "file_size": len(content or b""), "status": "indexed", "metadata": { "source_type": "judicial-interpretation", "authority": "最高人民法院", "effective_date": "2024-02-01", "region": "CN", }, } client = RemoteSharedAIPlatformClient( base_url="http://testserver", app_id="test-app", api_key="test-key", timeout_seconds=30.0, ) client.sdk = StubSdk() current_user = {"id": "user-1", "tenant_id": "tenant-1", "role": "admin"} listed = client.list_documents(current_user, "kb-1", 1, 20) uploaded = client.upload_document(current_user, "kb-1", "policy.md", "text/markdown", b"hello") retried = client.retry_document(current_user, "kb-1", "doc-3", file_name="retry.md", mime_type="text/markdown", content=b"world") self.assertEqual(listed["items"][0]["source_type"], "statute") self.assertEqual(listed["items"][0]["authority"], "全国人民代表大会") self.assertEqual(listed["items"][0]["effective_date"], "2021-01-01") self.assertEqual(listed["items"][0]["region"], "CN") self.assertEqual(listed["items"][0]["tags"], ["民法典", "合同"]) self.assertEqual(uploaded["authority"], "企业法务部") self.assertEqual(uploaded["region"], "CN-SH") self.assertEqual(uploaded["tags"], ["制度"]) self.assertEqual(retried["source_type"], "judicial-interpretation") self.assertEqual(retried["authority"], "最高人民法院") def test_remote_knowledge_bases_are_sorted_latest_first(self): class StubSdk: def list_knowledge_bases(self, context, *, page=1, page_size=20): return { "items": [ { "id": "kb-old", "name": "old", "document_count": 1, "created_at": "2026-05-12T00:00:00Z", "updated_at": "2026-05-12T00:00:00Z", }, { "id": "kb-new", "name": "new", "document_count": 2, "created_at": "2026-05-13T00:00:00Z", "updated_at": "2026-05-13T00:00:00Z", }, ], "page": page, "page_size": page_size, "total": 2, } client = RemoteSharedAIPlatformClient( base_url="http://testserver", app_id="test-app", api_key="test-key", timeout_seconds=30.0, ) client.sdk = StubSdk() payload = client.list_knowledge_bases({"id": "user-1", "tenant_id": "tenant-1", "role": "admin"}, 1, 100) self.assertEqual([item["id"] for item in payload["items"]], ["kb-new", "kb-old"]) self.assertEqual(payload["items"][0]["retrieval_mode"], "hybrid") def test_remote_knowledge_base_publish_chain_calls_sdk(self): class StubSdk: def list_knowledge_base_versions(self, context, knowledge_base_id): return [{"id": "kbver-1", "knowledge_base_id": knowledge_base_id, "version_label": "v1"}] def publish_knowledge_base(self, context, knowledge_base_id): return {"id": knowledge_base_id, "status": "active", "current_version_label": "v1"} def unpublish_knowledge_base(self, context, knowledge_base_id): return {"id": knowledge_base_id, "status": "archived", "current_version_label": "v1"} def rollback_knowledge_base(self, context, knowledge_base_id, *, target_version_label, change_summary=None): return { "id": knowledge_base_id, "status": "active", "current_version_label": target_version_label, "change_summary": change_summary, } client = RemoteSharedAIPlatformClient( base_url="http://testserver", app_id="test-app", api_key="test-key", timeout_seconds=30.0, ) client.sdk = StubSdk() current_user = {"id": "user-1", "tenant_id": "tenant-1", "role": "admin"} versions = client.list_knowledge_base_versions(current_user, "kb-1") published = client.publish_knowledge_base(current_user, "kb-1") unpublished = client.unpublish_knowledge_base(current_user, "kb-1") rolled_back = client.rollback_knowledge_base(current_user, "kb-1", "v1", "rollback") self.assertEqual(versions[0]["version_label"], "v1") self.assertEqual(published["current_version_label"], "v1") self.assertEqual(unpublished["status"], "archived") self.assertEqual(rolled_back["current_version_label"], "v1") if __name__ == "__main__": unittest.main()