import asyncio import json import httpx import quart import sqlalchemy from ... import group from ....authz import Permission from ....context import ExecutionContext, RequestContext from ......core import taskmgr from ......entity.persistence import metadata as persistence_metadata from ......workspace.errors import WorkspaceError, WorkspaceNotFoundError from ......utils import httpclient from langbot_plugin.runtime.plugin.mgr import PluginInstallSource LANGRAG_PLUGIN_AUTHOR = 'langbot-team' LANGRAG_PLUGIN_NAME = 'LangRAG' LANGRAG_PLUGIN_ID = f'{LANGRAG_PLUGIN_AUTHOR}/{LANGRAG_PLUGIN_NAME}' DEFAULT_SPACE_URL = 'https://space.langbot.app' # Old Retriever plugin_name -> New Connector plugin_name EXTERNAL_PLUGIN_NAME_MAPPING = { 'DifyDatasetsRetriever': 'DifyDatasetsConnector', 'RAGFlowRetriever': 'RAGFlowConnector', 'FastGPTRetriever': 'FastGPTConnector', } # Per-plugin: which old retriever_config fields belong to creation_settings. # Remaining fields go to retrieval_settings. # None means ALL fields go to creation_settings (no retrieval_schema). EXTERNAL_PLUGIN_CREATION_FIELDS: dict[str, set[str] | None] = { 'langbot-team/DifyDatasetsConnector': {'api_base_url', 'dify_apikey', 'dataset_id'}, 'langbot-team/RAGFlowConnector': {'api_base_url', 'api_key', 'dataset_ids'}, 'langbot-team/FastGPTConnector': None, # all fields -> creation_settings } _INFORMATION_SCHEMA_TABLES = sqlalchemy.table( 'tables', sqlalchemy.column('table_schema'), sqlalchemy.column('table_name'), schema='information_schema', ) _SQLITE_MASTER = sqlalchemy.table( 'sqlite_master', sqlalchemy.column('type'), sqlalchemy.column('name'), ) _LEGACY_KNOWLEDGE_BASE_BACKUP = sqlalchemy.table( 'knowledge_bases_backup', sqlalchemy.column('uuid'), sqlalchemy.column('name'), sqlalchemy.column('description'), sqlalchemy.column('emoji'), sqlalchemy.column('embedding_model_uuid'), sqlalchemy.column('top_k'), sqlalchemy.column('created_at'), sqlalchemy.column('updated_at'), ) _LEGACY_EXTERNAL_KNOWLEDGE_BASE = sqlalchemy.table( 'external_knowledge_bases', sqlalchemy.column('uuid'), sqlalchemy.column('name'), sqlalchemy.column('description'), sqlalchemy.column('emoji'), sqlalchemy.column('plugin_author'), sqlalchemy.column('plugin_name'), sqlalchemy.column('retriever_config'), sqlalchemy.column('created_at'), ) _CURRENT_KNOWLEDGE_BASE = sqlalchemy.table( 'knowledge_bases', sqlalchemy.column('uuid'), sqlalchemy.column('workspace_uuid'), sqlalchemy.column('name'), sqlalchemy.column('description'), sqlalchemy.column('emoji'), sqlalchemy.column('created_at'), sqlalchemy.column('updated_at'), sqlalchemy.column('knowledge_engine_plugin_id'), sqlalchemy.column('collection_id'), sqlalchemy.column('creation_settings'), sqlalchemy.column('retrieval_settings'), ) @group.group_class('knowledge/migration', '/api/v1/knowledge/migration') class KnowledgeMigrationRouterGroup(group.RouterGroup): async def _require_local_migration_context( self, execution_context: ExecutionContext, ) -> ExecutionContext: """Fence legacy-table migration to the OSS singleton Workspace. The backup tables predate Workspace scoping and are deliberately instance-global. A cloud projection must therefore never be allowed to inspect or restore them, even when it has a valid execution lease. """ try: binding = await self.ap.workspace_service.get_local_execution_binding( execution_context.workspace_uuid, expected_generation=execution_context.placement_generation, ) except WorkspaceNotFoundError: raise except WorkspaceError as exc: raise WorkspaceNotFoundError('RAG migration is unavailable') from exc if binding.instance_uuid != execution_context.instance_uuid: raise WorkspaceNotFoundError('RAG migration is unavailable') return ExecutionContext( instance_uuid=binding.instance_uuid, workspace_uuid=binding.workspace_uuid, placement_generation=binding.placement_generation, ) async def _get_migration_flag(self, execution_context: ExecutionContext) -> bool: """Check if rag_plugin_migration_needed flag is set.""" result = await self.ap.persistence_mgr.execute_async( sqlalchemy.select(persistence_metadata.WorkspaceMetadata.value) .where(persistence_metadata.WorkspaceMetadata.workspace_uuid == execution_context.workspace_uuid) .where(persistence_metadata.WorkspaceMetadata.key == 'rag_plugin_migration_needed') ) return result.scalar_one_or_none() == 'true' async def _set_migration_flag(self, execution_context: ExecutionContext, value: str): """Set rag_plugin_migration_needed flag.""" await self.ap.persistence_mgr.execute_async( sqlalchemy.update(persistence_metadata.WorkspaceMetadata) .where(persistence_metadata.WorkspaceMetadata.workspace_uuid == execution_context.workspace_uuid) .where(persistence_metadata.WorkspaceMetadata.key == 'rag_plugin_migration_needed') .values(value=value) ) async def _table_exists(self, table_name: str) -> bool: """Check if a table exists.""" if self.ap.persistence_mgr.db.name == 'postgresql': result = await self.ap.persistence_mgr.execute_async( sqlalchemy.select(_INFORMATION_SCHEMA_TABLES.c.table_name) .where(_INFORMATION_SCHEMA_TABLES.c.table_schema == 'public') .where(_INFORMATION_SCHEMA_TABLES.c.table_name == table_name) .limit(1) ) return result.first() is not None else: result = await self.ap.persistence_mgr.execute_async( sqlalchemy.select(_SQLITE_MASTER.c.name) .where(_SQLITE_MASTER.c.type == 'table') .where(_SQLITE_MASTER.c.name == table_name) .limit(1) ) return result.first() is not None async def _install_plugin_from_marketplace( self, execution_context: ExecutionContext, plugin_id: str, task_context: taskmgr.TaskContext, space_url: str, ) -> None: """Install a single plugin from the marketplace.""" p_author, p_name = plugin_id.split('/', 1) self.ap.logger.info(f'RAG migration: installing plugin {plugin_id} from marketplace...') task_context.trace(f'Installing plugin {plugin_id} from marketplace...') async with httpx.AsyncClient( trust_env=True, timeout=15, event_hooks=httpclient.httpx_response_limit_hooks(), ) as client: resp = await client.get(f'{space_url}/api/v1/marketplace/plugins/{p_author}/{p_name}') resp.raise_for_status() response_data = await httpclient.parse_json_response(resp) p_data = response_data.get('data', {}).get('plugin', {}) p_version = p_data.get('latest_version') if not p_version: raise Exception(f'Could not determine latest version for {plugin_id}') await self.ap.plugin_connector.require_workspace_context(execution_context) await self.ap.plugin_connector.install_plugin( PluginInstallSource.MARKETPLACE, { 'plugin_author': p_author, 'plugin_name': p_name, 'plugin_version': p_version, }, task_context=task_context, ) self.ap.logger.info(f'RAG migration: plugin {plugin_id} install request sent.') async def _execute_rag_migration( self, execution_context: ExecutionContext, task_context: taskmgr.TaskContext, install_plugin: bool = True, ): """Execute RAG migration: install required plugins and restore backup data.""" execution_context = await self._require_local_migration_context(execution_context) execution_context = await self.ap.plugin_connector.require_workspace_context(execution_context) warnings = [] # Collect all plugins we need: LangRAG (always) + connector plugins (from external KBs) needed_plugins: dict[str, str] = { LANGRAG_PLUGIN_ID: LANGRAG_PLUGIN_NAME, } has_external = await self._table_exists('external_knowledge_bases') if has_external: result = await self.ap.persistence_mgr.execute_async( sqlalchemy.select( _LEGACY_EXTERNAL_KNOWLEDGE_BASE.c.plugin_author, _LEGACY_EXTERNAL_KNOWLEDGE_BASE.c.plugin_name, ).distinct() ) for row in result.fetchall(): plugin_author = row[0] or '' plugin_name = row[1] or '' mapped_name = EXTERNAL_PLUGIN_NAME_MAPPING.get(plugin_name, plugin_name) plugin_id = f'{plugin_author}/{mapped_name}' if plugin_id not in needed_plugins: needed_plugins[plugin_id] = mapped_name self.ap.logger.info(f'RAG migration: plugins needed: {list(needed_plugins.keys())}') if install_plugin: # Step 1: Install all required plugins from marketplace task_context.trace('Installing required plugins...', action='install-plugin') space_url = self.ap.instance_config.data.get('space', {}).get('url', DEFAULT_SPACE_URL).rstrip('/') for plugin_id in needed_plugins: try: await self._install_plugin_from_marketplace( execution_context, plugin_id, task_context, space_url, ) except WorkspaceNotFoundError: raise except Exception as e: self.ap.logger.warning(f'RAG migration: plugin {plugin_id} install returned: {e}') task_context.trace(f'Plugin install note ({plugin_id}): {e}') # Step 2: Wait for all plugins to become available as knowledge engines task_context.trace( f'Waiting for plugins to become available: {list(needed_plugins.keys())}...', action='wait-plugin', ) max_retries = 30 engine_id_set: set[str] = set() for i in range(max_retries): try: await self.ap.plugin_connector.require_workspace_context(execution_context) engines = await self.ap.plugin_connector.list_knowledge_engines() engine_id_set = {e.get('plugin_id') for e in engines} except WorkspaceNotFoundError: raise except Exception: pass if all(pid in engine_id_set for pid in needed_plugins): self.ap.logger.info(f'RAG migration: all plugins ready: {engine_id_set}') task_context.trace('All required plugins are ready.') break if i == max_retries - 1: still_missing = [pid for pid in needed_plugins if pid not in engine_id_set] warning = f'Plugin(s) {still_missing} did not become available after {max_retries} retries' self.ap.logger.warning(f'RAG migration: {warning}') warnings.append(warning) task_context.trace(warning) await asyncio.sleep(2) else: try: await self.ap.plugin_connector.require_workspace_context(execution_context) engines = await self.ap.plugin_connector.list_knowledge_engines() engine_id_set = {e.get('plugin_id') for e in engines} except WorkspaceNotFoundError: raise except Exception: engine_id_set = set() # Step 3: Restore internal knowledge bases from backup task_context.trace('Restoring internal knowledge bases...', action='restore-internal') if await self._table_exists('knowledge_bases_backup'): result = await self.ap.persistence_mgr.execute_async(sqlalchemy.select(_LEGACY_KNOWLEDGE_BASE_BACKUP)) rows = result.fetchall() columns = result.keys() for row in rows: row_dict = dict(zip(columns, row)) kb_uuid = row_dict.get('uuid') name = row_dict.get('name', 'Untitled') description = row_dict.get('description', '') emoji = row_dict.get('emoji', '\U0001f4da') embedding_model_uuid = row_dict.get('embedding_model_uuid', '') top_k = row_dict.get('top_k', 5) created_at = row_dict.get('created_at') updated_at = row_dict.get('updated_at') # DB migration 20 created these columns as TEXT, while a fresh # schema uses SQLAlchemy JSON. Keep the statement structured, # but retain untyped bound values so both physical schemas and # SQLite's string-valued legacy DATETIME rows remain valid. creation_settings = json.dumps({'embedding_model_uuid': embedding_model_uuid}) retrieval_settings = json.dumps({'top_k': top_k}) await self.ap.persistence_mgr.execute_async( sqlalchemy.insert(_CURRENT_KNOWLEDGE_BASE).values( uuid=kb_uuid, workspace_uuid=execution_context.workspace_uuid, name=name, description=description, emoji=emoji, created_at=created_at, updated_at=updated_at, knowledge_engine_plugin_id=LANGRAG_PLUGIN_ID, collection_id=kb_uuid, creation_settings=creation_settings, retrieval_settings=retrieval_settings, ) ) await self.ap.plugin_connector.require_workspace_context(execution_context) try: config = {'embedding_model_uuid': embedding_model_uuid} await self.ap.plugin_connector.rag_on_kb_create(LANGRAG_PLUGIN_ID, kb_uuid, config) task_context.trace(f'Restored internal KB: {name} ({kb_uuid})') except Exception as e: warning = f'Failed to notify plugin for KB {name} ({kb_uuid}): {e}' warnings.append(warning) task_context.trace(warning) await self.ap.rag_mgr.load_knowledge_bases_from_db() # Step 4: Restore external knowledge bases task_context.trace('Restoring external knowledge bases...', action='restore-external') if has_external: result = await self.ap.persistence_mgr.execute_async(sqlalchemy.select(_LEGACY_EXTERNAL_KNOWLEDGE_BASE)) rows = result.fetchall() columns = result.keys() self.ap.logger.info( f'RAG migration: {len(rows)} external KB(s) to restore. Available engines: {engine_id_set}' ) task_context.trace(f'Found {len(rows)} external KB(s). Available engines: {engine_id_set}') for row in rows: row_dict = dict(zip(columns, row)) kb_uuid = row_dict.get('uuid') name = row_dict.get('name', 'Untitled') description = row_dict.get('description', '') emoji = row_dict.get('emoji', '\U0001f517') plugin_author = row_dict.get('plugin_author', '') plugin_name = row_dict.get('plugin_name', '') retriever_config = row_dict.get('retriever_config', {}) created_at = row_dict.get('created_at') mapped_plugin_name = EXTERNAL_PLUGIN_NAME_MAPPING.get(plugin_name, plugin_name) external_plugin_id = f'{plugin_author}/{mapped_plugin_name}' self.ap.logger.info( f'RAG migration: processing external KB "{name}" ({kb_uuid}), ' f'plugin: {plugin_author}/{plugin_name} -> {external_plugin_id}' ) if isinstance(retriever_config, str): try: retriever_config = json.loads(retriever_config) except (json.JSONDecodeError, TypeError): retriever_config = {} creation_fields = EXTERNAL_PLUGIN_CREATION_FIELDS.get(external_plugin_id) if creation_fields is None: creation_settings_dict = retriever_config retrieval_settings_dict = {} else: creation_settings_dict = {k: v for k, v in retriever_config.items() if k in creation_fields} retrieval_settings_dict = {k: v for k, v in retriever_config.items() if k not in creation_fields} await self.ap.persistence_mgr.execute_async( sqlalchemy.insert(_CURRENT_KNOWLEDGE_BASE).values( uuid=kb_uuid, workspace_uuid=execution_context.workspace_uuid, name=name, description=description, emoji=emoji, created_at=created_at, updated_at=created_at, knowledge_engine_plugin_id=external_plugin_id, collection_id=kb_uuid, creation_settings=json.dumps(creation_settings_dict), retrieval_settings=json.dumps(retrieval_settings_dict), ) ) if external_plugin_id not in engine_id_set: warning = ( f'External KB "{name}" ({kb_uuid}) record saved, but plugin {external_plugin_id} ' f'is not installed yet. Install the connector plugin to use it.' ) warnings.append(warning) task_context.trace(warning) else: await self.ap.plugin_connector.require_workspace_context(execution_context) try: await self.ap.plugin_connector.rag_on_kb_create( external_plugin_id, kb_uuid, creation_settings_dict ) task_context.trace(f'Restored external KB: {name} ({kb_uuid})') except Exception as e: warning = f'Failed to notify plugin for external KB {name} ({kb_uuid}): {e}' warnings.append(warning) task_context.trace(warning) await self.ap.rag_mgr.load_knowledge_bases_from_db() # Step 5: Clear migration flag await self._set_migration_flag(execution_context, 'false') task_context.trace('RAG migration completed.', action='done') if warnings: task_context.trace(f'Completed with {len(warnings)} warning(s).') async def initialize(self) -> None: @self.route( '/status', methods=['GET'], auth_type=group.AuthType.USER_TOKEN, permission=Permission.RESOURCE_VIEW, ) async def _(request_context: RequestContext) -> str: execution_context = ExecutionContext.from_request(request_context) execution_context = await self._require_local_migration_context(execution_context) needed = await self._get_migration_flag(execution_context) internal_kb_count = 0 external_kb_count = 0 if needed: if await self._table_exists('knowledge_bases_backup'): result = await self.ap.persistence_mgr.execute_async( sqlalchemy.select(sqlalchemy.func.count()).select_from(_LEGACY_KNOWLEDGE_BASE_BACKUP) ) internal_kb_count = result.scalar() or 0 if await self._table_exists('external_knowledge_bases'): result = await self.ap.persistence_mgr.execute_async( sqlalchemy.select(sqlalchemy.func.count()).select_from(_LEGACY_EXTERNAL_KNOWLEDGE_BASE) ) external_kb_count = result.scalar() or 0 return self.success( data={ 'needed': needed, 'internal_kb_count': internal_kb_count, 'external_kb_count': external_kb_count, } ) @self.route( '/execute', methods=['POST'], auth_type=group.AuthType.USER_TOKEN, permission=Permission.RESOURCE_MANAGE, ) async def _(request_context: RequestContext) -> str: execution_context = ExecutionContext.from_request(request_context) execution_context = await self._require_local_migration_context(execution_context) needed = await self._get_migration_flag(execution_context) if not needed: return self.http_status(400, -1, 'RAG migration is not needed') data = await quart.request.get_json(silent=True) or {} install_plugin = data.get('install_plugin', True) ctx = taskmgr.TaskContext.new() wrapper = self.ap.task_mgr.create_user_task( self._execute_rag_migration( execution_context, task_context=ctx, install_plugin=install_plugin, ), kind='rag-migration', name='rag-migration-execute', label='Migrating knowledge bases to plugin architecture', context=ctx, instance_uuid=execution_context.instance_uuid, workspace_uuid=execution_context.workspace_uuid, placement_generation=execution_context.placement_generation, ) return self.success(data={'task_id': wrapper.id}) @self.route( '/dismiss', methods=['POST'], auth_type=group.AuthType.USER_TOKEN, permission=Permission.RESOURCE_MANAGE, ) async def _(request_context: RequestContext) -> str: execution_context = ExecutionContext.from_request(request_context) execution_context = await self._require_local_migration_context(execution_context) needed = await self._get_migration_flag(execution_context) if not needed: return self.http_status(400, -1, 'RAG migration is not needed') await self._set_migration_flag(execution_context, 'false') return self.success()