From c1e3a1b3e78d6a98b93c35c178bd4c37af48b039 Mon Sep 17 00:00:00 2001
From: pablodanswer <pablo@danswer.ai>
Date: Thu, 8 Aug 2024 14:02:11 -0700
Subject: [PATCH] Select proper assistant override (#2068)

* encode images properly

* proper assistant default model updates

* remove now unneeded image encoding update

* update naming of persona llm option gathering
---
 backend/danswer/llm/utils.py          |  1 +
 web/src/app/chat/ChatPage.tsx         | 32 ++++++++++++++++++++++-----
 web/src/app/chat/message/Messages.tsx |  1 +
 web/src/lib/llm/utils.ts              | 28 +++++++++++++++++++++++
 4 files changed, 56 insertions(+), 6 deletions(-)

diff --git a/backend/danswer/llm/utils.py b/backend/danswer/llm/utils.py
index c1d463e2a070..2a464d61c5ff 100644
--- a/backend/danswer/llm/utils.py
+++ b/backend/danswer/llm/utils.py
@@ -29,6 +29,7 @@ from danswer.prompts.constants import CODE_BLOCK_PAT
 from danswer.utils.logger import setup_logger
 from shared_configs.configs import LOG_LEVEL
 
+
 if TYPE_CHECKING:
     from danswer.llm.answering.models import PreviousMessage
 
diff --git a/web/src/app/chat/ChatPage.tsx b/web/src/app/chat/ChatPage.tsx
index 0e25fbd837fd..accc4eb26e2d 100644
--- a/web/src/app/chat/ChatPage.tsx
+++ b/web/src/app/chat/ChatPage.tsx
@@ -45,7 +45,7 @@ import { useContext, useEffect, useRef, useState } from "react";
 import { usePopup } from "@/components/admin/connectors/Popup";
 import { SEARCH_PARAM_NAMES, shouldSubmitOnLoad } from "./searchParams";
 import { useDocumentSelection } from "./useDocumentSelection";
-import { useFilters, useLlmOverride } from "@/lib/hooks";
+import { LlmOverride, useFilters, useLlmOverride } from "@/lib/hooks";
 import { computeAvailableFilters } from "@/lib/filters";
 import { FeedbackType } from "./types";
 import { DocumentSidebar } from "./documentSidebar/DocumentSidebar";
@@ -60,7 +60,11 @@ import { AnswerPiecePacket, DanswerDocument } from "@/lib/search/interfaces";
 import { buildFilters } from "@/lib/search/utils";
 import { SettingsContext } from "@/components/settings/SettingsProvider";
 import Dropzone from "react-dropzone";
-import { checkLLMSupportsImageInput, getFinalLLM } from "@/lib/llm/utils";
+import {
+  checkLLMSupportsImageInput,
+  getFinalLLM,
+  getLLMProviderOverrideForPersona,
+} from "@/lib/llm/utils";
 import { ChatInputBar } from "./input/ChatInputBar";
 import { useChatContext } from "@/components/context/ChatContext";
 import { v4 as uuidv4 } from "uuid";
@@ -72,6 +76,10 @@ import { useSidebarVisibility } from "@/components/chat_search/hooks";
 import { SIDEBAR_TOGGLED_COOKIE_NAME } from "@/components/resizable/constants";
 import FixedLogo from "./shared_chat_search/FixedLogo";
 import { getSecondsUntilExpiration } from "@/lib/time";
+import {
+  FullLLMProvider,
+  LLMProviderDescriptor,
+} from "../admin/models/llm/interfaces";
 
 const TEMP_USER_MESSAGE_ID = -1;
 const TEMP_ASSISTANT_MESSAGE_ID = -2;
@@ -152,10 +160,6 @@ export function ChatPage({
         )
       ? 0
       : 0.7;
-  const llmOverrideManager = useLlmOverride(
-    selectedChatSession,
-    defaultTemperature
-  );
 
   const setSelectedAssistantFromId = (assistantId: number) => {
     // NOTE: also intentionally look through available assistants here, so that
@@ -168,6 +172,22 @@ export function ChatPage({
   const liveAssistant =
     selectedAssistant || filteredAssistants[0] || availableAssistants[0];
 
+  const llmOverrideManager = useLlmOverride(
+    selectedChatSession,
+    defaultTemperature
+  );
+
+  useEffect(() => {
+    const personaDefault = getLLMProviderOverrideForPersona(
+      liveAssistant,
+      llmProviders
+    );
+
+    if (personaDefault) {
+      llmOverrideManager.setLlmOverride(personaDefault);
+    }
+  }, [liveAssistant]);
+
   // this is for "@"ing assistants
   const [alternativeAssistant, setAlternativeAssistant] =
     useState<Persona | null>(null);
diff --git a/web/src/app/chat/message/Messages.tsx b/web/src/app/chat/message/Messages.tsx
index ff4f86e2d686..ba9b98ec78f0 100644
--- a/web/src/app/chat/message/Messages.tsx
+++ b/web/src/app/chat/message/Messages.tsx
@@ -332,6 +332,7 @@ export const AIMessage = ({
                       )}
 
                     {toolCall &&
+                      (!files || files.length == 0) &&
                       toolCall.tool_name === IMAGE_GENERATION_TOOL_NAME &&
                       !toolCall.tool_result && <GeneratingImageDisplay />}
 
diff --git a/web/src/lib/llm/utils.ts b/web/src/lib/llm/utils.ts
index 2d454235072b..c2a5a1ce8b40 100644
--- a/web/src/lib/llm/utils.ts
+++ b/web/src/lib/llm/utils.ts
@@ -34,6 +34,34 @@ export function getFinalLLM(
   return [provider, model];
 }
 
+export function getLLMProviderOverrideForPersona(
+  liveAssistant: Persona,
+  llmProviders: LLMProviderDescriptor[]
+): LlmOverride | null {
+  const overrideProvider = liveAssistant.llm_model_provider_override;
+  const overrideModel = liveAssistant.llm_model_version_override;
+
+  if (!overrideModel) {
+    return null;
+  }
+
+  const matchingProvider = llmProviders.find(
+    (provider) =>
+      (overrideProvider ? provider.name === overrideProvider : true) &&
+      provider.model_names.includes(overrideModel)
+  );
+
+  if (matchingProvider) {
+    return {
+      name: matchingProvider.name,
+      provider: matchingProvider.provider,
+      modelName: overrideModel,
+    };
+  }
+
+  return null;
+}
+
 const MODELS_SUPPORTING_IMAGES = [
   ["openai", "gpt-4o"],
   ["openai", "gpt-4o-mini"],