Shared filter utility for clarity (#3270)

* shared filter util * clearer comment
onyx-dot-app · Dec 3, 2024 · 66f47d2 · 66f47d2
1 parent 0a685bd
commit 66f47d2
Show file tree

Hide file tree

Showing 3 changed files with 42 additions and 41 deletions.
diff --git a/web/src/components/context/AssistantsContext.tsx b/web/src/components/context/AssistantsContext.tsx
@@ -11,6 +11,7 @@ import {
   classifyAssistants,
   orderAssistantsForUser,
   getUserCreatedAssistants,
+  filterAssistants,
 } from "@/lib/assistants/utils";
 import { useUser } from "../user/UserProvider";
 
@@ -145,22 +146,13 @@ export const AssistantsProvider: React.FC<{
       if (!response.ok) throw new Error("Failed to fetch assistants");
       let assistants: Persona[] = await response.json();
 
-      if (!hasImageCompatibleModel) {
-        assistants = assistants.filter(
-          (assistant) =>
-            !assistant.tools.some(
-              (tool) => tool.in_code_tool_id === "ImageGenerationTool"
-            )
-        );
-      }
-
-      if (!hasAnyConnectors) {
-        assistants = assistants.filter(
-          (assistant) => assistant.num_chunks === 0
-        );
-      }
+      let filteredAssistants = filterAssistants(
+        assistants,
+        hasAnyConnectors,
+        hasImageCompatibleModel
+      );
 
-      setAssistants(assistants);
+      setAssistants(filteredAssistants);
 
       // Fetch and update allAssistants for admins and curators
       await fetchPersonas();

diff --git a/web/src/lib/assistants/utils.ts b/web/src/lib/assistants/utils.ts
@@ -1,6 +1,7 @@
 import { Persona } from "@/app/admin/assistants/interfaces";
 import { User } from "../types";
 import { checkUserIsNoAuthUser } from "../user";
+import { personaComparator } from "@/app/admin/assistants/lib";
 
 export function checkUserOwnsAssistant(user: User | null, assistant: Persona) {
   return checkUserIdOwnsAssistant(user?.id, assistant);
@@ -117,3 +118,31 @@ export function getUserCreatedAssistants(
     checkUserOwnsAssistant(user, assistant)
   );
 }
+
+// Filter assistants based on connector status, image compatibility and visibility
+export function filterAssistants(
+  assistants: Persona[],
+  hasAnyConnectors: boolean,
+  hasImageCompatibleModel: boolean
+): Persona[] {
+  let filteredAssistants = assistants.filter(
+    (assistant) => assistant.is_visible
+  );
+
+  if (!hasAnyConnectors) {
+    filteredAssistants = filteredAssistants.filter(
+      (assistant) => assistant.num_chunks === 0
+    );
+  }
+
+  if (!hasImageCompatibleModel) {
+    filteredAssistants = filteredAssistants.filter(
+      (assistant) =>
+        !assistant.tools.some(
+          (tool) => tool.in_code_tool_id === "ImageGenerationTool"
+        )
+    );
+  }
+
+  return filteredAssistants.sort(personaComparator);
+}
diff --git a/web/src/lib/chat/fetchAssistantdata.ts b/web/src/lib/chat/fetchAssistantdata.ts
@@ -4,6 +4,7 @@ import { fetchLLMProvidersSS } from "@/lib/llm/fetchLLMs";
 import { personaComparator } from "@/app/admin/assistants/lib";
 import { fetchAssistantsSS } from "../assistants/fetchAssistantsSS";
 import { checkLLMSupportsImageInput } from "../llm/utils";
+import { filterAssistants } from "../assistants/utils";
 
 interface AssistantData {
   assistants: Persona[];
@@ -39,42 +40,21 @@ export async function fetchAssistantData(): Promise<AssistantData> {
       }),
     ]);
 
-    // Process visible assistants
-    let filteredAssistants = assistants.filter(
-      (assistant) => assistant.is_visible
-    );
-
-    // Process connector status
     const hasAnyConnectors = ccPairsResponse?.ok
       ? (await ccPairsResponse.json()).length > 0
       : false;
 
-    // Filter assistants based on connector status
-    if (!hasAnyConnectors) {
-      filteredAssistants = filteredAssistants.filter(
-        (assistant) => assistant.num_chunks === 0
-      );
-    }
-
-    // Sort assistants
-    filteredAssistants.sort(personaComparator);
-
-    // Check for image-compatible models
     const hasImageCompatibleModel = llmProviders.some(
       (provider) =>
         provider.provider === "openai" ||
         provider.model_names.some((model) => checkLLMSupportsImageInput(model))
     );
 
-    // Filter out image generation tools if no compatible model
-    if (!hasImageCompatibleModel) {
-      filteredAssistants = filteredAssistants.filter(
-        (assistant) =>
-          !assistant.tools.some(
-            (tool) => tool.in_code_tool_id === "ImageGenerationTool"
-          )
-      );
-    }
+    let filteredAssistants = filterAssistants(
+      assistants,
+      hasAnyConnectors,
+      hasImageCompatibleModel
+    );
 
     return {
       assistants: filteredAssistants,