huihuiSquare/digital-avatar-app/backend/services/knowledge_vectorizer.py

"""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 Avatar, KnowledgeChunk, KnowledgeDoc
from services.pdf_ocr_service import extract_scanned_pdf_text
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("原文件不可用，请重新上传")

            self._set_progress(db, doc, "extracting", 8)
            text = embeddings.extract_text(path, f".{doc.file_type}")
            if doc.file_type == "pdf" and not text.strip():
                avatar = db.get(Avatar, doc.avatar_id)
                if not avatar:
                    raise ValueError("文档所属分身不存在")

                def ocr_progress(done: int, total: int):
                    percent = 8 + int((done / max(1, total)) * 20)
                    self._set_progress(db, doc, "ocr", min(percent, 28))

                self._set_progress(db, doc, "ocr", 8)
                text = extract_scanned_pdf_text(
                    db,
                    avatar,
                    path,
                    on_progress=ocr_progress,
                )
            self._set_progress(db, doc, "chunking", 29)
            chunks = embeddings.chunk_text(text)
            if not chunks:
                raise ValueError("文档没有可建立索引的文字内容")
            self._set_progress(db, doc, "embedding", 30)

            def embedding_progress(done: int, total: int):
                percent = 30 + int((done / max(1, total)) * 65)
                self._set_progress(db, doc, "embedding", min(percent, 95))

            vectors = embeddings.embed(chunks, on_progress=embedding_progress)
            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 = ""
            doc.index_stage = "ready"
            doc.index_progress = 100
            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 "建立知识索引失败"
                failed_doc.index_stage = "failed"
                failed_doc.index_progress = 0
                db.commit()
            logger.exception("Knowledge vectorization failed for %s: %s", doc_id, exc)
        finally:
            db.close()

    @staticmethod
    def _set_progress(db, doc, stage: str, progress: int):
        doc.index_stage = stage
        doc.index_progress = progress
        db.commit()


knowledge_vectorizer = KnowledgeVectorizer()