"""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: from .service import AdaptiveRAGService as AdaptiveRAGServiceRuntime service = AdaptiveRAGServiceRuntime(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: from .service import AdaptiveRAGService as AdaptiveRAGServiceRuntime service = AdaptiveRAGServiceRuntime(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: from .service import AdaptiveRAGService as AdaptiveRAGServiceRuntime service = AdaptiveRAGServiceRuntime(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: from .service import AdaptiveRAGService as AdaptiveRAGServiceRuntime service = AdaptiveRAGServiceRuntime(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()