Chat: add vision support, upload images

Previously, it was only possible to send text to Cody via Chat. This PR
adds new support to additionally send Cody images. This can be helpful,
for example, when you want to write a basic HTML structure based on
a Figma design.

lib/shared/src/chat/chat.ts

@@ -57,12 +57,13 @@ export class ChatClient {
 
         // We only want to send up the speaker and prompt text, regardless of whatever other fields
         // might be on the messages objects (`file`, `displayText`, `contextFiles`, etc.).
-        const messagesToSend = augmentedMessages.map(({ speaker, text }) => ({
+        const messagesToSend = augmentedMessages.map<Message>(({ speaker, text, content }) => ({
             text,
             speaker,
+            content,
         }))
 
-        const completionParams = {
+        const completionParams: CompletionParameters = {
             ...DEFAULT_CHAT_COMPLETION_PARAMETERS,
             ...params,
             messages: messagesToSend,
@@ -107,8 +108,8 @@ export function sanitizeMessages(messages: Message[]): Message[] {
         // the next one
         const nextMessage = sanitizedMessages[index + 1]
         if (
-            (nextMessage.speaker === 'assistant' && !nextMessage.text?.length) ||
-            (message.speaker === 'assistant' && !message.text?.length)
+            (nextMessage.speaker === 'assistant' && !nextMessage.text?.length && !nextMessage.content) ||
+            (message.speaker === 'assistant' && !message.text?.length && !message.content)
         ) {
             return false
         }

lib/shared/src/chat/transcript/messages.ts

@@ -22,6 +22,7 @@ export interface SubMessage {
 
 export interface ChatMessage extends Message {
     contextFiles?: ContextItem[]
+    base64Image?: string
 
     contextAlternatives?: RankedContext[]
 

lib/shared/src/sourcegraph-api/completions/types.ts

@@ -5,27 +5,28 @@ interface DoneEvent {
     type: 'done'
 }
 
-interface CompletionEvent extends CompletionResponse {
-    type: 'completion'
-}
+// interface CompletionEvent extends CompletionResponse {
+//     type: 'completion'
+// }
 
 interface ErrorEvent {
     type: 'error'
     error: string
 }
 
-export type Event = DoneEvent | CompletionEvent | ErrorEvent
+export type Event = DoneEvent | ErrorEvent
 
 export interface Message {
     // Note: The unified API only supports one system message passed as the first message
     speaker: 'human' | 'assistant' | 'system'
     text?: PromptString
+    content?: string | MessagePart[]
+    base64Image?: string
 }
 
-export interface CompletionResponse {
-    completion: string
-    stopReason?: string
-}
+type MessagePart =
+    | { type: 'text'; text: string } // a normal text message
+    | { type: 'image_url'; image_url: { url: string } } // image message, per https://platform.openai.com/docs/guides/vision
 
 export interface CompletionParameters {
     fast?: boolean
@@ -45,6 +46,7 @@ export interface CompletionParameters {
         type: 'content'
         content: string
     }
+    base64Image?: string
 }
 
 export interface SerializedCompletionParameters extends Omit<CompletionParameters, 'messages'> {

vscode/src/chat/chat-view/ChatBuilder.ts

@@ -157,7 +157,7 @@ export class ChatBuilder {
         if (this.messages.at(-1)?.speaker === 'human') {
             throw new Error('Cannot add a user message after a user message')
         }
-        this.messages.push({ ...message, speaker: 'human' })
+        this.messages.push({ ...message, speaker: 'human', base64Image: this.getAndResetImage() })
         this.changeNotifications.next()
     }
 
@@ -322,6 +322,31 @@ export class ChatBuilder {
         }
         return result
     }
+
+    /**
+     * Store the base64-encoded image uploaded by user to a multi-modal model.
+     * Requires vision support in the model, added in the PR
+     * https://github.com/sourcegraph/sourcegraph/pull/546
+     */
+    private image: string | undefined = undefined
+
+    /**
+     * Sets the base64-encoded image for the chat model.
+     * @param base64Image - The base64-encoded image data to set.
+     */
+    public setImage(base64Image: string): void {
+        this.image = base64Image
+    }
+
+    /**
+     * Gets the base64-encoded image for the chat model and resets the internal image property to undefined.
+     * @returns The base64-encoded image, or undefined if no image has been set.
+     */
+    public getAndResetImage(): string | undefined {
+        const image = this.image
+        this.image = undefined
+        return image
+    }
 }
 
 function messageToSerializedChatInteraction(

vscode/src/chat/chat-view/ChatController.ts

@@ -513,6 +513,11 @@ export class ChatController implements vscode.Disposable, vscode.WebviewViewProv
                 }
                 break
             }
+
+            case 'chat/upload-file': {
+                this.chatBuilder.setImage(message.base64)
+                break
+            }
             case 'log': {
                 const logger = message.level === 'debug' ? logDebug : logError
                 logger(message.filterLabel, message.message)

vscode/src/chat/chat-view/prompt.ts

@@ -89,6 +89,9 @@ export class DefaultPrompter {
                     `Ignored ${messagesIgnored} chat messages due to context limit`
                 )
             }
+            for (const message of reverseTranscript) {
+                promptBuilder.tryAddImage(message.base64Image)
+            }
             // Counter for context items categorized by source
             const ignoredContext = { user: 0, corpus: 0, transcript: 0 }
 

vscode/src/chat/protocol.ts

@@ -156,6 +156,8 @@ export type WebviewMessage =
           selectedFilters: NLSSearchDynamicFilter[]
       }
     | { command: 'action/confirmation'; id: string; response: boolean }
+    | { command: 'log'; level: 'debug' | 'error'; filterLabel: string; message: string }
+    | { command: 'chat/upload-file'; base64: string }
 
 export interface SmartApplyResult {
     taskId: FixupTaskID

vscode/src/completions/nodeClient.ts

@@ -9,7 +9,6 @@ import {
     type CompletionCallbacks,
     type CompletionParameters,
     type CompletionRequestParameters,
-    type CompletionResponse,
     NetworkError,
     RateLimitError,
     SourcegraphCompletionsClient,
@@ -21,6 +20,7 @@ import {
     getTraceparentHeaders,
     globalAgentRef,
     isError,
+    logDebug,
     logError,
     onAbort,
     parseEvents,
@@ -38,6 +38,10 @@ export class SourcegraphNodeCompletionsClient extends SourcegraphCompletionsClie
         signal?: AbortSignal
     ): Promise<void> {
         const { apiVersion, interactionId } = requestParams
+        for (const message of params.messages) {
+            logDebug('apiVersion', JSON.stringify(apiVersion, null, 2))
+            logDebug('base64Image', JSON.stringify(message, null, 2))
+        }
 
         const url = new URL(await this.completionsEndpoint())
         if (apiVersion >= 1) {
@@ -326,7 +330,7 @@ export class SourcegraphNodeCompletionsClient extends SourcegraphCompletionsClie
                         getActiveTraceAndSpanId()?.traceId
                     )
                 }
-                const json = (await response.json()) as CompletionResponse
+                const json = await response.json()
                 if (typeof json?.completion === 'string') {
                     cb.onChange(json.completion)
                     cb.onComplete()

vscode/src/prompt-builder/index.ts

@@ -33,6 +33,7 @@ export class PromptBuilder {
      * A list of context items that are used to build context messages.
      */
     public contextItems: ContextItem[] = []
+    public images: string[] = []
 
     /**
      * Convenience constructor because loading the tokenizer is async due to its large size.
@@ -47,10 +48,28 @@ export class PromptBuilder {
         if (this.contextItems.length > 0) {
             this.buildContextMessages()
         }
-
+        this.buildImageMessages()
         return this.prefixMessages.concat([...this.reverseMessages].reverse())
     }
 
+    private buildImageMessages(): void {
+        for (const image of this.images) {
+            const imageMessage: Message = {
+                speaker: 'human',
+                content: [
+                    {
+                        type: 'image_url',
+                        image_url: {
+                            // TODO: Handle PNG/JPEG, don't hardcode to JPEG
+                            url: `data:image/jpeg;base64,${image}`,
+                        },
+                    },
+                ],
+            }
+            this.reverseMessages.push(...[ASSISTANT_MESSAGE, imageMessage])
+        }
+    }
+
     private buildContextMessages(): void {
         for (const item of this.contextItems) {
             // Create context messages for each context item, where
@@ -108,6 +127,12 @@ export class PromptBuilder {
         return undefined
     }
 
+    public tryAddImage(base64Image: string | undefined): void {
+        if (base64Image) {
+            this.images.push(base64Image)
+        }
+    }
+
     public async tryAddContext(
         type: ContextTokenUsageType | 'history',
         contextItems: ContextItem[]

vscode/webviews/chat/cells/messageCell/human/editor/HumanMessageEditor.tsx

@@ -31,6 +31,7 @@ import {
 import type { UserAccountInfo } from '../../../../../Chat'
 import { type ClientActionListener, useClientActionListener } from '../../../../../client/clientState'
 import { promptModeToIntent } from '../../../../../prompts/PromptsTab'
+import { getVSCodeAPI } from '../../../../../utils/VSCodeApi'
 import { useTelemetryRecorder } from '../../../../../utils/telemetry'
 import { useFeatureFlag } from '../../../../../utils/useFeatureFlags'
 import { useLinkOpener } from '../../../../../utils/useLinkOpener'
@@ -99,6 +100,8 @@ export const HumanMessageEditor: FunctionComponent<{
 }) => {
     const telemetryRecorder = useTelemetryRecorder()
 
+    const [imageFile, setImageFile] = useState<File | undefined>(undefined)
+
     const editorRef = useRef<PromptEditorRefAPI>(null)
     useImperativeHandle(parentEditorRef, (): PromptEditorRefAPI | null => editorRef.current, [])
 
@@ -126,7 +129,7 @@ export const HumanMessageEditor: FunctionComponent<{
     const experimentalPromptEditorEnabled = useFeatureFlag(FeatureFlag.CodyExperimentalPromptEditor)
 
     const onSubmitClick = useCallback(
-        (intent?: ChatMessage['intent'], forceSubmit?: boolean): void => {
+        async (intent?: ChatMessage['intent'], forceSubmit?: boolean): Promise<void> => {
             if (!forceSubmit && submitState === 'emptyEditorValue') {
                 return
             }
@@ -142,6 +145,28 @@ export const HumanMessageEditor: FunctionComponent<{
 
             const value = editorRef.current.getSerializedValue()
             parentOnSubmit(intent)
+            if (imageFile) {
+                const readFileGetBase64String = (file: File): Promise<string> => {
+                    return new Promise((resolve, reject) => {
+                        const reader = new FileReader()
+                        reader.onload = () => {
+                            const base64 = reader.result
+                            if (base64 && typeof base64 === 'string') {
+                                resolve(base64.split(',')[1])
+                            } else {
+                                reject(new Error('Failed to read file'))
+                            }
+                        }
+                        reader.onerror = () => reject(new Error('Failed to read file'))
+                        reader.readAsDataURL(file)
+                    })
+                }
+
+                const base64 = await readFileGetBase64String(imageFile)
+                getVSCodeAPI().postMessage({ command: 'chat/upload-file', base64 })
+                setImageFile(undefined)
+            }
+            parentOnSubmit(intent)
 
             telemetryRecorder.recordEvent('cody.humanMessageEditor', 'submit', {
                 metadata: {
@@ -157,7 +182,15 @@ export const HumanMessageEditor: FunctionComponent<{
                 },
             })
         },
-        [submitState, parentOnSubmit, onStop, telemetryRecorder.recordEvent, isFirstMessage, isSent]
+        [
+            submitState,
+            parentOnSubmit,
+            onStop,
+            telemetryRecorder.recordEvent,
+            isFirstMessage,
+            isSent,
+            imageFile,
+        ]
     )
 
     const onEditorEnterKey = useCallback(
@@ -423,6 +456,7 @@ export const HumanMessageEditor: FunctionComponent<{
     )
 
     const Editor = experimentalPromptEditorEnabled ? PromptEditorV2 : PromptEditor
+    const experimentalOneBoxEnabled = useFeatureFlag(FeatureFlag.CodyExperimentalOneBoxDebug)
 
     return (
         // biome-ignore lint/a11y/useKeyWithClickEvents: only relevant to click areas
@@ -470,6 +504,9 @@ export const HumanMessageEditor: FunctionComponent<{
                     hidden={!focused && isSent}
                     className={styles.toolbar}
                     intent={intent}
+                    imageFile={imageFile}
+                    setImageFile={setImageFile}
+                    experimentalOneBoxEnabled={experimentalOneBoxEnabled}
                 />
             )}
         </div>

vscode/webviews/chat/cells/messageCell/human/editor/toolbar/Toolbar.tsx

@@ -10,6 +10,7 @@ import { useActionSelect } from '../../../../../../prompts/PromptsTab'
 import { useClientConfig } from '../../../../../../utils/useClientConfig'
 import { AddContextButton } from './AddContextButton'
 import { SubmitButton, type SubmitButtonState } from './SubmitButton'
+import { UploadImageButton } from './UploadImageButton'
 
 /**
  * The toolbar for the human message editor.
@@ -35,6 +36,10 @@ export const Toolbar: FunctionComponent<{
     intent?: ChatMessage['intent']
 
     manuallySelectIntent: (intent: ChatMessage['intent']) => void
+    experimentalOneBoxEnabled?: boolean
+
+    imageFile?: File
+    setImageFile: (file: File | undefined) => void
 }> = ({
     userInfo,
     isEditorFocused,
@@ -48,6 +53,9 @@ export const Toolbar: FunctionComponent<{
     models,
     intent,
     manuallySelectIntent,
+    experimentalOneBoxEnabled,
+    imageFile,
+    setImageFile,
 }) => {
     /**
      * If the user clicks in a gap or on the toolbar outside of any of its buttons, report back to
@@ -88,6 +96,14 @@ export const Toolbar: FunctionComponent<{
                     />
                 )}
                 <PromptSelectFieldToolbarItem focusEditor={focusEditor} className="tw-ml-1 tw-mr-1" />
+                {
+                    <UploadImageButton
+                        className="tw-opacity-60"
+                        imageFile={imageFile}
+                        onClick={setImageFile}
+                    />
+                }
+                <PromptSelectFieldToolbarItem focusEditor={focusEditor} className="tw-ml-1 tw-mr-1" />
                 <ModelSelectFieldToolbarItem
                     models={models}
                     userInfo={userInfo}