Files
LangBot/src/langbot/pkg/rag/knowledge/kbmgr.py
2026-03-10 12:49:24 +00:00

566 lines
22 KiB
Python

from __future__ import annotations
import mimetypes
import os.path
import traceback
import uuid
import zipfile
import io
from typing import Any
from langbot.pkg.core import app
import sqlalchemy
from langbot.pkg.entity.persistence import rag as persistence_rag
from langbot.pkg.core import taskmgr
from langbot_plugin.api.entities.builtin.rag import context as rag_context
from .base import KnowledgeBaseInterface
class RuntimeKnowledgeBase(KnowledgeBaseInterface):
ap: app.Application
knowledge_base_entity: persistence_rag.KnowledgeBase
def __init__(self, ap: app.Application, knowledge_base_entity: persistence_rag.KnowledgeBase):
super().__init__(ap)
self.knowledge_base_entity = knowledge_base_entity
async def initialize(self):
pass
async def _store_file_task(
self, file: persistence_rag.File, task_context: taskmgr.TaskContext, parser_plugin_id: str | None = None
):
try:
# set file status to processing
await self.ap.persistence_mgr.execute_async(
sqlalchemy.update(persistence_rag.File)
.where(persistence_rag.File.uuid == file.uuid)
.values(status='processing')
)
task_context.set_current_action('Processing file')
# Get file size from storage
file_size = await self.ap.storage_mgr.storage_provider.size(file.file_name)
# Detect MIME type from extension
mime_type, _ = mimetypes.guess_type(file.file_name)
if mime_type is None:
mime_type = 'application/octet-stream'
# If a parser plugin is specified, call it before ingestion
parsed_content = None
if parser_plugin_id:
task_context.set_current_action('Parsing file')
file_bytes = await self.ap.storage_mgr.storage_provider.load(file.file_name)
parse_context = {
'mime_type': mime_type,
'filename': file.file_name,
'metadata': {},
}
parsed_content = await self.ap.plugin_connector.call_parser(parser_plugin_id, parse_context, file_bytes)
# Call plugin to ingest document
result = await self._ingest_document(
{
'document_id': file.uuid,
'filename': file.file_name,
'extension': file.extension,
'file_size': file_size,
'mime_type': mime_type,
},
file.file_name, # storage path
parsed_content=parsed_content,
)
# Check plugin result status
if result.get('status') == 'failed':
error_msg = result.get('error_message', 'Plugin ingestion returned failed status')
raise Exception(error_msg)
# set file status to completed
await self.ap.persistence_mgr.execute_async(
sqlalchemy.update(persistence_rag.File)
.where(persistence_rag.File.uuid == file.uuid)
.values(status='completed')
)
except Exception as e:
self.ap.logger.error(f'Error storing file {file.uuid}: {e}')
traceback.print_exc()
# set file status to failed
await self.ap.persistence_mgr.execute_async(
sqlalchemy.update(persistence_rag.File)
.where(persistence_rag.File.uuid == file.uuid)
.values(status='failed')
)
raise
finally:
# delete file from storage
await self.ap.storage_mgr.storage_provider.delete(file.file_name)
async def store_file(self, file_id: str, parser_plugin_id: str | None = None) -> str:
# pre checking
if not await self.ap.storage_mgr.storage_provider.exists(file_id):
raise Exception(f'File {file_id} not found')
file_name = file_id
_, ext = os.path.splitext(file_name)
extension = ext.lstrip('.').lower() if ext else ''
if extension == 'zip':
return await self._store_zip_file(file_id, parser_plugin_id=parser_plugin_id)
file_uuid = str(uuid.uuid4())
kb_id = self.knowledge_base_entity.uuid
file_obj_data = {
'uuid': file_uuid,
'kb_id': kb_id,
'file_name': file_name,
'extension': extension,
'status': 'pending',
}
file_obj = persistence_rag.File(**file_obj_data)
await self.ap.persistence_mgr.execute_async(sqlalchemy.insert(persistence_rag.File).values(file_obj_data))
# run background task asynchronously
ctx = taskmgr.TaskContext.new()
wrapper = self.ap.task_mgr.create_user_task(
self._store_file_task(file_obj, task_context=ctx, parser_plugin_id=parser_plugin_id),
kind='knowledge-operation',
name=f'knowledge-store-file-{file_id}',
label=f'Store file {file_id}',
context=ctx,
)
return wrapper.id
async def _store_zip_file(self, zip_file_id: str, parser_plugin_id: str | None = None) -> str:
"""Handle ZIP file by extracting each document and storing them separately."""
self.ap.logger.info(f'Processing ZIP file: {zip_file_id}')
zip_bytes = await self.ap.storage_mgr.storage_provider.load(zip_file_id)
supported_extensions = {'txt', 'pdf', 'docx', 'md', 'html'}
stored_file_tasks = []
# use utf-8 encoding
with zipfile.ZipFile(io.BytesIO(zip_bytes), 'r', metadata_encoding='utf-8') as zip_ref:
for file_info in zip_ref.filelist:
# skip directories and hidden files
if file_info.is_dir() or file_info.filename.startswith('.'):
continue
_, file_ext = os.path.splitext(file_info.filename)
file_extension = file_ext.lstrip('.').lower()
if file_extension not in supported_extensions:
self.ap.logger.debug(f'Skipping unsupported file in ZIP: {file_info.filename}')
continue
try:
file_content = zip_ref.read(file_info.filename)
base_name = file_info.filename.replace('/', '_').replace('\\', '_')
file_stem, file_ext = os.path.splitext(base_name)
extension = file_ext.lstrip('.')
if file_stem.startswith('__MACOSX'):
continue
extracted_file_id = file_stem + '_' + str(uuid.uuid4())[:8] + '.' + extension
# save file to storage
await self.ap.storage_mgr.storage_provider.save(extracted_file_id, file_content)
task_id = await self.store_file(extracted_file_id, parser_plugin_id=parser_plugin_id)
stored_file_tasks.append(task_id)
self.ap.logger.info(
f'Extracted and stored file from ZIP: {file_info.filename} -> {extracted_file_id}'
)
except Exception as e:
self.ap.logger.warning(f'Failed to extract file {file_info.filename} from ZIP: {e}')
continue
if not stored_file_tasks:
raise Exception('No supported files found in ZIP archive')
self.ap.logger.info(f'Successfully processed ZIP file {zip_file_id}, extracted {len(stored_file_tasks)} files')
await self.ap.storage_mgr.storage_provider.delete(zip_file_id)
return stored_file_tasks[0] if stored_file_tasks else ''
async def retrieve(self, query: str, settings: dict | None = None) -> list[rag_context.RetrievalResultEntry]:
# Merge stored retrieval_settings with per-request overrides
stored = self.knowledge_base_entity.retrieval_settings or {}
merged = {**stored, **(settings or {})}
if 'top_k' not in merged:
merged['top_k'] = 5 # fallback default
response = await self._retrieve(query, merged)
results_data = response.get('results', [])
entries = []
for r in results_data:
if isinstance(r, dict):
entries.append(rag_context.RetrievalResultEntry(**r))
elif isinstance(r, rag_context.RetrievalResultEntry):
entries.append(r)
return entries
async def delete_file(self, file_id: str):
await self._delete_document(file_id)
# Also cleanup DB record
await self.ap.persistence_mgr.execute_async(
sqlalchemy.delete(persistence_rag.File).where(persistence_rag.File.uuid == file_id)
)
def get_uuid(self) -> str:
"""Get the UUID of the knowledge base"""
return self.knowledge_base_entity.uuid
def get_name(self) -> str:
"""Get the name of the knowledge base"""
return self.knowledge_base_entity.name
def get_knowledge_engine_plugin_id(self) -> str:
"""Get the Knowledge Engine plugin ID"""
return self.knowledge_base_entity.knowledge_engine_plugin_id or ''
async def dispose(self):
"""Dispose the knowledge base, notifying the plugin to cleanup."""
await self._on_kb_delete()
# ========== Plugin Communication Methods ==========
async def _on_kb_create(self) -> None:
"""Notify plugin about KB creation."""
plugin_id = self.knowledge_base_entity.knowledge_engine_plugin_id
if not plugin_id:
return
try:
config = self.knowledge_base_entity.creation_settings or {}
self.ap.logger.info(
f'Calling RAG plugin {plugin_id}: on_knowledge_base_create(kb_id={self.knowledge_base_entity.uuid})'
)
await self.ap.plugin_connector.rag_on_kb_create(plugin_id, self.knowledge_base_entity.uuid, config)
except Exception as e:
self.ap.logger.error(f'Failed to notify plugin {plugin_id} on KB create: {e}')
raise
async def _on_kb_delete(self) -> None:
"""Notify plugin about KB deletion."""
plugin_id = self.knowledge_base_entity.knowledge_engine_plugin_id
if not plugin_id:
return
try:
self.ap.logger.info(
f'Calling RAG plugin {plugin_id}: on_knowledge_base_delete(kb_id={self.knowledge_base_entity.uuid})'
)
await self.ap.plugin_connector.rag_on_kb_delete(plugin_id, self.knowledge_base_entity.uuid)
except Exception as e:
self.ap.logger.error(f'Failed to notify plugin {plugin_id} on KB delete: {e}')
async def _ingest_document(
self,
file_metadata: dict[str, Any],
storage_path: str,
parsed_content: dict[str, Any] | None = None,
) -> dict[str, Any]:
"""Call plugin to ingest document."""
kb = self.knowledge_base_entity
plugin_id = kb.knowledge_engine_plugin_id
if not plugin_id:
self.ap.logger.error(f'No RAG plugin ID configured for KB {kb.uuid}. Ingestion failed.')
raise ValueError('RAG Plugin ID required')
self.ap.logger.info(f'Calling RAG plugin {plugin_id}: ingest(doc={file_metadata.get("filename")})')
# Inject knowledge_base_id into file metadata as required by SDK schema
file_metadata['knowledge_base_id'] = kb.uuid
context_data = {
'file_object': {
'metadata': file_metadata,
'storage_path': storage_path,
},
'knowledge_base_id': kb.uuid,
'collection_id': kb.collection_id or kb.uuid,
'creation_settings': kb.creation_settings or {},
'parsed_content': parsed_content,
}
try:
result = await self.ap.plugin_connector.call_rag_ingest(plugin_id, context_data)
return result
except Exception as e:
self.ap.logger.error(f'Plugin ingestion failed: {e}')
raise
async def _retrieve(
self,
query: str,
settings: dict[str, Any],
) -> dict[str, Any]:
"""Call plugin to retrieve documents.
Raises:
ValueError: If no RAG plugin is configured for this KB.
Exception: If the plugin retrieval call fails.
"""
kb = self.knowledge_base_entity
plugin_id = kb.knowledge_engine_plugin_id
if not plugin_id:
raise ValueError(f'No RAG plugin ID configured for KB {kb.uuid}. Retrieval failed.')
# Session context (e.g. session_name) stays in retrieval_settings
# for plugins that need it. Do NOT move them into filters, as filters
# are passed directly to vector_search by some plugins (e.g. LangRAG)
# and would cause empty results when the metadata field doesn't exist.
filters = settings.pop('filters', {})
retrieval_context = {
'query': query,
'knowledge_base_id': kb.uuid,
'collection_id': kb.collection_id or kb.uuid,
'retrieval_settings': settings,
'creation_settings': kb.creation_settings or {},
'filters': filters,
}
result = await self.ap.plugin_connector.call_rag_retrieve(
plugin_id,
retrieval_context,
)
return result
async def _delete_document(self, document_id: str) -> bool:
"""Call plugin to delete document."""
kb = self.knowledge_base_entity
plugin_id = kb.knowledge_engine_plugin_id
if not plugin_id:
return False
self.ap.logger.info(f'Calling RAG plugin {plugin_id}: delete_document(doc_id={document_id})')
try:
return await self.ap.plugin_connector.call_rag_delete_document(plugin_id, document_id, kb.uuid)
except Exception as e:
self.ap.logger.error(f'Plugin document deletion failed: {e}')
return False
class RAGManager:
ap: app.Application
knowledge_bases: dict[str, KnowledgeBaseInterface]
def __init__(self, ap: app.Application):
self.ap = ap
self.knowledge_bases = {}
async def initialize(self):
await self.load_knowledge_bases_from_db()
async def get_all_knowledge_base_details(self) -> list[dict]:
"""Get all knowledge bases with enriched Knowledge Engine details."""
# 1. Get raw KBs from DB
result = await self.ap.persistence_mgr.execute_async(sqlalchemy.select(persistence_rag.KnowledgeBase))
knowledge_bases = result.all()
# 2. Get all available Knowledge Engines for enrichment
engine_map = {}
if self.ap.plugin_connector.is_enable_plugin:
try:
engines = await self.ap.plugin_connector.list_knowledge_engines()
engine_map = {e['plugin_id']: e for e in engines}
except Exception as e:
self.ap.logger.warning(f'Failed to list Knowledge Engines: {e}')
# 3. Serialize and enrich
kb_list = []
for kb in knowledge_bases:
kb_dict = self.ap.persistence_mgr.serialize_model(persistence_rag.KnowledgeBase, kb)
self._enrich_kb_dict(kb_dict, engine_map)
kb_list.append(kb_dict)
return kb_list
async def get_knowledge_base_details(self, kb_uuid: str) -> dict | None:
"""Get specific knowledge base with enriched Knowledge Engine details."""
result = await self.ap.persistence_mgr.execute_async(
sqlalchemy.select(persistence_rag.KnowledgeBase).where(persistence_rag.KnowledgeBase.uuid == kb_uuid)
)
kb = result.first()
if not kb:
return None
kb_dict = self.ap.persistence_mgr.serialize_model(persistence_rag.KnowledgeBase, kb)
# Fetch engines
engine_map = {}
if self.ap.plugin_connector.is_enable_plugin:
try:
engines = await self.ap.plugin_connector.list_knowledge_engines()
engine_map = {e['plugin_id']: e for e in engines}
except Exception as e:
self.ap.logger.warning(f'Failed to list Knowledge Engines: {e}')
self._enrich_kb_dict(kb_dict, engine_map)
return kb_dict
@staticmethod
def _to_i18n_name(name) -> dict:
"""Ensure name is always an I18nObject-compatible dict.
If *name* is already a dict (with ``en_US`` / ``zh_Hans`` keys) it is
returned as-is. A plain string is wrapped into an I18nObject so the
frontend ``extractI18nObject`` helper never receives an unexpected type.
"""
if isinstance(name, dict):
return name
return {'en_US': str(name), 'zh_Hans': str(name)}
def _enrich_kb_dict(self, kb_dict: dict, engine_map: dict) -> None:
"""Helper to inject engine info into KB dict."""
plugin_id = kb_dict.get('knowledge_engine_plugin_id')
# Default fallback structure — name must be I18nObject for frontend compatibility
fallback_name = self._to_i18n_name(plugin_id or 'Internal (Legacy)')
fallback_info = {
'plugin_id': plugin_id,
'name': fallback_name,
'capabilities': [],
}
if not plugin_id:
kb_dict['knowledge_engine'] = fallback_info
return
engine_info = engine_map.get(plugin_id)
if engine_info:
kb_dict['knowledge_engine'] = {
'plugin_id': plugin_id,
'name': self._to_i18n_name(engine_info.get('name', plugin_id)),
'capabilities': engine_info.get('capabilities', []),
}
else:
kb_dict['knowledge_engine'] = fallback_info
async def create_knowledge_base(
self,
name: str,
knowledge_engine_plugin_id: str,
creation_settings: dict,
retrieval_settings: dict | None = None,
description: str = '',
) -> persistence_rag.KnowledgeBase:
"""Create a new knowledge base using a RAG plugin."""
# Validate that the Knowledge Engine plugin exists
if self.ap.plugin_connector.is_enable_plugin:
try:
engines = await self.ap.plugin_connector.list_knowledge_engines()
engine_ids = [e.get('plugin_id') for e in engines]
if knowledge_engine_plugin_id not in engine_ids:
raise ValueError(f'Knowledge Engine plugin {knowledge_engine_plugin_id} not found')
except ValueError:
raise
except Exception as e:
self.ap.logger.warning(f'Failed to validate Knowledge Engine plugin existence: {e}')
kb_uuid = str(uuid.uuid4())
# Use UUID as collection ID by default for isolation
collection_id = kb_uuid
kb_data = {
'uuid': kb_uuid,
'name': name,
'description': description,
'knowledge_engine_plugin_id': knowledge_engine_plugin_id,
'collection_id': collection_id,
'creation_settings': creation_settings,
'retrieval_settings': retrieval_settings or {},
}
# Create Entity
kb = persistence_rag.KnowledgeBase(**kb_data)
# Persist
await self.ap.persistence_mgr.execute_async(sqlalchemy.insert(persistence_rag.KnowledgeBase).values(kb_data))
# Load into Runtime
runtime_kb = await self.load_knowledge_base(kb)
# Notify Plugin — rollback DB record and runtime entry on failure
try:
await runtime_kb._on_kb_create()
except Exception:
self.knowledge_bases.pop(kb_uuid, None)
await self.ap.persistence_mgr.execute_async(
sqlalchemy.delete(persistence_rag.KnowledgeBase).where(persistence_rag.KnowledgeBase.uuid == kb_uuid)
)
raise
self.ap.logger.info(f'Created new Knowledge Base {name} ({kb_uuid}) using plugin {knowledge_engine_plugin_id}')
return kb
async def load_knowledge_bases_from_db(self):
self.ap.logger.info('Loading knowledge bases from db...')
self.knowledge_bases = {}
# Load knowledge bases
result = await self.ap.persistence_mgr.execute_async(sqlalchemy.select(persistence_rag.KnowledgeBase))
knowledge_bases = result.all()
for knowledge_base in knowledge_bases:
try:
await self.load_knowledge_base(knowledge_base)
except Exception as e:
self.ap.logger.error(
f'Error loading knowledge base {knowledge_base.uuid}: {e}\n{traceback.format_exc()}'
)
async def load_knowledge_base(
self,
knowledge_base_entity: persistence_rag.KnowledgeBase | sqlalchemy.Row | dict,
) -> RuntimeKnowledgeBase:
if isinstance(knowledge_base_entity, sqlalchemy.Row):
# Safe access to _mapping for SQLAlchemy 1.4+
knowledge_base_entity = persistence_rag.KnowledgeBase(**knowledge_base_entity._mapping)
elif isinstance(knowledge_base_entity, dict):
# Filter out non-database fields (like knowledge_engine which is computed)
filtered_dict = {
k: v for k, v in knowledge_base_entity.items() if k in persistence_rag.KnowledgeBase.ALL_DB_FIELDS
}
knowledge_base_entity = persistence_rag.KnowledgeBase(**filtered_dict)
runtime_knowledge_base = RuntimeKnowledgeBase(ap=self.ap, knowledge_base_entity=knowledge_base_entity)
await runtime_knowledge_base.initialize()
self.knowledge_bases[runtime_knowledge_base.get_uuid()] = runtime_knowledge_base
return runtime_knowledge_base
async def get_knowledge_base_by_uuid(self, kb_uuid: str) -> KnowledgeBaseInterface | None:
return self.knowledge_bases.get(kb_uuid)
async def remove_knowledge_base_from_runtime(self, kb_uuid: str):
self.knowledge_bases.pop(kb_uuid, None)
async def delete_knowledge_base(self, kb_uuid: str):
kb = self.knowledge_bases.pop(kb_uuid, None)
if kb is not None:
await kb.dispose()
else:
self.ap.logger.warning(f'Knowledge base {kb_uuid} not found in runtime, skipping plugin notification')