bhetherman/open-webui
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

refac

Browse Source
This commit is contained in:
Timothy Jaeryang Baek
2026-08-24 19:15:10 -04:00
parent dc03e7e595
commit 0169b10979
1 changed files with 41 additions and 39 deletions
Download Patch File Download Diff File Expand all files Collapse all files
backend/open_webui/routers/knowledge.py
+41 -39
View File
@@ -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