fix: by review

Signed-off-by: yihong0618 <zouzou0208@gmail.com> Signed-off-by: Fangyin Cheng <staneyffer@gmail.com>
eosphoros-ai · Dec 20, 2023 · 0c21315 · 0c21315
1 parent 4c0bab6
commit 0c21315
Show file tree

Hide file tree

Showing 2 changed files with 11 additions and 1 deletion.
diff --git a/dbgpt/config.py → dbgpt/_private/config.py b/dbgpt/config.py → dbgpt/_private/config.py
diff --git a/dbgpt/model/proxy/llms/gemini.py b/dbgpt/model/proxy/llms/gemini.py
@@ -47,7 +47,17 @@ def gemini_generate_stream(
 
     import google.generativeai as genai
 
-    genai.configure(api_key=proxy_api_key)
+    if model_params.proxy_api_base:
+        from google.api_core import client_options
+
+        client_opts = client_options.ClientOptions(
+            api_endpoint=model_params.proxy_api_base
+        )
+        genai.configure(
+            api_key=proxy_api_key, transport="rest", client_options=client_opts
+        )
+    else:
+        genai.configure(api_key=proxy_api_key)
     model = genai.GenerativeModel(
         model_name=proxyllm_backend,
         generation_config=generation_config,