PMFlow/backend/app/services/adaptive_rag_service/mixin_knowledge_docs.py

"""Adaptive RAG mixin module."""
from __future__ import annotations

from .shared import *


class AdaptiveRAGKnowledgeDocsMixin:
    async def create_knowledge_doc(
        self,
        *,
        title: str,
        file_type: str,
        raw_content: str,
        uploaded_by: Optional[int],
        source_type: str = "file",
        file_url: Optional[str] = None,
        project_id: Optional[int] = None,
    ) -> dict[str, Any]:
        doc = await self.create_knowledge_doc_placeholder(
            title=title,
            file_type=file_type,
            uploaded_by=uploaded_by,
            source_type=source_type,
            file_url=file_url,
        )
        await self._index_doc_content(
            doc=doc,
            raw_content=raw_content,
            project_id=project_id,
            source_type=source_type,
        )
        return self._serialize_doc(doc)

    async def create_knowledge_doc_placeholder(
        self,
        *,
        title: str,
        file_type: str,
        uploaded_by: Optional[int],
        source_type: str = "file",
        file_url: Optional[str] = None,
    ) -> KnowledgeBaseDoc:
        doc = KnowledgeBaseDoc(
            tenant_id=self.tenant_id,
            title=(title or "未命名文档")[:255],
            file_type=(file_type or "txt")[:32],
            source_type=(source_type or "file")[:32],
            file_url=(file_url or "")[:1024] or None,
            uploaded_by=uploaded_by,
            status="processing",
            last_fetched_at=now_cn_naive() if source_type == "url" else None,
        )
        self.db.add(doc)
        await self.db.commit()
        await self.db.refresh(doc)
        return doc

    async def process_knowledge_doc_content(
        self,
        *,
        doc_id: int,
        raw_content: str,
        project_id: Optional[int] = None,
    ) -> dict[str, Any]:
        doc = (
            await self.db.execute(
                select(KnowledgeBaseDoc).where(
                    and_(
                        KnowledgeBaseDoc.id == int(doc_id),
                        KnowledgeBaseDoc.tenant_id == self.tenant_id,
                    )
                )
            )
        ).scalar_one_or_none()
        if not doc:
            raise ValueError("doc_not_found")
        await self._index_doc_content(
            doc=doc,
            raw_content=raw_content,
            project_id=project_id,
            source_type=str(doc.source_type or "file"),
        )
        return self._serialize_doc(doc)

    async def _index_doc_content(
        self,
        *,
        doc: KnowledgeBaseDoc,
        raw_content: str,
        project_id: Optional[int],
        source_type: str,
    ) -> None:
        clean_content = str(raw_content or "").strip()
        if not clean_content:
            doc.status = "failed"
            doc.error_message = "文档内容为空，无法向量化"
            await self.db.commit()
            return

        try:
            await self.db.execute(
                delete(KnowledgeBaseChunk).where(
                    and_(
                        KnowledgeBaseChunk.tenant_id == self.tenant_id,
                        KnowledgeBaseChunk.doc_id == int(doc.id),
                    )
                )
            )

            chunk_items = self._chunk_text_with_metadata(clean_content, chunk_size=500, overlap=50)
            for idx, (part, extra_meta) in enumerate(chunk_items):
                emb = await self.llm.get_embedding_async(part, db=self.db, tenant_id=self.tenant_id)
                chunk_metadata = {
                    "title": doc.title,
                    "file_type": doc.file_type,
                    "source_type": source_type,
                    "file_url": doc.file_url,
                }
                chunk_metadata.update(extra_meta)
                self.db.add(
                    KnowledgeBaseChunk(
                        tenant_id=self.tenant_id,
                        doc_id=int(doc.id),
                        source_type="url" if source_type == "url" else "doc",
                        project_id=project_id,
                        content=part,
                        embedding=emb if emb else None,
                        chunk_index=idx,
                        chunk_metadata=chunk_metadata,
                    )
                )
            doc.status = "complete"
            doc.error_message = None
            if source_type == "url":
                doc.last_fetched_at = now_cn_naive()
            await self.db.commit()
        except Exception as exc:
            doc.status = "failed"
            doc.error_message = str(exc)[:500]
            await self.db.commit()
            raise

    async def list_knowledge_docs(self) -> list[dict[str, Any]]:
        docs = (
            await self.db.execute(
                select(KnowledgeBaseDoc)
                .where(KnowledgeBaseDoc.tenant_id == self.tenant_id)
                .order_by(KnowledgeBaseDoc.created_at.desc())
                .limit(200)
            )
        ).scalars().all()
        return [
            self._serialize_doc(d)
            for d in docs
        ]

    async def get_knowledge_doc(self, *, doc_id: int) -> dict[str, Any]:
        doc = (
            await self.db.execute(
                select(KnowledgeBaseDoc).where(
                    and_(
                        KnowledgeBaseDoc.id == int(doc_id),
                        KnowledgeBaseDoc.tenant_id == self.tenant_id,
                    )
                )
            )
        ).scalar_one_or_none()
        if not doc:
            raise ValueError("doc_not_found")
        return self._serialize_doc(doc)

    async def delete_knowledge_doc(self, *, doc_id: int) -> bool:
        doc = (
            await self.db.execute(
                select(KnowledgeBaseDoc).where(
                    and_(
                        KnowledgeBaseDoc.id == int(doc_id),
                        KnowledgeBaseDoc.tenant_id == self.tenant_id,
                    )
                )
            )
        ).scalar_one_or_none()
        if not doc:
            return False

        await self.db.execute(
            delete(KnowledgeBaseChunk).where(
                and_(
                    KnowledgeBaseChunk.doc_id == int(doc_id),
                    KnowledgeBaseChunk.tenant_id == self.tenant_id,
                )
            )
        )
        await self.db.delete(doc)
        await self.db.commit()
        return True