feat: add knowledge reindex workflow

2026-06-22 23:08:59 +08:00
parent a94597f659
commit cf03b0eafd
10 changed files with 333 additions and 8 deletions
@@ -37,6 +37,13 @@
   - OpenAPI `KnowledgeChunk` 新增 `embedding_provider`。
   - Python SDK `KnowledgeChunk` 新增 `embedding_provider`。

+6. 重建索引闭环。
+   - 新增 `POST /api/v1/knowledge-documents/reindex`。
+   - 支持按 `document_id`、`knowledge_base_id`、`import_batch_id` 过滤。
+   - 使用当前 embedding provider 重新切分和生成 chunk。
+   - PostgreSQL + pgvector 环境会重新写入 `knowledge_chunks.embedding`。
+   - Python SDK 新增 `reindex_knowledge_documents(...)`。
+
 ## 生产配置建议

 ```powershell
@@ -55,6 +62,7 @@ python -m alembic upgrade head

 - pgvector 是数据库扩展，需要 PostgreSQL 实例支持 `CREATE EXTENSION vector`。
 - 迁移只新增列，不会自动把历史 JSON embedding 转成真实向量；启用真实 embedding 后需要重导入或重建索引。
+- 启用真实 embedding 后优先调用 `POST /api/v1/knowledge-documents/reindex`，避免历史 chunk 继续停留在 `local-hash-v1`。
 - `local-hash-v1` 继续作为 fallback，不作为最终生产语义向量模型。
 - `AI_PLATFORM_VECTOR_SEARCH_BACKEND=json` 可强制关闭 pgvector，用于排障。
 - `AI_PLATFORM_EMBEDDING_FALLBACK_ENABLED=false` 会在 embedding 参数缺失或 provider 失败时返回空向量，更适合严格生产验收；默认 `true` 更适合平滑上线。
@@ -221,6 +221,23 @@ paths:
            application/json:
              schema:
                $ref: "#/components/schemas/KnowledgeDocument"
+  /api/v1/knowledge-documents/reindex:
+    post:
+      operationId: reindexKnowledgeDocuments
+      summary: Rebuild knowledge document chunks with the current embedding provider
+      requestBody:
+        required: true
+        content:
+          application/json:
+            schema:
+              $ref: "#/components/schemas/KnowledgeReindexRequest"
+      responses:
+        "200":
+          description: Reindex result
+          content:
+            application/json:
+              schema:
+                $ref: "#/components/schemas/KnowledgeReindexResult"
  /api/v1/knowledge-import-batches:
    get:
      operationId: listKnowledgeImportBatches
@@ -796,6 +813,47 @@ components:
        index_status:
          type: string
          default: indexed
+    KnowledgeReindexRequest:
+      type: object
+      properties:
+        document_id:
+          type: string
+          nullable: true
+        knowledge_base_id:
+          type: string
+          nullable: true
+        import_batch_id:
+          type: string
+          nullable: true
+        limit:
+          type: integer
+          default: 100
+          minimum: 1
+          maximum: 1000
+    KnowledgeReindexResult:
+      type: object
+      required:
+        - documents_count
+        - chunks_count
+        - embedding_provider
+        - embedding_model
+        - embedding_dimension
+        - document_ids
+      properties:
+        documents_count:
+          type: integer
+        chunks_count:
+          type: integer
+        embedding_provider:
+          type: string
+        embedding_model:
+          type: string
+        embedding_dimension:
+          type: integer
+        document_ids:
+          type: array
+          items:
+            type: string
    KnowledgeImportBatchCreate:
      type: object
      required:
@@ -67,6 +67,7 @@ Core MVP endpoints:
 - `GET|POST /api/v1/prompt-templates`
 - `GET|POST /api/v1/knowledge-import-batches`
 - `POST /api/v1/knowledge-documents`
+- `POST /api/v1/knowledge-documents/reindex`
 - `GET /api/v1/knowledge-documents/{document_id}/chunks`
 - `GET /api/v1/knowledge-documents/search`
 - `GET /api/v1/knowledge-chunks/search?strategy=keyword|vector|hybrid`
@@ -77,6 +78,11 @@ 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.

+After switching from `local-hash-v1` to a real embedding model, call
+`POST /api/v1/knowledge-documents/reindex` with a `document_id`,
+`knowledge_base_id`, or `import_batch_id` filter to rebuild historical chunks
+with the current embedding provider and refresh pgvector values.
+
 ## Test

 From repository root:
@@ -86,6 +86,7 @@ class OpenAICompatibleEmbeddingProvider:
        self._transport = transport
        self._fallback_provider = fallback_provider or LocalHashEmbeddingProvider()
        self.model = settings.embedding_model
+        self.dimension = settings.embedding_dimension

    def embed_texts(self, texts: list[str]) -> list[EmbeddingVector]:
        if not texts:
@@ -19,6 +19,8 @@ from yuqei_ai_platform_api.repository import (
    KnowledgeDocumentCreate,
    KnowledgeImportBatch,
    KnowledgeImportBatchCreate,
+    KnowledgeReindexRequest,
+    KnowledgeReindexResult,
    KnowledgeSearchResult,
    PromptTemplate,
    PromptTemplateCreate,
@@ -109,6 +111,14 @@ def create_app(
    def add_knowledge_document(payload: KnowledgeDocumentCreate) -> KnowledgeDocument:
        return store.add_knowledge_document(payload)

+    @app.post(
+        f"{resolved_settings.api_prefix}/knowledge-documents/reindex",
+        response_model=KnowledgeReindexResult,
+        tags=["knowledge"],
+    )
+    def reindex_knowledge_documents(payload: KnowledgeReindexRequest) -> KnowledgeReindexResult:
+        return store.reindex_knowledge_documents(payload)
+
    @app.post(
        f"{resolved_settings.api_prefix}/knowledge-import-batches",
        response_model=KnowledgeImportBatch,
@@ -8,7 +8,7 @@ from typing import Protocol
 from uuid import uuid4

 from pydantic import BaseModel, Field
-from sqlalchemy import create_engine, select, text
+from sqlalchemy import create_engine, delete, select, text
 from sqlalchemy.engine import Engine
 from sqlalchemy.orm import Session, sessionmaker
 from sqlalchemy.pool import StaticPool
@@ -124,6 +124,22 @@ class KnowledgeSearchResult(BaseModel):
    matched_terms: list[str] = Field(default_factory=list)


+class KnowledgeReindexRequest(BaseModel):
+    document_id: str | None = None
+    knowledge_base_id: str | None = None
+    import_batch_id: str | None = None
+    limit: int = Field(default=100, ge=1, le=1000)
+
+
+class KnowledgeReindexResult(BaseModel):
+    documents_count: int = 0
+    chunks_count: int = 0
+    embedding_provider: str
+    embedding_model: str
+    embedding_dimension: int = 0
+    document_ids: list[str] = Field(default_factory=list)
+
+
 class RetrievalSource(BaseModel):
    chunk_id: str | None = None
    document_id: str | None = None
@@ -200,6 +216,8 @@ class AiPlatformRepository(Protocol):
        strategy: str = "hybrid",
    ) -> list[KnowledgeSearchResult]: ...

+    def reindex_knowledge_documents(self, payload: KnowledgeReindexRequest) -> KnowledgeReindexResult: ...
+
    def add_ai_run_audit(self, audit: AiRunAudit) -> AiRunAudit: ...

    def list_ai_run_audits(self, *, limit: int = 20) -> list[AiRunAudit]: ...
@@ -379,6 +397,53 @@ class InMemoryAiPlatformRepository:
            embedding_provider=self.embedding_provider,
        )

+    def reindex_knowledge_documents(self, payload: KnowledgeReindexRequest) -> KnowledgeReindexResult:
+        with self._lock:
+            documents = [
+                document
+                for document in self.knowledge_documents.values()
+                if _matches_reindex_request(document, payload)
+            ][: payload.limit]
+            document_ids = {document.id for document in documents}
+            affected_batch_ids = {
+                document.import_batch_id for document in documents if document.import_batch_id
+            }
+            if document_ids:
+                self.knowledge_chunks = {
+                    chunk_id: chunk
+                    for chunk_id, chunk in self.knowledge_chunks.items()
+                    if chunk.document_id not in document_ids
+                }
+
+            reindexed_chunks: list[KnowledgeChunk] = []
+            for document in documents:
+                indexed_document = document.model_copy(update={"index_status": "indexed"})
+                self.knowledge_documents[document.id] = indexed_document
+                chunks = split_knowledge_document(indexed_document, embedding_provider=self.embedding_provider)
+                for chunk in chunks:
+                    self.knowledge_chunks[chunk.id] = chunk
+                reindexed_chunks.extend(chunks)
+
+            self._refresh_import_batch_counts(affected_batch_ids)
+            return _build_reindex_result(documents, reindexed_chunks, self.embedding_provider)
+
+    def _refresh_import_batch_counts(self, batch_ids: set[str]) -> None:
+        for batch_id in batch_ids:
+            batch = self.knowledge_import_batches.get(batch_id)
+            if not batch:
+                continue
+            documents_count = sum(
+                1 for document in self.knowledge_documents.values() if document.import_batch_id == batch_id
+            )
+            chunks_count = sum(1 for chunk in self.knowledge_chunks.values() if chunk.import_batch_id == batch_id)
+            self.knowledge_import_batches[batch_id] = batch.model_copy(
+                update={
+                    "documents_count": documents_count,
+                    "chunks_count": chunks_count,
+                    "updated_at": datetime.now(UTC),
+                }
+            )
+
    def add_ai_run_audit(self, audit: AiRunAudit) -> AiRunAudit:
        with self._lock:
            self.ai_run_audits.insert(0, audit)
@@ -517,13 +582,7 @@ class SqlAlchemyAiPlatformRepository:
            chunks = [_knowledge_chunk_model_from_dto(chunk) for chunk in chunk_dtos]
            session.add_all(chunks)
            session.flush()
-            if self._engine.dialect.name == "postgresql" and self._has_pgvector_embedding_column(session):
-                for chunk in chunk_dtos:
-                    if chunk.embedding_vector:
-                        session.execute(
-                            text("UPDATE knowledge_chunks SET embedding = CAST(:embedding AS vector) WHERE id = :id"),
-                            {"embedding": _to_pgvector_literal(chunk.embedding_vector), "id": chunk.id},
-                        )
+            self._write_pgvector_embeddings(session, chunk_dtos)
            if document.import_batch_id:
                batch = session.get(KnowledgeImportBatchModel, document.import_batch_id)
                if batch:
@@ -624,6 +683,37 @@ class SqlAlchemyAiPlatformRepository:
            embedding_provider=self._embedding_provider,
        )

+    def reindex_knowledge_documents(self, payload: KnowledgeReindexRequest) -> KnowledgeReindexResult:
+        with self._session_factory.begin() as session:
+            statement = select(KnowledgeDocumentModel).order_by(KnowledgeDocumentModel.created_at.desc())
+            if payload.document_id:
+                statement = statement.where(KnowledgeDocumentModel.id == payload.document_id)
+            if payload.knowledge_base_id:
+                statement = statement.where(KnowledgeDocumentModel.knowledge_base_id == payload.knowledge_base_id)
+            if payload.import_batch_id:
+                statement = statement.where(KnowledgeDocumentModel.import_batch_id == payload.import_batch_id)
+            rows = session.scalars(statement.limit(payload.limit)).all()
+            documents = [_knowledge_document_from_model(row) for row in rows]
+            document_ids = [document.id for document in documents]
+            affected_batch_ids = {
+                document.import_batch_id for document in documents if document.import_batch_id
+            }
+            if document_ids:
+                session.execute(
+                    delete(KnowledgeChunkModel).where(KnowledgeChunkModel.document_id.in_(document_ids))
+                )
+
+            reindexed_chunks: list[KnowledgeChunk] = []
+            for row, document in zip(rows, documents):
+                row.index_status = "indexed"
+                chunks = split_knowledge_document(document, embedding_provider=self._embedding_provider)
+                reindexed_chunks.extend(chunks)
+                session.add_all(_knowledge_chunk_model_from_dto(chunk) for chunk in chunks)
+            session.flush()
+            self._write_pgvector_embeddings(session, reindexed_chunks)
+            self._refresh_import_batch_counts(session, affected_batch_ids)
+            return _build_reindex_result(documents, reindexed_chunks, self._embedding_provider)
+
    def _should_use_pgvector(self, strategy: str) -> bool:
        if strategy not in {"vector", "hybrid"}:
            return False
@@ -633,6 +723,35 @@ class SqlAlchemyAiPlatformRepository:
            return self._engine.dialect.name == "postgresql"
        return self._engine.dialect.name == "postgresql"

+    def _write_pgvector_embeddings(self, session: Session, chunks: list[KnowledgeChunk]) -> None:
+        if self._engine.dialect.name != "postgresql" or not self._has_pgvector_embedding_column(session):
+            return
+        for chunk in chunks:
+            if chunk.embedding_vector:
+                session.execute(
+                    text("UPDATE knowledge_chunks SET embedding = CAST(:embedding AS vector) WHERE id = :id"),
+                    {"embedding": _to_pgvector_literal(chunk.embedding_vector), "id": chunk.id},
+                )
+
+    def _refresh_import_batch_counts(self, session: Session, batch_ids: set[str]) -> None:
+        for batch_id in batch_ids:
+            batch = session.get(KnowledgeImportBatchModel, batch_id)
+            if not batch:
+                continue
+            documents_count = len(
+                session.scalars(
+                    select(KnowledgeDocumentModel.id).where(KnowledgeDocumentModel.import_batch_id == batch_id)
+                ).all()
+            )
+            chunks_count = len(
+                session.scalars(
+                    select(KnowledgeChunkModel.id).where(KnowledgeChunkModel.import_batch_id == batch_id)
+                ).all()
+            )
+            batch.documents_count = documents_count
+            batch.chunks_count = chunks_count
+            batch.updated_at = datetime.now(UTC)
+
    def _has_pgvector_embedding_column(self, session: Session) -> bool:
        try:
            return bool(
@@ -818,6 +937,32 @@ def search_result_to_retrieval_source(result: KnowledgeSearchResult) -> Retrieva
    )


+def _matches_reindex_request(document: KnowledgeDocument, payload: KnowledgeReindexRequest) -> bool:
+    if payload.document_id and document.id != payload.document_id:
+        return False
+    if payload.knowledge_base_id and document.knowledge_base_id != payload.knowledge_base_id:
+        return False
+    if payload.import_batch_id and document.import_batch_id != payload.import_batch_id:
+        return False
+    return True
+
+
+def _build_reindex_result(
+    documents: list[KnowledgeDocument],
+    chunks: list[KnowledgeChunk],
+    embedding_provider: EmbeddingProvider,
+) -> KnowledgeReindexResult:
+    first_chunk = chunks[0] if chunks else None
+    return KnowledgeReindexResult(
+        documents_count=len(documents),
+        chunks_count=len(chunks),
+        embedding_provider=first_chunk.embedding_provider if first_chunk else embedding_provider.provider_name,
+        embedding_model=first_chunk.embedding_model if first_chunk else embedding_provider.model,
+        embedding_dimension=first_chunk.embedding_dimension if first_chunk else int(getattr(embedding_provider, "dimension", 0) or 0),
+        document_ids=[document.id for document in documents],
+    )
+
+
 def _provider_config_from_model(row: AiProviderConfigModel) -> ProviderConfig:
    return ProviderConfig(
        id=row.id,
@@ -8,6 +8,7 @@ from yuqei_ai_platform_api.repository import (
    InMemoryAiPlatformRepository,
    KnowledgeDocumentCreate,
    KnowledgeImportBatchCreate,
+    KnowledgeReindexRequest,
    ProviderConfigCreate,
    PromptTemplateCreate,
    SqlAlchemyAiPlatformRepository,
@@ -193,6 +194,44 @@ def test_repository_uses_injected_embedding_provider_for_chunks_and_search() ->
    assert result.vector_score >= 0.35


+def test_knowledge_reindex_rebuilds_chunks_with_current_embedding_provider() -> None:
+    repository = InMemoryAiPlatformRepository()
+    app = create_app(AiPlatformSettings(environment="test"), repository=repository)
+    client = TestClient(app)
+
+    document_response = client.post(
+        "/api/v1/knowledge-documents",
+        json={
+            "knowledge_base_id": "laws-cn",
+            "title": "Lease Policy",
+            "source_type": "policy",
+            "reference": "Lease-001",
+            "content": "Deposit refund must be clearly written.",
+        },
+    )
+    assert document_response.status_code == 200
+    document = document_response.json()
+    assert client.get(f"/api/v1/knowledge-documents/{document['id']}/chunks").json()[0]["embedding_provider"] == "local-hash"
+
+    repository.embedding_provider = TestEmbeddingProvider()
+    reindex_response = client.post(
+        "/api/v1/knowledge-documents/reindex",
+        json={"document_id": document["id"]},
+    )
+
+    assert reindex_response.status_code == 200
+    result = reindex_response.json()
+    assert result["documents_count"] == 1
+    assert result["chunks_count"] == 1
+    assert result["embedding_provider"] == "test-embedding"
+    assert result["embedding_model"] == "test-embedding-v1"
+    assert result["embedding_dimension"] == 3
+    assert result["document_ids"] == [document["id"]]
+    rebuilt_chunk = client.get(f"/api/v1/knowledge-documents/{document['id']}/chunks").json()[0]
+    assert rebuilt_chunk["embedding_provider"] == "test-embedding"
+    assert rebuilt_chunk["embedding_dimension"] == 3
+
+
 def test_legal_qa_uses_chunk_search_and_records_audit() -> None:
    client = make_client()
    client.post(
@@ -316,6 +355,10 @@ 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
+    reindex_result = reloaded.reindex_knowledge_documents(KnowledgeReindexRequest(import_batch_id=batch.id))
+    assert reindex_result.documents_count == 1
+    assert reindex_result.chunks_count >= 1
+    assert reindex_result.embedding_provider == "local-hash"
    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
@@ -6,6 +6,8 @@ from yuqei_sdk.ai_platform import (
    KnowledgeDocumentCreate,
    KnowledgeImportBatch,
    KnowledgeImportBatchCreate,
+    KnowledgeReindexRequest,
+    KnowledgeReindexResult,
    KnowledgeSearchResult,
    LegalCitation,
    LegalQaRequest,
@@ -26,6 +28,8 @@ __all__ = [
    "KnowledgeDocumentCreate",
    "KnowledgeImportBatch",
    "KnowledgeImportBatchCreate",
+    "KnowledgeReindexRequest",
+    "KnowledgeReindexResult",
    "KnowledgeSearchResult",
    "LegalCitation",
    "LegalQaRequest",
@@ -123,6 +123,22 @@ class KnowledgeSearchResult(BaseModel):
    matched_terms: list[str] = Field(default_factory=list)


+class KnowledgeReindexRequest(BaseModel):
+    document_id: str | None = None
+    knowledge_base_id: str | None = None
+    import_batch_id: str | None = None
+    limit: int = 100
+
+
+class KnowledgeReindexResult(BaseModel):
+    documents_count: int = 0
+    chunks_count: int = 0
+    embedding_provider: str
+    embedding_model: str
+    embedding_dimension: int = 0
+    document_ids: list[str] = Field(default_factory=list)
+
+
 class RetrievalSource(BaseModel):
    chunk_id: str | None = None
    document_id: str | None = None
@@ -254,6 +270,20 @@ class AiPlatformClient:
        response.raise_for_status()
        return KnowledgeDocument.model_validate(response.json())

+    def reindex_knowledge_documents(
+        self,
+        request: KnowledgeReindexRequest | None = None,
+        *,
+        context: RequestContext | None = None,
+    ) -> KnowledgeReindexResult:
+        response = self._client.post(
+            f"{self._api_prefix}/knowledge-documents/reindex",
+            json=(request or KnowledgeReindexRequest()).model_dump(exclude_none=True),
+            headers=(context or RequestContext()).to_headers(),
+        )
+        response.raise_for_status()
+        return KnowledgeReindexResult.model_validate(response.json())
+
    def create_knowledge_import_batch(
        self,
        batch: KnowledgeImportBatchCreate,
@@ -1,3 +1,4 @@
+import json
 from urllib.parse import parse_qs

 import httpx
@@ -6,6 +7,7 @@ from yuqei_sdk import (
    AiPlatformClient,
    KnowledgeDocumentCreate,
    KnowledgeImportBatchCreate,
+    KnowledgeReindexRequest,
    LegalQaRequest,
    PromptTemplateCreate,
    ProviderConfigCreate,
@@ -238,6 +240,20 @@ def test_ai_platform_client_covers_admin_data_endpoints() -> None:
                    }
                ],
            )
+        if request.url.path == "/api/v1/knowledge-documents/reindex":
+            payload = json.loads(request.read().decode("utf-8"))
+            assert payload["knowledge_base_id"] == "laws-cn"
+            return httpx.Response(
+                200,
+                json={
+                    "documents_count": 1,
+                    "chunks_count": 1,
+                    "embedding_provider": "openai-compatible",
+                    "embedding_model": "text-embedding-prod",
+                    "embedding_dimension": 1536,
+                    "document_ids": ["doc-1"],
+                },
+            )
        if request.url.path == "/api/v1/knowledge-chunks/search":
            return httpx.Response(
                200,
@@ -346,6 +362,10 @@ def test_ai_platform_client_covers_admin_data_endpoints() -> None:
    assert client.list_knowledge_chunks("doc-1")[0].embedding_provider == "local-hash"
    assert client.list_knowledge_chunks("doc-1")[0].embedding_model == "local-hash-v1"
    assert client.search_knowledge_documents("lease", limit=3)[0].title == "Civil Code"
+    reindex_result = client.reindex_knowledge_documents(KnowledgeReindexRequest(knowledge_base_id="laws-cn"))
+    assert reindex_result.embedding_provider == "openai-compatible"
+    assert reindex_result.embedding_dimension == 1536
+    assert reindex_result.document_ids == ["doc-1"]
    chunk_result = client.search_knowledge_chunks("lease", limit=3)
    assert chunk_result[0].score == 0.82
    assert chunk_result[0].hybrid_score == 0.82