From 1b6b1347222b79e218560cfa892a6ca443a2d4fa Mon Sep 17 00:00:00 2001
From: pablodanswer <pablo@danswer.ai>
Date: Thu, 24 Oct 2024 10:29:36 -0700
Subject: [PATCH] Clearer azure models (#2898)

* clear up llm

* remove logs
---
 backend/danswer/llm/factory.py              | 4 +++-
 backend/danswer/server/manage/llm/api.py    | 2 ++
 backend/danswer/server/manage/llm/models.py | 1 +
 backend/requirements/model_server.txt       | 4 ++--
 4 files changed, 8 insertions(+), 3 deletions(-)

diff --git a/backend/danswer/llm/factory.py b/backend/danswer/llm/factory.py
index f930c3d33..eedf7ccc7 100644
--- a/backend/danswer/llm/factory.py
+++ b/backend/danswer/llm/factory.py
@@ -51,6 +51,7 @@ def get_llms_for_persona(
         return get_llm(
             provider=llm_provider.provider,
             model=model,
+            deployment_name=llm_provider.deployment_name,
             api_key=llm_provider.api_key,
             api_base=llm_provider.api_base,
             api_version=llm_provider.api_version,
@@ -104,7 +105,7 @@ def get_default_llms(
 def get_llm(
     provider: str,
     model: str,
-    deployment_name: str | None = None,
+    deployment_name: str | None,
     api_key: str | None = None,
     api_base: str | None = None,
     api_version: str | None = None,
@@ -116,6 +117,7 @@ def get_llm(
     return DefaultMultiLLM(
         model_provider=provider,
         model_name=model,
+        deployment_name=deployment_name,
         api_key=api_key,
         api_base=api_base,
         api_version=api_version,
diff --git a/backend/danswer/server/manage/llm/api.py b/backend/danswer/server/manage/llm/api.py
index 06501d683..9cac96236 100644
--- a/backend/danswer/server/manage/llm/api.py
+++ b/backend/danswer/server/manage/llm/api.py
@@ -54,6 +54,7 @@ def test_llm_configuration(
         api_base=test_llm_request.api_base,
         api_version=test_llm_request.api_version,
         custom_config=test_llm_request.custom_config,
+        deployment_name=test_llm_request.deployment_name,
     )
 
     functions_with_args: list[tuple[Callable, tuple]] = [(test_llm, (llm,))]
@@ -70,6 +71,7 @@ def test_llm_configuration(
             api_base=test_llm_request.api_base,
             api_version=test_llm_request.api_version,
             custom_config=test_llm_request.custom_config,
+            deployment_name=test_llm_request.deployment_name,
         )
         functions_with_args.append((test_llm, (fast_llm,)))
 
diff --git a/backend/danswer/server/manage/llm/models.py b/backend/danswer/server/manage/llm/models.py
index 2e3b38448..9b371099c 100644
--- a/backend/danswer/server/manage/llm/models.py
+++ b/backend/danswer/server/manage/llm/models.py
@@ -21,6 +21,7 @@ class TestLLMRequest(BaseModel):
     # model level
     default_model_name: str
     fast_default_model_name: str | None = None
+    deployment_name: str | None = None
 
 
 class LLMProviderDescriptor(BaseModel):
diff --git a/backend/requirements/model_server.txt b/backend/requirements/model_server.txt
index 2ea6df66d..6160555f7 100644
--- a/backend/requirements/model_server.txt
+++ b/backend/requirements/model_server.txt
@@ -1,5 +1,5 @@
-cohere==5.6.1
 einops==0.8.0
+cohere==5.6.1
 fastapi==0.109.2
 google-cloud-aiplatform==1.58.0
 numpy==1.26.4
@@ -13,4 +13,4 @@ transformers==4.39.2
 uvicorn==0.21.1
 voyageai==0.2.3
 litellm==1.49.5
-sentry-sdk[fastapi,celery,starlette]==2.14.0
+sentry-sdk[fastapi,celery,starlette]==2.14.0
\ No newline at end of file