[WIP] Readdding image support

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}