mirror of
https://github.com/langbot-app/LangBot.git
synced 2026-06-04 21:06:03 +00:00
81 lines
3.0 KiB
Python
81 lines
3.0 KiB
Python
from __future__ import annotations
|
|
|
|
from ....core import app
|
|
import sqlalchemy
|
|
from langbot.pkg.entity.persistence import rag as persistence_rag
|
|
import uuid
|
|
|
|
|
|
class ExternalKBService:
|
|
"""External KB service"""
|
|
|
|
ap: app.Application
|
|
|
|
def __init__(self, ap: app.Application) -> None:
|
|
self.ap = ap
|
|
|
|
# External Knowledge Base methods
|
|
async def get_external_knowledge_bases(self) -> list[dict]:
|
|
result = await self.ap.persistence_mgr.execute_async(sqlalchemy.select(persistence_rag.ExternalKnowledgeBase))
|
|
external_kbs = result.all()
|
|
return [
|
|
self.ap.persistence_mgr.serialize_model(persistence_rag.ExternalKnowledgeBase, external_kb)
|
|
for external_kb in external_kbs
|
|
]
|
|
|
|
async def get_external_knowledge_base(self, kb_uuid: str) -> dict | None:
|
|
result = await self.ap.persistence_mgr.execute_async(
|
|
sqlalchemy.select(persistence_rag.ExternalKnowledgeBase).where(
|
|
persistence_rag.ExternalKnowledgeBase.uuid == kb_uuid
|
|
)
|
|
)
|
|
external_kb = result.first()
|
|
if external_kb is None:
|
|
return None
|
|
return self.ap.persistence_mgr.serialize_model(persistence_rag.ExternalKnowledgeBase, external_kb)
|
|
|
|
async def create_external_knowledge_base(self, kb_data: dict) -> str:
|
|
kb_data['uuid'] = str(uuid.uuid4())
|
|
await self.ap.persistence_mgr.execute_async(
|
|
sqlalchemy.insert(persistence_rag.ExternalKnowledgeBase).values(kb_data)
|
|
)
|
|
|
|
kb = await self.get_external_knowledge_base(kb_data['uuid'])
|
|
|
|
await self.ap.rag_mgr.load_external_knowledge_base(kb)
|
|
|
|
return kb_data['uuid']
|
|
|
|
async def retrieve_external_knowledge_base(self, kb_uuid: str, query: str) -> list[dict]:
|
|
"""Retrieve external knowledge base"""
|
|
runtime_kb = await self.ap.rag_mgr.get_knowledge_base_by_uuid(kb_uuid)
|
|
if runtime_kb is None:
|
|
raise Exception('Knowledge base not found')
|
|
return [
|
|
result.model_dump() for result in await runtime_kb.retrieve(query, 5)
|
|
] # top_k is just a placeholder for external knowledge base
|
|
|
|
async def update_external_knowledge_base(self, kb_uuid: str, kb_data: dict) -> None:
|
|
if 'uuid' in kb_data:
|
|
del kb_data['uuid']
|
|
|
|
await self.ap.persistence_mgr.execute_async(
|
|
sqlalchemy.update(persistence_rag.ExternalKnowledgeBase)
|
|
.values(kb_data)
|
|
.where(persistence_rag.ExternalKnowledgeBase.uuid == kb_uuid)
|
|
)
|
|
await self.ap.rag_mgr.remove_knowledge_base_from_runtime(kb_uuid)
|
|
|
|
kb = await self.get_external_knowledge_base(kb_uuid)
|
|
|
|
await self.ap.rag_mgr.load_external_knowledge_base(kb)
|
|
|
|
async def delete_external_knowledge_base(self, kb_uuid: str) -> None:
|
|
await self.ap.persistence_mgr.execute_async(
|
|
sqlalchemy.delete(persistence_rag.ExternalKnowledgeBase).where(
|
|
persistence_rag.ExternalKnowledgeBase.uuid == kb_uuid
|
|
)
|
|
)
|
|
|
|
await self.ap.rag_mgr.delete_knowledge_base(kb_uuid)
|