feat: show generated embedding results in messages

gregnr · gregnr · commit 7c26bc0ae54c · 2024-07-30T15:49:41.000-05:00
diff --git a/apps/postgres-new/components/markdown-accordion.tsx b/apps/postgres-new/components/markdown-accordion.tsx
@@ -0,0 +1,53 @@
+import { DatabaseZap } from 'lucide-react'
+import ReactMarkdown from 'react-markdown'
+import remarkGfm from 'remark-gfm'
+import {
+  Accordion,
+  AccordionContent,
+  AccordionItem,
+  AccordionTrigger,
+} from '~/components/ui/accordion'
+import { cn } from '~/lib/utils'
+
+export type MarkdownAccordionProps = {
+  title: string
+  content: string
+  error?: string
+  className?: string
+}
+
+export default function MarkdownAccordion({
+  title,
+  content,
+  error,
+  className,
+}: MarkdownAccordionProps) {
+  return (
+    <Accordion type="single" collapsible>
+      <AccordionItem
+        value="item-1"
+        className={cn(
+          'border-2 border-neutral-100 bg-neutral-50 px-3 py-2 rounded-md',
+          error ? 'bg-destructive-300' : undefined,
+          className
+        )}
+      >
+        <AccordionTrigger className="p-0 gap-2">
+          <div className="flex gap-2 items-center font-normal text-lighter text-sm">
+            <DatabaseZap size={14} />
+            {title}
+          </div>
+        </AccordionTrigger>
+        <AccordionContent className="py-2 [&_>div]:pb-0 flex flex-col gap-2">
+          <ReactMarkdown
+            remarkPlugins={[remarkGfm]}
+            className="prose text-xs mt-2 [&_ul>li::before]:top-2 [&_ol>li::before]:top-0"
+          >
+            {content}
+          </ReactMarkdown>
+          {error && <div className="text-destructive-600 text-xs">{error}</div>}
+        </AccordionContent>
+      </AccordionItem>
+    </Accordion>
+  )
+}
diff --git a/apps/postgres-new/components/tools/generated-embedding.tsx b/apps/postgres-new/components/tools/generated-embedding.tsx
@@ -0,0 +1,39 @@
+import { codeBlock } from 'common-tags'
+import { ToolInvocation } from '~/lib/tools'
+import MarkdownAccordion from '../markdown-accordion'
+
+export type GeneratedEmbeddingProps = {
+  toolInvocation: ToolInvocation<'embed'>
+}
+
+export default function GeneratedEmbedding({ toolInvocation }: GeneratedEmbeddingProps) {
+  const { texts } = toolInvocation.args
+
+  if (!('result' in toolInvocation)) {
+    return null
+  }
+
+  if (!toolInvocation.result.success) {
+    const content = codeBlock`
+      ${texts.map((text) => `> ${text}`).join('\n')}
+    `
+
+    return (
+      <MarkdownAccordion
+        title="Error generating embedding"
+        content={content}
+        error={toolInvocation.result.error}
+      />
+    )
+  }
+
+  const { ids } = toolInvocation.result
+
+  const content = codeBlock`
+    Results stored in \`meta.embeddings\` table:
+
+    ${texts.map((text, i) => `- **\`id\`:** ${ids[i]}\n\n  **\`content\`:** ${text}`).join('\n')}
+  `
+
+  return <MarkdownAccordion title="Generated embedding" content={content} />
+}
diff --git a/apps/postgres-new/components/tools/index.tsx b/apps/postgres-new/components/tools/index.tsx
@@ -5,6 +5,7 @@ import CsvImport from './csv-import'
 import CsvRequest from './csv-request'
 import ExecutedSql from './executed-sql'
 import GeneratedChart from './generated-chart'
+import GeneratedEmbedding from './generated-embedding'
 
 export type ToolUiProps = {
   toolInvocation: ToolInvocation
@@ -24,6 +25,8 @@ export function ToolUi({ toolInvocation }: ToolUiProps) {
       return <CsvExport toolInvocation={toolInvocation} />
     case 'renameConversation':
       return <ConversationRename toolInvocation={toolInvocation} />
+    case 'embed':
+      return <GeneratedEmbedding toolInvocation={toolInvocation} />
   }
   return null
 }
diff --git a/apps/postgres-new/lib/hooks.ts b/apps/postgres-new/lib/hooks.ts
@@ -83,7 +83,7 @@ export function useLocalStorage<T>(key: string, initialValue: T) {
   const queryKey = ['local-storage', key]
 
   const currentValue =
-    typeof window !== 'undefined' ? window.localStorage.getItem(key) ?? undefined : undefined
+    typeof window !== 'undefined' ? (window.localStorage.getItem(key) ?? undefined) : undefined
 
   const { data: storedValue = currentValue ? (JSON.parse(currentValue) as T) : initialValue } =
     useQuery({
