chore: Set GPT_MODEL environment variable in mini_llm.py

microsoft · Jul 25, 2024 · 0b2ff45 · 0b2ff45
1 parent afc3352
commit 0b2ff45
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 1 deletion.
diff --git a/env.template b/env.template
@@ -1,4 +1,5 @@
 OPENAI_API_KEY=
+MODEL_NAME=
 
 #AZURE_OPENAI_API_KEY=
 #AZURE_OPENAI_ENDPOINT=

diff --git a/nl2dsl/utils/mini_llm.py b/nl2dsl/utils/mini_llm.py
@@ -15,8 +15,10 @@
 else:
     client = OpenAI()
 
-GPT_MODEL = "gpt-4-turbo"
+GPT_MODEL = os.getenv('MODEL_NAME')
 
+if GPT_MODEL is None:
+    raise EnvironmentError("MODEL_NAME environment variable is not set.")
 
 @retry(wait=wait_random_exponential(multiplier=1, max=40), stop=stop_after_attempt(10))
 def mini_llm(utterance, model="gpt-4-turbo", temperature=0.3):