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

refac

Browse Source 
Co-Authored-By: Jacob Leksan <63938553+jmleksan@users.noreply.github.com>
This commit is contained in:
Timothy Jaeryang Baek
2026-07-27 02:17:11 -04:00
co-authored by Jacob Leksan
parent 3fff80ad2f
commit 312d8a8e7f
1 changed files with 11 additions and 4 deletions
Download Patch File Download Diff File Expand all files Collapse all files
backend/open_webui/tools/builtin.py
+11 -4
View File
@@ -3045,9 +3045,10 @@ async def query_knowledge_files(
user_role = __user__.get('role', 'user') user_role = __user__.get('role', 'user')
user_group_ids = [group.id for group in await Groups.get_groups_by_member_id(user_id)] user_group_ids = [group.id for group in await Groups.get_groups_by_member_id(user_id)]
embedding_function = __request__.app.state.EMBEDDING_FUNCTION embedding_function = getattr(__request__.app.state, 'EMBEDDING_FUNCTION', None)
if not embedding_function: if not embedding_function:
return json.dumps({'error': 'Embedding function not configured'}) return json.dumps({'error': 'Embedding function not configured'})
user_model = UserModel.model_construct(id=user_id, role=user_role)
collection_names = [] collection_names = []
external_knowledges = [] external_knowledges = []
@@ -3155,7 +3156,9 @@ async def query_knowledge_files(
__request__, __request__,
collection_names=collection_names, collection_names=collection_names,
queries=[query], queries=[query],
embedding_function=embedding_function, embedding_function=lambda queries, prefix: embedding_function(
queries, prefix=prefix, user=user_model
),
k=count, k=count,
) )
@@ -3180,7 +3183,7 @@ async def query_knowledge_files(
knowledge, knowledge,
queries=[query], queries=[query],
count=count, count=count,
user=type('UserContext', (), {'id': user_id, 'role': user_role})(), user=user_model,
) )
documents = query_results.get('documents', [[]])[0] documents = query_results.get('documents', [[]])[0]
metadatas = query_results.get('metadatas', [[]])[0] metadatas = query_results.get('metadatas', [[]])[0]
@@ -3238,7 +3241,11 @@ async def query_knowledge_bases(
user_id = __user__.get('id') user_id = __user__.get('id')
user_group_ids = [group.id for group in await Groups.get_groups_by_member_id(user_id)] user_group_ids = [group.id for group in await Groups.get_groups_by_member_id(user_id)]
query_embedding = await __request__.app.state.EMBEDDING_FUNCTION(query, prefix=RAG_EMBEDDING_QUERY_PREFIX) embedding_function = getattr(__request__.app.state, 'EMBEDDING_FUNCTION', None)
if not embedding_function:
return json.dumps({'error': 'Embedding function not configured'})
user_model = UserModel.model_construct(id=user_id, role=__user__.get('role', 'user'))
query_embedding = await embedding_function(query, prefix=RAG_EMBEDDING_QUERY_PREFIX, user=user_model)
# Min-heap of (distance, knowledge_base_id) - only holds top `count` results # Min-heap of (distance, knowledge_base_id) - only holds top `count` results
top_results_heap = [] top_results_heap = []