Fix azure

backend/alembic/versions/e4334d5b33ba_add_deployment_name_to_llmprovider.py

@@ -0,0 +1,26 @@
+"""add_deployment_name_to_llmprovider
+
+Revision ID: e4334d5b33ba
+Revises: 46b7a812670f
+Create Date: 2024-10-04 09:52:34.896867
+
+"""
+from alembic import op
+import sqlalchemy as sa
+
+
+# revision identifiers, used by Alembic.
+revision = "e4334d5b33ba"
+down_revision = "46b7a812670f"
+branch_labels = None
+depends_on = None
+
+
+def upgrade() -> None:
+    op.add_column(
+        "llm_provider", sa.Column("deployment_name", sa.String(), nullable=True)
+    )
+
+
+def downgrade() -> None:
+    op.drop_column("llm_provider", "deployment_name")

backend/danswer/chat/process_message.py

@@ -18,6 +18,10 @@
 from danswer.chat.models import MessageSpecificCitations
 from danswer.chat.models import QADocsResponse
 from danswer.chat.models import StreamingError
+from danswer.configs.app_configs import AZURE_DALLE_API_BASE
+from danswer.configs.app_configs import AZURE_DALLE_API_KEY
+from danswer.configs.app_configs import AZURE_DALLE_API_VERSION
+from danswer.configs.app_configs import AZURE_DEPLOYMENT_NAME
 from danswer.configs.chat_configs import BING_API_KEY
 from danswer.configs.chat_configs import CHAT_TARGET_CHUNK_PERCENTAGE
 from danswer.configs.chat_configs import DISABLE_LLM_CHOOSE_SEARCH
@@ -560,7 +564,26 @@ def stream_chat_message_objects(
                         and llm.config.api_key
                         and llm.config.model_provider == "openai"
                     ):
-                        img_generation_llm_config = llm.config
+                        img_generation_llm_config = LLMConfig(
+                            model_provider=llm.config.model_provider,
+                            model_name="dall-e-3",
+                            temperature=GEN_AI_TEMPERATURE,
+                            api_key=llm.config.api_key,
+                            api_base=llm.config.api_base,
+                            api_version=llm.config.api_version,
+                        )
+                    elif (
+                        llm.config.model_provider == "azure"
+                        and AZURE_DALLE_API_KEY is not None
+                    ):
+                        img_generation_llm_config = LLMConfig(
+                            model_provider="azure",
+                            model_name=f"azure/{AZURE_DEPLOYMENT_NAME}",
+                            temperature=GEN_AI_TEMPERATURE,
+                            api_key=AZURE_DALLE_API_KEY,
+                            api_base=AZURE_DALLE_API_BASE,
+                            api_version=AZURE_DALLE_API_VERSION,
+                        )
                     else:
                         llm_providers = fetch_existing_llm_providers(db_session)
                         openai_provider = next(
@@ -579,7 +602,7 @@ def stream_chat_message_objects(
                             )
                         img_generation_llm_config = LLMConfig(
                             model_provider=openai_provider.provider,
-                            model_name=openai_provider.default_model_name,
+                            model_name="dall-e-3",
                             temperature=GEN_AI_TEMPERATURE,
                             api_key=openai_provider.api_key,
                             api_base=openai_provider.api_base,
@@ -591,6 +614,7 @@ def stream_chat_message_objects(
                             api_base=img_generation_llm_config.api_base,
                             api_version=img_generation_llm_config.api_version,
                             additional_headers=litellm_additional_headers,
+                            model=img_generation_llm_config.model_name,
                         )
                     ]
                 elif tool_cls.__name__ == InternetSearchTool.__name__:

backend/danswer/configs/app_configs.py

@@ -135,7 +135,7 @@
     os.environ.get("POSTGRES_PASSWORD") or "password"
 )
 POSTGRES_HOST = os.environ.get("POSTGRES_HOST") or "localhost"
-POSTGRES_PORT = os.environ.get("POSTGRES_PORT") or "5432"
+POSTGRES_PORT = os.environ.get("POSTGRES_PORT") or "5433"
 POSTGRES_DB = os.environ.get("POSTGRES_DB") or "postgres"
 
 POSTGRES_API_SERVER_POOL_SIZE = int(
@@ -413,6 +413,12 @@
     os.environ.get("ENABLE_PAID_ENTERPRISE_EDITION_FEATURES", "").lower() == "true"
 )
 
+# Azure DALL-E Configurations
+AZURE_DALLE_API_VERSION = os.environ.get("AZURE_DALLE_API_VERSION")
+AZURE_DALLE_API_KEY = os.environ.get("AZURE_DALLE_API_KEY")
+AZURE_DALLE_API_BASE = os.environ.get("AZURE_DALLE_API_BASE")
+AZURE_DEPLOYMENT_NAME = os.environ.get("AZURE_DEPLOYMENT_NAME")
+
 
 MULTI_TENANT = os.environ.get("MULTI_TENANT", "").lower() == "true"
 SECRET_JWT_KEY = os.environ.get("SECRET_JWT_KEY", "")

backend/danswer/db/llm.py

@@ -83,6 +83,7 @@ def upsert_llm_provider(
     existing_llm_provider.model_names = llm_provider.model_names
     existing_llm_provider.is_public = llm_provider.is_public
     existing_llm_provider.display_model_names = llm_provider.display_model_names
+    existing_llm_provider.deployment_name = llm_provider.deployment_name
 
     if not existing_llm_provider.id:
         # If its not already in the db, we need to generate an ID by flushing

backend/danswer/db/models.py

@@ -1143,6 +1143,8 @@ class LLMProvider(Base):
         postgresql.ARRAY(String), nullable=True
     )
 
+    deployment_name: Mapped[str | None] = mapped_column(String, nullable=True)
+
     # should only be set for a single provider
     is_default_provider: Mapped[bool | None] = mapped_column(Boolean, unique=True)
     # EE only

backend/danswer/llm/chat_llm.py

@@ -204,6 +204,7 @@ def __init__(
         model_name: str,
         api_base: str | None = None,
         api_version: str | None = None,
+        deployment_name: str | None = None,
         max_output_tokens: int | None = None,
         custom_llm_provider: str | None = None,
         temperature: float = GEN_AI_TEMPERATURE,
@@ -215,6 +216,7 @@ def __init__(
         self._model_version = model_name
         self._temperature = temperature
         self._api_key = api_key
+        self._deployment_name = deployment_name
         self._api_base = api_base
         self._api_version = api_version
         self._custom_llm_provider = custom_llm_provider
@@ -290,7 +292,7 @@ def _completion(
         try:
             return litellm.completion(
                 # model choice
-                model=f"{self.config.model_provider}/{self.config.model_name}",
+                model=f"{self.config.model_provider}/{self.config.deployment_name or self.config.model_name}",
                 # NOTE: have to pass in None instead of empty string for these
                 # otherwise litellm can have some issues with bedrock
                 api_key=self._api_key or None,
@@ -325,6 +327,7 @@ def config(self) -> LLMConfig:
             api_key=self._api_key,
             api_base=self._api_base,
             api_version=self._api_version,
+            deployment_name=self._deployment_name,
         )
 
     def _invoke_implementation(

backend/danswer/llm/factory.py

@@ -88,6 +88,7 @@ def _create_llm(model: str) -> LLM:
         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,
@@ -103,6 +104,7 @@ def _create_llm(model: str) -> LLM:
 def get_llm(
     provider: str,
     model: str,
+    deployment_name: str | None = None,
     api_key: str | None = None,
     api_base: str | None = None,
     api_version: str | None = None,

backend/danswer/llm/interfaces.py

@@ -24,7 +24,7 @@ class LLMConfig(BaseModel):
     api_key: str | None = None
     api_base: str | None = None
     api_version: str | None = None
-
+    deployment_name: str | None = None
     # This disables the "model_" protected namespace for pydantic
     model_config = {"protected_namespaces": ()}
 

backend/danswer/llm/llm_provider_options.py

@@ -16,10 +16,12 @@ class WellKnownLLMProviderDescriptor(BaseModel):
     api_base_required: bool
     api_version_required: bool
     custom_config_keys: list[CustomConfigKey] | None = None
+    single_model_supported: bool = False
 
     llm_names: list[str]
     default_model: str | None = None
     default_fast_model: str | None = None
+    deployment_name_required: bool = False
 
 
 OPENAI_PROVIDER_NAME = "openai"
@@ -108,6 +110,8 @@ def fetch_available_well_known_llms() -> list[WellKnownLLMProviderDescriptor]:
             api_version_required=True,
             custom_config_keys=[],
             llm_names=fetch_models_for_provider(AZURE_PROVIDER_NAME),
+            deployment_name_required=True,
+            single_model_supported=True,
         ),
         WellKnownLLMProviderDescriptor(
             name=BEDROCK_PROVIDER_NAME,

backend/danswer/server/features/tool/api.py

@@ -21,6 +21,8 @@
 from danswer.tools.custom.openapi_parsing import MethodSpec
 from danswer.tools.custom.openapi_parsing import openapi_to_method_specs
 from danswer.tools.custom.openapi_parsing import validate_openapi_schema
+from danswer.tools.images.image_generation_tool import ImageGenerationTool
+from danswer.tools.utils import is_image_generation_available
 
 router = APIRouter(prefix="/tool")
 admin_router = APIRouter(prefix="/admin/tool")
@@ -127,4 +129,9 @@ def list_tools(
     _: User | None = Depends(current_user),
 ) -> list[ToolSnapshot]:
     tools = get_tools(db_session)
-    return [ToolSnapshot.from_model(tool) for tool in tools]
+    return [
+        ToolSnapshot.from_model(tool)
+        for tool in tools
+        if tool.in_code_tool_id != ImageGenerationTool.name
+        or is_image_generation_available(db_session=db_session)
+    ]

backend/danswer/server/manage/llm/api.py

@@ -55,6 +55,7 @@ def test_llm_configuration(
         api_version=test_llm_request.api_version,
         custom_config=test_llm_request.custom_config,
     )
+
     functions_with_args: list[tuple[Callable, tuple]] = [(test_llm, (llm,))]
 
     if (
@@ -141,6 +142,20 @@ def put_llm_provider(
             detail=f"LLM Provider with name {llm_provider.name} already exists",
         )
 
+    # Ensure default_model_name and fast_default_model_name are in display_model_names
+    # This is necessary for custom models and Bedrock/Azure models
+    if llm_provider.display_model_names is None:
+        llm_provider.display_model_names = []
+
+    if llm_provider.default_model_name not in llm_provider.display_model_names:
+        llm_provider.display_model_names.append(llm_provider.default_model_name)
+
+    if (
+        llm_provider.fast_default_model_name
+        and llm_provider.fast_default_model_name not in llm_provider.display_model_names
+    ):
+        llm_provider.display_model_names.append(llm_provider.fast_default_model_name)
+
     try:
         return upsert_llm_provider(
             llm_provider=llm_provider,

backend/danswer/server/manage/llm/models.py

@@ -66,6 +66,7 @@ class LLMProvider(BaseModel):
     is_public: bool = True
     groups: list[int] = Field(default_factory=list)
     display_model_names: list[str] | None = None
+    deployment_name: str | None = None
 
 
 class LLMProviderUpsertRequest(LLMProvider):
@@ -100,4 +101,5 @@ def from_model(cls, llm_provider_model: "LLMProviderModel") -> "FullLLMProvider"
             ),
             is_public=llm_provider_model.is_public,
             groups=[group.id for group in llm_provider_model.groups],
+            deployment_name=llm_provider_model.deployment_name,
         )

backend/danswer/tools/utils.py

@@ -1,5 +1,11 @@
 import json
 
+from sqlalchemy.orm import Session
+
+from danswer.configs.app_configs import AZURE_DALLE_API_KEY
+from danswer.db.connector import check_connectors_exist
+from danswer.db.document import check_docs_exist
+from danswer.db.models import LLMProvider
 from danswer.natural_language_processing.utils import BaseTokenizer
 from danswer.tools.tool import Tool
 
@@ -26,3 +32,18 @@ def compute_tool_tokens(tool: Tool, llm_tokenizer: BaseTokenizer) -> int:
 
 def compute_all_tool_tokens(tools: list[Tool], llm_tokenizer: BaseTokenizer) -> int:
     return sum(compute_tool_tokens(tool, llm_tokenizer) for tool in tools)
+
+
+def is_image_generation_available(db_session: Session) -> bool:
+    providers = db_session.query(LLMProvider).all()
+    for provider in providers:
+        if provider.name == "OpenAI":
+            return True
+
+    return bool(AZURE_DALLE_API_KEY)
+
+
+def is_document_search_available(db_session: Session) -> bool:
+    docs_exist = check_docs_exist(db_session)
+    connectors_exist = check_connectors_exist(db_session)
+    return docs_exist or connectors_exist

deployment/docker_compose/docker-compose.dev.yml

@@ -298,7 +298,7 @@ services:
       - POSTGRES_USER=${POSTGRES_USER:-postgres}
       - POSTGRES_PASSWORD=${POSTGRES_PASSWORD:-password}
     ports:
-      - "5432:5432"
+      - "5433:5432"
     volumes:
       - db_volume:/var/lib/postgresql/data
 

deployment/docker_compose/docker-compose.gpu-dev.yml

@@ -308,7 +308,7 @@ services:
       - POSTGRES_USER=${POSTGRES_USER:-postgres}
       - POSTGRES_PASSWORD=${POSTGRES_PASSWORD:-password}
     ports:
-      - "5432:5432"
+      - "5433:5432"
     volumes:
       - db_volume:/var/lib/postgresql/data
 

deployment/docker_compose/docker-compose.search-testing.yml

@@ -157,7 +157,7 @@ services:
       - POSTGRES_USER=${POSTGRES_USER:-postgres}
       - POSTGRES_PASSWORD=${POSTGRES_PASSWORD:-password}
     ports:
-      - "5432"
+      - "5433"
     volumes:
       - db_volume:/var/lib/postgresql/data
 

web/src/app/admin/assistants/AssistantEditor.tsx

@@ -192,15 +192,11 @@ export function AssistantEditor({
     modelOptionsByProvider.set(llmProvider.name, providerOptions);
   });
 
-  const providerSupportingImageGenerationExists =
-    providersContainImageGeneratingSupport(llmProviders);
-
   const personaCurrentToolIds =
     existingPersona?.tools.map((tool) => tool.id) || [];
+
   const searchTool = findSearchTool(tools);
-  const imageGenerationTool = providerSupportingImageGenerationExists
-    ? findImageGenerationTool(tools)
-    : undefined;
+  const imageGenerationTool = findImageGenerationTool(tools);
   const internetSearchTool = findInternetSearchTool(tools);
 
   const customTools = tools.filter(
@@ -997,7 +993,7 @@ export function AssistantEditor({
                           alignTop={tool.description != null}
                           key={tool.id}
                           name={`enabled_tools_map.${tool.id}`}
-                          label={tool.name}
+                          label={tool.display_name}
                           subtext={tool.description}
                           onChange={() => {
                             toggleToolInValues(tool.id);

web/src/app/admin/configuration/llm/ConfiguredLLMProviderDisplay.tsx

@@ -133,6 +133,7 @@ function LLMProviderDisplay({
           </Button>
         </div>
       </div>
+
       {formIsVisible && (
         <LLMProviderUpdateModal
           llmProviderDescriptor={llmProviderDescriptor}