238 lines
8.2 KiB
Python
238 lines
8.2 KiB
Python
"""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()
|