fix parameter position

openchatai · Dec 12, 2023 · 3d9cbe1 · 3d9cbe1
1 parent bea0c85
commit 3d9cbe1
Show file tree

Hide file tree

Showing 2 changed files with 11 additions and 9 deletions.
diff --git a/llm-server/routes/workflow/extractors/convert_json_to_text.py b/llm-server/routes/workflow/extractors/convert_json_to_text.py
@@ -11,7 +11,7 @@
 openai_api_key = os.getenv("OPENAI_API_KEY")
 logger = CustomLogger(module_name=__name__)
 
-chat = get_chat_model(CHAT_MODELS.gpt_3_5_turbo_16k)
+chat = get_chat_model()
 
 
 def convert_json_to_text(

diff --git a/llm-server/utils/get_chat_model.py b/llm-server/utils/get_chat_model.py
@@ -1,21 +1,23 @@
-from typing import Dict
+import os
+from functools import lru_cache
+
+from langchain.chat_models import ChatOllama, ChatAnthropic
 from langchain.chat_models import ChatOpenAI
 from langchain.chat_models.base import BaseChatModel
-from langchain.chat_models import ChatOllama, ChatAnthropic
+
 from utils import llm_consts
 from .chat_models import CHAT_MODELS
-from functools import lru_cache
-import os
 
 localip = os.getenv("LOCAL_IP", "localhost")
 
 model_name = os.getenv(llm_consts.model_env_var, CHAT_MODELS.gpt_4_32k)
 
+
 @lru_cache(maxsize=1)
 def get_chat_model() -> BaseChatModel:
     if model_name == CHAT_MODELS.gpt_3_5_turbo:
         model = ChatOpenAI(
-            model=CHAT_MODELS.gpt_3_5_turbo, 
+            model=CHAT_MODELS.gpt_3_5_turbo,
             temperature=0
         )
     elif model_name == CHAT_MODELS.gpt_4_32k:
@@ -30,12 +32,12 @@ def get_chat_model() -> BaseChatModel:
         )
     elif model_name == CHAT_MODELS.gpt_3_5_turbo_16k:
         model = ChatOpenAI(
-            model=CHAT_MODELS.gpt_3_5_turbo_16k, 
+            model=CHAT_MODELS.gpt_3_5_turbo_16k,
             temperature=0
         )
     elif model_name == "claude":
         model = ChatAnthropic(
-           anthropic_api_key=os.getenv("CLAUDE_API_KEY"),
+            anthropic_api_key=os.getenv("CLAUDE_API_KEY"),
         )
     elif model_name == "openchat":
         model = ChatOllama(
@@ -45,4 +47,4 @@ def get_chat_model() -> BaseChatModel:
         )
     else:
         raise ValueError(f"Unsupported model: {model_name}")
-    return model
+    return model