Merge pull request transformerlab#399 from transformerlab/add/supports-filtering-interact

Add filtering based on the supports field for interact in loader plugins

Author: deep1401

src/renderer/components/Experiment/Interact/Interact.tsx

@@ -40,13 +40,25 @@ import VisualizeGeneration from './VisualizeGeneration';
 import ModelLayerVisualization from './ModelLayerVisualization';
 
 const fetcher = (url) => fetch(url).then((res) => res.json());
+// const supports = [
+//   'chat',
+//   'completion',
+//   'rag',
+//   'tools',
+//   'template',
+//   'embeddings',
+//   'tokenize',
+//   'logprobs',
+//   'batched',
+// ];
 
 export default function Chat({
   experimentInfo,
   experimentInfoMutate,
   setRagEngine,
   mode,
   setMode,
+  supports,
 }) {
   const { models } = chatAPI.useModelStatus();
   const [conversationId, setConversationId] = React.useState(null);
@@ -764,17 +776,51 @@ export default function Chat({
               }
             }
           >
-            <Option value="chat">Chat</Option>
-            <Option value="completion">Completion</Option>
-            <Option value="visualize_model">Model Activations</Option>
-            <Option value="model_layers">Model Architecture</Option>
-            <Option value="rag">Query Docs (RAG)</Option>
-            <Option value="tools">Tool Calling</Option>
-            <Option value="template">Templated Prompt</Option>
-            <Option value="embeddings">Embeddings</Option>
-            <Option value="tokenize">Tokenize</Option>
-            <Option value="logprobs">Visualize Logprobs</Option>
-            <Option value="batched">Batched Query</Option>
+            <Option value="chat" disabled={!supports.includes('chat')}>
+              Chat
+            </Option>
+            <Option
+              value="completion"
+              disabled={!supports.includes('completion')}
+            >
+              Completion
+            </Option>
+            <Option
+              value="visualize_model"
+              disabled={!supports.includes('visualize_model')}
+            >
+              Model Activations
+            </Option>
+            <Option
+              value="model_layers"
+              disabled={!supports.includes('model_layers')}
+            >
+              Model Architecture
+            </Option>
+            <Option value="rag" disabled={!supports.includes('rag')}>
+              Query Docs (RAG)
+            </Option>
+            <Option value="tools" disabled={!supports.includes('tools')}>
+              Tool Calling
+            </Option>
+            <Option value="template" disabled={!supports.includes('template')}>
+              Templated Prompt
+            </Option>
+            <Option
+              value="embeddings"
+              disabled={!supports.includes('embeddings')}
+            >
+              Embeddings
+            </Option>
+            <Option value="tokenize" disabled={!supports.includes('tokenize')}>
+              Tokenize
+            </Option>
+            <Option value="logprobs" disabled={!supports.includes('logprobs')}>
+              Visualize Logprobs
+            </Option>
+            <Option value="batched" disabled={!supports.includes('batched')}>
+              Batched Query
+            </Option>
           </Select>
         </FormControl>
         <Typography level="title-md">

src/renderer/components/MainAppPanel.tsx

@@ -16,6 +16,8 @@ import {
   useEffect,
   useState,
 } from 'react';
+import useSWR from 'swr';
+
 import { AnalyticsBrowser } from '@segment/analytics-next';
 import Data from './Data/Data';
 import Interact from './Experiment/Interact/Interact';
@@ -117,6 +119,49 @@ export default function MainAppPanel({
   const [selectedInteractSubpage, setSelectedInteractSubpage] =
     useState('chat');
 
+  const fetcher = (url) => fetch(url).then((res) => res.json());
+
+  // Extract pluginId at the top level
+  const inferenceParams = experimentInfo?.config?.inferenceParams;
+  const pluginId = inferenceParams
+    ? JSON.parse(inferenceParams)?.inferenceEngine
+    : null;
+
+  // Use SWR at the top level, not inside useEffect
+  const { data: modelData } = useSWR(
+    experimentInfo?.id && pluginId
+      ? chatAPI.Endpoints.Experiment.ScriptGetFile(
+          experimentInfo.id,
+          pluginId,
+          'index.json',
+        )
+      : null,
+    fetcher,
+  );
+
+  let modelSupports = [
+    'chat',
+    'completion',
+    'rag',
+    'tools',
+    'template',
+    'embeddings',
+    'tokenize',
+    'batched',
+  ];
+
+  if (modelData && modelData !== 'null' && modelData !== 'undefined') {
+    modelSupports = JSON.parse(modelData)?.supports || [
+      'chat',
+      'completion',
+      'rag',
+      'tools',
+      'template',
+      'embeddings',
+      'tokenize',
+      'batched',
+    ];
+  }
   const setFoundation = useCallback(
     (model) => {
       let model_name = '';
@@ -351,6 +396,7 @@ export default function MainAppPanel({
               setRagEngine={setRagEngine}
               mode={selectedInteractSubpage}
               setMode={setSelectedInteractSubpage}
+              supports={modelSupports}
             />
           }
         />
@@ -363,6 +409,7 @@ export default function MainAppPanel({
               setRagEngine={setRagEngine}
               mode={'model_layers'}
               setMode={setSelectedInteractSubpage}
+              supports={modelSupports}
             />
           }
         />