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

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

from .shared import *


class AdaptiveRAGIngestMixin:
    async def ingest_feishu_message(
        self,
        *,
        content: str,
        sender: str,
        chat_id: str,
        message_id: Optional[str],
        captured_at: Optional[datetime] = None,
        project_id: Optional[int] = None,
    ) -> None:
        text = str(content or "").strip()
        if not text:
            return
        chunks = self._chunk_text(text, chunk_size=320, overlap=40)
        for idx, part in enumerate(chunks):
            emb = await self.llm.get_embedding_async(part, db=self.db, tenant_id=self.tenant_id)
            self.db.add(
                KnowledgeBaseChunk(
                    tenant_id=self.tenant_id,
                    doc_id=None,
                    source_type="feishu",
                    project_id=project_id,
                    content=part,
                    embedding=emb if emb else None,
                    chunk_index=idx,
                    chunk_metadata={
                        "sender": sender,
                        "chat_id": chat_id,
                        "message_id": message_id,
                        "captured_at": (captured_at or now_cn_aware()).isoformat(),
                    },
                )
            )
        await self.db.commit()

    async def mark_doc_processing(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")
        doc.status = "processing"
        doc.error_message = None
        await self.db.commit()
        await self.db.refresh(doc)
        return self._serialize_doc(doc)

    @staticmethod
    async def ingest_feishu_message_async(
        *,
        tenant_id: int,
        content: str,
        sender: str,
        chat_id: str,
        message_id: Optional[str],
        captured_at: Optional[datetime] = None,
        project_id: Optional[int] = None,
    ) -> None:
        async with AsyncSessionLocal() as session:
            service = AdaptiveRAGService(session, tenant_id=tenant_id)
            try:
                await service.ingest_feishu_message(
                    content=content,
                    sender=sender,
                    chat_id=chat_id,
                    message_id=message_id,
                    captured_at=captured_at,
                    project_id=project_id,
                )
            except Exception:
                await session.rollback()

    @staticmethod
    async def process_knowledge_doc_content_async(
        *,
        tenant_id: int,
        doc_id: int,
        raw_content: str,
        project_id: Optional[int] = None,
    ) -> None:
        async with AsyncSessionLocal() as session:
            service = AdaptiveRAGService(session, tenant_id=tenant_id)
            try:
                await service.process_knowledge_doc_content(
                    doc_id=doc_id,
                    raw_content=raw_content,
                    project_id=project_id,
                )
            except Exception:
                await session.rollback()

    async def refresh_url_doc(self, *, doc_id: int, 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")
        if str(doc.source_type or "").lower() != "url":
            raise ValueError("doc_is_not_url_source")
        if not str(doc.file_url or "").strip():
            raise ValueError("doc_url_missing")
        text = await self._fetch_url_text(str(doc.file_url))
        await self._index_doc_content(
            doc=doc,
            raw_content=text,
            project_id=project_id,
            source_type="url",
        )
        await self.db.refresh(doc)
        return self._serialize_doc(doc)

    @staticmethod
    async def refresh_url_doc_async(
        *,
        tenant_id: int,
        doc_id: int,
        project_id: Optional[int] = None,
    ) -> None:
        async with AsyncSessionLocal() as session:
            service = AdaptiveRAGService(session, tenant_id=tenant_id)
            try:
                await service.refresh_url_doc(doc_id=doc_id, project_id=project_id)
            except Exception:
                await session.rollback()

    async def upsert_pmflow_record(
        self,
        *,
        record_type: str,
        record_id: int,
        project_id: Optional[int],
        title: str,
        content: str,
    ) -> None:
        safe_type = str(record_type or "").strip().lower()[:32]
        safe_id = int(record_id)
        if not safe_type or safe_id <= 0:
            return

        stmt = select(KnowledgeBaseChunk).where(
            and_(
                KnowledgeBaseChunk.tenant_id == self.tenant_id,
                KnowledgeBaseChunk.source_type == "pmflow",
            )
        )
        if project_id is not None:
            stmt = stmt.where(KnowledgeBaseChunk.project_id == int(project_id))
        existing = (await self.db.execute(stmt.limit(2000))).scalars().all()
        for row in existing:
            meta = self._safe_dict(row.chunk_metadata)
            if str(meta.get("record_type") or "").strip().lower() != safe_type:
                continue
            if str(meta.get("record_id") or "").strip() != str(safe_id):
                continue
            await self.db.delete(row)

        text = f"{title}\n\n{content}".strip()
        chunks = self._chunk_text(text, chunk_size=500, overlap=50)
        for idx, part in enumerate(chunks):
            emb = await self.llm.get_embedding_async(part, db=self.db, tenant_id=self.tenant_id)
            target_url = self._build_pmflow_target_url(
                record_type=safe_type,
                record_id=safe_id,
                project_id=project_id,
            )
            self.db.add(
                KnowledgeBaseChunk(
                    tenant_id=self.tenant_id,
                    doc_id=None,
                    source_type="pmflow",
                    project_id=project_id,
                    content=part,
                    embedding=emb if emb else None,
                    chunk_index=idx,
                    chunk_metadata={
                        "record_type": safe_type,
                        "record_id": safe_id,
                        "project_id": int(project_id) if project_id is not None else None,
                        "target_url": target_url,
                        "title": title[:255],
                        "updated_at": now_cn_aware().isoformat(),
                    },
                )
            )
        await self.db.commit()

    @staticmethod
    async def upsert_pmflow_record_async(
        *,
        tenant_id: int,
        record_type: str,
        record_id: int,
        project_id: Optional[int],
        title: str,
        content: str,
    ) -> None:
        async with AsyncSessionLocal() as session:
            service = AdaptiveRAGService(session, tenant_id=tenant_id)
            try:
                await service.upsert_pmflow_record(
                    record_type=record_type,
                    record_id=record_id,
                    project_id=project_id,
                    title=title,
                    content=content,
                )
            except Exception:
                await session.rollback()