Update embedding.py (#431)

dataelement · Mar 25, 2024 · 5e4ea52 · 5e4ea52
2 parents a4fd107 + 8e33533
commit 5e4ea52
Showing 1 changed file with 4 additions and 3 deletions.
diff --git a/src/backend/bisheng/utils/embedding.py b/src/backend/bisheng/utils/embedding.py
@@ -12,15 +12,16 @@ def decide_embeddings(model: str) -> Embeddings:
     """embed method"""
     model_list = settings.get_knowledge().get('embeddings')
     params = model_list.get(model)
-    if model == 'text-embedding-ada-002' or params.get('component', '') == 'openai':
-        if is_openai_v1() and 'openai_proxy' in params:
+    component = params.pop('component', '')
+    if model == 'text-embedding-ada-002' or component == 'openai':
+        if is_openai_v1() and params.get('openai_proxy'):
             client_params = langchain_bug_openv1(params)
             client_params['http_client'] = httpx.Client(proxies=params['openai_proxy'])
             params['client'] = openai.OpenAI(**client_params).embeddings
             client_params['http_client'] = httpx.AsyncClient(proxies=params['openai_proxy'])
             params['async_client'] = openai.AsyncOpenAI(**client_params).embeddings
         return OpenAIEmbeddings(**params)
-    elif params.get('component', '') == 'custom':
+    elif component == 'custom':
         return CustomHostEmbedding(**params)
     else:
         return HostEmbeddings(**model_list.get(model))