From b319d41787da38c91163ff4e4851ceaf54857973 Mon Sep 17 00:00:00 2001 From: yuqei Date: Mon, 22 Jun 2026 21:20:09 +0800 Subject: [PATCH] feat: enhance ai knowledge retrieval --- ...wledge-retrieval-enhancement-2026-06-22.md | 109 ++++++ .../packages/ai-openapi/openapi.yaml | 226 +++++++++++++ .../services/ai-platform-api/README.md | 3 + ...2_0002_add_knowledge_batches_and_chunks.py | 78 +++++ .../src/yuqei_ai_platform_api/legal_qa.py | 8 +- .../src/yuqei_ai_platform_api/main.py | 114 ++++++- .../src/yuqei_ai_platform_api/models.py | 41 +++ .../src/yuqei_ai_platform_api/repository.py | 319 ++++++++++++++++++ .../tests/test_platform_data.py | 109 ++++-- yuqei-ai-sdk-python/README.md | 2 +- yuqei-ai-sdk-python/src/yuqei_sdk/__init__.py | 8 + .../src/yuqei_sdk/ai_platform.py | 113 ++++++- yuqei-ai-sdk-python/tests/test_ai_platform.py | 155 +++++++-- 13 files changed, 1223 insertions(+), 62 deletions(-) create mode 100644 project-docs/V2-ai-platform-knowledge-retrieval-enhancement-2026-06-22.md create mode 100644 yuqei-ai-platform/services/ai-platform-api/migrations/versions/20260622_0002_add_knowledge_batches_and_chunks.py diff --git a/project-docs/V2-ai-platform-knowledge-retrieval-enhancement-2026-06-22.md b/project-docs/V2-ai-platform-knowledge-retrieval-enhancement-2026-06-22.md new file mode 100644 index 00000000..86d603a1 --- /dev/null +++ b/project-docs/V2-ai-platform-knowledge-retrieval-enhancement-2026-06-22.md @@ -0,0 +1,109 @@ +# V2 AI Platform Knowledge Retrieval Enhancement + +日期：2026-06-22 + +## 本轮目标 + +本轮推进 AI 中台知识库检索增强，把上一轮“整篇知识文档关键词检索”升级为更接近 RAG 的基础结构： + +- 导入批次。 +- 文档切分 chunk。 +- chunk 级检索结果。 +- 检索命中分数。 +- 来源定位。 +- 法律问答返回结构化主要依据。 + +## 已完成 + +1. 数据模型 + - 新增 `KnowledgeImportBatchModel`。 + - 新增 `KnowledgeChunkModel`。 + - `KnowledgeDocumentModel` 增加： + - `import_batch_id` + - `index_status` + +2. 数据库迁移 + - 新增 Alembic migration： + - `20260622_0002_add_knowledge_batches_and_chunks.py` + - 支持升级/降级： + - `knowledge_import_batches` + - `knowledge_chunks` + - `knowledge_documents.import_batch_id` + - `knowledge_documents.index_status` + +3. Repository 能力 + - 创建知识导入批次。 + - 列出知识导入批次。 + - 新增知识文档时自动切分 chunk。 + - 列出文档 chunk。 + - chunk 级关键词检索。 + - 返回检索分数、命中词、locator、document_id、chunk_id。 + +4. API 能力 + - `GET /api/v1/knowledge-import-batches` + - `POST /api/v1/knowledge-import-batches` + - `GET /api/v1/knowledge-documents/{document_id}/chunks` + - `GET /api/v1/knowledge-chunks/search` + - 旧接口 `GET /api/v1/knowledge-documents/search` 保持兼容。 + +5. 法律问答增强 + - 法律问答优先使用 chunk 检索。 + - `LegalCitation` 增加： + - `score` + - `locator` + - `document_id` + - `chunk_id` + - `knowledge_base_id` + - `LegalQaResponse` 增加 `primary_sources`。 + - 旧字段 `citations` 保持兼容。 + +6. SDK 能力 + - 新增知识导入批次类型和方法。 + - 新增 chunk 类型和方法。 + - 新增 scored chunk search 方法。 + - 法律问答响应支持 `primary_sources`。 + +## 验证结果 + +已通过： + +```powershell +python -m pytest +``` + +当前结果： + +```text +21 passed +``` + +覆盖点： + +- 知识导入批次创建和统计更新。 +- 文档自动切分 chunk。 +- chunk 搜索返回 score 和 locator。 +- SQLite 持久化仓库可跨实例读取批次、chunk 和审计。 +- 法律问答返回 `primary_sources`。 +- SDK 覆盖新增知识库检索能力。 + +## 当前边界 + +本轮仍是关键词检索，不是正式向量检索。 + +尚未完成： + +- embedding 生成。 +- 向量库或 pgvector。 +- 混合检索排序。 +- 文档上传解析、OCR、复杂 DOCX/PDF 切分。 +- AI 审计记录命中文档 ID 列表、token、费用、provider 错误详情。 + +## 下一轮建议 + +下一轮进入知识库索引质量增强： + +1. 增加 embedding 字段和向量检索适配层。 +2. 引入 `pgvector` 或独立向量库的技术选择和迁移。 +3. 增加检索 explain 信息：关键词分、向量分、最终分。 +4. AI 审计记录 `hit_chunk_ids`、provider 错误、token 和成本。 +5. 法律问答输出按“结论、主要依据、风险提示、建议动作”结构化。 diff --git a/yuqei-ai-platform/packages/ai-openapi/openapi.yaml b/yuqei-ai-platform/packages/ai-openapi/openapi.yaml index 83181a79..2410d04c 100644 --- a/yuqei-ai-platform/packages/ai-openapi/openapi.yaml +++ b/yuqei-ai-platform/packages/ai-openapi/openapi.yaml @@ -221,6 +221,68 @@ paths: application/json: schema: $ref: "#/components/schemas/KnowledgeDocument" + /api/v1/knowledge-import-batches: + get: + operationId: listKnowledgeImportBatches + summary: List knowledge import batches + parameters: + - name: knowledge_base_id + in: query + required: false + schema: + type: string + - name: limit + in: query + required: false + schema: + type: integer + default: 20 + minimum: 1 + maximum: 100 + responses: + "200": + description: Knowledge import batches + content: + application/json: + schema: + type: array + items: + $ref: "#/components/schemas/KnowledgeImportBatch" + post: + operationId: createKnowledgeImportBatch + summary: Create a knowledge import batch + requestBody: + required: true + content: + application/json: + schema: + $ref: "#/components/schemas/KnowledgeImportBatchCreate" + responses: + "200": + description: Knowledge import batch + content: + application/json: + schema: + $ref: "#/components/schemas/KnowledgeImportBatch" + /api/v1/knowledge-documents/{document_id}/chunks: + get: + operationId: listKnowledgeChunks + summary: List chunks for a knowledge document + parameters: + - name: document_id + in: path + required: true + schema: + type: string + responses: + "200": + description: Knowledge chunks + content: + application/json: + schema: + type: array + items: + $ref: "#/components/schemas/KnowledgeChunk" /api/v1/knowledge-documents/search: get: operationId: searchKnowledgeDocuments @@ -240,6 +302,11 @@ paths: default: 5 minimum: 1 maximum: 20 + - name: knowledge_base_id + in: query + required: false + schema: + type: string responses: "200": description: Knowledge document matches @@ -249,6 +316,39 @@ paths: type: array items: $ref: "#/components/schemas/KnowledgeDocument" + /api/v1/knowledge-chunks/search: + get: + operationId: searchKnowledgeChunks + summary: Search knowledge chunks with relevance score and locator + parameters: + - name: keyword + in: query + required: false + schema: + type: string + default: "" + - name: knowledge_base_id + in: query + required: false + schema: + type: string + - name: limit + in: query + required: false + schema: + type: integer + default: 5 + minimum: 1 + maximum: 20 + responses: + "200": + description: Knowledge chunk matches + content: + application/json: + schema: + type: array + items: + $ref: "#/components/schemas/KnowledgeSearchResult" /api/v1/audit/ai-runs: get: operationId: listAiRuns @@ -350,6 +450,22 @@ components: url: type: string nullable: true + score: + type: number + format: float + nullable: true + locator: + type: string + nullable: true + document_id: + type: string + nullable: true + chunk_id: + type: string + nullable: true + knowledge_base_id: + type: string + nullable: true LegalQaRequest: type: object required: @@ -376,6 +492,10 @@ components: type: array items: $ref: "#/components/schemas/LegalCitation" + primary_sources: + type: array + items: + $ref: "#/components/schemas/LegalCitation" confidence: type: number format: float @@ -600,6 +720,9 @@ components: properties: id: type: string + import_batch_id: + type: string + nullable: true knowledge_base_id: type: string title: @@ -611,6 +734,10 @@ components: type: string content: type: string + index_status: + type: string + enum: [pending, indexing, indexed, failed] + default: indexed created_at: type: string format: date-time @@ -620,6 +747,9 @@ components: - title - content properties: + import_batch_id: + type: string + nullable: true knowledge_base_id: type: string default: default @@ -633,6 +763,102 @@ components: default: "" content: type: string + index_status: + type: string + default: indexed + KnowledgeImportBatchCreate: + type: object + required: + - source_name + properties: + knowledge_base_id: + type: string + default: default + source_name: + type: string + source_type: + type: string + default: manual_text + status: + type: string + default: indexed + KnowledgeImportBatch: + allOf: + - $ref: "#/components/schemas/KnowledgeImportBatchCreate" + - type: object + required: + - id + - documents_count + - chunks_count + - created_at + - updated_at + properties: + id: + type: string + documents_count: + type: integer + chunks_count: + type: integer + created_at: + type: string + format: date-time + updated_at: + type: string + format: date-time + KnowledgeChunk: + type: object + required: + - id + - document_id + - knowledge_base_id + - title + - source_type + - reference + - content + - chunk_index + - locator + - created_at + properties: + id: + type: string + document_id: + type: string + import_batch_id: + type: string + nullable: true + knowledge_base_id: + type: string + title: + type: string + source_type: + type: string + reference: + type: string + content: + type: string + chunk_index: + type: integer + locator: + type: string + created_at: + type: string + format: date-time + KnowledgeSearchResult: + type: object + required: + - chunk + - score + - matched_terms + properties: + chunk: + $ref: "#/components/schemas/KnowledgeChunk" + score: + type: number + format: float + matched_terms: + type: array + items: + type: string AiRun: type: object required: diff --git a/yuqei-ai-platform/services/ai-platform-api/README.md b/yuqei-ai-platform/services/ai-platform-api/README.md index 636c0fc6..8602c2ef 100644 --- a/yuqei-ai-platform/services/ai-platform-api/README.md +++ b/yuqei-ai-platform/services/ai-platform-api/README.md @@ -43,8 +43,11 @@ Core MVP endpoints: - `GET|POST /api/v1/model-configs` - `GET|POST /api/v1/prompt-templates` +- `GET|POST /api/v1/knowledge-import-batches` - `POST /api/v1/knowledge-documents` +- `GET /api/v1/knowledge-documents/{document_id}/chunks` - `GET /api/v1/knowledge-documents/search` +- `GET /api/v1/knowledge-chunks/search` - `POST /api/v1/legal/qa` - `GET /api/v1/audit/ai-runs` diff --git a/yuqei-ai-platform/services/ai-platform-api/migrations/versions/20260622_0002_add_knowledge_batches_and_chunks.py b/yuqei-ai-platform/services/ai-platform-api/migrations/versions/20260622_0002_add_knowledge_batches_and_chunks.py new file mode 100644 index 00000000..0527be61 --- /dev/null +++ b/yuqei-ai-platform/services/ai-platform-api/migrations/versions/20260622_0002_add_knowledge_batches_and_chunks.py @@ -0,0 +1,78 @@ +"""add knowledge batches and chunks + +Revision ID: 20260622_0002 +Revises: 20260622_0001 +Create Date: 2026-06-22 + +""" +from typing import Sequence, Union + +from alembic import op +import sqlalchemy as sa + + +revision: str = "20260622_0002" +down_revision: Union[str, None] = "20260622_0001" +branch_labels: Union[str, Sequence[str], None] = None +depends_on: Union[str, Sequence[str], None] = None + + +def upgrade() -> None: + op.add_column("knowledge_documents", sa.Column("import_batch_id", sa.String(length=64), nullable=True)) + op.add_column( + "knowledge_documents", + sa.Column("index_status", sa.String(length=32), nullable=False, server_default="indexed"), + ) + op.create_index("ix_knowledge_documents_import_batch_id", "knowledge_documents", ["import_batch_id"]) + + op.create_table( + "knowledge_import_batches", + sa.Column("id", sa.String(length=64), nullable=False), + sa.Column("knowledge_base_id", sa.String(length=64), nullable=False), + sa.Column("source_name", sa.String(length=256), nullable=False), + sa.Column("source_type", sa.String(length=64), nullable=False), + sa.Column("status", sa.String(length=32), nullable=False), + sa.Column("documents_count", sa.Integer(), nullable=False), + sa.Column("chunks_count", sa.Integer(), nullable=False), + sa.Column("created_at", sa.DateTime(timezone=True), nullable=False), + sa.Column("updated_at", sa.DateTime(timezone=True), nullable=False), + sa.PrimaryKeyConstraint("id"), + ) + op.create_index( + "ix_knowledge_import_batches_knowledge_base_id", + "knowledge_import_batches", + ["knowledge_base_id"], + ) + + op.create_table( + "knowledge_chunks", + sa.Column("id", sa.String(length=64), nullable=False), + sa.Column("document_id", sa.String(length=64), nullable=False), + sa.Column("import_batch_id", sa.String(length=64), nullable=True), + sa.Column("knowledge_base_id", sa.String(length=64), nullable=False), + sa.Column("title", sa.String(length=256), nullable=False), + sa.Column("source_type", sa.String(length=64), nullable=False), + sa.Column("reference", sa.String(length=256), nullable=False), + sa.Column("content", sa.Text(), nullable=False), + sa.Column("chunk_index", sa.Integer(), nullable=False), + sa.Column("locator", sa.String(length=128), nullable=False), + sa.Column("created_at", sa.DateTime(timezone=True), nullable=False), + sa.PrimaryKeyConstraint("id"), + ) + op.create_index("ix_knowledge_chunks_document_id", "knowledge_chunks", ["document_id"]) + op.create_index("ix_knowledge_chunks_import_batch_id", "knowledge_chunks", ["import_batch_id"]) + op.create_index("ix_knowledge_chunks_knowledge_base_id", "knowledge_chunks", ["knowledge_base_id"]) + + +def downgrade() -> None: + op.drop_index("ix_knowledge_chunks_knowledge_base_id", table_name="knowledge_chunks") + op.drop_index("ix_knowledge_chunks_import_batch_id", table_name="knowledge_chunks") + op.drop_index("ix_knowledge_chunks_document_id", table_name="knowledge_chunks") + op.drop_table("knowledge_chunks") + + op.drop_index("ix_knowledge_import_batches_knowledge_base_id", table_name="knowledge_import_batches") + op.drop_table("knowledge_import_batches") + + op.drop_index("ix_knowledge_documents_import_batch_id", table_name="knowledge_documents") + op.drop_column("knowledge_documents", "index_status") + op.drop_column("knowledge_documents", "import_batch_id") diff --git a/yuqei-ai-platform/services/ai-platform-api/src/yuqei_ai_platform_api/legal_qa.py b/yuqei-ai-platform/services/ai-platform-api/src/yuqei_ai_platform_api/legal_qa.py index 249acd57..1cf1f93e 100644 --- a/yuqei-ai-platform/services/ai-platform-api/src/yuqei_ai_platform_api/legal_qa.py +++ b/yuqei-ai-platform/services/ai-platform-api/src/yuqei_ai_platform_api/legal_qa.py @@ -13,11 +13,17 @@ class LegalCitation(BaseModel): reference: str quote: str url: str | None = None + score: float | None = None + locator: str | None = None + document_id: str | None = None + chunk_id: str | None = None + knowledge_base_id: str | None = None class LegalQaResponse(BaseModel): answer: str citations: list[LegalCitation] + primary_sources: list[LegalCitation] = Field(default_factory=list) confidence: float | None = None @@ -128,4 +134,4 @@ def _with_retrieved_citations( seen.add(key) deduped.append(citation) - return response.model_copy(update={"citations": deduped}) + return response.model_copy(update={"citations": deduped, "primary_sources": deduped[:3]}) diff --git a/yuqei-ai-platform/services/ai-platform-api/src/yuqei_ai_platform_api/main.py b/yuqei-ai-platform/services/ai-platform-api/src/yuqei_ai_platform_api/main.py index 2d101885..c3d5f36b 100644 --- a/yuqei-ai-platform/services/ai-platform-api/src/yuqei_ai_platform_api/main.py +++ b/yuqei-ai-platform/services/ai-platform-api/src/yuqei_ai_platform_api/main.py @@ -1,5 +1,5 @@ -from typing import Literal from time import perf_counter +from typing import Literal from uuid import uuid4 from fastapi import FastAPI, Query @@ -7,19 +7,24 @@ from pydantic import BaseModel from yuqei_ai_platform_api import __version__ from yuqei_ai_platform_api.config import AiPlatformSettings, get_settings -from yuqei_ai_platform_api.legal_qa import LegalQaRequest, LegalQaResponse, answer_legal_question +from yuqei_ai_platform_api.legal_qa import LegalCitation, LegalQaRequest, LegalQaResponse, answer_legal_question from yuqei_ai_platform_api.logging import configure_logging from yuqei_ai_platform_api.prompts import render_legal_qa_prompt from yuqei_ai_platform_api.providers import ProviderRequest, build_provider_adapter from yuqei_ai_platform_api.repository import ( AiPlatformRepository, AiRunAudit, + KnowledgeChunk, KnowledgeDocument, KnowledgeDocumentCreate, + KnowledgeImportBatch, + KnowledgeImportBatchCreate, + KnowledgeSearchResult, PromptTemplate, PromptTemplateCreate, ProviderConfig, ProviderConfigCreate, + chunk_to_citation, document_to_citation, ) from yuqei_ai_platform_api.state import get_repository @@ -103,13 +108,64 @@ def create_app( def add_knowledge_document(payload: KnowledgeDocumentCreate) -> KnowledgeDocument: return store.add_knowledge_document(payload) + @app.post( + f"{resolved_settings.api_prefix}/knowledge-import-batches", + response_model=KnowledgeImportBatch, + tags=["knowledge"], + ) + def create_knowledge_import_batch(payload: KnowledgeImportBatchCreate) -> KnowledgeImportBatch: + return store.create_knowledge_import_batch(payload) + + @app.get( + f"{resolved_settings.api_prefix}/knowledge-import-batches", + response_model=list[KnowledgeImportBatch], + tags=["knowledge"], + ) + def list_knowledge_import_batches( + knowledge_base_id: str | None = None, + limit: int = Query(20, ge=1, le=100), + ) -> list[KnowledgeImportBatch]: + return store.list_knowledge_import_batches(knowledge_base_id=knowledge_base_id, limit=limit) + + @app.get( + f"{resolved_settings.api_prefix}/knowledge-documents/{{document_id}}/chunks", + response_model=list[KnowledgeChunk], + tags=["knowledge"], + ) + def list_knowledge_chunks(document_id: str) -> list[KnowledgeChunk]: + return store.list_knowledge_chunks(document_id) + @app.get( f"{resolved_settings.api_prefix}/knowledge-documents/search", response_model=list[KnowledgeDocument], tags=["knowledge"], ) - def search_knowledge_documents(keyword: str = "", limit: int = Query(5, ge=1, le=20)) -> list[KnowledgeDocument]: - return store.search_knowledge(keyword, limit=limit) + def search_knowledge_documents( + keyword: str = "", + knowledge_base_id: str | None = None, + limit: int = Query(5, ge=1, le=20), + ) -> list[KnowledgeDocument]: + return store.search_knowledge( + keyword, + knowledge_base_ids=[knowledge_base_id] if knowledge_base_id else None, + limit=limit, + ) + + @app.get( + f"{resolved_settings.api_prefix}/knowledge-chunks/search", + response_model=list[KnowledgeSearchResult], + tags=["knowledge"], + ) + def search_knowledge_chunks( + keyword: str = "", + knowledge_base_id: str | None = None, + limit: int = Query(5, ge=1, le=20), + ) -> list[KnowledgeSearchResult]: + return store.search_knowledge_chunks( + keyword, + knowledge_base_ids=[knowledge_base_id] if knowledge_base_id else None, + limit=limit, + ) @app.get( f"{resolved_settings.api_prefix}/audit/ai-runs", @@ -126,8 +182,7 @@ def create_app( ) def run_legal_qa(request: LegalQaRequest) -> LegalQaResponse: started_at = perf_counter() - documents = _search_citation_documents(store, request) - citations = [document_to_citation(document) for document in documents] + citations = _search_citations(store, request) provider_config = store.get_default_provider_config() prompt_template = store.get_prompt_template("legal_qa.default") prompt = render_legal_qa_prompt( @@ -150,8 +205,9 @@ def create_app( audit_status = "fallback" else: response = LegalQaResponse( - answer="AI 服务暂时不可用，请稍后重试。", + answer="AI service is temporarily unavailable. Please try again later.", citations=citations, + primary_sources=citations[:3], confidence=0.0, ) audit_status = "failed" @@ -189,13 +245,7 @@ def _search_citation_documents( return documents question_text = f"{request.question}\n{request.matter_context or ''}" - fallback_keywords: list[str] = [] - if any(keyword in question_text for keyword in ["试用期", "劳动合同", "工资", "解除", "竞业"]): - fallback_keywords.extend(["试用期", "劳动合同"]) - if any(keyword in question_text for keyword in ["租赁", "房屋", "租金", "押金", "保证金"]): - fallback_keywords.extend(["租赁", "民法典"]) - - for keyword in fallback_keywords: + for keyword in _fallback_keywords(question_text): documents = repository.search_knowledge( keyword, knowledge_base_ids=request.knowledge_base_ids, @@ -204,3 +254,39 @@ def _search_citation_documents( if documents: return documents return [] + + +def _search_citations( + repository: AiPlatformRepository, + request: LegalQaRequest, +) -> list[LegalCitation]: + results = repository.search_knowledge_chunks( + request.question, + knowledge_base_ids=request.knowledge_base_ids, + limit=3, + ) + if not results: + question_text = f"{request.question}\n{request.matter_context or ''}" + for keyword in _fallback_keywords(question_text): + results = repository.search_knowledge_chunks( + keyword, + knowledge_base_ids=request.knowledge_base_ids, + limit=3, + ) + if results: + break + if results: + return [chunk_to_citation(result) for result in results] + return [document_to_citation(document) for document in _search_citation_documents(repository, request)] + + +def _fallback_keywords(question_text: str) -> list[str]: + normalized = question_text.lower() + fallback_keywords: list[str] = [] + labor_terms = ["probation", "labor", "salary", "termination", "\u8bd5\u7528\u671f", "\u52b3\u52a8\u5408\u540c"] + lease_terms = ["lease", "rent", "deposit", "house", "\u79df\u8d41", "\u79df\u91d1", "\u62bc\u91d1"] + if any(keyword in normalized for keyword in labor_terms): + fallback_keywords.extend(["probation", "labor contract", "\u8bd5\u7528\u671f", "\u52b3\u52a8\u5408\u540c"]) + if any(keyword in normalized for keyword in lease_terms): + fallback_keywords.extend(["lease", "deposit", "\u79df\u8d41", "\u62bc\u91d1"]) + return fallback_keywords diff --git a/yuqei-ai-platform/services/ai-platform-api/src/yuqei_ai_platform_api/models.py b/yuqei-ai-platform/services/ai-platform-api/src/yuqei_ai_platform_api/models.py index 027db01f..debc9658 100644 --- a/yuqei-ai-platform/services/ai-platform-api/src/yuqei_ai_platform_api/models.py +++ b/yuqei-ai-platform/services/ai-platform-api/src/yuqei_ai_platform_api/models.py @@ -44,11 +44,52 @@ class KnowledgeDocumentModel(Base): __tablename__ = "knowledge_documents" id: Mapped[str] = mapped_column(String(64), primary_key=True) + import_batch_id: Mapped[str | None] = mapped_column(String(64), nullable=True, index=True) knowledge_base_id: Mapped[str] = mapped_column(String(64), index=True) title: Mapped[str] = mapped_column(String(256)) source_type: Mapped[str] = mapped_column(String(64), default="knowledge_document") reference: Mapped[str] = mapped_column(String(256)) content: Mapped[str] = mapped_column(Text) + index_status: Mapped[str] = mapped_column(String(32), default="indexed") + created_at: Mapped[datetime] = mapped_column( + DateTime(timezone=True), + default=lambda: datetime.now(UTC), + ) + + +class KnowledgeImportBatchModel(Base): + __tablename__ = "knowledge_import_batches" + + id: Mapped[str] = mapped_column(String(64), primary_key=True) + knowledge_base_id: Mapped[str] = mapped_column(String(64), index=True) + source_name: Mapped[str] = mapped_column(String(256)) + source_type: Mapped[str] = mapped_column(String(64), default="manual_text") + status: Mapped[str] = mapped_column(String(32), default="indexed") + documents_count: Mapped[int] = mapped_column(Integer, default=0) + chunks_count: Mapped[int] = mapped_column(Integer, default=0) + created_at: Mapped[datetime] = mapped_column( + DateTime(timezone=True), + default=lambda: datetime.now(UTC), + ) + updated_at: Mapped[datetime] = mapped_column( + DateTime(timezone=True), + default=lambda: datetime.now(UTC), + ) + + +class KnowledgeChunkModel(Base): + __tablename__ = "knowledge_chunks" + + id: Mapped[str] = mapped_column(String(64), primary_key=True) + document_id: Mapped[str] = mapped_column(String(64), index=True) + import_batch_id: Mapped[str | None] = mapped_column(String(64), nullable=True, index=True) + knowledge_base_id: Mapped[str] = mapped_column(String(64), index=True) + title: Mapped[str] = mapped_column(String(256)) + source_type: Mapped[str] = mapped_column(String(64), default="knowledge_document") + reference: Mapped[str] = mapped_column(String(256)) + content: Mapped[str] = mapped_column(Text) + chunk_index: Mapped[int] = mapped_column(Integer, default=0) + locator: Mapped[str] = mapped_column(String(128), default="") created_at: Mapped[datetime] = mapped_column( DateTime(timezone=True), default=lambda: datetime.now(UTC), diff --git a/yuqei-ai-platform/services/ai-platform-api/src/yuqei_ai_platform_api/repository.py b/yuqei-ai-platform/services/ai-platform-api/src/yuqei_ai_platform_api/repository.py index f2a65050..b58f1392 100644 --- a/yuqei-ai-platform/services/ai-platform-api/src/yuqei_ai_platform_api/repository.py +++ b/yuqei-ai-platform/services/ai-platform-api/src/yuqei_ai_platform_api/repository.py @@ -17,7 +17,9 @@ from yuqei_ai_platform_api.models import ( AiProviderConfigModel, AiRunAuditModel, Base, + KnowledgeChunkModel, KnowledgeDocumentModel, + KnowledgeImportBatchModel, PromptTemplateModel, ) @@ -51,11 +53,13 @@ class PromptTemplate(PromptTemplateCreate): class KnowledgeDocumentCreate(BaseModel): + import_batch_id: str | None = None knowledge_base_id: str = Field(default="default") title: str = Field(min_length=1) source_type: str = "knowledge_document" reference: str = Field(default="") content: str = Field(min_length=1) + index_status: str = "indexed" class KnowledgeDocument(KnowledgeDocumentCreate): @@ -63,6 +67,41 @@ class KnowledgeDocument(KnowledgeDocumentCreate): created_at: datetime +class KnowledgeImportBatchCreate(BaseModel): + knowledge_base_id: str = Field(default="default") + source_name: str = Field(min_length=1) + source_type: str = "manual_text" + status: str = "indexed" + + +class KnowledgeImportBatch(KnowledgeImportBatchCreate): + id: str + documents_count: int = 0 + chunks_count: int = 0 + created_at: datetime + updated_at: datetime + + +class KnowledgeChunk(BaseModel): + id: str + document_id: str + import_batch_id: str | None = None + knowledge_base_id: str + title: str + source_type: str + reference: str + content: str + chunk_index: int + locator: str + created_at: datetime + + +class KnowledgeSearchResult(BaseModel): + chunk: KnowledgeChunk + score: float + matched_terms: list[str] = Field(default_factory=list) + + class AiRunAudit(BaseModel): id: str operation: str @@ -91,6 +130,17 @@ class AiPlatformRepository(Protocol): def add_knowledge_document(self, payload: KnowledgeDocumentCreate) -> KnowledgeDocument: ... + def create_knowledge_import_batch(self, payload: KnowledgeImportBatchCreate) -> KnowledgeImportBatch: ... + + def list_knowledge_import_batches( + self, + *, + knowledge_base_id: str | None = None, + limit: int = 20, + ) -> list[KnowledgeImportBatch]: ... + + def list_knowledge_chunks(self, document_id: str) -> list[KnowledgeChunk]: ... + def search_knowledge( self, keyword: str, @@ -99,6 +149,14 @@ class AiPlatformRepository(Protocol): limit: int = 5, ) -> list[KnowledgeDocument]: ... + def search_knowledge_chunks( + self, + keyword: str, + *, + knowledge_base_ids: list[str] | None = None, + limit: int = 5, + ) -> list[KnowledgeSearchResult]: ... + def add_ai_run_audit(self, audit: AiRunAudit) -> AiRunAudit: ... def list_ai_run_audits(self, *, limit: int = 20) -> list[AiRunAudit]: ... @@ -110,7 +168,9 @@ class AiPlatformRepository(Protocol): class InMemoryAiPlatformRepository: provider_configs: dict[str, ProviderConfig] = field(default_factory=dict) prompt_templates: dict[str, PromptTemplate] = field(default_factory=dict) + knowledge_import_batches: dict[str, KnowledgeImportBatch] = field(default_factory=dict) knowledge_documents: dict[str, KnowledgeDocument] = field(default_factory=dict) + knowledge_chunks: dict[str, KnowledgeChunk] = field(default_factory=dict) ai_run_audits: list[AiRunAudit] = field(default_factory=list) _lock: Lock = field(default_factory=Lock) @@ -193,8 +253,50 @@ class InMemoryAiPlatformRepository: **payload.model_dump(), ) self.knowledge_documents[document.id] = document + chunks = split_knowledge_document(document) + for chunk in chunks: + self.knowledge_chunks[chunk.id] = chunk + if document.import_batch_id and document.import_batch_id in self.knowledge_import_batches: + batch = self.knowledge_import_batches[document.import_batch_id] + self.knowledge_import_batches[document.import_batch_id] = batch.model_copy( + update={ + "documents_count": batch.documents_count + 1, + "chunks_count": batch.chunks_count + len(chunks), + "updated_at": datetime.now(UTC), + } + ) return document + def create_knowledge_import_batch(self, payload: KnowledgeImportBatchCreate) -> KnowledgeImportBatch: + with self._lock: + batch = KnowledgeImportBatch( + id=f"kb-batch-{uuid4().hex[:12]}", + documents_count=0, + chunks_count=0, + created_at=datetime.now(UTC), + updated_at=datetime.now(UTC), + **payload.model_dump(), + ) + self.knowledge_import_batches[batch.id] = batch + return batch + + def list_knowledge_import_batches( + self, + *, + knowledge_base_id: str | None = None, + limit: int = 20, + ) -> list[KnowledgeImportBatch]: + with self._lock: + batches = list(self.knowledge_import_batches.values()) + if knowledge_base_id: + batches = [batch for batch in batches if batch.knowledge_base_id == knowledge_base_id] + return sorted(batches, key=lambda item: item.updated_at, reverse=True)[:limit] + + def list_knowledge_chunks(self, document_id: str) -> list[KnowledgeChunk]: + with self._lock: + chunks = [chunk for chunk in self.knowledge_chunks.values() if chunk.document_id == document_id] + return sorted(chunks, key=lambda item: item.chunk_index) + def search_knowledge( self, keyword: str, @@ -214,6 +316,17 @@ class InMemoryAiPlatformRepository: with self._lock: return [document for document in self.knowledge_documents.values() if matches(document)][:limit] + def search_knowledge_chunks( + self, + keyword: str, + *, + knowledge_base_ids: list[str] | None = None, + limit: int = 5, + ) -> list[KnowledgeSearchResult]: + with self._lock: + chunks = list(self.knowledge_chunks.values()) + return _rank_knowledge_chunks(chunks, keyword, knowledge_base_ids=knowledge_base_ids, limit=limit) + def add_ai_run_audit(self, audit: AiRunAudit) -> AiRunAudit: with self._lock: self.ai_run_audits.insert(0, audit) @@ -341,9 +454,57 @@ class SqlAlchemyAiPlatformRepository: ) with self._session_factory.begin() as session: session.add(document) + chunks = [ + _knowledge_chunk_model_from_dto(chunk) + for chunk in split_knowledge_document(_knowledge_document_from_model(document)) + ] + session.add_all(chunks) + if document.import_batch_id: + batch = session.get(KnowledgeImportBatchModel, document.import_batch_id) + if batch: + batch.documents_count += 1 + batch.chunks_count += len(chunks) + batch.updated_at = datetime.now(UTC) session.flush() return _knowledge_document_from_model(document) + def create_knowledge_import_batch(self, payload: KnowledgeImportBatchCreate) -> KnowledgeImportBatch: + now = datetime.now(UTC) + batch = KnowledgeImportBatchModel( + id=f"kb-batch-{uuid4().hex[:12]}", + documents_count=0, + chunks_count=0, + created_at=now, + updated_at=now, + **payload.model_dump(), + ) + with self._session_factory.begin() as session: + session.add(batch) + session.flush() + return _knowledge_import_batch_from_model(batch) + + def list_knowledge_import_batches( + self, + *, + knowledge_base_id: str | None = None, + limit: int = 20, + ) -> list[KnowledgeImportBatch]: + with self._session_factory() as session: + statement = select(KnowledgeImportBatchModel).order_by(KnowledgeImportBatchModel.updated_at.desc()) + if knowledge_base_id: + statement = statement.where(KnowledgeImportBatchModel.knowledge_base_id == knowledge_base_id) + rows = session.scalars(statement.limit(limit)) + return [_knowledge_import_batch_from_model(row) for row in rows] + + def list_knowledge_chunks(self, document_id: str) -> list[KnowledgeChunk]: + with self._session_factory() as session: + rows = session.scalars( + select(KnowledgeChunkModel) + .where(KnowledgeChunkModel.document_id == document_id) + .order_by(KnowledgeChunkModel.chunk_index) + ) + return [_knowledge_chunk_from_model(row) for row in rows] + def search_knowledge( self, keyword: str, @@ -368,6 +529,25 @@ class SqlAlchemyAiPlatformRepository: break return matches + def search_knowledge_chunks( + self, + keyword: str, + *, + knowledge_base_ids: list[str] | None = None, + limit: int = 5, + ) -> list[KnowledgeSearchResult]: + with self._session_factory() as session: + statement = select(KnowledgeChunkModel).order_by(KnowledgeChunkModel.created_at.desc()) + if knowledge_base_ids: + statement = statement.where(KnowledgeChunkModel.knowledge_base_id.in_(knowledge_base_ids)) + rows = session.scalars(statement).all() + return _rank_knowledge_chunks( + [_knowledge_chunk_from_model(row) for row in rows], + keyword, + knowledge_base_ids=knowledge_base_ids, + limit=limit, + ) + def add_ai_run_audit(self, audit: AiRunAudit) -> AiRunAudit: row = AiRunAuditModel(**audit.model_dump()) with self._session_factory.begin() as session: @@ -410,6 +590,21 @@ def document_to_citation(document: KnowledgeDocument) -> LegalCitation: ) +def chunk_to_citation(result: KnowledgeSearchResult) -> LegalCitation: + chunk = result.chunk + return LegalCitation( + source_type=chunk.source_type, + title=chunk.title, + reference=f"{chunk.reference} {chunk.locator}".strip(), + quote=chunk.content, + score=result.score, + locator=chunk.locator, + document_id=chunk.document_id, + chunk_id=chunk.id, + knowledge_base_id=chunk.knowledge_base_id, + ) + + def _provider_config_from_model(row: AiProviderConfigModel) -> ProviderConfig: return ProviderConfig( id=row.id, @@ -439,15 +634,63 @@ def _prompt_template_from_model(row: PromptTemplateModel) -> PromptTemplate: def _knowledge_document_from_model(row: KnowledgeDocumentModel) -> KnowledgeDocument: return KnowledgeDocument( id=row.id, + import_batch_id=row.import_batch_id, knowledge_base_id=row.knowledge_base_id, title=row.title, source_type=row.source_type, reference=row.reference, content=row.content, + index_status=row.index_status, created_at=_ensure_utc(row.created_at), ) +def _knowledge_import_batch_from_model(row: KnowledgeImportBatchModel) -> KnowledgeImportBatch: + return KnowledgeImportBatch( + id=row.id, + knowledge_base_id=row.knowledge_base_id, + source_name=row.source_name, + source_type=row.source_type, + status=row.status, + documents_count=row.documents_count, + chunks_count=row.chunks_count, + created_at=_ensure_utc(row.created_at), + updated_at=_ensure_utc(row.updated_at), + ) + + +def _knowledge_chunk_from_model(row: KnowledgeChunkModel) -> KnowledgeChunk: + return KnowledgeChunk( + id=row.id, + document_id=row.document_id, + import_batch_id=row.import_batch_id, + knowledge_base_id=row.knowledge_base_id, + title=row.title, + source_type=row.source_type, + reference=row.reference, + content=row.content, + chunk_index=row.chunk_index, + locator=row.locator, + created_at=_ensure_utc(row.created_at), + ) + + +def _knowledge_chunk_model_from_dto(chunk: KnowledgeChunk) -> KnowledgeChunkModel: + return KnowledgeChunkModel( + id=chunk.id, + document_id=chunk.document_id, + import_batch_id=chunk.import_batch_id, + knowledge_base_id=chunk.knowledge_base_id, + title=chunk.title, + source_type=chunk.source_type, + reference=chunk.reference, + content=chunk.content, + chunk_index=chunk.chunk_index, + locator=chunk.locator, + created_at=chunk.created_at, + ) + + def _ai_run_audit_from_model(row: AiRunAuditModel) -> AiRunAudit: return AiRunAudit( id=row.id, @@ -525,3 +768,79 @@ def _default_knowledge_documents() -> list[KnowledgeDocumentCreate]: ), ), ] + + +def split_knowledge_document(document: KnowledgeDocument, *, max_chars: int = 500) -> list[KnowledgeChunk]: + raw_segments = [segment.strip() for segment in document.content.splitlines() if segment.strip()] + segments: list[str] = [] + current = "" + for segment in raw_segments or [document.content]: + if len(current) + len(segment) + 1 <= max_chars: + current = f"{current}\n{segment}".strip() + continue + if current: + segments.append(current) + current = segment + if current: + segments.append(current) + + chunks: list[KnowledgeChunk] = [] + for index, content in enumerate(segments, start=1): + chunks.append( + KnowledgeChunk( + id=f"chunk-{uuid4().hex[:12]}", + document_id=document.id, + import_batch_id=document.import_batch_id, + knowledge_base_id=document.knowledge_base_id, + title=document.title, + source_type=document.source_type, + reference=document.reference, + content=content, + chunk_index=index, + locator=f"chunk-{index}", + created_at=datetime.now(UTC), + ) + ) + return chunks + + +def _rank_knowledge_chunks( + chunks: list[KnowledgeChunk], + keyword: str, + *, + knowledge_base_ids: list[str] | None = None, + limit: int = 5, +) -> list[KnowledgeSearchResult]: + allowed_base_ids = set(knowledge_base_ids or []) + terms = _tokenize(keyword) + results: list[KnowledgeSearchResult] = [] + for chunk in chunks: + if allowed_base_ids and chunk.knowledge_base_id not in allowed_base_ids: + continue + haystack = f"{chunk.title}\n{chunk.reference}\n{chunk.content}".lower() + matched_terms = [term for term in terms if term in haystack] + if terms and not matched_terms: + continue + score = _score_chunk(chunk, matched_terms, terms_count=len(terms)) + results.append(KnowledgeSearchResult(chunk=chunk, score=score, matched_terms=matched_terms)) + return sorted(results, key=lambda item: (item.score, item.chunk.created_at), reverse=True)[:limit] + + +def _tokenize(keyword: str) -> list[str]: + normalized = keyword.strip().lower() + if not normalized: + return [] + terms = [term for term in normalized.replace("\n", " ").split(" ") if term] + return terms or [normalized] + + +def _score_chunk(chunk: KnowledgeChunk, matched_terms: list[str], *, terms_count: int) -> float: + if terms_count == 0: + return 0.1 + haystack = f"{chunk.title}\n{chunk.reference}\n{chunk.content}".lower() + occurrences = sum(haystack.count(term) for term in matched_terms) + coverage = len(set(matched_terms)) / max(terms_count, 1) + density = min(occurrences / 5, 1.0) + title_bonus = 0.2 if any(term in chunk.title.lower() for term in matched_terms) else 0.0 + reference_bonus = 0.15 if any(term in chunk.reference.lower() for term in matched_terms) else 0.0 + return round(min(coverage * 0.65 + density * 0.2 + title_bonus + reference_bonus, 1.0), 4) diff --git a/yuqei-ai-platform/services/ai-platform-api/tests/test_platform_data.py b/yuqei-ai-platform/services/ai-platform-api/tests/test_platform_data.py index 055a8472..8e207646 100644 --- a/yuqei-ai-platform/services/ai-platform-api/tests/test_platform_data.py +++ b/yuqei-ai-platform/services/ai-platform-api/tests/test_platform_data.py @@ -6,6 +6,7 @@ from yuqei_ai_platform_api.repository import ( AiRunAudit, InMemoryAiPlatformRepository, KnowledgeDocumentCreate, + KnowledgeImportBatchCreate, ProviderConfigCreate, PromptTemplateCreate, SqlAlchemyAiPlatformRepository, @@ -56,7 +57,7 @@ def test_prompt_template_can_be_hot_updated_and_listed() -> None: json={ "name": "legal_qa.default", "version": "v2", - "content": "先检索依据，再给出结论。", + "content": "Search evidence first, then answer. Question: {question}. Evidence: {citations}", "enabled": True, }, ) @@ -79,43 +80,97 @@ def test_knowledge_document_can_be_added_and_searched() -> None: "/api/v1/knowledge-documents", json={ "knowledge_base_id": "enterprise", - "title": "公司租赁审查规范", + "title": "Lease Review Policy", "source_type": "internal_policy", - "reference": "租赁-001", - "content": "房屋租赁合同必须明确押金返还条件。", + "reference": "Lease-001", + "content": "A lease contract must clearly state the deposit refund condition.", }, ) assert add_response.status_code == 200 - search_response = client.get("/api/v1/knowledge-documents/search", params={"keyword": "押金"}) + search_response = client.get("/api/v1/knowledge-documents/search", params={"keyword": "deposit"}) assert search_response.status_code == 200 - assert search_response.json()[0]["title"] == "公司租赁审查规范" + assert search_response.json()[0]["title"] == "Lease Review Policy" -def test_legal_qa_uses_knowledge_and_records_audit() -> None: +def test_knowledge_import_batch_chunks_and_scored_search() -> None: + repository = InMemoryAiPlatformRepository() + repository.seed_defaults() + app = create_app(AiPlatformSettings(environment="test"), repository=repository) + client = TestClient(app) + + batch_response = client.post( + "/api/v1/knowledge-import-batches", + json={ + "knowledge_base_id": "laws-cn", + "source_name": "civil-code-lease", + "source_type": "law_import", + }, + ) + assert batch_response.status_code == 200 + batch = batch_response.json() + + document_response = client.post( + "/api/v1/knowledge-documents", + json={ + "import_batch_id": batch["id"], + "knowledge_base_id": "laws-cn", + "title": "Civil Code", + "source_type": "law", + "reference": "Article 703", + "content": "Lease contract deposit refund.\nTenant rent payment obligation.", + }, + ) + assert document_response.status_code == 200 + document = document_response.json() + + chunks_response = client.get(f"/api/v1/knowledge-documents/{document['id']}/chunks") + assert chunks_response.status_code == 200 + chunks = chunks_response.json() + assert chunks[0]["locator"] == "chunk-1" + + search_response = client.get( + "/api/v1/knowledge-chunks/search", + params={"keyword": "deposit refund", "knowledge_base_id": "laws-cn"}, + ) + assert search_response.status_code == 200 + result = search_response.json()[0] + assert result["score"] > 0 + assert result["chunk"]["document_id"] == document["id"] + + batches_response = client.get("/api/v1/knowledge-import-batches", params={"knowledge_base_id": "laws-cn"}) + assert batches_response.status_code == 200 + refreshed_batch = next(item for item in batches_response.json() if item["id"] == batch["id"]) + assert refreshed_batch["documents_count"] == 1 + assert refreshed_batch["chunks_count"] >= 1 + + +def test_legal_qa_uses_chunk_search_and_records_audit() -> None: client = make_client() client.post( "/api/v1/knowledge-documents", json={ - "knowledge_base_id": "laws-cn", - "title": "中华人民共和国劳动合同法", - "source_type": "law", - "reference": "第十九条", - "content": "试用期不得超过法定上限。", + "knowledge_base_id": "enterprise", + "title": "Lease Review Policy", + "source_type": "internal_policy", + "reference": "Lease-001", + "content": "Deposit refund conditions must be written into the lease contract.", }, ) response = client.post( "/api/v1/legal/qa", json={ - "question": "劳动合同试用期怎么约定？", - "knowledge_base_ids": ["laws-cn"], + "question": "How should deposit refund be handled?", + "knowledge_base_ids": ["enterprise"], }, ) assert response.status_code == 200 payload = response.json() - assert payload["citations"][0]["reference"] == "第十九条" + assert payload["citations"][0]["reference"].startswith("Lease-001") + assert payload["primary_sources"][0]["score"] > 0 + assert payload["primary_sources"][0]["chunk_id"] audit_response = client.get("/api/v1/audit/ai-runs") assert audit_response.status_code == 200 @@ -146,16 +201,24 @@ def test_sqlalchemy_repository_persists_platform_data(tmp_path) -> None: PromptTemplateCreate( name="legal_qa.default", version="v2", - content="先检索依据，再回答。", + content="Search evidence first, then answer.", + ) + ) + batch = repository.create_knowledge_import_batch( + KnowledgeImportBatchCreate( + knowledge_base_id="laws-cn", + source_name="civil-code-lease", + source_type="law_import", ) ) repository.add_knowledge_document( KnowledgeDocumentCreate( + import_batch_id=batch.id, knowledge_base_id="laws-cn", - title="中华人民共和国民法典", + title="Civil Code", source_type="law", - reference="第七百零三条", - content="租赁合同是出租人将租赁物交付承租人使用、收益，承租人支付租金的合同。", + reference="Article 703", + content="A lease contract is formed when the tenant pays rent for use of the leased property.", ) ) repository.add_ai_run_audit( @@ -164,8 +227,8 @@ def test_sqlalchemy_repository_persists_platform_data(tmp_path) -> None: operation="legal_qa", provider_config_id=provider.id, prompt_template_id=prompt.id, - question="租赁合同如何约定押金？", - answer="应明确押金返还条件。", + question="How should lease deposit refund be handled?", + answer="The contract should clearly state deposit refund conditions.", citations_count=1, latency_ms=12, status="succeeded", @@ -176,5 +239,7 @@ def test_sqlalchemy_repository_persists_platform_data(tmp_path) -> None: assert reloaded.get_default_provider_config().model == "deepseek-chat" # type: ignore[union-attr] assert reloaded.list_prompt_templates()[0].version == "v2" - assert reloaded.search_knowledge("租赁合同", knowledge_base_ids=["laws-cn"])[0].reference == "第七百零三条" + assert reloaded.search_knowledge("lease", knowledge_base_ids=["laws-cn"])[0].reference == "Article 703" + assert reloaded.list_knowledge_import_batches(knowledge_base_id="laws-cn")[0].documents_count == 1 + assert reloaded.search_knowledge_chunks("lease", knowledge_base_ids=["laws-cn"])[0].score > 0 assert reloaded.list_ai_run_audits()[0].id == "run-test-001" diff --git a/yuqei-ai-sdk-python/README.md b/yuqei-ai-sdk-python/README.md index 08c5b0e6..0d766f84 100644 --- a/yuqei-ai-sdk-python/README.md +++ b/yuqei-ai-sdk-python/README.md @@ -9,5 +9,5 @@ Initial scope: - Typed request and response models. - Model provider config management. - Prompt template management. -- Knowledge document insert/search. +- Knowledge import batches, document insert/search, chunk listing, and scored chunk search. - AI run audit listing. diff --git a/yuqei-ai-sdk-python/src/yuqei_sdk/__init__.py b/yuqei-ai-sdk-python/src/yuqei_sdk/__init__.py index dc5832e3..ab2925a0 100644 --- a/yuqei-ai-sdk-python/src/yuqei_sdk/__init__.py +++ b/yuqei-ai-sdk-python/src/yuqei_sdk/__init__.py @@ -1,8 +1,12 @@ from yuqei_sdk.ai_platform import ( AiPlatformClient, AiRunAudit, + KnowledgeChunk, KnowledgeDocument, KnowledgeDocumentCreate, + KnowledgeImportBatch, + KnowledgeImportBatchCreate, + KnowledgeSearchResult, LegalCitation, LegalQaRequest, LegalQaResponse, @@ -16,8 +20,12 @@ from yuqei_sdk.context import RequestContext __all__ = [ "AiPlatformClient", "AiRunAudit", + "KnowledgeChunk", "KnowledgeDocument", "KnowledgeDocumentCreate", + "KnowledgeImportBatch", + "KnowledgeImportBatchCreate", + "KnowledgeSearchResult", "LegalCitation", "LegalQaRequest", "LegalQaResponse", diff --git a/yuqei-ai-sdk-python/src/yuqei_sdk/ai_platform.py b/yuqei-ai-sdk-python/src/yuqei_sdk/ai_platform.py index 533781d5..1bef77d4 100644 --- a/yuqei-ai-sdk-python/src/yuqei_sdk/ai_platform.py +++ b/yuqei-ai-sdk-python/src/yuqei_sdk/ai_platform.py @@ -18,11 +18,17 @@ class LegalCitation(BaseModel): reference: str quote: str url: str | None = None + score: float | None = None + locator: str | None = None + document_id: str | None = None + chunk_id: str | None = None + knowledge_base_id: str | None = None class LegalQaResponse(BaseModel): answer: str citations: list[LegalCitation] + primary_sources: list[LegalCitation] = Field(default_factory=list) confidence: float | None = None @@ -55,11 +61,13 @@ class PromptTemplate(PromptTemplateCreate): class KnowledgeDocumentCreate(BaseModel): + import_batch_id: str | None = None knowledge_base_id: str = "default" title: str = Field(min_length=1) source_type: str = "knowledge_document" reference: str = "" content: str = Field(min_length=1) + index_status: str = "indexed" class KnowledgeDocument(KnowledgeDocumentCreate): @@ -67,6 +75,41 @@ class KnowledgeDocument(KnowledgeDocumentCreate): created_at: str +class KnowledgeImportBatchCreate(BaseModel): + knowledge_base_id: str = "default" + source_name: str = Field(min_length=1) + source_type: str = "manual_text" + status: str = "indexed" + + +class KnowledgeImportBatch(KnowledgeImportBatchCreate): + id: str + documents_count: int = 0 + chunks_count: int = 0 + created_at: str + updated_at: str + + +class KnowledgeChunk(BaseModel): + id: str + document_id: str + import_batch_id: str | None = None + knowledge_base_id: str + title: str + source_type: str + reference: str + content: str + chunk_index: int + locator: str + created_at: str + + +class KnowledgeSearchResult(BaseModel): + chunk: KnowledgeChunk + score: float + matched_terms: list[str] = Field(default_factory=list) + + class AiRunAudit(BaseModel): id: str operation: str @@ -178,21 +221,89 @@ class AiPlatformClient: response.raise_for_status() return KnowledgeDocument.model_validate(response.json()) + def create_knowledge_import_batch( + self, + batch: KnowledgeImportBatchCreate, + *, + context: RequestContext | None = None, + ) -> KnowledgeImportBatch: + response = self._client.post( + f"{self._api_prefix}/knowledge-import-batches", + json=batch.model_dump(exclude_none=True), + headers=(context or RequestContext()).to_headers(), + ) + response.raise_for_status() + return KnowledgeImportBatch.model_validate(response.json()) + + def list_knowledge_import_batches( + self, + *, + knowledge_base_id: str | None = None, + limit: int = 20, + context: RequestContext | None = None, + ) -> list[KnowledgeImportBatch]: + params: dict[str, Any] = {"limit": limit} + if knowledge_base_id: + params["knowledge_base_id"] = knowledge_base_id + response = self._client.get( + f"{self._api_prefix}/knowledge-import-batches", + params=params, + headers=(context or RequestContext()).to_headers(), + ) + response.raise_for_status() + return [KnowledgeImportBatch.model_validate(item) for item in response.json()] + + def list_knowledge_chunks( + self, + document_id: str, + *, + context: RequestContext | None = None, + ) -> list[KnowledgeChunk]: + response = self._client.get( + f"{self._api_prefix}/knowledge-documents/{document_id}/chunks", + headers=(context or RequestContext()).to_headers(), + ) + response.raise_for_status() + return [KnowledgeChunk.model_validate(item) for item in response.json()] + def search_knowledge_documents( self, keyword: str = "", *, + knowledge_base_id: str | None = None, limit: int = 5, context: RequestContext | None = None, ) -> list[KnowledgeDocument]: + params: dict[str, Any] = {"keyword": keyword, "limit": limit} + if knowledge_base_id: + params["knowledge_base_id"] = knowledge_base_id response = self._client.get( f"{self._api_prefix}/knowledge-documents/search", - params={"keyword": keyword, "limit": limit}, + params=params, headers=(context or RequestContext()).to_headers(), ) response.raise_for_status() return [KnowledgeDocument.model_validate(item) for item in response.json()] + def search_knowledge_chunks( + self, + keyword: str = "", + *, + knowledge_base_id: str | None = None, + limit: int = 5, + context: RequestContext | None = None, + ) -> list[KnowledgeSearchResult]: + params: dict[str, Any] = {"keyword": keyword, "limit": limit} + if knowledge_base_id: + params["knowledge_base_id"] = knowledge_base_id + response = self._client.get( + f"{self._api_prefix}/knowledge-chunks/search", + params=params, + headers=(context or RequestContext()).to_headers(), + ) + response.raise_for_status() + return [KnowledgeSearchResult.model_validate(item) for item in response.json()] + def list_ai_run_audits( self, *, diff --git a/yuqei-ai-sdk-python/tests/test_ai_platform.py b/yuqei-ai-sdk-python/tests/test_ai_platform.py index 3af2c730..94c4054d 100644 --- a/yuqei-ai-sdk-python/tests/test_ai_platform.py +++ b/yuqei-ai-sdk-python/tests/test_ai_platform.py @@ -5,6 +5,7 @@ import httpx from yuqei_sdk import ( AiPlatformClient, KnowledgeDocumentCreate, + KnowledgeImportBatchCreate, LegalQaRequest, PromptTemplateCreate, ProviderConfigCreate, @@ -21,13 +22,31 @@ def test_ai_platform_client_posts_legal_qa_with_context_headers() -> None: return httpx.Response( 200, json={ - "answer": "依据劳动合同法第十九条，需要核对合同期限。", + "answer": "Use Article 19 and verify the contract term.", "citations": [ { "source_type": "law", - "title": "中华人民共和国劳动合同法", - "reference": "第十九条", - "quote": "劳动合同期限三个月以上不满一年的，试用期不得超过一个月。", + "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", + } + ], + "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", } ], "confidence": 0.72, @@ -40,11 +59,12 @@ def test_ai_platform_client_posts_legal_qa_with_context_headers() -> None: ) response = client.run_legal_qa( - LegalQaRequest(question="试用期如何约定？"), + 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 == "第十九条" + assert response.citations[0].reference == "Article 19 chunk-1" + assert response.primary_sources[0].score == 0.82 assert captured_request is not None assert captured_request.url.path == "/api/v1/legal/qa" assert captured_request.headers["X-Trace-Id"] == "trace-1" @@ -98,7 +118,7 @@ def test_ai_platform_client_covers_admin_data_endpoints() -> None: "id": "prompt-1", "name": "legal_qa.default", "version": "v2", - "content": "先检索依据，再回答。", + "content": "Search evidence first, then answer.", "enabled": True, "updated_at": "2026-06-22T00:00:00Z", }, @@ -111,43 +131,121 @@ def test_ai_platform_client_covers_admin_data_endpoints() -> None: "id": "prompt-1", "name": "legal_qa.default", "version": "v2", - "content": "先检索依据，再回答。", + "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": "中华人民共和国民法典", + "title": "Civil Code", "source_type": "law", - "reference": "第七百零三条", - "content": "租赁合同是出租人将租赁物交付承租人使用、收益，承租人支付租金的合同。", + "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", + "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["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": "中华人民共和国民法典", + "title": "Civil Code", "source_type": "law", - "reference": "第七百零三条", - "content": "租赁合同是出租人将租赁物交付承租人使用、收益，承租人支付租金的合同。", + "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-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", + "created_at": "2026-06-22T00:00:00Z", + }, + "score": 0.82, + "matched_terms": ["lease"], + } + ], + ) if request.url.path == "/api/v1/audit/ai-runs": return httpx.Response( 200, @@ -157,8 +255,8 @@ def test_ai_platform_client_covers_admin_data_endpoints() -> None: "operation": "legal_qa", "provider_config_id": "provider-1", "prompt_template_id": "prompt-1", - "question": "租赁合同如何约定押金？", - "answer": "应明确押金返还条件。", + "question": "How should lease deposit refund be handled?", + "answer": "The contract should clearly state deposit refund conditions.", "citations_count": 1, "latency_ms": 12, "status": "succeeded", @@ -181,18 +279,29 @@ def test_ai_platform_client_covers_admin_data_endpoints() -> None: ).model == "deepseek-chat" assert client.list_model_configs()[0].provider == "deepseek" assert client.upsert_prompt_template( - PromptTemplateCreate(name="legal_qa.default", version="v2", content="先检索依据，再回答。") + 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="中华人民共和国民法典", + title="Civil Code", source_type="law", - reference="第七百零三条", - content="租赁合同是出租人将租赁物交付承租人使用、收益，承租人支付租金的合同。", + reference="Article 703", + content="A lease contract defines rent and use of leased property.", ) - ).reference == "第七百零三条" - assert client.search_knowledge_documents("租赁", limit=3)[0].title == "中华人民共和国民法典" + ).reference == "Article 703" + assert client.list_knowledge_chunks("doc-1")[0].locator == "chunk-1" + assert client.search_knowledge_documents("lease", limit=3)[0].title == "Civil Code" + assert client.search_knowledge_chunks("lease", limit=3)[0].score == 0.82 assert client.list_ai_run_audits(limit=5)[0].operation == "legal_qa" assert "/api/v1/audit/ai-runs" in seen_paths