From 95796a84f2f095d3ada4d59abdb0584457329173 Mon Sep 17 00:00:00 2001
From: Benjamin Shafii <ben@prologe.io>
Date: Wed, 1 Jan 2025 16:09:49 +0100
Subject: [PATCH] fix: a lot of stuff

---
 .../plugin/views/assistant/ai-chat/chat.tsx   | 131 +++++++++++++++---
 .../ai-chat/context-limit-indicator.tsx       |  75 ++++++----
 .../views/assistant/ai-chat/submit-button.tsx |   2 +-
 .../ai-chat/tool-handlers/screenpipe-utils.ts |   1 +
 .../assistant/ai-chat/types/annotations.ts    |  32 ++++-
 5 files changed, 184 insertions(+), 57 deletions(-)

diff --git a/packages/plugin/views/assistant/ai-chat/chat.tsx b/packages/plugin/views/assistant/ai-chat/chat.tsx
index 57271d0b..30bdabba 100644
--- a/packages/plugin/views/assistant/ai-chat/chat.tsx
+++ b/packages/plugin/views/assistant/ai-chat/chat.tsx
@@ -24,7 +24,10 @@ import { ContextItems } from "./components/context-items";
 import { ClearAllButton } from "./components/clear-all-button";
 import { useCurrentFile } from "./hooks/use-current-file";
 import { SearchAnnotationHandler } from "./tool-handlers/search-annotation-handler";
-import { isSearchResultsAnnotation, SearchResultsAnnotation } from "./types/annotations";
+import {
+  isSearchResultsAnnotation,
+  SearchResultsAnnotation,
+} from "./types/annotations";
 import { ExamplePrompts } from "./components/example-prompts";
 
 interface ChatComponentProps {
@@ -79,29 +82,35 @@ export const ChatComponent: React.FC<ChatComponentProps> = ({
         files: Object.fromEntries(
           Object.entries(files).map(([id, file]) => [
             id,
-            { ...file, content: '' }
+            { ...file, content: "" },
           ])
         ),
         folders: Object.fromEntries(
           Object.entries(folders).map(([id, folder]) => [
             id,
-            { ...folder, files: folder.files.map(f => ({ ...f, content: '' })) }
+            {
+              ...folder,
+              files: folder.files.map(f => ({ ...f, content: "" })),
+            },
           ])
         ),
         tags: Object.fromEntries(
           Object.entries(tags).map(([id, tag]) => [
             id,
-            { ...tag, files: tag.files.map(f => ({ ...f, content: '' })) }
+            { ...tag, files: tag.files.map(f => ({ ...f, content: "" })) },
           ])
         ),
         searchResults: Object.fromEntries(
           Object.entries(searchResults).map(([id, search]) => [
             id,
-            { ...search, results: search.results.map(r => ({ ...r, content: '' })) }
+            {
+              ...search,
+              results: search.results.map(r => ({ ...r, content: "" })),
+            },
           ])
         ),
         // Keep these as is
-        currentFile: currentFile ? { ...currentFile, content: '' } : null,
+        currentFile: currentFile ? { ...currentFile, content: "" } : null,
         screenpipe,
         textSelections,
       };
@@ -129,6 +138,8 @@ export const ChatComponent: React.FC<ChatComponentProps> = ({
     handleSubmit,
     stop,
     addToolResult,
+    error,
+    reload,
   } = useChat({
     onDataChunk: (chunk: DataChunk) => {
       if (chunk.type === "metadata" && chunk.data?.groundingMetadata) {
@@ -183,7 +194,8 @@ export const ChatComponent: React.FC<ChatComponentProps> = ({
     onError: error => {
       logger.error(error.message);
       setErrorMessage(
-        "Connection failed. If the problem persists, please check your internet connection or VPN."
+        error.message ||
+          "Connection failed. If the problem persists, please check your internet connection or VPN."
       );
     },
     onFinish: () => {
@@ -253,10 +265,55 @@ export const ChatComponent: React.FC<ChatComponentProps> = ({
     } as React.ChangeEvent<HTMLInputElement>);
   };
 
+  const handleRetry = () => {
+    setErrorMessage(null);
+    reload();
+  };
+
   return (
     <div className="flex flex-col h-full">
       <div className="flex-grow overflow-y-auto p-4 h-full">
         <div className="flex flex-col min-h-min-content">
+          {errorMessage && (
+            <div className="bg-[--background-modifier-error] bg-opacity-10 text-[--text-error] p-4 rounded-md mb-4 flex items-center justify-between">
+              <span className="flex items-center">
+                <svg
+                  className="w-5 h-5 mr-2"
+                  fill="none"
+                  stroke="currentColor"
+                  viewBox="0 0 24 24"
+                >
+                  <path
+                    strokeLinecap="round"
+                    strokeLinejoin="round"
+                    strokeWidth={2}
+                    d="M12 8v4m0 4h.01M21 12a9 9 0 11-18 0 9 9 0 0118 0z"
+                  />
+                </svg>
+                {errorMessage}
+              </span>
+              <button
+                onClick={handleRetry}
+                className="px-3 py-1 bg-[--interactive-accent] hover:bg-[--interactive-accent-hover] text-[--text-on-accent] rounded-md text-sm flex items-center"
+              >
+                <svg
+                  className="w-4 h-4 mr-1"
+                  fill="none"
+                  stroke="currentColor"
+                  viewBox="0 0 24 24"
+                >
+                  <path
+                    strokeLinecap="round"
+                    strokeLinejoin="round"
+                    strokeWidth={2}
+                    d="M4 4v5h.582m15.356 2A8.001 8.001 0 004.582 9m0 0H9m11 11v-5h-.581m0 0a8.003 8.003 0 01-15.357-2m15.357 2H15"
+                  />
+                </svg>
+                Retry
+              </button>
+            </div>
+          )}
+
           {messages.length === 0 ? (
             <div className="flex flex-col items-center justify-center h-full">
               <h3 className="text-[--text-normal] mb-4">Try these examples</h3>
@@ -271,7 +328,7 @@ export const ChatComponent: React.FC<ChatComponentProps> = ({
                     return (
                       <SearchAnnotationHandler
                         key={`${message.id}-annotation-${index}`}
-                        annotation={annotation as SearchResultsAnnotation}
+                        annotation={annotation}
                       />
                     );
                   }
@@ -292,6 +349,33 @@ export const ChatComponent: React.FC<ChatComponentProps> = ({
               </React.Fragment>
             ))
           )}
+
+          {isGenerating && (
+            <div className="flex items-center text-[--text-muted] text-sm mt-4">
+              <svg
+                className="animate-spin -ml-1 mr-3 h-4 w-4"
+                xmlns="http://www.w3.org/2000/svg"
+                fill="none"
+                viewBox="0 0 24 24"
+              >
+                <circle
+                  className="opacity-25"
+                  cx="12"
+                  cy="12"
+                  r="10"
+                  stroke="currentColor"
+                  strokeWidth="4"
+                ></circle>
+                <path
+                  className="opacity-75"
+                  fill="currentColor"
+                  d="M4 12a8 8 0 018-8V0C5.373 0 0 5.373 0 12h4zm2 5.291A7.962 7.962 0 014 12H0c0 3.042 1.135 5.824 3 7.938l3-2.647z"
+                ></path>
+              </svg>
+              Generating response...
+            </div>
+          )}
+
           <div ref={messagesEndRef} />
           {groundingMetadata && (
             <SourcesSection groundingMetadata={groundingMetadata} />
@@ -302,25 +386,30 @@ export const ChatComponent: React.FC<ChatComponentProps> = ({
       <div className="border-t border-[--background-modifier-border] p-4">
         <div className="flex items-center space-x-2 mb-4">
           <ContextItems />
-
           <ClearAllButton />
         </div>
 
-        <form onSubmit={handleSendMessage} className="flex items-end">
-          <div className="flex-grow overflow-y-auto relative" ref={inputRef}>
-            <Tiptap
-              value={input}
-              onChange={handleTiptapChange}
-              onKeyDown={handleKeyDown}
-            />
-
-            <div className="absolute bottom-0 right-0 h-full flex items-center">
-              <AudioRecorder
-                onTranscriptionComplete={handleTranscriptionComplete}
+        <form onSubmit={handleSendMessage} className="flex items-stretch min-h-min max-h-full">
+          <div
+            className={`flex flex-grow ${
+              error ? "opacity-50 pointer-events-none" : ""
+            }`}
+          >
+            <div className="overflow-y-auto relative w-full " ref={inputRef}>
+              <Tiptap
+                value={input}
+                onChange={handleTiptapChange}
+                onKeyDown={handleKeyDown}
               />
+
+              <div className="absolute bottom-0 right-0 h-full flex items-center">
+                <AudioRecorder
+                  onTranscriptionComplete={handleTranscriptionComplete}
+                />
+              </div>
             </div>
+            <SubmitButton isGenerating={isGenerating} />
           </div>
-          <SubmitButton isGenerating={isGenerating} />
         </form>
 
         <div className="flex items-center justify-between">
diff --git a/packages/plugin/views/assistant/ai-chat/context-limit-indicator.tsx b/packages/plugin/views/assistant/ai-chat/context-limit-indicator.tsx
index 11e71deb..cff91f63 100644
--- a/packages/plugin/views/assistant/ai-chat/context-limit-indicator.tsx
+++ b/packages/plugin/views/assistant/ai-chat/context-limit-indicator.tsx
@@ -70,36 +70,55 @@ export function ContextLimitIndicator({
   const shouldWarn = stats.percentUsed > 80;
   
   return (
-    <div className="mt-2 space-y-2">
-      <div className={`p-2 rounded text-xs flex gap-1 items-center justify-between
-        ${isOverLimit 
-          ? "border border-[--text-error] text-[--text-error]" 
-          : shouldWarn
-            ? "border border-[--text-warning] text-[--text-warning]"
-            : "bg-[--background-modifier-border] text-[--text-muted]"}`}>
-        <span>
-          {isOverLimit 
-            ? "Context size exceeds maximum" 
+    <div className="mt-2 space-y-2 flex">
+      <div className="relative group">
+        <div className={`p-2 min-w-max rounded text-xs flex gap-1 items-center justify-between cursor-pointer hover:bg-[--background-modifier-hover] transition-colors
+          ${isOverLimit 
+            ? "border border-[--text-error] text-[--text-error]" 
             : shouldWarn
-              ? "Context size nearing limit"
-              : "Context used"}
-        </span>
-        <span className="font-mono">{stats.percentUsed.toFixed(0)}%</span>
-      </div>
-      
-      <button
-        onClick={toggleLightweightMode}
-        title="Enable when processing multiple files to remove content from context. Useful for batch operations like moving, renaming, or tagging files."
-        className={`w-full p-2 rounded text-xs flex items-center justify-between group relative
-          ${isLightweightMode 
-            ? "bg-[--interactive-accent] text-[--text-on-accent]" 
-            : "bg-[--background-modifier-border] text-[--text-muted] hover:bg-[--background-modifier-border-hover]"}`}
-      >
-        <span>Disable Context</span>
-        <div className="absolute bottom-full left-0 w-full p-2 bg-[--background-secondary] rounded shadow-lg text-[--text-normal] opacity-0 group-hover:opacity-100 transition-opacity mb-2 text-left pointer-events-none">
-          Use this when processing multiple files to reduce context size. Removes file content from context while preserving metadata for batch operations.
+              ? "border border-[--text-warning] text-[--text-warning]"
+              : "border border-[--background-modifier-border] text-[--text-muted]"}`}>
+          <span>
+            {isOverLimit 
+              ? "Context size exceeds maximum" 
+              : shouldWarn
+                ? "Context size nearing limit"
+                : "Context used"}
+          </span>
+          <span className="font-mono">{stats.percentUsed.toFixed(0)}%</span>
+        </div>
+
+        {/* Enhanced menu-style tooltip */}
+        <div className="absolute left-0 top-full mt-1 w-72 bg-[--background-secondary] rounded-md shadow-lg border border-[--background-modifier-border] opacity-0 group-hover:opacity-100 transition-opacity z-10">
+          <div
+            onClick={toggleLightweightMode}
+            className={`w-full px-4 py-3.5 text-left text-xs flex items-start justify-between gap-4 hover:bg-[--background-modifier-hover] rounded-md cursor-pointer
+              ${isLightweightMode 
+                ? "text-[--interactive-accent]" 
+                : "text-[--text-normal]"}`}
+          >
+            <div className="space-y-1.5 flex-1">
+              <div className="font-medium flex items-center gap-2">
+                <div className={`w-4 h-4 rounded border flex items-center justify-center
+                  ${isLightweightMode 
+                    ? "border-[--interactive-accent] bg-[--interactive-accent]" 
+                    : "border-[--background-modifier-border]"}`}
+                >
+                  {isLightweightMode && (
+                    <svg className="w-3 h-3 text-[--text-on-accent]" viewBox="0 0 14 14" fill="none">
+                      <path d="M11.6666 3.5L5.24992 9.91667L2.33325 7" stroke="currentColor" strokeWidth="2" strokeLinecap="round" strokeLinejoin="round"/>
+                    </svg>
+                  )}
+                </div>
+                <span>Disable Context</span>
+              </div>
+              <div className="text-[--text-muted] text-[11px] leading-relaxed">
+                Removes file content from context while preserving metadata. Useful for batch operations like moving, renaming, or tagging files.
+              </div>
+            </div>
+          </div>
         </div>
-      </button>
+      </div>
     </div>
   );
 }
\ No newline at end of file
diff --git a/packages/plugin/views/assistant/ai-chat/submit-button.tsx b/packages/plugin/views/assistant/ai-chat/submit-button.tsx
index 518d977a..19ead3f7 100644
--- a/packages/plugin/views/assistant/ai-chat/submit-button.tsx
+++ b/packages/plugin/views/assistant/ai-chat/submit-button.tsx
@@ -13,7 +13,7 @@ export const SubmitButton: React.FC<SubmitButtonProps> = ({
   return (
     <Button
       type="submit"
-      className={`h-full ml-2 font-bold py-2 px-4 rounded ${
+      className={`flex-none ml-2 font-bold px-4 flex items-center justify-center h-full ${
         isGenerating
           ? "bg-[--background-modifier-form-field] text-[--text-muted] cursor-not-allowed"
           : "bg-[--interactive-accent] hover:bg-[--interactive-accent-hover] text-[--text-on-accent]"
diff --git a/packages/plugin/views/assistant/ai-chat/tool-handlers/screenpipe-utils.ts b/packages/plugin/views/assistant/ai-chat/tool-handlers/screenpipe-utils.ts
index 0ab7cd09..f7f0a9e5 100644
--- a/packages/plugin/views/assistant/ai-chat/tool-handlers/screenpipe-utils.ts
+++ b/packages/plugin/views/assistant/ai-chat/tool-handlers/screenpipe-utils.ts
@@ -135,6 +135,7 @@ export async function getDailyInformation({
       endTime: endTime.toISOString(),
       contentType: "ocr",
       limit: plugin.settings.queryScreenpipeLimit,
+      appName: "Arc"
     }),
     queryScreenpipe({
       startTime: startTime.toISOString(),
diff --git a/packages/plugin/views/assistant/ai-chat/types/annotations.ts b/packages/plugin/views/assistant/ai-chat/types/annotations.ts
index fd19427a..c48e75bd 100644
--- a/packages/plugin/views/assistant/ai-chat/types/annotations.ts
+++ b/packages/plugin/views/assistant/ai-chat/types/annotations.ts
@@ -1,20 +1,38 @@
-import { MessageAnnotation } from 'ai';
+import type { Message } from 'ai';
 
 export interface SearchResult {
-  content: string;
-  score: number;
-  filePath: string;
+  segment: {
+    text: string;
+    startIndex: number;
+    endIndex: number;
+  };
+  groundingChunkIndices: number[];
+  confidenceScores: number[];
 }
 
-export interface SearchResultsAnnotation extends MessageAnnotation {
+export interface WebSource {
+  web: {
+    uri: string;
+    title: string;
+  };
+}
+
+export interface SearchResultsAnnotation {
   type: 'search-results';
-  results: SearchResult[];
+  groundingMetadata: {
+    webSearchQueries: string[];
+    searchEntryPoint: {
+      renderedContent: string;
+    };
+    groundingChunks: WebSource[];
+    groundingSupports: SearchResult[];
+  };
 }
 
 export type CustomAnnotation = SearchResultsAnnotation;
 
 export function isSearchResultsAnnotation(
-  annotation: MessageAnnotation
+  annotation: any
 ): annotation is SearchResultsAnnotation {
   return annotation.type === 'search-results';
 } 
\ No newline at end of file