feat: enhance ai knowledge retrieval

This commit is contained in:
2026-06-22 21:20:09 +08:00
parent ccf35ca586
commit b319d41787
13 changed files with 1223 additions and 62 deletions
@@ -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. 法律问答输出按“结论、主要依据、风险提示、建议动作”结构化。
@@ -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:
@@ -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`
@@ -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")
@@ -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]})
@@ -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
@@ -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),
@@ -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)
@@ -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"
+1 -1
View File
@@ -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.
@@ -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",
@@ -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,
*,
+132 -23
View File
@@ -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