From 0169b109793c17b66fa1f81ca90a3224c4043823 Mon Sep 17 00:00:00 2001 From: Timothy Jaeryang Baek Date: Mon, 24 Aug 2026 19:15:10 -0400 Subject: [PATCH] refac --- backend/open_webui/routers/knowledge.py | 80 +++++++++++++------------ 1 file changed, 41 insertions(+), 39 deletions(-) diff --git a/backend/open_webui/routers/knowledge.py b/backend/open_webui/routers/knowledge.py index 31865bce63..5c259e8c44 100644 --- a/backend/open_webui/routers/knowledge.py +++ b/backend/open_webui/routers/knowledge.py @@ -541,7 +541,7 @@ EXTERNAL_KNOWLEDGE_CONNECTIONS_CONFIG_KEY = 'external_knowledge.connections' EXTERNAL_KNOWLEDGE_PROVIDERS = {'qdrant', 'milvus', 'pgvector'} -def _validate_external_connection_form(form_data: ExternalKnowledgeConnectionForm) -> tuple[str, dict]: +def _get_external_connection_provider_and_config(form_data: ExternalKnowledgeConnectionForm) -> tuple[str, dict]: provider = form_data.provider.lower().strip() if provider not in EXTERNAL_KNOWLEDGE_PROVIDERS: raise HTTPException( @@ -563,13 +563,13 @@ def _validate_external_connection_form(form_data: ExternalKnowledgeConnectionFor return provider, {key: value for key, value in config.items() if key in allowed_config_keys} -def _external_auth_config(provider: str, incoming: Optional[dict], existing: Optional[dict] = None) -> dict: +def _get_external_auth_config(provider: str, incoming: Optional[dict], existing: Optional[dict] = None) -> dict: if provider == 'pgvector': return {} return existing if incoming is None else incoming or {} -def _normalize_external_source(source: ExternalKnowledgeSourceForm, provider: str) -> ExternalKnowledgeSourceForm: +def _get_normalized_external_source(source: ExternalKnowledgeSourceForm, provider: str) -> ExternalKnowledgeSourceForm: source.type = (source.type or 'collection').strip() source.name = source.name.strip() @@ -600,7 +600,7 @@ def _normalize_external_source(source: ExternalKnowledgeSourceForm, provider: st return source -def _sanitize_external_connection(connection: dict) -> dict: +def _get_sanitized_external_connection(connection: dict) -> dict: sanitized = {**connection} sanitized.pop('auth_config', None) sanitized['auth_configured'] = bool(connection.get('auth_config')) @@ -615,17 +615,17 @@ async def _set_external_connections(connections: list[dict]) -> None: await Config.upsert({EXTERNAL_KNOWLEDGE_CONNECTIONS_CONFIG_KEY: connections}) -def _external_connection_dict( +def _get_external_connection_from_form( form_data: ExternalKnowledgeConnectionForm, user_id: str, id: Optional[str] = None ) -> dict: - provider, config = _validate_external_connection_form(form_data) + provider, config = _get_external_connection_provider_and_config(form_data) now = int(time.time()) return { 'id': id or str(uuid.uuid4()), 'name': form_data.name.strip(), 'provider': provider, 'endpoint': form_data.endpoint.strip(), - 'auth_config': _external_auth_config(provider, form_data.auth_config), + 'auth_config': _get_external_auth_config(provider, form_data.auth_config), 'config': config, 'capabilities': form_data.capabilities or {'retrieve': True}, 'health': None, @@ -636,17 +636,17 @@ def _external_connection_dict( } -def _external_connection_update_dict( +def _get_external_connection_update_from_form( form_data: ExternalKnowledgeConnectionForm, existing: dict, ) -> dict: - provider, config = _validate_external_connection_form(form_data) + provider, config = _get_external_connection_provider_and_config(form_data) return { **existing, 'name': form_data.name.strip(), 'provider': provider, 'endpoint': form_data.endpoint.strip(), - 'auth_config': _external_auth_config(provider, form_data.auth_config, existing.get('auth_config')) or {}, + 'auth_config': _get_external_auth_config(provider, form_data.auth_config, existing.get('auth_config')) or {}, 'config': config, 'capabilities': form_data.capabilities or {'retrieve': True}, 'enabled': form_data.enabled, @@ -654,12 +654,14 @@ def _external_connection_update_dict( } -async def _get_external_connection(id: str) -> Optional[dict]: +async def _get_external_connection_by_id(id: str) -> Optional[dict]: connections = await _get_external_connections() return next((connection for connection in connections if connection.get('id') == id), None) -async def _count_external_connection_mappings(connection_id: str, db: Optional[AsyncSession] = None) -> int: +async def _get_knowledge_base_count_for_external_connection( + connection_id: str, db: Optional[AsyncSession] = None +) -> int: count = 0 for knowledge in await Knowledges.get_knowledge_bases(db=db): if (knowledge.meta or {}).get('external', {}).get('connection_id') == connection_id: @@ -669,7 +671,7 @@ async def _count_external_connection_mappings(connection_id: str, db: Optional[A @router.get('/external/connections', response_model=ExternalKnowledgeConnectionListResponse) async def get_external_knowledge_connections(user=Depends(get_admin_user)): - connections = [_sanitize_external_connection(connection) for connection in await _get_external_connections()] + connections = [_get_sanitized_external_connection(connection) for connection in await _get_external_connections()] return ExternalKnowledgeConnectionListResponse(items=connections, total=len(connections)) @@ -680,10 +682,10 @@ async def create_external_knowledge_connection( user=Depends(get_admin_user), ): connections = await _get_external_connections() - connection = _external_connection_dict(form_data, user.id) + connection = _get_external_connection_from_form(form_data, user.id) connections.append(connection) await _set_external_connections(connections) - sanitized = _sanitize_external_connection(connection) + sanitized = _get_sanitized_external_connection(connection) await publish_event( request, EVENTS.KNOWLEDGE_EXTERNAL_CONNECTION_CREATED, @@ -699,10 +701,10 @@ async def get_external_knowledge_connection( id: str, user=Depends(get_admin_user), ): - connection = await _get_external_connection(id) + connection = await _get_external_connection_by_id(id) if not connection: raise HTTPException(status_code=status.HTTP_404_NOT_FOUND, detail=ERROR_MESSAGES.NOT_FOUND) - return _sanitize_external_connection(connection) + return _get_sanitized_external_connection(connection) @router.patch('/external/connections/{id}', response_model=dict) @@ -717,10 +719,10 @@ async def update_external_knowledge_connection( if idx is None: raise HTTPException(status_code=status.HTTP_404_NOT_FOUND, detail=ERROR_MESSAGES.NOT_FOUND) - connection = _external_connection_update_dict(form_data, connections[idx]) + connection = _get_external_connection_update_from_form(form_data, connections[idx]) connections[idx] = connection await _set_external_connections(connections) - sanitized = _sanitize_external_connection(connection) + sanitized = _get_sanitized_external_connection(connection) await publish_event( request, EVENTS.KNOWLEDGE_EXTERNAL_CONNECTION_UPDATED, @@ -738,11 +740,11 @@ async def delete_external_knowledge_connection( user=Depends(get_admin_user), db: AsyncSession = Depends(get_async_session), ): - connection = await _get_external_connection(id) + connection = await _get_external_connection_by_id(id) if not connection: raise HTTPException(status_code=status.HTTP_404_NOT_FOUND, detail=ERROR_MESSAGES.NOT_FOUND) - if await _count_external_connection_mappings(id, db=db) > 0: + if await _get_knowledge_base_count_for_external_connection(id, db=db) > 0: raise HTTPException( status_code=status.HTTP_400_BAD_REQUEST, detail='External connection is still used by knowledge bases.', @@ -765,7 +767,7 @@ async def test_external_knowledge_connection( id: str, user=Depends(get_admin_user), ): - connection = await _get_external_connection(id) + connection = await _get_external_connection_by_id(id) if not connection: raise HTTPException(status_code=status.HTTP_404_NOT_FOUND, detail=ERROR_MESSAGES.NOT_FOUND) @@ -784,7 +786,7 @@ async def test_external_knowledge_connection( return health -async def _test_external_source_definition( +async def _get_external_source_test_result( request: Request, connection: dict, source: ExternalKnowledgeSourceForm, @@ -795,7 +797,7 @@ async def _test_external_source_definition( if not query.strip(): raise HTTPException(status_code=status.HTTP_400_BAD_REQUEST, detail='Test query is required.') - source = _normalize_external_source(source, connection.get('provider')) + source = _get_normalized_external_source(source, connection.get('provider')) test_knowledge = KnowledgeResponse( id='external-test', user_id=user.id, @@ -838,14 +840,14 @@ async def test_external_knowledge_source( user=Depends(get_admin_user), ): if form_data.connection_id: - existing_connection = await _get_external_connection(form_data.connection_id) + existing_connection = await _get_external_connection_by_id(form_data.connection_id) if not existing_connection: raise HTTPException(status_code=status.HTTP_404_NOT_FOUND, detail='External connection not found.') - connection = _external_connection_update_dict(form_data.connection, existing_connection) + connection = _get_external_connection_update_from_form(form_data.connection, existing_connection) else: - connection = _external_connection_dict(form_data.connection, user.id, id='external-test') + connection = _get_external_connection_from_form(form_data.connection, user.id, id='external-test') - return await _test_external_source_definition( + return await _get_external_source_test_result( request, connection, form_data.source, @@ -862,12 +864,12 @@ async def test_external_knowledge_retrieval( form_data: ExternalKnowledgeRetrieveTestForm, user=Depends(get_admin_user), ): - connection = await _get_external_connection(id) + connection = await _get_external_connection_by_id(id) if not connection: raise HTTPException(status_code=status.HTTP_404_NOT_FOUND, detail=ERROR_MESSAGES.NOT_FOUND) source = form_data.source or ExternalKnowledgeSourceForm(name='test', config={'content_field': 'payload.text'}) - return await _test_external_source_definition(request, connection, source, form_data.query, form_data.count, user) + return await _get_external_source_test_result(request, connection, source, form_data.query, form_data.count, user) @router.post('/external/knowledge/create', response_model=KnowledgeResponse | None) @@ -877,12 +879,12 @@ async def create_external_knowledge( user=Depends(get_admin_user), db: AsyncSession = Depends(get_async_session), ): - connection = await _get_external_connection(form_data.connection_id) + connection = await _get_external_connection_by_id(form_data.connection_id) if not connection: raise HTTPException(status_code=status.HTTP_404_NOT_FOUND, detail=ERROR_MESSAGES.NOT_FOUND) if not form_data.name.strip(): raise HTTPException(status_code=status.HTTP_400_BAD_REQUEST, detail='Knowledge name is required.') - source = _normalize_external_source(form_data.source, connection.get('provider')) + source = _get_normalized_external_source(form_data.source, connection.get('provider')) form_data.access_grants = await filter_allowed_access_grants( await Config.get('user.permissions'), @@ -930,9 +932,9 @@ async def create_external_knowledge_source( if not form_data.name.strip(): raise HTTPException(status_code=status.HTTP_400_BAD_REQUEST, detail='Knowledge name is required.') - connection = _external_connection_dict(form_data.connection, user.id) - source = _normalize_external_source(form_data.source, connection.get('provider')) - test_result = await _test_external_source_definition( + connection = _get_external_connection_from_form(form_data.connection, user.id) + source = _get_normalized_external_source(form_data.source, connection.get('provider')) + test_result = await _get_external_source_test_result( request, connection, source, @@ -1006,9 +1008,9 @@ async def update_external_knowledge_source( raise HTTPException(status_code=status.HTTP_404_NOT_FOUND, detail='External connection not found.') existing_connection = connections[idx] - connection = _external_connection_update_dict(form_data.connection, existing_connection) - source = _normalize_external_source(form_data.source, connection.get('provider')) - test_result = await _test_external_source_definition( + connection = _get_external_connection_update_from_form(form_data.connection, existing_connection) + source = _get_normalized_external_source(form_data.source, connection.get('provider')) + test_result = await _get_external_source_test_result( request, connection, source, @@ -1737,7 +1739,7 @@ async def delete_knowledge_by_id( if ( connection_id and user.role == 'admin' - and await _count_external_connection_mappings(connection_id, db=db) <= 1 + and await _get_knowledge_base_count_for_external_connection(connection_id, db=db) <= 1 ): connections = [ connection for connection in await _get_external_connections() if connection.get('id') != connection_id