Merge pull request #63 from SciPhi-AI/Nolan/ImproveChat

Nolan/improve chat

package.json

@@ -30,8 +30,8 @@
   "license": "MIT",
   "dependencies": {
     "@headlessui/react": "^2.1.2",
-    "@heroicons/react": "^2.1.5",
     "@nextui-org/react": "^2.4.6",
+    "@radix-ui/react-accordion": "^1.2.0",
     "@radix-ui/react-alert-dialog": "^1.1.1",
     "@radix-ui/react-checkbox": "^1.1.1",
     "@radix-ui/react-collapsible": "^1.1.0",

src/components/ChatDemo/ConfigurationSheet.tsx

@@ -1,4 +1,4 @@
-import { AdjustmentsHorizontalIcon } from '@heroicons/react/24/outline';
+import { SlidersHorizontal } from 'lucide-react';
 import React from 'react';
 
 import { Input } from '@/components/ui/input';
@@ -12,48 +12,7 @@ import {
   SheetTrigger,
 } from '@/components/ui/sheet';
 import { Slider } from '@/components/ui/slider';
-
-interface GenerationConfig {
-  // RAG Configuration
-  temperature?: number;
-  setTemperature?: (value: number) => void;
-  top_p?: number;
-  setTopP?: (value: number) => void;
-  top_k?: number;
-  setTop_k?: (value: number) => void;
-  max_tokens_to_sample?: number;
-  setMax_tokens_to_sample?: (value: number) => void;
-  model?: string;
-  setModel?: (value: string) => void;
-  stream?: boolean;
-  setStream?: (value: boolean) => void;
-  functions?: Array<Record<string, any>> | null;
-  setFunctions?: (value: Array<Record<string, any>> | null) => void;
-  skip_special_tokens?: boolean;
-  setSkipSpecialTokens?: (value: boolean) => void;
-  stop_token?: string | null;
-  setStopToken?: (value: string | null) => void;
-  num_beams?: number;
-  setNumBeams?: (value: number) => void;
-  do_sample?: boolean;
-  setDoSample?: (value: boolean) => void;
-  generate_with_chat?: boolean;
-  setGenerateWithChat?: (value: boolean) => void;
-  add_generation_kwargs?: Record<string, any>;
-  setAddGenerationKwargs?: (value: Record<string, any>) => void;
-  api_base?: string | null;
-  setApiBase?: (value: string | null) => void;
-
-  // Knowledge Graph Configuration
-  kg_temperature?: number;
-  setKgTemperature?: (value: number) => void;
-  kg_top_p?: number;
-  setKgTopP?: (value: number) => void;
-  kg_top_k?: number;
-  setKgTop_k?: (value: number) => void;
-  kg_max_tokens_to_sample?: number;
-  setKgMax_tokens_to_sample?: (value: number) => void;
-}
+import { GenerationConfig } from '@/types';
 
 const ConfigurationSheet: React.FC<GenerationConfig> = ({
   temperature,
@@ -96,7 +55,7 @@ const ConfigurationSheet: React.FC<GenerationConfig> = ({
   return (
     <Sheet>
       <SheetTrigger>
-        <AdjustmentsHorizontalIcon className="h-8 w-8 mb-3" />
+        <SlidersHorizontal className="h-8 w-8 mb-3" />
       </SheetTrigger>
       <SheetContent side="left">
         <SheetHeader>

src/components/ChatDemo/DefaultQueries.tsx

@@ -1,34 +1,44 @@
-import { LightBulbIcon } from '@heroicons/react/24/outline';
+import { Lightbulb, FlaskConical, Flame, Earth } from 'lucide-react';
 import { FC } from 'react';
 
 import { Logo } from '@/components/shared/Logo';
-import { Alert, AlertDescription, AlertTitle } from '@/components/ui/alert';
-
-interface DefaultQueriesProps {
-  setQuery: (query: string) => void;
-}
+import { Alert, AlertDescription } from '@/components/ui/alert';
+import { DefaultQueriesProps } from '@/types';
 
 export const DefaultQueries: FC<DefaultQueriesProps> = ({ setQuery }) => {
   const defaultQueries = [
-    'What is the main topic of the uploaded documents?',
-    'Can you summarize the key points from the documents?',
-    'Are there any common themes across the documents?',
-    'What are the most frequently mentioned entities in the documents?',
+    {
+      query: 'What is the main topic of the uploaded documents?',
+      icon: <Lightbulb className="h-6 w-6 text-yellow-400" />,
+    },
+    {
+      query: 'Summarize key points for me.',
+      icon: <FlaskConical className="h-6 w-6 text-purple-400" />,
+    },
+    {
+      query: 'What issues do you see with the documents?',
+      icon: <Flame className="h-6 w-6 text-red-400" />,
+    },
+    {
+      query: 'How are these documents interrelated?',
+      icon: <Earth className="h-6 w-6 text-green-400" />,
+    },
   ];
 
   return (
     <div className="flex flex-col items-center justify-center h-full space-y-8">
       <Logo width={150} height={150} disableLink={true} />
-      <div className="grid grid-cols-2 gap-4 w-full max-w-2xl">
-        {defaultQueries.map((query, index) => (
+      <div className="grid grid-cols-1 sm:grid-cols-2 md:grid-cols-3 lg:grid-cols-4 gap-4 w-full max-w-4xl px-4">
+        {defaultQueries.map(({ query, icon }, index) => (
           <Alert
             key={index}
-            className="cursor-pointer hover:bg-zinc-700"
+            className="cursor-pointer hover:bg-zinc-700 flex flex-col items-start p-3 h-[100px]"
             onClick={() => setQuery(query)}
           >
-            <LightBulbIcon className="h-4 w-4" />
-            <AlertTitle>Sample Query {index + 1}</AlertTitle>
-            <AlertDescription>{query}</AlertDescription>
+            <div className="mb-2">{icon}</div>
+            <AlertDescription className="text-sm text-left">
+              {query}
+            </AlertDescription>
           </Alert>
         ))}
       </div>

src/components/ChatDemo/MessageBubble.tsx

@@ -0,0 +1,32 @@
+import React from 'react';
+
+import { Message } from '@/types';
+
+const MessageBubble: React.FC<{ message: Message; isStreaming?: boolean }> = ({
+  message,
+  isStreaming,
+}) => {
+  if (message.sender === 'user') {
+    return (
+      <div className="flex justify-end mb-4">
+        <div className="bg-zinc-800 text-white rounded-lg p-3 max-w-xs lg:max-w-md">
+          {message.content}
+        </div>
+      </div>
+    );
+  } else if (message.sender === 'assistant') {
+    return (
+      <div className="flex justify-start mb-4">
+        <div
+          className={`bg-gray-200 rounded-lg p-3 max-w-xs lg:max-w-md ${isStreaming ? 'animate-pulse' : ''}`}
+        >
+          {message.content}
+          {isStreaming && <span className="inline-block animate-pulse">▋</span>}
+        </div>
+      </div>
+    );
+  }
+  return null;
+};
+
+export default MessageBubble;

src/components/ChatDemo/PipelineStatus.tsx

@@ -1,64 +1,83 @@
 import { r2rClient } from 'r2r-js';
-import React, { useState, useEffect } from 'react';
+import React, { useState, useEffect, useCallback } from 'react';
+
+import { useUserContext } from '@/context/UserContext';
+import { PipelineStatusProps } from '@/types';
+
+const MAX_RETRIES = 3;
+const RETRY_DELAY = 2000;
 
 async function checkPipelineStatus(
-  deploymentUrl: string | undefined
+  deploymentUrl: string | undefined,
+  getClient: () => Promise<r2rClient | null>
 ): Promise<'Connected' | 'No Connection'> {
   if (!deploymentUrl) {
     return 'No Connection';
   }
 
-  try {
-    const client = new r2rClient(deploymentUrl);
-    await client.health();
-    return 'Connected';
-  } catch (error) {
-    console.error('Health check failed:', error);
-    return 'No Connection';
+  for (let attempt = 0; attempt < MAX_RETRIES; attempt++) {
+    try {
+      const client = await getClient();
+      if (!client) {
+        return 'No Connection';
+      }
+      await client.health();
+      return 'Connected';
+    } catch (error) {
+      console.warn(`Health check attempt ${attempt + 1} failed:`, error);
+      if (attempt < MAX_RETRIES - 1) {
+        await new Promise((resolve) => setTimeout(resolve, RETRY_DELAY));
+      }
+    }
   }
+
+  console.error('Health check failed after multiple attempts');
+  return 'No Connection';
 }
 
 export function useConnectionStatus(
   deploymentUrl?: string,
   onStatusChange?: (isConnected: boolean) => void
 ) {
+  const { getClient, refreshAuth } = useUserContext();
   const [isConnected, setIsConnected] = useState(false);
 
-  useEffect(() => {
-    const checkStatus = async () => {
-      if (deploymentUrl) {
-        const status = await checkPipelineStatus(deploymentUrl);
-        const newConnectionStatus = status === 'Connected';
-        setIsConnected(newConnectionStatus);
-        onStatusChange?.(newConnectionStatus);
-      } else {
-        setIsConnected(false);
-        onStatusChange?.(false);
+  const checkStatus = useCallback(async () => {
+    if (deploymentUrl) {
+      const status = await checkPipelineStatus(deploymentUrl, getClient);
+      const newConnectionStatus = status === 'Connected';
+      setIsConnected(newConnectionStatus);
+      onStatusChange?.(newConnectionStatus);
+
+      if (!newConnectionStatus) {
+        // Attempt to refresh authentication if connection is lost
+        try {
+          await refreshAuth();
+        } catch (error) {
+          console.error('Failed to refresh authentication:', error);
+        }
       }
-    };
+    } else {
+      setIsConnected(false);
+      onStatusChange?.(false);
+    }
+  }, [deploymentUrl, getClient, onStatusChange, refreshAuth]);
 
+  useEffect(() => {
     checkStatus();
     const interval = setInterval(checkStatus, 10000);
 
     return () => clearInterval(interval);
-  }, [deploymentUrl, onStatusChange]);
+  }, [checkStatus]);
 
   return isConnected;
 }
 
-interface PipelineStatusProps {
-  pipeline: {
-    deploymentUrl: string;
-  };
-  className?: string;
-  onStatusChange?: (isConnected: boolean) => void;
-}
-
 export function PipelineStatus({
-  pipeline,
   className = '',
   onStatusChange,
 }: PipelineStatusProps) {
+  const { pipeline } = useUserContext();
   const isConnected = useConnectionStatus(
     pipeline?.deploymentUrl,
     onStatusChange

src/components/ChatDemo/SingleSwitch.tsx

@@ -7,14 +7,7 @@ import {
   TooltipProvider,
   TooltipTrigger,
 } from '@/components/ui/tooltip';
-
-interface SingleSwitchProps {
-  id: string;
-  initialChecked: boolean;
-  onChange: (id: string, checked: boolean) => void;
-  label: string;
-  tooltipText: string;
-}
+import { SingleSwitchProps } from '@/types';
 
 const SingleSwitch: React.FC<SingleSwitchProps> = ({
   id,