From 82d2a9f0aa7c4cfb138f278ec9b52362154b9081 Mon Sep 17 00:00:00 2001 From: yuqei Date: Mon, 22 Jun 2026 21:42:12 +0800 Subject: [PATCH] feat: add hybrid rag retrieval audit --- ...wledge-retrieval-enhancement-2026-06-22.md | 94 ++++--- .../packages/ai-openapi/openapi.yaml | 106 ++++++++ .../services/ai-platform-api/README.md | 6 +- ...0260622_0003_add_hybrid_retrieval_audit.py | 51 ++++ .../src/yuqei_ai_platform_api/legal_qa.py | 31 ++- .../src/yuqei_ai_platform_api/main.py | 29 ++- .../src/yuqei_ai_platform_api/models.py | 7 + .../src/yuqei_ai_platform_api/repository.py | 246 +++++++++++++++++- .../tests/test_platform_data.py | 45 +++- yuqei-ai-sdk-python/README.md | 4 +- yuqei-ai-sdk-python/src/yuqei_sdk/__init__.py | 2 + .../src/yuqei_sdk/ai_platform.py | 35 ++- yuqei-ai-sdk-python/tests/test_ai_platform.py | 50 +++- 13 files changed, 625 insertions(+), 81 deletions(-) create mode 100644 yuqei-ai-platform/services/ai-platform-api/migrations/versions/20260622_0003_add_hybrid_retrieval_audit.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 index 86d603a1..b8ebcf60 100644 --- 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 @@ -4,106 +4,102 @@ ## 本轮目标 -本轮推进 AI 中台知识库检索增强，把上一轮“整篇知识文档关键词检索”升级为更接近 RAG 的基础结构： +本轮推进 AI 中台知识库检索增强，把早期“整篇知识文档关键词检索”升级为更接近生产 RAG 的基础结构： - 导入批次。 - 文档切分 chunk。 - chunk 级检索结果。 -- 检索命中分数。 +- 关键词分、向量分、混合分。 - 来源定位。 - 法律问答返回结构化主要依据。 +- AI 审计记录命中的 chunk ids、来源、分数和检索策略。 ## 已完成 1. 数据模型 - 新增 `KnowledgeImportBatchModel`。 - 新增 `KnowledgeChunkModel`。 - - `KnowledgeDocumentModel` 增加： - - `import_batch_id` - - `index_status` + - `KnowledgeDocumentModel` 增加 `import_batch_id`、`index_status`。 + - `KnowledgeChunkModel` 增加本地 embedding 字段：`embedding_model`、`embedding_dimension`、`embedding_json`。 + - `AiRunAuditModel` 增加检索追踪字段：`retrieval_query`、`retrieval_strategy`、`retrieved_chunk_ids_json`、`retrieved_sources_json`。 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` + - `20260622_0002_add_knowledge_batches_and_chunks.py` + - `20260622_0003_add_hybrid_retrieval_audit.py` + - 支持升级/回滚。 3. Repository 能力 - - 创建知识导入批次。 - - 列出知识导入批次。 + - 创建和列出知识导入批次。 - 新增知识文档时自动切分 chunk。 - - 列出文档 chunk。 - - chunk 级关键词检索。 - - 返回检索分数、命中词、locator、document_id、chunk_id。 + - chunk 持久化本地 hash embedding。 + - 支持 `keyword`、`vector`、`hybrid` 三种检索策略。 + - 搜索结果返回 `score`、`keyword_score`、`vector_score`、`hybrid_score`、`matched_terms`、`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` 保持兼容。 + - `GET /api/v1/knowledge-chunks/search?strategy=keyword|vector|hybrid` + - `GET /api/v1/knowledge-documents/search` + - `POST /api/v1/legal/qa` + - `GET /api/v1/audit/ai-runs` 5. 法律问答增强 - - 法律问答优先使用 chunk 检索。 - - `LegalCitation` 增加： - - `score` - - `locator` - - `document_id` - - `chunk_id` - - `knowledge_base_id` - - `LegalQaResponse` 增加 `primary_sources`。 - - 旧字段 `citations` 保持兼容。 + - 默认使用 `hybrid` 检索。 + - `LegalQaResponse.primary_sources` 返回结构化主要依据。 + - `LegalCitation` 增加检索分数、定位、chunk id、document id、knowledge base id、命中词和检索策略。 + - AI 审计记录保留本次命中的 chunk ids 和结构化来源。 + - 清理规则兜底回答里的历史乱码文本，改为可读 UTF-8 中文。 6. SDK 能力 - - 新增知识导入批次类型和方法。 - - 新增 chunk 类型和方法。 - - 新增 scored chunk search 方法。 - - 法律问答响应支持 `primary_sources`。 + - Python SDK 同步知识批次、chunk、混合检索和审计来源字段。 + - `search_knowledge_chunks(..., strategy="hybrid")` 支持策略参数。 + - `AiRunAudit` 支持读取 `retrieved_chunk_ids` 和 `retrieved_sources`。 ## 验证结果 已通过： ```powershell -python -m pytest +python -m pytest yuqei-ai-platform/services/ai-platform-api/tests yuqei-ai-sdk-python/tests ``` 当前结果： ```text -21 passed +16 passed ``` 覆盖点： - 知识导入批次创建和统计更新。 - 文档自动切分 chunk。 -- chunk 搜索返回 score 和 locator。 -- SQLite 持久化仓库可跨实例读取批次、chunk 和审计。 +- chunk 返回 embedding 元数据。 +- chunk 搜索返回关键词分、向量分、混合分和命中词。 - 法律问答返回 `primary_sources`。 -- SDK 覆盖新增知识库检索能力。 +- AI 审计记录命中的 chunk ids 和结构化来源。 +- SQLite 持久化后可重新读取 batch、chunk、audit。 +- SDK 覆盖新增检索和审计字段。 ## 当前边界 -本轮仍是关键词检索，不是正式向量检索。 +本轮采用确定性的本地 hash embedding，目的是先稳定 API、数据结构、审计链路和测试闭环。它不是最终生产级语义向量模型。 尚未完成： -- embedding 生成。 -- 向量库或 pgvector。 -- 混合检索排序。 -- 文档上传解析、OCR、复杂 DOCX/PDF 切分。 -- AI 审计记录命中文档 ID 列表、token、费用、provider 错误详情。 +- 接入真实 embedding 模型。 +- 引入 pgvector 或独立向量库。 +- 大规模文档解析、OCR、DOCX/PDF 智能切分。 +- provider token、费用、错误详情的完整审计。 +- 检索 explain 面板和质量评测集。 ## 下一轮建议 -下一轮进入知识库索引质量增强： +下一步进入生产级 RAG 基础设施： -1. 增加 embedding 字段和向量检索适配层。 -2. 引入 `pgvector` 或独立向量库的技术选择和迁移。 -3. 增加检索 explain 信息：关键词分、向量分、最终分。 -4. AI 审计记录 `hit_chunk_ids`、provider 错误、token 和成本。 -5. 法律问答输出按“结论、主要依据、风险提示、建议动作”结构化。 +1. 选择并接入 embedding provider。 +2. 增加 pgvector 迁移和向量索引。 +3. 把当前 hash embedding 作为测试/离线 fallback，生产使用真实向量。 +4. 增加检索 explain API，展示关键词分、向量分、混合分、过滤原因。 +5. 建立法律问答评测集，检查“回答是否有依据、依据是否匹配问题、是否误引法条”。 diff --git a/yuqei-ai-platform/packages/ai-openapi/openapi.yaml b/yuqei-ai-platform/packages/ai-openapi/openapi.yaml index 2410d04c..143c5b7d 100644 --- a/yuqei-ai-platform/packages/ai-openapi/openapi.yaml +++ b/yuqei-ai-platform/packages/ai-openapi/openapi.yaml @@ -340,6 +340,13 @@ paths: default: 5 minimum: 1 maximum: 20 + - name: strategy + in: query + required: false + schema: + type: string + enum: [keyword, vector, hybrid] + default: hybrid responses: "200": description: Knowledge chunk matches @@ -466,6 +473,25 @@ components: knowledge_base_id: type: string nullable: true + keyword_score: + type: number + format: float + nullable: true + vector_score: + type: number + format: float + nullable: true + hybrid_score: + type: number + format: float + nullable: true + retrieval_strategy: + type: string + nullable: true + matched_terms: + type: array + items: + type: string LegalQaRequest: type: object required: @@ -480,6 +506,10 @@ components: type: array items: type: string + retrieval_strategy: + type: string + enum: [keyword, vector, hybrid] + default: hybrid LegalQaResponse: type: object required: @@ -840,6 +870,12 @@ components: type: integer locator: type: string + embedding_model: + type: string + default: local-hash-v1 + embedding_dimension: + type: integer + default: 64 created_at: type: string format: date-time @@ -855,6 +891,62 @@ components: score: type: number format: float + keyword_score: + type: number + format: float + vector_score: + type: number + format: float + hybrid_score: + type: number + format: float + retrieval_strategy: + type: string + enum: [keyword, vector, hybrid] + matched_terms: + type: array + items: + type: string + RetrievalSource: + type: object + required: + - title + - source_type + - reference + properties: + chunk_id: + type: string + nullable: true + document_id: + type: string + nullable: true + knowledge_base_id: + type: string + nullable: true + title: + type: string + source_type: + type: string + reference: + type: string + locator: + type: string + nullable: true + score: + type: number + format: float + keyword_score: + type: number + format: float + vector_score: + type: number + format: float + hybrid_score: + type: number + format: float + retrieval_strategy: + type: string + enum: [keyword, vector, hybrid] matched_terms: type: array items: @@ -885,6 +977,20 @@ components: nullable: true citations_count: type: integer + retrieval_query: + type: string + nullable: true + retrieval_strategy: + type: string + nullable: true + retrieved_chunk_ids: + type: array + items: + type: string + retrieved_sources: + type: array + items: + $ref: "#/components/schemas/RetrievalSource" latency_ms: type: integer status: diff --git a/yuqei-ai-platform/services/ai-platform-api/README.md b/yuqei-ai-platform/services/ai-platform-api/README.md index 8602c2ef..95c6ae6c 100644 --- a/yuqei-ai-platform/services/ai-platform-api/README.md +++ b/yuqei-ai-platform/services/ai-platform-api/README.md @@ -47,10 +47,14 @@ Core MVP endpoints: - `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` +- `GET /api/v1/knowledge-chunks/search?strategy=keyword|vector|hybrid` - `POST /api/v1/legal/qa` - `GET /api/v1/audit/ai-runs` +Knowledge chunk search returns keyword, vector, and hybrid scores. Legal QA uses +hybrid retrieval by default and records retrieved chunk ids plus structured source +metadata in AI run audits. + ## Test From repository root: diff --git a/yuqei-ai-platform/services/ai-platform-api/migrations/versions/20260622_0003_add_hybrid_retrieval_audit.py b/yuqei-ai-platform/services/ai-platform-api/migrations/versions/20260622_0003_add_hybrid_retrieval_audit.py new file mode 100644 index 00000000..7e5e0af9 --- /dev/null +++ b/yuqei-ai-platform/services/ai-platform-api/migrations/versions/20260622_0003_add_hybrid_retrieval_audit.py @@ -0,0 +1,51 @@ +"""add hybrid retrieval audit fields + +Revision ID: 20260622_0003 +Revises: 20260622_0002 +Create Date: 2026-06-22 + +""" +from typing import Sequence, Union + +from alembic import op +import sqlalchemy as sa + + +revision: str = "20260622_0003" +down_revision: Union[str, None] = "20260622_0002" +branch_labels: Union[str, Sequence[str], None] = None +depends_on: Union[str, Sequence[str], None] = None + + +def upgrade() -> None: + op.add_column( + "knowledge_chunks", + sa.Column("embedding_model", sa.String(length=64), nullable=False, server_default="local-hash-v1"), + ) + op.add_column( + "knowledge_chunks", + sa.Column("embedding_dimension", sa.Integer(), nullable=False, server_default="64"), + ) + op.add_column("knowledge_chunks", sa.Column("embedding_json", sa.Text(), nullable=True)) + + op.add_column("ai_run_audits", sa.Column("retrieval_query", sa.Text(), nullable=True)) + op.add_column("ai_run_audits", sa.Column("retrieval_strategy", sa.String(length=32), nullable=True)) + op.add_column( + "ai_run_audits", + sa.Column("retrieved_chunk_ids_json", sa.Text(), nullable=False, server_default="[]"), + ) + op.add_column( + "ai_run_audits", + sa.Column("retrieved_sources_json", sa.Text(), nullable=False, server_default="[]"), + ) + + +def downgrade() -> None: + op.drop_column("ai_run_audits", "retrieved_sources_json") + op.drop_column("ai_run_audits", "retrieved_chunk_ids_json") + op.drop_column("ai_run_audits", "retrieval_strategy") + op.drop_column("ai_run_audits", "retrieval_query") + + op.drop_column("knowledge_chunks", "embedding_json") + op.drop_column("knowledge_chunks", "embedding_dimension") + op.drop_column("knowledge_chunks", "embedding_model") 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 1cf1f93e..67780747 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 @@ -5,6 +5,7 @@ class LegalQaRequest(BaseModel): question: str = Field(min_length=1) matter_context: str | None = None knowledge_base_ids: list[str] = Field(default_factory=list) + retrieval_strategy: str = "hybrid" class LegalCitation(BaseModel): @@ -18,6 +19,11 @@ class LegalCitation(BaseModel): document_id: str | None = None chunk_id: str | None = None knowledge_base_id: str | None = None + keyword_score: float | None = None + vector_score: float | None = None + hybrid_score: float | None = None + retrieval_strategy: str | None = None + matched_terms: list[str] = Field(default_factory=list) class LegalQaResponse(BaseModel): @@ -34,20 +40,23 @@ def answer_legal_question( ) -> LegalQaResponse: question = request.question.strip() context = request.matter_context.strip() if request.matter_context else "" - question_text = question + "\n" + context + question_text = f"{question}\n{context}" citations_from_search = retrieved_citations or [] - if _contains_any(question_text, ["试用期", "劳动合同", "工资", "解除", "竞业"]): + if _contains_any(question_text, ["劳动合同", "试用期", "工资", "解除劳动", "竞业限制"]): return _with_retrieved_citations(_labor_answer(question), citations_from_search) - if _contains_any(question_text, ["租赁", "房屋", "租金", "押金", "保证金"]): + if _contains_any(question_text, ["租赁", "房屋", "租金", "押金", "保证金", "出租", "承租"]): return _with_retrieved_citations(_lease_answer(question), citations_from_search) - return _with_retrieved_citations( - LegalQaResponse( + return _with_retrieved_citations(_general_contract_answer(), citations_from_search) + + +def _general_contract_answer() -> LegalQaResponse: + return LegalQaResponse( answer=( - "当前问题需要结合合同类型、交易背景和适用地区进一步判断。" - "V2 MVP 会先返回可审计的基础意见；正式版本会先检索企业知识库和法规库，" - "再给出更精确的法律依据、条文引用和处理建议。" + "当前问题需要结合合同类型、交易背景、主体身份和履行地点进一步判断。" + "建议先核对合同主体、标的、价款、履行期限、违约责任、解除条件、争议解决和证据留存安排。" + "如果需要更精确的结论，应补充合同正文或关键条款后再进行法律检索和风险判断。" ), citations=[ LegalCitation( @@ -58,8 +67,6 @@ def answer_legal_question( ) ], confidence=0.45, - ), - citations_from_search, ) @@ -67,7 +74,7 @@ def _labor_answer(question: str) -> LegalQaResponse: return LegalQaResponse( answer=( f"针对“{question}”，应优先核对劳动合同期限、试用期约定、工资支付记录和解除理由。" - "如果试用期超过法定上限，或者解除缺少事实和制度依据，企业侧会面临补足工资、" + "如果试用期超过法定上限，或者解除缺少事实和制度依据，用人单位可能面临补足工资、" "违法解除赔偿或继续履行等风险。建议在处理前形成证据清单，并把结论写入审批意见。" ), citations=[ @@ -92,7 +99,7 @@ def _lease_answer(question: str) -> LegalQaResponse: return LegalQaResponse( answer=( f"针对“{question}”，应先确认租赁标的、租赁期限、租金支付、押金返还、维修义务和解除条件。" - "如果合同没有明确交付状态、押金扣除条件或提前解除责任，后续很容易产生举证争议。" + "如果合同没有明确交付状态、押金扣除条件或提前解除责任，后续容易产生举证争议。" "建议在起草或审查时把交付清单、付款节点和违约责任写成可执行条款。" ), citations=[ 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 c3d5f36b..d1628a56 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 @@ -26,6 +26,7 @@ from yuqei_ai_platform_api.repository import ( ProviderConfigCreate, chunk_to_citation, document_to_citation, + search_result_to_retrieval_source, ) from yuqei_ai_platform_api.state import get_repository @@ -160,11 +161,13 @@ def create_app( keyword: str = "", knowledge_base_id: str | None = None, limit: int = Query(5, ge=1, le=20), + strategy: str = Query("hybrid", pattern="^(keyword|vector|hybrid)$"), ) -> list[KnowledgeSearchResult]: return store.search_knowledge_chunks( keyword, knowledge_base_ids=[knowledge_base_id] if knowledge_base_id else None, limit=limit, + strategy=strategy, ) @app.get( @@ -182,7 +185,11 @@ def create_app( ) def run_legal_qa(request: LegalQaRequest) -> LegalQaResponse: started_at = perf_counter() - citations = _search_citations(store, request) + search_results = _search_chunk_results(store, request) + citations = [chunk_to_citation(result) for result in search_results] + if not citations: + citations = [document_to_citation(document) for document in _search_citation_documents(store, request)] + retrieval_sources = [search_result_to_retrieval_source(result) for result in search_results] provider_config = store.get_default_provider_config() prompt_template = store.get_prompt_template("legal_qa.default") prompt = render_legal_qa_prompt( @@ -220,6 +227,10 @@ def create_app( question=request.question, answer=response.answer, citations_count=len(response.citations), + retrieval_query=request.question, + retrieval_strategy=request.retrieval_strategy, + retrieved_chunk_ids=[source.chunk_id for source in retrieval_sources if source.chunk_id], + retrieved_sources=retrieval_sources, latency_ms=max(0, int((perf_counter() - started_at) * 1000)), status=audit_status, ) @@ -260,10 +271,21 @@ def _search_citations( repository: AiPlatformRepository, request: LegalQaRequest, ) -> list[LegalCitation]: + results = _search_chunk_results(repository, request) + 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 _search_chunk_results( + repository: AiPlatformRepository, + request: LegalQaRequest, +) -> list[KnowledgeSearchResult]: results = repository.search_knowledge_chunks( request.question, knowledge_base_ids=request.knowledge_base_ids, limit=3, + strategy=request.retrieval_strategy, ) if not results: question_text = f"{request.question}\n{request.matter_context or ''}" @@ -272,12 +294,11 @@ def _search_citations( keyword, knowledge_base_ids=request.knowledge_base_ids, limit=3, + strategy=request.retrieval_strategy, ) 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)] + return results def _fallback_keywords(question_text: str) -> list[str]: 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 debc9658..62f2e446 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 @@ -90,6 +90,9 @@ class KnowledgeChunkModel(Base): content: Mapped[str] = mapped_column(Text) chunk_index: Mapped[int] = mapped_column(Integer, default=0) locator: Mapped[str] = mapped_column(String(128), default="") + embedding_model: Mapped[str] = mapped_column(String(64), default="local-hash-v1") + embedding_dimension: Mapped[int] = mapped_column(Integer, default=64) + embedding_json: Mapped[str | None] = mapped_column(Text, nullable=True) created_at: Mapped[datetime] = mapped_column( DateTime(timezone=True), default=lambda: datetime.now(UTC), @@ -106,6 +109,10 @@ class AiRunAuditModel(Base): question: Mapped[str | None] = mapped_column(Text, nullable=True) answer: Mapped[str | None] = mapped_column(Text, nullable=True) citations_count: Mapped[int] = mapped_column(Integer, default=0) + retrieval_query: Mapped[str | None] = mapped_column(Text, nullable=True) + retrieval_strategy: Mapped[str | None] = mapped_column(String(32), nullable=True) + retrieved_chunk_ids_json: Mapped[str] = mapped_column(Text, default="[]") + retrieved_sources_json: Mapped[str] = mapped_column(Text, default="[]") latency_ms: Mapped[int] = mapped_column(Integer, default=0) status: Mapped[str] = mapped_column(String(32), default="succeeded") created_at: Mapped[datetime] = mapped_column( 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 b58f1392..d09ff45c 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 @@ -1,5 +1,9 @@ from __future__ import annotations +import hashlib +import json +import math +import re from dataclasses import dataclass, field from datetime import UTC, datetime from threading import Lock @@ -24,6 +28,29 @@ from yuqei_ai_platform_api.models import ( ) +EMBEDDING_MODEL = "local-hash-v1" +EMBEDDING_DIMENSION = 64 +RETRIEVAL_STRATEGIES = {"keyword", "vector", "hybrid"} +MIN_VECTOR_RECALL_SCORE = 0.35 +GENERIC_RETRIEVAL_TERMS = { + "合同", + "这个", + "需要", + "注意", + "什么", + "事项", + "问题", + "风险", + "如何", + "怎么", + "是否", + "可以", + "应该", + "处理", + "相关", +} + + class ProviderConfigCreate(BaseModel): provider: str = Field(min_length=1) model: str = Field(min_length=1) @@ -93,12 +120,35 @@ class KnowledgeChunk(BaseModel): content: str chunk_index: int locator: str + embedding_model: str = EMBEDDING_MODEL + embedding_dimension: int = EMBEDDING_DIMENSION + embedding_vector: list[float] = Field(default_factory=list, exclude=True) created_at: datetime class KnowledgeSearchResult(BaseModel): chunk: KnowledgeChunk score: float + keyword_score: float = 0.0 + vector_score: float = 0.0 + hybrid_score: float = 0.0 + retrieval_strategy: str = "hybrid" + matched_terms: list[str] = Field(default_factory=list) + + +class RetrievalSource(BaseModel): + chunk_id: str | None = None + document_id: str | None = None + knowledge_base_id: str | None = None + title: str + source_type: str + reference: str + locator: str | None = None + score: float = 0.0 + keyword_score: float = 0.0 + vector_score: float = 0.0 + hybrid_score: float = 0.0 + retrieval_strategy: str = "hybrid" matched_terms: list[str] = Field(default_factory=list) @@ -110,6 +160,10 @@ class AiRunAudit(BaseModel): question: str | None = None answer: str | None = None citations_count: int = 0 + retrieval_query: str | None = None + retrieval_strategy: str | None = None + retrieved_chunk_ids: list[str] = Field(default_factory=list) + retrieved_sources: list[RetrievalSource] = Field(default_factory=list) latency_ms: int = 0 status: str = "succeeded" created_at: datetime = Field(default_factory=lambda: datetime.now(UTC)) @@ -155,6 +209,7 @@ class AiPlatformRepository(Protocol): *, knowledge_base_ids: list[str] | None = None, limit: int = 5, + strategy: str = "hybrid", ) -> list[KnowledgeSearchResult]: ... def add_ai_run_audit(self, audit: AiRunAudit) -> AiRunAudit: ... @@ -322,10 +377,17 @@ class InMemoryAiPlatformRepository: *, knowledge_base_ids: list[str] | None = None, limit: int = 5, + strategy: str = "hybrid", ) -> 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) + return _rank_knowledge_chunks( + chunks, + keyword, + knowledge_base_ids=knowledge_base_ids, + limit=limit, + strategy=strategy, + ) def add_ai_run_audit(self, audit: AiRunAudit) -> AiRunAudit: with self._lock: @@ -535,6 +597,7 @@ class SqlAlchemyAiPlatformRepository: *, knowledge_base_ids: list[str] | None = None, limit: int = 5, + strategy: str = "hybrid", ) -> list[KnowledgeSearchResult]: with self._session_factory() as session: statement = select(KnowledgeChunkModel).order_by(KnowledgeChunkModel.created_at.desc()) @@ -546,10 +609,28 @@ class SqlAlchemyAiPlatformRepository: keyword, knowledge_base_ids=knowledge_base_ids, limit=limit, + strategy=strategy, ) def add_ai_run_audit(self, audit: AiRunAudit) -> AiRunAudit: - row = AiRunAuditModel(**audit.model_dump()) + row = AiRunAuditModel( + id=audit.id, + operation=audit.operation, + provider_config_id=audit.provider_config_id, + prompt_template_id=audit.prompt_template_id, + question=audit.question, + answer=audit.answer, + citations_count=audit.citations_count, + retrieval_query=audit.retrieval_query, + retrieval_strategy=audit.retrieval_strategy, + retrieved_chunk_ids_json=_encode_json_list(audit.retrieved_chunk_ids), + retrieved_sources_json=_encode_json_list( + [source.model_dump() for source in audit.retrieved_sources] + ), + latency_ms=audit.latency_ms, + status=audit.status, + created_at=audit.created_at, + ) with self._session_factory.begin() as session: session.add(row) session.flush() @@ -602,6 +683,30 @@ def chunk_to_citation(result: KnowledgeSearchResult) -> LegalCitation: document_id=chunk.document_id, chunk_id=chunk.id, knowledge_base_id=chunk.knowledge_base_id, + keyword_score=result.keyword_score, + vector_score=result.vector_score, + hybrid_score=result.hybrid_score, + retrieval_strategy=result.retrieval_strategy, + matched_terms=result.matched_terms, + ) + + +def search_result_to_retrieval_source(result: KnowledgeSearchResult) -> RetrievalSource: + chunk = result.chunk + return RetrievalSource( + chunk_id=chunk.id, + document_id=chunk.document_id, + knowledge_base_id=chunk.knowledge_base_id, + title=chunk.title, + source_type=chunk.source_type, + reference=chunk.reference, + locator=chunk.locator, + score=result.score, + keyword_score=result.keyword_score, + vector_score=result.vector_score, + hybrid_score=result.hybrid_score, + retrieval_strategy=result.retrieval_strategy, + matched_terms=result.matched_terms, ) @@ -671,6 +776,9 @@ def _knowledge_chunk_from_model(row: KnowledgeChunkModel) -> KnowledgeChunk: content=row.content, chunk_index=row.chunk_index, locator=row.locator, + embedding_model=row.embedding_model, + embedding_dimension=row.embedding_dimension, + embedding_vector=_decode_embedding(row.embedding_json), created_at=_ensure_utc(row.created_at), ) @@ -687,6 +795,9 @@ def _knowledge_chunk_model_from_dto(chunk: KnowledgeChunk) -> KnowledgeChunkMode content=chunk.content, chunk_index=chunk.chunk_index, locator=chunk.locator, + embedding_model=chunk.embedding_model, + embedding_dimension=chunk.embedding_dimension, + embedding_json=_encode_embedding(chunk.embedding_vector), created_at=chunk.created_at, ) @@ -700,6 +811,10 @@ def _ai_run_audit_from_model(row: AiRunAuditModel) -> AiRunAudit: question=row.question, answer=row.answer, citations_count=row.citations_count, + retrieval_query=row.retrieval_query, + retrieval_strategy=row.retrieval_strategy, + retrieved_chunk_ids=_decode_json_list(row.retrieved_chunk_ids_json), + retrieved_sources=[RetrievalSource.model_validate(item) for item in _decode_json_list(row.retrieved_sources_json)], latency_ms=row.latency_ms, status=row.status, created_at=_ensure_utc(row.created_at), @@ -786,6 +901,7 @@ def split_knowledge_document(document: KnowledgeDocument, *, max_chars: int = 50 chunks: list[KnowledgeChunk] = [] for index, content in enumerate(segments, start=1): + embedding_text = f"{document.title}\n{document.reference}\n{content}" chunks.append( KnowledgeChunk( id=f"chunk-{uuid4().hex[:12]}", @@ -798,6 +914,9 @@ def split_knowledge_document(document: KnowledgeDocument, *, max_chars: int = 50 content=content, chunk_index=index, locator=f"chunk-{index}", + embedding_model=EMBEDDING_MODEL, + embedding_dimension=EMBEDDING_DIMENSION, + embedding_vector=_build_embedding(embedding_text), created_at=datetime.now(UTC), ) ) @@ -810,19 +929,48 @@ def _rank_knowledge_chunks( *, knowledge_base_ids: list[str] | None = None, limit: int = 5, + strategy: str = "hybrid", ) -> list[KnowledgeSearchResult]: allowed_base_ids = set(knowledge_base_ids or []) + retrieval_strategy = strategy if strategy in RETRIEVAL_STRATEGIES else "hybrid" terms = _tokenize(keyword) + if keyword.strip() and not terms: + return [] + query_vector = _build_embedding(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: + keyword_score = _score_chunk(chunk, matched_terms, terms_count=len(terms)) + vector_score = _vector_score(query_vector, _chunk_embedding(chunk)) + if terms and not _passes_retrieval_threshold( + keyword_score=keyword_score, + vector_score=vector_score, + strategy=retrieval_strategy, + ): continue - score = _score_chunk(chunk, matched_terms, terms_count=len(terms)) - results.append(KnowledgeSearchResult(chunk=chunk, score=score, matched_terms=matched_terms)) + hybrid_score = round(keyword_score * 0.55 + vector_score * 0.45, 4) + if retrieval_strategy == "keyword": + score = keyword_score + elif retrieval_strategy == "vector": + score = vector_score + else: + score = hybrid_score + if terms and score <= 0: + continue + results.append( + KnowledgeSearchResult( + chunk=chunk, + score=score, + keyword_score=keyword_score, + vector_score=vector_score, + hybrid_score=hybrid_score, + retrieval_strategy=retrieval_strategy, + matched_terms=matched_terms, + ) + ) return sorted(results, key=lambda item: (item.score, item.chunk.created_at), reverse=True)[:limit] @@ -830,8 +978,13 @@ 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] + latin_terms = re.findall(r"[a-z0-9]+", normalized) + cjk_chars = re.findall(r"[\u4e00-\u9fff]", normalized) + cjk_bigrams = [f"{cjk_chars[index]}{cjk_chars[index + 1]}" for index in range(len(cjk_chars) - 1)] + terms = [*latin_terms, *cjk_bigrams] + if not terms: + terms = [term for term in normalized.replace("\n", " ").split(" ") if term] + return [term for term in dict.fromkeys(terms or [normalized]) if term not in GENERIC_RETRIEVAL_TERMS] def _score_chunk(chunk: KnowledgeChunk, matched_terms: list[str], *, terms_count: int) -> float: @@ -844,3 +997,82 @@ def _score_chunk(chunk: KnowledgeChunk, matched_terms: list[str], *, terms_count 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) + + +def _passes_retrieval_threshold( + *, + keyword_score: float, + vector_score: float, + strategy: str, +) -> bool: + if strategy == "keyword": + return keyword_score > 0 + if strategy == "vector": + return vector_score >= MIN_VECTOR_RECALL_SCORE + return keyword_score > 0 or vector_score >= MIN_VECTOR_RECALL_SCORE + + +def _build_embedding(text: str, *, dimension: int = EMBEDDING_DIMENSION) -> list[float]: + terms = _tokenize(text) + if not terms: + return [] + vector = [0.0 for _ in range(dimension)] + for term in terms: + digest = hashlib.sha256(term.encode("utf-8")).digest() + bucket = int.from_bytes(digest[:4], "big") % dimension + sign = 1.0 if digest[4] % 2 == 0 else -1.0 + vector[bucket] += sign + norm = math.sqrt(sum(value * value for value in vector)) + if norm <= 0: + return [] + return [round(value / norm, 6) for value in vector] + + +def _chunk_embedding(chunk: KnowledgeChunk) -> list[float]: + if chunk.embedding_vector: + return chunk.embedding_vector + return _build_embedding(f"{chunk.title}\n{chunk.reference}\n{chunk.content}") + + +def _vector_score(query_vector: list[float], chunk_vector: list[float]) -> float: + if not query_vector or not chunk_vector: + return 0.0 + length = min(len(query_vector), len(chunk_vector)) + dot = sum(query_vector[index] * chunk_vector[index] for index in range(length)) + return round(max(dot, 0.0), 4) + + +def _encode_embedding(vector: list[float]) -> str | None: + if not vector: + return None + return json.dumps(vector, ensure_ascii=False, separators=(",", ":")) + + +def _decode_embedding(payload: str | None) -> list[float]: + if not payload: + return [] + try: + parsed = json.loads(payload) + except json.JSONDecodeError: + return [] + if not isinstance(parsed, list): + return [] + values: list[float] = [] + for item in parsed: + if isinstance(item, int | float): + values.append(float(item)) + return values + + +def _encode_json_list(values: list[object]) -> str: + return json.dumps(values, ensure_ascii=False, separators=(",", ":")) + + +def _decode_json_list(payload: str | None) -> list[object]: + if not payload: + return [] + try: + parsed = json.loads(payload) + except json.JSONDecodeError: + return [] + return parsed if isinstance(parsed, list) else [] 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 8e207646..b6158237 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 @@ -128,14 +128,20 @@ def test_knowledge_import_batch_chunks_and_scored_search() -> None: assert chunks_response.status_code == 200 chunks = chunks_response.json() assert chunks[0]["locator"] == "chunk-1" + assert chunks[0]["embedding_model"] == "local-hash-v1" + assert chunks[0]["embedding_dimension"] == 64 search_response = client.get( "/api/v1/knowledge-chunks/search", - params={"keyword": "deposit refund", "knowledge_base_id": "laws-cn"}, + params={"keyword": "deposit refund", "knowledge_base_id": "laws-cn", "strategy": "hybrid"}, ) assert search_response.status_code == 200 result = search_response.json()[0] assert result["score"] > 0 + assert result["retrieval_strategy"] == "hybrid" + assert result["keyword_score"] > 0 + assert result["vector_score"] >= 0 + assert result["hybrid_score"] == result["score"] assert result["chunk"]["document_id"] == document["id"] batches_response = client.get("/api/v1/knowledge-import-batches", params={"knowledge_base_id": "laws-cn"}) @@ -171,6 +177,8 @@ def test_legal_qa_uses_chunk_search_and_records_audit() -> None: assert payload["citations"][0]["reference"].startswith("Lease-001") assert payload["primary_sources"][0]["score"] > 0 assert payload["primary_sources"][0]["chunk_id"] + assert payload["primary_sources"][0]["retrieval_strategy"] == "hybrid" + assert payload["primary_sources"][0]["hybrid_score"] == payload["primary_sources"][0]["score"] audit_response = client.get("/api/v1/audit/ai-runs") assert audit_response.status_code == 200 @@ -178,6 +186,11 @@ def test_legal_qa_uses_chunk_search_and_records_audit() -> None: assert audits[0]["operation"] == "legal_qa" assert audits[0]["citations_count"] >= 1 assert audits[0]["prompt_template_id"] is not None + assert audits[0]["retrieval_query"] == "How should deposit refund be handled?" + assert audits[0]["retrieval_strategy"] == "hybrid" + assert audits[0]["retrieved_chunk_ids"] + assert audits[0]["retrieved_sources"][0]["reference"] == "Lease-001" + assert audits[0]["retrieved_sources"][0]["score"] > 0 assert audits[0]["status"] == "fallback" @@ -230,6 +243,26 @@ def test_sqlalchemy_repository_persists_platform_data(tmp_path) -> None: question="How should lease deposit refund be handled?", answer="The contract should clearly state deposit refund conditions.", citations_count=1, + retrieval_query="How should lease deposit refund be handled?", + retrieval_strategy="hybrid", + retrieved_chunk_ids=["chunk-test-001"], + retrieved_sources=[ + { + "chunk_id": "chunk-test-001", + "document_id": "doc-test-001", + "knowledge_base_id": "laws-cn", + "title": "Civil Code", + "source_type": "law", + "reference": "Article 703", + "locator": "chunk-1", + "score": 0.8, + "keyword_score": 0.7, + "vector_score": 0.9, + "hybrid_score": 0.8, + "retrieval_strategy": "hybrid", + "matched_terms": ["lease"], + } + ], latency_ms=12, status="succeeded", ) @@ -241,5 +274,11 @@ def test_sqlalchemy_repository_persists_platform_data(tmp_path) -> None: assert reloaded.list_prompt_templates()[0].version == "v2" 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" + chunk_result = reloaded.search_knowledge_chunks("lease", knowledge_base_ids=["laws-cn"])[0] + assert chunk_result.score > 0 + assert chunk_result.hybrid_score == chunk_result.score + assert chunk_result.chunk.embedding_vector + audit = reloaded.list_ai_run_audits()[0] + assert audit.id == "run-test-001" + assert audit.retrieved_chunk_ids == ["chunk-test-001"] + assert audit.retrieved_sources[0].reference == "Article 703" diff --git a/yuqei-ai-sdk-python/README.md b/yuqei-ai-sdk-python/README.md index 0d766f84..d7e90307 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 import batches, document insert/search, chunk listing, and scored chunk search. -- AI run audit listing. +- Knowledge import batches, document insert/search, chunk listing, and keyword/vector/hybrid chunk search. +- AI run audit listing with retrieved chunk ids and structured source metadata. diff --git a/yuqei-ai-sdk-python/src/yuqei_sdk/__init__.py b/yuqei-ai-sdk-python/src/yuqei_sdk/__init__.py index ab2925a0..93fd68cb 100644 --- a/yuqei-ai-sdk-python/src/yuqei_sdk/__init__.py +++ b/yuqei-ai-sdk-python/src/yuqei_sdk/__init__.py @@ -14,6 +14,7 @@ from yuqei_sdk.ai_platform import ( PromptTemplateCreate, ProviderConfig, ProviderConfigCreate, + RetrievalSource, ) from yuqei_sdk.context import RequestContext @@ -34,4 +35,5 @@ __all__ = [ "ProviderConfig", "ProviderConfigCreate", "RequestContext", + "RetrievalSource", ] 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 1bef77d4..89e368ec 100644 --- a/yuqei-ai-sdk-python/src/yuqei_sdk/ai_platform.py +++ b/yuqei-ai-sdk-python/src/yuqei_sdk/ai_platform.py @@ -10,6 +10,7 @@ class LegalQaRequest(BaseModel): question: str = Field(min_length=1) matter_context: str | None = None knowledge_base_ids: list[str] = Field(default_factory=list) + retrieval_strategy: str = "hybrid" class LegalCitation(BaseModel): @@ -23,6 +24,11 @@ class LegalCitation(BaseModel): document_id: str | None = None chunk_id: str | None = None knowledge_base_id: str | None = None + keyword_score: float | None = None + vector_score: float | None = None + hybrid_score: float | None = None + retrieval_strategy: str | None = None + matched_terms: list[str] = Field(default_factory=list) class LegalQaResponse(BaseModel): @@ -101,12 +107,34 @@ class KnowledgeChunk(BaseModel): content: str chunk_index: int locator: str + embedding_model: str = "local-hash-v1" + embedding_dimension: int = 64 created_at: str class KnowledgeSearchResult(BaseModel): chunk: KnowledgeChunk score: float + keyword_score: float = 0.0 + vector_score: float = 0.0 + hybrid_score: float = 0.0 + retrieval_strategy: str = "hybrid" + matched_terms: list[str] = Field(default_factory=list) + + +class RetrievalSource(BaseModel): + chunk_id: str | None = None + document_id: str | None = None + knowledge_base_id: str | None = None + title: str + source_type: str + reference: str + locator: str | None = None + score: float = 0.0 + keyword_score: float = 0.0 + vector_score: float = 0.0 + hybrid_score: float = 0.0 + retrieval_strategy: str = "hybrid" matched_terms: list[str] = Field(default_factory=list) @@ -118,6 +146,10 @@ class AiRunAudit(BaseModel): question: str | None = None answer: str | None = None citations_count: int = 0 + retrieval_query: str | None = None + retrieval_strategy: str | None = None + retrieved_chunk_ids: list[str] = Field(default_factory=list) + retrieved_sources: list[RetrievalSource] = Field(default_factory=list) latency_ms: int = 0 status: str created_at: str @@ -291,9 +323,10 @@ class AiPlatformClient: *, knowledge_base_id: str | None = None, limit: int = 5, + strategy: str = "hybrid", context: RequestContext | None = None, ) -> list[KnowledgeSearchResult]: - params: dict[str, Any] = {"keyword": keyword, "limit": limit} + params: dict[str, Any] = {"keyword": keyword, "limit": limit, "strategy": strategy} if knowledge_base_id: params["knowledge_base_id"] = knowledge_base_id response = self._client.get( diff --git a/yuqei-ai-sdk-python/tests/test_ai_platform.py b/yuqei-ai-sdk-python/tests/test_ai_platform.py index 94c4054d..7c4bf270 100644 --- a/yuqei-ai-sdk-python/tests/test_ai_platform.py +++ b/yuqei-ai-sdk-python/tests/test_ai_platform.py @@ -34,6 +34,11 @@ def test_ai_platform_client_posts_legal_qa_with_context_headers() -> None: "document_id": "doc-1", "chunk_id": "chunk-1", "knowledge_base_id": "laws-cn", + "keyword_score": 0.78, + "vector_score": 0.87, + "hybrid_score": 0.82, + "retrieval_strategy": "hybrid", + "matched_terms": ["probation"], } ], "primary_sources": [ @@ -47,6 +52,11 @@ def test_ai_platform_client_posts_legal_qa_with_context_headers() -> None: "document_id": "doc-1", "chunk_id": "chunk-1", "knowledge_base_id": "laws-cn", + "keyword_score": 0.78, + "vector_score": 0.87, + "hybrid_score": 0.82, + "retrieval_strategy": "hybrid", + "matched_terms": ["probation"], } ], "confidence": 0.72, @@ -65,6 +75,8 @@ def test_ai_platform_client_posts_legal_qa_with_context_headers() -> None: assert response.citations[0].reference == "Article 19 chunk-1" assert response.primary_sources[0].score == 0.82 + assert response.primary_sources[0].hybrid_score == 0.82 + assert response.primary_sources[0].retrieval_strategy == "hybrid" assert captured_request is not None assert captured_request.url.path == "/api/v1/legal/qa" assert captured_request.headers["X-Trace-Id"] == "trace-1" @@ -199,6 +211,8 @@ def test_ai_platform_client_covers_admin_data_endpoints() -> None: "content": "A lease contract defines rent and use of leased property.", "chunk_index": 1, "locator": "chunk-1", + "embedding_model": "local-hash-v1", + "embedding_dimension": 64, "created_at": "2026-06-22T00:00:00Z", } ], @@ -239,9 +253,15 @@ def test_ai_platform_client_covers_admin_data_endpoints() -> None: "content": "A lease contract defines rent and use of leased property.", "chunk_index": 1, "locator": "chunk-1", + "embedding_model": "local-hash-v1", + "embedding_dimension": 64, "created_at": "2026-06-22T00:00:00Z", }, "score": 0.82, + "keyword_score": 0.76, + "vector_score": 0.89, + "hybrid_score": 0.82, + "retrieval_strategy": "hybrid", "matched_terms": ["lease"], } ], @@ -258,6 +278,26 @@ def test_ai_platform_client_covers_admin_data_endpoints() -> None: "question": "How should lease deposit refund be handled?", "answer": "The contract should clearly state deposit refund conditions.", "citations_count": 1, + "retrieval_query": "How should lease deposit refund be handled?", + "retrieval_strategy": "hybrid", + "retrieved_chunk_ids": ["chunk-1"], + "retrieved_sources": [ + { + "chunk_id": "chunk-1", + "document_id": "doc-1", + "knowledge_base_id": "laws-cn", + "title": "Civil Code", + "source_type": "law", + "reference": "Article 703", + "locator": "chunk-1", + "score": 0.82, + "keyword_score": 0.76, + "vector_score": 0.89, + "hybrid_score": 0.82, + "retrieval_strategy": "hybrid", + "matched_terms": ["lease"], + } + ], "latency_ms": 12, "status": "succeeded", "created_at": "2026-06-22T00:00:00Z", @@ -301,7 +341,13 @@ def test_ai_platform_client_covers_admin_data_endpoints() -> None: ) ).reference == "Article 703" assert client.list_knowledge_chunks("doc-1")[0].locator == "chunk-1" + assert client.list_knowledge_chunks("doc-1")[0].embedding_model == "local-hash-v1" assert client.search_knowledge_documents("lease", limit=3)[0].title == "Civil Code" - assert client.search_knowledge_chunks("lease", limit=3)[0].score == 0.82 - assert client.list_ai_run_audits(limit=5)[0].operation == "legal_qa" + chunk_result = client.search_knowledge_chunks("lease", limit=3) + assert chunk_result[0].score == 0.82 + assert chunk_result[0].hybrid_score == 0.82 + audit = client.list_ai_run_audits(limit=5)[0] + assert audit.operation == "legal_qa" + assert audit.retrieved_chunk_ids == ["chunk-1"] + assert audit.retrieved_sources[0].reference == "Article 703" assert "/api/v1/audit/ai-runs" in seen_paths