From c157004e07980de2ac969d01a5381e29b26c6547 Mon Sep 17 00:00:00 2001 From: "Timothy J. Baek" Date: Wed, 9 Oct 2024 12:05:16 -0700 Subject: [PATCH] refac --- backend/open_webui/apps/retrieval/utils.py | 19 +++++++------------ 1 file changed, 7 insertions(+), 12 deletions(-) diff --git a/backend/open_webui/apps/retrieval/utils.py b/backend/open_webui/apps/retrieval/utils.py index 8364728d5..027101d1e 100644 --- a/backend/open_webui/apps/retrieval/utils.py +++ b/backend/open_webui/apps/retrieval/utils.py @@ -271,18 +271,13 @@ def get_embedding_function( if embedding_engine == "": return lambda query: embedding_function.encode(query).tolist() elif embedding_engine in ["ollama", "openai"]: - if embedding_engine == "ollama": - func = lambda query: generate_embeddings( - model=embedding_model, - text=query, - ) - elif embedding_engine == "openai": - func = lambda query: generate_embeddings( - model=embedding_model, - text=query, - key=openai_key, - url=openai_url, - ) + func = lambda query: generate_embeddings( + engine=embedding_engine, + model=embedding_model, + text=query, + key=openai_key if embedding_engine == "openai" else "", + url=openai_url if embedding_engine == "openai" else "", + ) def generate_multiple(query, func): if isinstance(query, list):