Multimodal output support (#1826)

Author: neoxelox

apps/web/src/components/ChatWrapper/Message/index.tsx

@@ -9,7 +9,12 @@ import {
   ToolContent,
 } from '@latitude-data/constants/legacyCompiler'
 
-import { AgentToolsMap, isSafeUrl } from '@latitude-data/constants'
+import {
+  AgentToolsMap,
+  isEncodedImage,
+  isInlineImage,
+  isSafeUrl,
+} from '@latitude-data/constants'
 import { Badge, BadgeProps } from '@latitude-data/web-ui/atoms/Badge'
 import { Button } from '@latitude-data/web-ui/atoms/Button'
 import { CodeBlock } from '@latitude-data/web-ui/atoms/CodeBlock'
@@ -648,13 +653,30 @@ const ContentImage = memo(
     collapseParameters?: boolean
     sourceMap?: PromptlSourceRef[]
   }) => {
+    const TextComponent = size === 'small' ? Text.H6 : Text.H5
     const segment = useMemo(
       () => computeSegments('image', image.toString(), sourceMap, parameters),
       [image, sourceMap, parameters],
     )[0]
 
     if (!isSafeUrl(image)) {
-      const TextComponent = size === 'small' ? Text.H6 : Text.H5
+      if (isInlineImage(image)) {
+        return (
+          <Image
+            src={image.toString()}
+            className='max-h-72 rounded-xl w-fit object-contain'
+          />
+        )
+      }
+
+      if (isEncodedImage(image)) {
+        return (
+          <Image
+            src={`data:image/png;base64,${image.toString()}`}
+            className='max-h-72 rounded-xl w-fit object-contain'
+          />
+        )
+      }
 
       return (
         <div className='flex flex-row p-4 gap-2 bg-muted rounded-xl w-fit items-center'>
@@ -670,8 +692,6 @@ const ContentImage = memo(
       )
     }
 
-    const TextComponent = size === 'small' ? Text.H6 : Text.H5
-
     if (!segment || typeof segment === 'string') {
       return (
         <Image
@@ -719,6 +739,7 @@ const ContentFile = memo(
     collapseParameters?: boolean
     sourceMap?: PromptlSourceRef[]
   }) => {
+    const TextComponent = size === 'small' ? Text.H6 : Text.H5
     const segment = useMemo(
       () => computeSegments('file', file.toString(), sourceMap, parameters),
       [file, sourceMap, parameters],
@@ -735,8 +756,6 @@ const ContentFile = memo(
       )
     }
 
-    const TextComponent = size === 'small' ? Text.H6 : Text.H5
-
     if (!segment || typeof segment === 'string') {
       return <FileComponent src={file.toString()} />
     }

apps/web/src/hooks/playgroundChat/useProviderEventHandler.ts

@@ -1,15 +1,18 @@
 import { tokenizeText } from '$/lib/tokenize'
 import { ChainEvent } from '@latitude-data/constants'
 import {
+  AssistantMessage,
+  FileContent,
+  ImageContent,
   Message,
   MessageContent,
   MessageRole,
   ToolCall,
   ToolRequestContent,
 } from '@latitude-data/constants/legacyCompiler'
+import { StreamEventTypes } from '@latitude-data/core/constants'
 import { ParsedEvent } from 'eventsource-parser/stream'
 import React, { useCallback } from 'react'
-import { StreamEventTypes } from '@latitude-data/core/constants'
 
 type SetMessagesFunction = React.Dispatch<React.SetStateAction<Message[]>>
 type SetUnrespondedToolCallsFunction = React.Dispatch<
@@ -183,17 +186,13 @@ export function useProviderEventHandler({
 
   // Helper function to handle tool-result events
   const handleToolResult = useCallback(
-    (
-      data: {
-        type: 'tool-result'
-      } & {
-        type: 'tool-result'
-        toolCallId: string
-        toolName: string
-        args: any
-        result: any
-      },
-    ) => {
+    (data: {
+      type: 'tool-result'
+      toolCallId: string
+      toolName: string
+      args: any
+      result: any
+    }) => {
       setMessages((messages) => {
         const lastMessage = messages.at(-1)!
         if (lastMessage.role === MessageRole.assistant) {
@@ -293,6 +292,54 @@ export function useProviderEventHandler({
     [setMessages, incrementUsageDelta],
   )
 
+  const handleFile = useCallback(
+    ({
+      file: { mediaType, base64Data, base64, uint8Array },
+    }: {
+      type: 'file'
+      file: {
+        mediaType: string
+        base64Data?: string
+        base64?: string
+        uint8Array?: Uint8Array
+      }
+    }) => {
+      const mediaData = base64Data ?? base64 ?? uint8Array?.toString() ?? ''
+
+      const content = mediaType?.startsWith('image/')
+        ? ({
+            type: 'image',
+            image: mediaData,
+          } as ImageContent)
+        : ({
+            type: 'file',
+            file: mediaData,
+            mimeType: mediaType,
+          } as FileContent)
+
+      setMessages((messages) => {
+        const lastMessage = messages.at(-1)!
+        if (lastMessage.role === MessageRole.assistant) {
+          return [
+            ...messages.slice(0, -1),
+            {
+              ...lastMessage,
+              content: [...(lastMessage.content ?? []), content],
+            } as AssistantMessage,
+          ]
+        } else {
+          // Should not be possible
+          throw new Error('Expected assistant message')
+        }
+      })
+
+      incrementUsageDelta({
+        completionTokens: tokenizeText(mediaData),
+      })
+    },
+    [setMessages, incrementUsageDelta],
+  )
+
   // Main handler that delegates to the appropriate helper based on event type
   const handleProviderEvent = useCallback(
     (parsedEvent: ParsedEvent, data: ChainEvent['data']) => {
@@ -314,6 +361,9 @@ export function useProviderEventHandler({
         case 'reasoning':
           handleReasoning(data)
           break
+        case 'file':
+          handleFile(data)
+          break
       }
     },
     [
@@ -322,6 +372,7 @@ export function useProviderEventHandler({
       handleToolCall,
       handleToolResult,
       handleReasoning,
+      handleFile,
     ],
   )
 

packages/constants/src/helpers.ts

@@ -66,6 +66,27 @@ export function simplifyDocument(
   }
 }
 
+const INLINE_IMAGE_REGEX = /^data:image\/[^;]+;base64,/i
+export function isInlineImage(src: unknown): boolean {
+  if (!src) return false
+
+  return INLINE_IMAGE_REGEX.test(src.toString().trim())
+}
+
+export function isEncodedImage(src: unknown): boolean {
+  if (!src) return false
+
+  const clean = src.toString().trim().replace(/\s+/g, '')
+  const pad = clean.length % 4
+  const decoded = pad ? clean + '==='.slice(pad) : clean
+
+  try {
+    return !!atob(decoded)
+  } catch {
+    return false
+  }
+}
+
 export function isSafeUrl(url: unknown): url is string | URL {
   const isUrl =
     url instanceof URL ||