add variable LLM_MODEL_NAME

ittia-research · Aug 2, 2024 · 6ef1589 · 6ef1589
1 parent 98701b8
commit 6ef1589
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 3 deletions.
diff --git a/.env b/.env
@@ -1,4 +1,5 @@
 LLM_LOCAL_BASE_URL=http://xinference:9997/v1
+LLM_MODEL_NAME=google/gemma-2-27b-it
 OLLAMA_BASE_URL=http://ollama:11434
 HOSTING_CHECK_BASE_URL=http://127.0.0.1:8000
 SEARCH_BASE_URL=https://s.jina.ai

diff --git a/src/llm.py b/src/llm.py
@@ -6,8 +6,6 @@
 
 logger = logging.getLogger(__name__)
 
-config_list_openai = [{"model": "gpt-4o-mini", "api_key": os.environ.get("OPENAI_API_KEY")}]
-
 """
 About models:
   - Gemma 2 does not support system rule
@@ -18,9 +16,10 @@
 Todo:
   - With xinference + Gemma 2 + AutoGen, why 'system message' does not work well
 """
+LLM_MODEL_NAME = os.environ.get("LLM_MODEL_NAME") or "google/gemma-2-27b-it"
 config_list_local = [
     # set prices, otherwise there will be warnings
-    {"model": "gemma2:9b-instruct-q6_K", "base_url": os.environ.get("OLLAMA_BASE_URL") + "/v1", "tags": ["gemma", "local"], "price": [0, 0]},
+    {"model": LLM_MODEL_NAME, "base_url": os.environ.get("OLLAMA_BASE_URL") + "/v1", "tags": ["gemma", "local"], "price": [0, 0]},
 ]
 
 llm_config = {"config_list": config_list_local}