From ef78c36e988d03c79b8a953414a1c9f7af78daa2 Mon Sep 17 00:00:00 2001
From: nzlov <wfhtqp@gmail.com>
Date: Wed, 26 Feb 2025 17:35:49 +0800
Subject: [PATCH] fix: openai env (#1400)

---
 lua/avante/rag_service.lua | 3 ++-
 py/rag-service/src/main.py | 3 +++
 2 files changed, 5 insertions(+), 1 deletion(-)

diff --git a/lua/avante/rag_service.lua b/lua/avante/rag_service.lua
index 55f36aa..a97324a 100644
--- a/lua/avante/rag_service.lua
+++ b/lua/avante/rag_service.lua
@@ -63,12 +63,13 @@ function M.launch_rag_service(cb)
     Utils.debug(string.format("container %s not found, starting...", container_name))
   end
   local cmd_ = string.format(
-    "docker run -d -p %d:8000 --name %s -v %s:/data -v /:/host -e DATA_DIR=/data -e OPENAI_API_KEY=%s -e OPENAI_BASE_URL=%s %s",
+    "docker run -d -p %d:8000 --name %s -v %s:/data -v /:/host -e DATA_DIR=/data -e OPENAI_API_KEY=%s -e OPENAI_API_BASE=%s -e OPENAI_EMBED_MODEL=%s %s",
     port,
     container_name,
     data_path,
     openai_api_key,
     openai_base_url,
+    os.getenv("OPENAI_EMBED_MODEL"),
     image
   )
   vim.fn.jobstart(cmd_, {
diff --git a/py/rag-service/src/main.py b/py/rag-service/src/main.py
index fa59d60..9e52e6f 100644
--- a/py/rag-service/src/main.py
+++ b/py/rag-service/src/main.py
@@ -315,6 +315,9 @@ chroma_collection = chroma_client.get_or_create_collection("documents")
 vector_store = ChromaVectorStore(chroma_collection=chroma_collection)
 storage_context = StorageContext.from_defaults(vector_store=vector_store)
 embed_model = OpenAIEmbedding()
+model = os.getenv("OPENAI_EMBED_MODEL", "")
+if model:
+    embed_model = OpenAIEmbedding(model=model)
 Settings.embed_model = embed_model