Merge pull request #217 from leeeizhang/lei/openai-base-url

[MRG] add manual openai base_url
MLSysOps · Sep 20, 2024 · f377b78 · f377b78
2 parents c802a3c + 24fb280
commit f377b78
Showing 1 changed file with 5 additions and 1 deletion.
diff --git a/mle/model.py b/mle/model.py
@@ -1,3 +1,4 @@
+import os
 import copy
 import importlib.util
 import json
@@ -107,7 +108,10 @@ def __init__(self, api_key, model, temperature=0.7):
         self.model = model if model else 'gpt-4o-2024-08-06'
         self.model_type = MODEL_OPENAI
         self.temperature = temperature
-        self.client = self.openai(api_key=api_key)
+        self.client = self.openai(
+            api_key=api_key,
+            base_url=os.getenv("OPENAI_BASE_URL", "https://api.openai.com/v1"),
+        )
         self.func_call_history = []
 
     def query(self, chat_history, **kwargs):