Update models.py
Add OLLAMA_BASE_URL and LM_STUDIO_BASE_URL to the models.py file so both instances can be easily changed when using an installation on another or another docker container.
MasterX committed
Aug 17, 2024 at 17:55 UTC
cf34240e47f2de6139c64b38deb962afcd7c8456
1 file changed
+3
-3
models.py
+3
-3
@@ -22,7 +22,7 @@ def get_api_key(service):
22
23
24
# Ollama models
25
-def get_ollama_chat(model_name:str, temperature=DEFAULT_TEMPERATURE, base_url="http://localhost:11434"):
25
+def get_ollama_chat(model_name:str, temperature=DEFAULT_TEMPERATURE, base_url=os.getenv("OLLAMA_BASE_URL")):
26
return Ollama(model=model_name,temperature=temperature, base_url=base_url)
27
28
def get_ollama_embedding(model_name:str, temperature=DEFAULT_TEMPERATURE):
@@ -34,10 +34,10 @@ def get_huggingface_embedding(model_name:str):
34
return HuggingFaceEmbeddings(model_name=model_name)
35
36
# LM Studio and other OpenAI compatible interfaces
37
-def get_lmstudio_chat(model_name:str, base_url="http://localhost:1234/v1", temperature=DEFAULT_TEMPERATURE):
37
+def get_lmstudio_chat(model_name:str, base_url=os.getenv("LM_STUDIO_BASE_URL"), temperature=DEFAULT_TEMPERATURE):
38
return ChatOpenAI(model_name=model_name, base_url=base_url, temperature=temperature, api_key="none") # type: ignore
39
40
-def get_lmstudio_embedding(model_name:str, base_url="http://localhost:1234/v1"):
40
+def get_lmstudio_embedding(model_name:str, base_url=os.getenv("LM_STUDIO_BASE_URL")):
41
return OpenAIEmbeddings(model_name=model_name, base_url=base_url) # type: ignore
42
43
# Anthropic models