fix(avatar): index knowledge documents asynchronously

2026-09-04 11:53:56 +08:00
parent 59350fb41d
commit b98a2b9507
8 changed files with 312 additions and 69 deletions
@@ -0,0 +1,125 @@
+"""Durable, serial knowledge-document indexing for the avatar knowledge base."""
+
+import json
+import logging
+import os
+import queue
+import threading
+from datetime import datetime, timezone
+
+from database import SessionLocal
+from models import KnowledgeChunk, KnowledgeDoc
+import embeddings
+
+logger = logging.getLogger(__name__)
+
+BACKEND_DIR = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
+UPLOAD_DIR = os.path.abspath(
+    os.getenv("UPLOAD_DIR", os.path.join(BACKEND_DIR, "routers", "uploads"))
+)
+
+
+class KnowledgeVectorizer:
+    """Indexes one document at a time so slow providers cannot block uploads."""
+
+    def __init__(self):
+        self._queue: queue.Queue[str] = queue.Queue()
+        self._queued: set[str] = set()
+        self._lock = threading.Lock()
+        self._thread: threading.Thread | None = None
+
+    def start(self):
+        if self._thread and self._thread.is_alive():
+            return
+        self._thread = threading.Thread(
+            target=self._run, name="knowledge-vectorizer", daemon=True
+        )
+        self._thread.start()
+        db = SessionLocal()
+        try:
+            # A process restart must not abandon documents already accepted by upload.
+            for (doc_id,) in db.query(KnowledgeDoc.id).filter(KnowledgeDoc.status == "parsing"):
+                self.enqueue(doc_id)
+        finally:
+            db.close()
+
+    def enqueue(self, doc_id: str):
+        with self._lock:
+            if doc_id in self._queued:
+                return
+            self._queued.add(doc_id)
+        self._queue.put(doc_id)
+
+    def _run(self):
+        while True:
+            doc_id = self._queue.get()
+            try:
+                self.vectorize_document(doc_id)
+            except Exception:
+                logger.exception("Unexpected knowledge vectorizer failure for %s", doc_id)
+            finally:
+                with self._lock:
+                    self._queued.discard(doc_id)
+                self._queue.task_done()
+
+    def vectorize_document(self, doc_id: str):
+        db = SessionLocal()
+        try:
+            doc = db.get(KnowledgeDoc, doc_id)
+            if not doc or doc.status != "parsing":
+                return
+
+            stored_name = os.path.basename(doc.file_url or "")
+            path = os.path.join(UPLOAD_DIR, doc.avatar_id, stored_name)
+            if not stored_name or not os.path.isfile(path):
+                raise FileNotFoundError("原文件不可用，请重新上传")
+
+            text = embeddings.extract_text(path, f".{doc.file_type}")
+            chunks = embeddings.chunk_text(text)
+            if not chunks:
+                raise ValueError("文档没有可建立索引的文字内容")
+            vectors = embeddings.embed(chunks)
+            if len(vectors) != len(chunks):
+                raise ValueError("向量服务返回数量与文档分段不一致")
+
+            # Commit the document and every chunk together. Chat only sees complete indexes.
+            db.query(KnowledgeChunk).filter(KnowledgeChunk.doc_id == doc.id).delete()
+            db.add_all(
+                [
+                    KnowledgeChunk(
+                        doc_id=doc.id,
+                        avatar_id=doc.avatar_id,
+                        content=chunk,
+                        vector=json.dumps(vector),
+                        chunk_index=index,
+                        embedding_model=embeddings.MODEL,
+                    )
+                    for index, (chunk, vector) in enumerate(zip(chunks, vectors))
+                ]
+            )
+            doc.vectorized = True
+            doc.embedding_model = embeddings.MODEL
+            doc.chunk_count = len(chunks)
+            doc.vectorized_at = datetime.now(timezone.utc)
+            doc.status = "ready"
+            doc.error_message = ""
+            db.commit()
+            logger.info("Knowledge document %s indexed with %s chunks", doc.id, len(chunks))
+        except Exception as exc:
+            db.rollback()
+            failed_doc = db.get(KnowledgeDoc, doc_id)
+            if failed_doc:
+                db.query(KnowledgeChunk).filter(KnowledgeChunk.doc_id == failed_doc.id).delete()
+                failed_doc.status = "failed"
+                failed_doc.vectorized = False
+                failed_doc.embedding_model = ""
+                failed_doc.chunk_count = 0
+                failed_doc.vectorized_at = None
+                failed_doc.error_message = str(exc)[:300] or "建立知识索引失败"
+                db.commit()
+            logger.exception("Knowledge vectorization failed for %s: %s", doc_id, exc)
+        finally:
+            db.close()
+
+
+knowledge_vectorizer = KnowledgeVectorizer()