from __future__ import annotations import asyncio import io import mimetypes import os.path import traceback import uuid import zipfile from typing import Any import sqlalchemy from langbot_plugin.api.entities.builtin.rag import context as rag_context from langbot.pkg.api.http.authz import WorkspaceRequiredError from langbot.pkg.api.http.context import ExecutionContext, RequestContext from langbot.pkg.api.http.service.tenant import TenantContext, require_workspace_uuid from langbot.pkg.core import app, taskmgr from langbot.pkg.core.task_boundary import run_in_workspace_uow from langbot.pkg.entity.persistence import rag as persistence_rag from langbot.pkg.workspace.errors import WorkspaceInvariantError, WorkspaceNotFoundError from .base import KnowledgeBaseInterface _MAX_ZIP_ARCHIVE_ENTRIES = 1024 _MAX_ZIP_DOCUMENTS = 8 _MAX_ZIP_FILE_BYTES = 10 * 1024 * 1024 _MAX_ZIP_UNCOMPRESSED_BYTES = 40 * 1024 * 1024 _MAX_ZIP_COMPRESSION_RATIO = 100 class RuntimeKnowledgeBase(KnowledgeBaseInterface): ap: app.Application knowledge_base_entity: persistence_rag.KnowledgeBase def __init__( self, ap: app.Application, knowledge_base_entity: persistence_rag.KnowledgeBase, execution_context: ExecutionContext, ): super().__init__(ap) self.knowledge_base_entity = knowledge_base_entity self.execution_context = execution_context async def initialize(self): pass async def _assert_execution_context(self, execution_context: ExecutionContext) -> None: """Reject stale or cross-Workspace runtime access.""" if not isinstance(execution_context, ExecutionContext): raise WorkspaceRequiredError('ExecutionContext is required for knowledge runtime access') if ( execution_context.instance_uuid != self.execution_context.instance_uuid or execution_context.workspace_uuid != self.execution_context.workspace_uuid or execution_context.placement_generation != self.execution_context.placement_generation ): raise WorkspaceNotFoundError('Knowledge base not found') if self.knowledge_base_entity.workspace_uuid != execution_context.workspace_uuid: raise WorkspaceNotFoundError('Knowledge base not found') binding = await self.ap.workspace_service.get_execution_binding( execution_context.workspace_uuid, expected_generation=execution_context.placement_generation, ) if binding.instance_uuid != execution_context.instance_uuid: raise WorkspaceNotFoundError('Knowledge base not found') async def _require_plugin_runtime_context( self, execution_context: ExecutionContext, ) -> ExecutionContext: """Fence every singleton Plugin Runtime call to this runtime KB.""" await self._assert_execution_context(execution_context) return await self.ap.plugin_connector.require_workspace_context(execution_context) def _require_upload_object_key( self, execution_context: ExecutionContext, object_key: str, ) -> None: """Reject raw, cross-Workspace, stale, or non-upload object keys.""" try: self.ap.storage_mgr.require_scoped_object_key( execution_context, object_key, expected_owner_type='upload_document', ) except (WorkspaceRequiredError, ValueError) as exc: raise WorkspaceNotFoundError('Upload not found') from exc async def _store_file_task( self, execution_context: ExecutionContext, file: persistence_rag.File, task_context: taskmgr.TaskContext, parser_plugin_id: str | None = None, ): await run_in_workspace_uow( self.ap, execution_context.workspace_uuid, lambda: self._assert_execution_context(execution_context), ) self._require_upload_object_key(execution_context, file.file_name) try: # set file status to processing status_visible = False for retry_delay in (0.0, 0.01, 0.05, 0.1): if retry_delay: await asyncio.sleep(retry_delay) if await self._set_file_status(execution_context, file.uuid, 'processing'): status_visible = True break if not status_visible: raise WorkspaceNotFoundError('Knowledge file was not committed before its background task started') task_context.set_current_action('Processing file') # Get file size from storage file_size = await self.ap.storage_mgr.size_scoped_object_key( execution_context, file.file_name, expected_owner_type='upload_document', ) # 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.load_scoped_object_key( execution_context, file.file_name, expected_owner_type='upload_document', ) parse_context = { 'mime_type': mime_type, 'filename': file.file_name, 'metadata': {}, } await self._require_plugin_runtime_context(execution_context) 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( execution_context, { '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 if not await self._set_file_status(execution_context, file.uuid, 'completed'): raise WorkspaceNotFoundError('Knowledge file not found') except Exception as e: self.ap.logger.error(f'Error storing file {file.uuid}: {e}') traceback.print_exc() # A stale placement is fenced from all writes, including failure # status updates from an old background task. try: if not await self._set_file_status(execution_context, file.uuid, 'failed'): raise WorkspaceNotFoundError('Knowledge file not found') except Exception: self.ap.logger.warning(f'Skipping stale RAG task status update for file {file.uuid}') raise finally: # An old background task must not touch an upload after its # placement generation has been fenced off. try: await self._assert_execution_context(execution_context) await self.ap.storage_mgr.delete_scoped_object_key( execution_context, file.file_name, expected_owner_type='upload_document', ) except (WorkspaceRequiredError, WorkspaceNotFoundError): self.ap.logger.warning(f'Skipping stale RAG upload cleanup for file {file.uuid}') async def _set_file_status( self, execution_context: ExecutionContext, file_uuid: str, status: str, ) -> bool: """Commit one detached-task status transition in its own tenant UoW.""" async def update() -> bool: await self._assert_execution_context(execution_context) result = await self.ap.persistence_mgr.execute_async( sqlalchemy.update(persistence_rag.File) .where(persistence_rag.File.workspace_uuid == execution_context.workspace_uuid) .where(persistence_rag.File.uuid == file_uuid) .values(status=status) ) return getattr(result, 'rowcount', 0) > 0 persistence_mgr = self.ap.persistence_mgr managed_mode = getattr(getattr(persistence_mgr, 'mode', None), 'value', None) in { 'cloud_runtime', 'oss_compat', } tenant_uow = getattr(persistence_mgr, 'tenant_uow', None) if managed_mode: if not callable(tenant_uow): raise RuntimeError('Knowledge tasks require an explicit tenant UoW') async with tenant_uow(execution_context.workspace_uuid): return await update() return await update() async def store_file( self, execution_context: ExecutionContext, file_id: str, parser_plugin_id: str | None = None, ) -> str: await self._assert_execution_context(execution_context) self._require_upload_object_key(execution_context, file_id) # pre checking if not await self.ap.storage_mgr.exists_scoped_object_key( execution_context, file_id, expected_owner_type='upload_document', ): raise WorkspaceNotFoundError('Upload 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(execution_context, 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, 'workspace_uuid': execution_context.workspace_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() try: wrapper = self.ap.task_mgr.create_user_task( self._store_file_task( execution_context, 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, instance_uuid=execution_context.instance_uuid, workspace_uuid=execution_context.workspace_uuid, placement_generation=execution_context.placement_generation, ) except taskmgr.TaskCapacityError: await self.ap.persistence_mgr.execute_async( sqlalchemy.delete(persistence_rag.File) .where(persistence_rag.File.workspace_uuid == execution_context.workspace_uuid) .where(persistence_rag.File.uuid == file_uuid) ) raise return wrapper.id async def _store_zip_file( self, execution_context: ExecutionContext, zip_file_id: str, parser_plugin_id: str | None = None, ) -> str: """Handle ZIP file by extracting each document and storing them separately.""" await self._assert_execution_context(execution_context) self._require_upload_object_key(execution_context, zip_file_id) self.ap.logger.info(f'Processing ZIP file: {zip_file_id}') zip_bytes = await self.ap.storage_mgr.load_scoped_object_key( execution_context, zip_file_id, expected_owner_type='upload_document', ) supported_extensions = {'txt', 'pdf', 'docx', 'md', 'html'} stored_file_tasks = [] try: # use utf-8 encoding with zipfile.ZipFile(io.BytesIO(zip_bytes), 'r', metadata_encoding='utf-8') as zip_ref: if len(zip_ref.filelist) > _MAX_ZIP_ARCHIVE_ENTRIES: raise ValueError('ZIP archive contains too many entries') supported_files: list[zipfile.ZipInfo] = [] total_uncompressed_bytes = 0 for file_info in zip_ref.filelist: # skip directories and hidden files normalized_name = file_info.filename.replace('\\', '/').strip('/') path_parts = normalized_name.split('/') if ( file_info.is_dir() or not normalized_name or any(part.startswith('.') for part in path_parts) or '__MACOSX' in path_parts ): 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 if file_info.flag_bits & 0x1: raise ValueError('Encrypted ZIP entries are not supported') if file_info.file_size > _MAX_ZIP_FILE_BYTES: raise ValueError(f'ZIP document exceeds the file size limit: {file_info.filename}') if ( file_info.file_size and file_info.file_size > max(file_info.compress_size, 1) * _MAX_ZIP_COMPRESSION_RATIO ): raise ValueError(f'ZIP document exceeds the compression-ratio limit: {file_info.filename}') total_uncompressed_bytes += file_info.file_size if total_uncompressed_bytes > _MAX_ZIP_UNCOMPRESSED_BYTES: raise ValueError('ZIP documents exceed the uncompressed size limit') supported_files.append(file_info) if len(supported_files) > _MAX_ZIP_DOCUMENTS: raise ValueError('ZIP archive contains too many supported documents') for file_info in supported_files: try: file_content = await asyncio.to_thread(zip_ref.read, file_info) base_name = file_info.filename.replace('/', '_').replace('\\', '_') file_stem, file_ext = os.path.splitext(base_name) extension = file_ext.lstrip('.') extracted_file_id = file_stem + '_' + str(uuid.uuid4())[:8] + '.' + extension extracted_object_key = await self.ap.storage_mgr.save_scoped( execution_context, owner_type='upload_document', owner=f'knowledge-base:{self.knowledge_base_entity.uuid}', key=extracted_file_id, value=file_content, ) try: task_id = await self.store_file( execution_context, extracted_object_key, parser_plugin_id=parser_plugin_id, ) except Exception: await self.ap.storage_mgr.delete_scoped_object_key( execution_context, extracted_object_key, expected_owner_type='upload_document', ) raise stored_file_tasks.append(task_id) self.ap.logger.info( f'Extracted and stored file from ZIP: {file_info.filename} -> {extracted_object_key}' ) except taskmgr.TaskCapacityError: raise 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' ) return stored_file_tasks[0] if stored_file_tasks else '' finally: try: await self._assert_execution_context(execution_context) await self.ap.storage_mgr.delete_scoped_object_key( execution_context, zip_file_id, expected_owner_type='upload_document', ) except FileNotFoundError: pass except (WorkspaceRequiredError, WorkspaceNotFoundError): self.ap.logger.warning(f'Skipping stale RAG ZIP cleanup for upload {zip_file_id}') except Exception as e: self.ap.logger.warning(f'Failed to cleanup ZIP file {zip_file_id}: {e}') async def retrieve( self, execution_context: ExecutionContext, query: str, settings: dict | None = None, ) -> list[rag_context.RetrievalResultEntry]: await self._assert_execution_context(execution_context) # 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(execution_context, 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, execution_context: ExecutionContext, file_id: str): await self._assert_execution_context(execution_context) result = await self.ap.persistence_mgr.execute_async( sqlalchemy.select(persistence_rag.File.uuid) .where(persistence_rag.File.workspace_uuid == execution_context.workspace_uuid) .where(persistence_rag.File.kb_id == self.knowledge_base_entity.uuid) .where(persistence_rag.File.uuid == file_id) .limit(1) ) if result.first() is None: raise WorkspaceNotFoundError('Knowledge file not found') await self._delete_document(execution_context, file_id) # Also cleanup DB record await self.ap.persistence_mgr.execute_async( sqlalchemy.delete(persistence_rag.File) .where(persistence_rag.File.workspace_uuid == execution_context.workspace_uuid) .where(persistence_rag.File.kb_id == self.knowledge_base_entity.uuid) .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, execution_context: ExecutionContext): """Dispose the knowledge base, notifying the plugin to cleanup.""" await self._assert_execution_context(execution_context) await self._on_kb_delete(execution_context) # ========== Plugin Communication Methods ========== async def _on_kb_create(self, execution_context: ExecutionContext) -> None: """Notify plugin about KB creation.""" await self._assert_execution_context(execution_context) 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._require_plugin_runtime_context(execution_context) 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, execution_context: ExecutionContext) -> None: """Notify plugin about KB deletion.""" await self._assert_execution_context(execution_context) plugin_id = self.knowledge_base_entity.knowledge_engine_plugin_id if not plugin_id: return await self._require_plugin_runtime_context(execution_context) 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, execution_context: ExecutionContext, file_metadata: dict[str, Any], storage_path: str, parsed_content: dict[str, Any] | None = None, ) -> dict[str, Any]: """Call plugin to ingest document.""" await self._assert_execution_context(execution_context) 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, } await self._require_plugin_runtime_context(execution_context) 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, execution_context: ExecutionContext, 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. """ await self._assert_execution_context(execution_context) 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. plugin_settings = dict(settings) filters = plugin_settings.pop('filters', {}) retrieval_context = { 'query': query, 'knowledge_base_id': kb.uuid, 'collection_id': kb.collection_id or kb.uuid, 'retrieval_settings': plugin_settings, 'creation_settings': kb.creation_settings or {}, 'filters': filters, } await self._require_plugin_runtime_context(execution_context) result = await self.ap.plugin_connector.call_rag_retrieve( plugin_id, retrieval_context, ) return result async def _delete_document(self, execution_context: ExecutionContext, document_id: str) -> bool: """Call plugin to delete document.""" await self._assert_execution_context(execution_context) 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})') await self._require_plugin_runtime_context(execution_context) 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[tuple[str, str], RuntimeKnowledgeBase] def __init__(self, ap: app.Application): self.ap = ap self.knowledge_bases = {} self._scope_generations: dict[tuple[str, str], int] = {} self._knowledge_keys_by_scope: dict[ tuple[str, str], set[tuple[str, str]], ] = {} def _cache_runtime( self, runtime: RuntimeKnowledgeBase, ) -> None: context = runtime.execution_context self._observe_execution_context(context) key = ( context.workspace_uuid, runtime.get_uuid(), ) self.knowledge_bases[key] = runtime scope = (context.instance_uuid, context.workspace_uuid) self._knowledge_keys_by_scope.setdefault(scope, set()).add(key) def _pop_runtime( self, context: ExecutionContext, kb_uuid: str, ) -> RuntimeKnowledgeBase | None: key = (context.workspace_uuid, kb_uuid) runtime = self.knowledge_bases.pop(key, None) scope = (context.instance_uuid, context.workspace_uuid) keys = self._knowledge_keys_by_scope.get(scope) if keys is not None: keys.discard(key) if not keys: self._knowledge_keys_by_scope.pop(scope, None) self._scope_generations.pop(scope, None) return runtime def _observe_execution_context(self, context: ExecutionContext) -> None: scope = (context.instance_uuid, context.workspace_uuid) previous_generation = self._scope_generations.get(scope) if previous_generation is not None and context.placement_generation < previous_generation: raise WorkspaceInvariantError('RAG runtime placement generation rolled back') if previous_generation == context.placement_generation: return if previous_generation is not None: for key in self._knowledge_keys_by_scope.pop(scope, ()): self.knowledge_bases.pop(key, None) self._scope_generations[scope] = context.placement_generation async def initialize(self): await self.load_knowledge_bases_from_db() async def _to_execution_context( self, context: RequestContext | ExecutionContext, *, _binding_validated: bool = False, ) -> ExecutionContext: if isinstance(context, RequestContext): execution_context = ExecutionContext.from_request(context) elif isinstance(context, ExecutionContext): execution_context = context else: raise WorkspaceRequiredError('RequestContext or ExecutionContext is required') if not _binding_validated: binding = await self.ap.workspace_service.get_execution_binding( execution_context.workspace_uuid, expected_generation=execution_context.placement_generation, ) if binding.instance_uuid != execution_context.instance_uuid: raise WorkspaceNotFoundError('Workspace not found') scope = ( execution_context.instance_uuid, execution_context.workspace_uuid, ) if scope in self._scope_generations: self._observe_execution_context(execution_context) return execution_context async def _get_engine_map(self, context: TenantContext) -> dict[str, dict]: engine_map: dict[str, dict] = {} connector = getattr(self.ap, 'plugin_connector', None) if connector is not None and connector.is_enable_plugin: await connector.require_workspace_context(context) try: engines = await connector.list_knowledge_engines() engine_map = {engine['plugin_id']: engine for engine in engines} except Exception as e: self.ap.logger.warning(f'Failed to list Knowledge Engines: {e}') return engine_map async def get_all_knowledge_base_details(self, context: TenantContext) -> list[dict]: """Get all knowledge bases with enriched Knowledge Engine details.""" workspace_uuid = require_workspace_uuid(context) # 1. Get raw KBs from DB result = await self.ap.persistence_mgr.execute_async( sqlalchemy.select(persistence_rag.KnowledgeBase).where( persistence_rag.KnowledgeBase.workspace_uuid == workspace_uuid ) ) knowledge_bases = result.all() # 2. Get all available Knowledge Engines for enrichment engine_map = await self._get_engine_map(context) # 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, context: TenantContext, kb_uuid: str) -> dict | None: """Get specific knowledge base with enriched Knowledge Engine details.""" workspace_uuid = require_workspace_uuid(context) result = await self.ap.persistence_mgr.execute_async( sqlalchemy.select(persistence_rag.KnowledgeBase) .where(persistence_rag.KnowledgeBase.workspace_uuid == workspace_uuid) .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 = await self._get_engine_map(context) 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, context: RequestContext | ExecutionContext, 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.""" execution_context = await self._to_execution_context(context) # Validate that the Knowledge Engine plugin exists if self.ap.plugin_connector.is_enable_plugin: await self.ap.plugin_connector.require_workspace_context(execution_context) 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, 'workspace_uuid': execution_context.workspace_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(execution_context, kb) # Notify Plugin — rollback DB record and runtime entry on failure try: await runtime_kb._on_kb_create(execution_context) except Exception: self._pop_runtime(execution_context, kb_uuid) await self.ap.persistence_mgr.execute_async( sqlalchemy.delete(persistence_rag.KnowledgeBase) .where(persistence_rag.KnowledgeBase.workspace_uuid == execution_context.workspace_uuid) .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 = {} self._scope_generations = {} self._knowledge_keys_by_scope = {} list_bindings = getattr(self.ap.workspace_service, 'list_active_execution_bindings', None) tenant_uow = getattr(self.ap.persistence_mgr, 'tenant_uow', None) cloud_runtime = getattr(getattr(self.ap.persistence_mgr, 'mode', None), 'value', None) == 'cloud_runtime' if cloud_runtime: if not callable(list_bindings) or not callable(tenant_uow): raise RuntimeError('Cloud knowledge loading requires explicit instance discovery and tenant UoWs') for binding in await list_bindings(): async with tenant_uow(binding.workspace_uuid): result = await self.ap.persistence_mgr.execute_async( sqlalchemy.select(persistence_rag.KnowledgeBase) .where(persistence_rag.KnowledgeBase.workspace_uuid == binding.workspace_uuid) .order_by(persistence_rag.KnowledgeBase.uuid) ) for knowledge_base in result.all(): try: await self.load_knowledge_base( ExecutionContext( instance_uuid=binding.instance_uuid, workspace_uuid=binding.workspace_uuid, placement_generation=binding.placement_generation, ), knowledge_base, _binding_validated=True, ) except Exception as e: self.ap.logger.error( f'Error loading knowledge base {knowledge_base.uuid}: {e}\n{traceback.format_exc()}' ) return # Compatibility path for isolated manager tests and older embedders. result = await self.ap.persistence_mgr.execute_async(sqlalchemy.select(persistence_rag.KnowledgeBase)) knowledge_bases = result.all() for knowledge_base in knowledge_bases: try: binding = await self.ap.workspace_service.get_execution_binding(knowledge_base.workspace_uuid) execution_context = ExecutionContext( instance_uuid=binding.instance_uuid, workspace_uuid=binding.workspace_uuid, placement_generation=binding.placement_generation, ) await self.load_knowledge_base(execution_context, 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, context: RequestContext | ExecutionContext, knowledge_base_entity: persistence_rag.KnowledgeBase | sqlalchemy.Row | dict, *, _binding_validated: bool = False, ) -> 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) execution_context = await self._to_execution_context( context, _binding_validated=_binding_validated, ) if knowledge_base_entity.workspace_uuid != execution_context.workspace_uuid: raise WorkspaceNotFoundError('Knowledge base not found') runtime_knowledge_base = RuntimeKnowledgeBase( ap=self.ap, knowledge_base_entity=knowledge_base_entity, execution_context=execution_context, ) await runtime_knowledge_base.initialize() self._cache_runtime(runtime_knowledge_base) return runtime_knowledge_base async def get_knowledge_base_by_uuid( self, context: RequestContext | ExecutionContext, kb_uuid: str, ) -> RuntimeKnowledgeBase | None: execution_context = await self._to_execution_context(context) return self.knowledge_bases.get((execution_context.workspace_uuid, kb_uuid)) async def remove_knowledge_base_from_runtime( self, context: RequestContext | ExecutionContext, kb_uuid: str, ) -> None: execution_context = await self._to_execution_context(context) self._pop_runtime(execution_context, kb_uuid) async def delete_knowledge_base( self, context: RequestContext | ExecutionContext, kb_uuid: str, ) -> None: execution_context = await self._to_execution_context(context) kb = self._pop_runtime(execution_context, kb_uuid) if kb is not None: await kb.dispose(execution_context) else: self.ap.logger.warning(f'Knowledge base {kb_uuid} not found in runtime, skipping plugin notification')