Implement edit/insert prompts (#5958)

closes: https://linear.app/sourcegraph/issue/SRCH-1172/prompts-can-output-into-3-places This PR integrates with the prompt mode (chat | edit | insert) from the SG API. Prompts with edit or insert mode are executed using the `executeEdit`, the same way the commands are executed. The prompt is used as the instruction text for `executeEdit`. First the chat transcript is created, context is fetched and then the `executeEdit` is called. The `FixupTask` from the execution is used to construct the response in chat. The response for now shows the diff only. In BG, the state for interactions mode is saved as intent. The dropdown and icon indicator for to allow users to manually change the mode and to see which mode is set, is only available behind the OneBox feature flag and will not be released yet. ![CleanShot 2024-10-21 at 18 28 08@2x](https://github.com/user-attachments/assets/5cb3e6b1-b53a-41b8-824b-a74f35d1a4b9) ## Test plan - create a edit/insert prompt - execute it from Cody - it should work as expected. ## Changelog - Add ability to execute prompts to perform edits or insert code.
sourcegraph · Oct 22, 2024 · 07a44d3 · 07a44d3
1 parent fdcc8a1
commit 07a44d3
Show file tree

Hide file tree

Showing 15 changed files with 316 additions and 169 deletions.
diff --git a/...ib/src/main/kotlin/com/sourcegraph/cody/agent/protocol_generated/SerializedChatMessage.kt b/...ib/src/main/kotlin/com/sourcegraph/cody/agent/protocol_generated/SerializedChatMessage.kt
@@ -10,7 +10,7 @@ data class SerializedChatMessage(
   val speaker: SpeakerEnum, // Oneof: human, assistant, system
   val text: String? = null,
   val model: String? = null,
-  val intent: IntentEnum? = null, // Oneof: search, chat
+  val intent: IntentEnum? = null, // Oneof: search, chat, edit, insert
 ) {
 
   enum class SpeakerEnum {
@@ -22,6 +22,8 @@ data class SerializedChatMessage(
   enum class IntentEnum {
     @SerializedName("search") Search,
     @SerializedName("chat") Chat,
+    @SerializedName("edit") Edit,
+    @SerializedName("insert") Insert,
   }
 }
 
diff --git a/lib/prompt-editor/src/nodes/ContextItemMentionNode.tsx b/lib/prompt-editor/src/nodes/ContextItemMentionNode.tsx
@@ -177,7 +177,9 @@ function iconForContextItem(contextItem: SerializedContextItem): React.Component
               ? SYMBOL_CONTEXT_MENTION_PROVIDER.id
               : contextItem.type === 'repository' || contextItem.type === 'tree'
                 ? REMOTE_REPOSITORY_PROVIDER_URI
-                : contextItem.providerUri
+                : contextItem.type === 'openctx'
+                  ? contextItem.providerUri
+                  : 'unknown'
     return iconForProvider[providerUri] ?? AtSignIcon
 }
 

diff --git a/lib/shared/src/chat/transcript/messages.ts b/lib/shared/src/chat/transcript/messages.ts
@@ -35,7 +35,7 @@ export interface ChatMessage extends Message {
     model?: string
 
     /* The detected intent of the message */
-    intent?: 'search' | 'chat' | undefined | null
+    intent?: 'search' | 'chat' | 'edit' | 'insert' | undefined | null
 }
 
 // An unsafe version of the {@link ChatMessage} that has the PromptString

diff --git a/lib/shared/src/sourcegraph-api/graphql/client.ts b/lib/shared/src/sourcegraph-api/graphql/client.ts
@@ -415,6 +415,7 @@ export interface Prompt {
     description?: string
     draft: boolean
     autoSubmit?: boolean
+    mode?: PromptMode
     definition: {
         text: string
     }
@@ -427,6 +428,8 @@ export interface Prompt {
     }
 }
 
+export type PromptMode = 'CHAT' | 'EDIT' | 'INSERT'
+
 interface ContextFiltersResponse {
     site: {
         codyContextFilters: {

diff --git a/lib/shared/src/sourcegraph-api/graphql/queries.ts b/lib/shared/src/sourcegraph-api/graphql/queries.ts
@@ -354,6 +354,7 @@ query ViewerPrompts($query: String!) {
             description
             draft
             autoSubmit
+            mode
             definition {
                 text
             }

diff --git a/lib/shared/src/telemetry-v2/events/chat-question.ts b/lib/shared/src/telemetry-v2/events/chat-question.ts
@@ -180,6 +180,8 @@ export const events = [
                 auto: 1,
                 chat: 2,
                 search: 3,
+                edit: 4,
+                insert: 5,
             } satisfies Record<
                 typeof fallbackValue | 'auto' | Exclude<ChatMessage['intent'], null | undefined>,
                 number
@@ -215,6 +217,7 @@ function publicContextSummary(globalPrefix: string, context: ContextItem[]) {
         openctx: cloneDeep(defaultByTypeCount),
         repository: cloneDeep(defaultByTypeCount),
         symbol: cloneDeep(defaultByTypeCount),
+        mode: cloneDeep(defaultByTypeCount),
         tree: {
             ...cloneDeep(defaultByTypeCount),
             isWorkspaceRoot: undefined as number | undefined,

diff --git a/vscode/src/chat/chat-view/ChatController.ts b/vscode/src/chat/chat-view/ChatController.ts
@@ -1,6 +1,7 @@
 import {
     type ChatModel,
     type CodyClientConfig,
+    DefaultEditCommands,
     cenv,
     clientCapabilities,
     currentSiteVersion,
@@ -89,6 +90,7 @@ import {
     startAuthProgressIndicator,
 } from '../../auth/auth-progress-indicator'
 import type { startTokenReceiver } from '../../auth/token-receiver'
+import { executeCodyCommand } from '../../commands/CommandsController'
 import { getContextFileFromUri } from '../../commands/context/file-path'
 import { getContextFileFromCursor } from '../../commands/context/selection'
 import { resolveContextItems } from '../../editor/utils/editor-context'
@@ -719,6 +721,12 @@ export class ChatController implements vscode.Disposable, vscode.WebviewViewProv
             signal.throwIfAborted()
             const corpusContext = contextAlternatives[0].items
 
+            const inputTextWithoutContextChips = editorState
+                ? PromptString.unsafe_fromUserQuery(
+                      inputTextWithoutContextChipsFromPromptEditorState(editorState)
+                  )
+                : inputText
+
             const repositoryMentioned = mentions.find(contextItem =>
                 ['repository', 'tree'].includes(contextItem.type)
             )
@@ -733,55 +741,64 @@ export class ChatController implements vscode.Disposable, vscode.WebviewViewProv
                 ? detectedIntentScores
                 : undefined
 
-            const userSpecifiedIntent = this.featureCodyExperimentalOneBox
-                ? manuallySelectedIntent && detectedIntent
+            const userSpecifiedIntent =
+                manuallySelectedIntent && detectedIntent
                     ? detectedIntent
-                    : 'auto'
-                : undefined
+                    : this.featureCodyExperimentalOneBox
+                      ? 'auto'
+                      : 'chat'
+
+            const finalIntentDetectionResponse = detectedIntent
+                ? { intent: detectedIntent, allScores: detectedIntentScores }
+                : this.featureCodyExperimentalOneBox && repositoryMentioned
+                  ? await this.detectChatIntent({
+                        requestID,
+                        text: inputTextWithoutContextChips.toString(),
+                    })
+                        .then(async response => {
+                            signal.throwIfAborted()
+                            this.chatBuilder.setLastMessageIntent(response?.intent)
+                            this.postEmptyMessageInProgress(model)
+                            return response
+                        })
+                        .catch(() => undefined)
+                  : undefined
 
-            if (this.featureCodyExperimentalOneBox && repositoryMentioned) {
-                const inputTextWithoutContextChips = editorState
-                    ? PromptString.unsafe_fromUserQuery(
-                          inputTextWithoutContextChipsFromPromptEditorState(editorState)
-                      )
-                    : inputText
-
-                const finalIntentDetectionResponse = detectedIntent
-                    ? { intent: detectedIntent, allScores: detectedIntentScores }
-                    : await this.detectChatIntent({
-                          requestID,
-                          text: inputTextWithoutContextChips.toString(),
-                      })
-                          .then(async response => {
-                              signal.throwIfAborted()
-                              this.chatBuilder.setLastMessageIntent(response?.intent)
-                              this.postEmptyMessageInProgress(model)
-                              return response
-                          })
-                          .catch(() => undefined)
-
-                intent = finalIntentDetectionResponse?.intent
-                intentScores = finalIntentDetectionResponse?.allScores
-                signal.throwIfAborted()
-                if (intent === 'search') {
-                    telemetryEvents['cody.chat-question/executed'].record(
-                        {
-                            ...telemetryProperties,
-                            context: corpusContext,
-                            userSpecifiedIntent,
-                            detectedIntent: intent,
-                            detectedIntentScores: intentScores,
-                        },
-                        { current: span, firstToken: firstTokenSpan, addMetadata: true },
-                        tokenCounterUtils
-                    )
+            intent = finalIntentDetectionResponse?.intent
+            intentScores = finalIntentDetectionResponse?.allScores
+            signal.throwIfAborted()
 
-                    return await this.handleSearchIntent({
+            if (['search', 'edit', 'insert'].includes(intent || '')) {
+                telemetryEvents['cody.chat-question/executed'].record(
+                    {
+                        ...telemetryProperties,
                         context: corpusContext,
-                        signal,
-                        contextAlternatives,
-                    })
-                }
+                        userSpecifiedIntent,
+                        detectedIntent: intent,
+                        detectedIntentScores: intentScores,
+                    },
+                    { current: span, firstToken: firstTokenSpan, addMetadata: true },
+                    tokenCounterUtils
+                )
+            }
+
+            if (intent === 'edit' || intent === 'insert') {
+                return await this.handleEditMode({
+                    requestID,
+                    mode: intent,
+                    instruction: inputTextWithoutContextChips,
+                    context: corpusContext,
+                    signal,
+                    contextAlternatives,
+                })
+            }
+
+            if (intent === 'search') {
+                return await this.handleSearchIntent({
+                    context: corpusContext,
+                    signal,
+                    contextAlternatives,
+                })
             }
 
             // Experimental Feature: Deep Cody
@@ -897,6 +914,84 @@ export class ChatController implements vscode.Disposable, vscode.WebviewViewProv
         this.postViewTranscript()
     }
 
+    private async handleEditMode({
+        requestID,
+        mode,
+        instruction,
+        context,
+        signal,
+        contextAlternatives,
+    }: {
+        requestID: string
+        instruction: PromptString
+        mode: 'edit' | 'insert'
+        context: ContextItem[]
+        signal: AbortSignal
+        contextAlternatives: RankedContext[]
+    }): Promise<void> {
+        signal.throwIfAborted()
+
+        this.chatBuilder.setLastMessageContext(context, contextAlternatives)
+        this.chatBuilder.setLastMessageIntent(mode)
+
+        const result = await executeCodyCommand(DefaultEditCommands.Edit, {
+            requestID,
+            runInChatMode: true,
+            userContextFiles: context,
+            configuration: {
+                instruction,
+                mode,
+                intent: mode === 'edit' ? 'edit' : 'add',
+            },
+        })
+
+        if (result?.type !== 'edit' || !result.task) {
+            this.postError(new Error('Failed to execute edit command'), 'transcript')
+            return
+        }
+
+        const task = result.task
+
+        let responseMessage = `Here is the response for the ${task.intent} instruction:\n`
+        task.diff?.map(diff => {
+            responseMessage += '\n```diff\n'
+            if (diff.type === 'deletion') {
+                responseMessage += task.document
+                    .getText(diff.range)
+                    .split('\n')
+                    .map(line => `- ${line}`)
+                    .join('\n')
+            }
+            if (diff.type === 'decoratedReplacement') {
+                responseMessage += diff.oldText
+                    .split('\n')
+                    .map(line => `- ${line}`)
+                    .join('\n')
+                responseMessage += diff.text
+                    .split('\n')
+                    .map(line => `+ ${line}`)
+                    .join('\n')
+            }
+            if (diff.type === 'insertion') {
+                responseMessage += diff.text
+                    .split('\n')
+                    .map(line => `+ ${line}`)
+                    .join('\n')
+            }
+            responseMessage += '\n```'
+        })
+
+        this.chatBuilder.addBotMessage(
+            {
+                text: ps`${PromptString.unsafe_fromLLMResponse(responseMessage)}`,
+            },
+            this.chatBuilder.selectedModel || ChatBuilder.NO_MODEL
+        )
+
+        void this.saveSession()
+        this.postViewTranscript()
+    }
+
     private async computeContext(
         { text, mentions }: HumanInput,
         requestID: string,

diff --git a/vscode/src/chat/protocol.ts b/vscode/src/chat/protocol.ts
@@ -169,6 +169,7 @@ export type ExtensionMessage =
           type: 'clientAction'
           addContextItemsToLastHumanInput?: ContextItem[] | null | undefined
           appendTextToLastPromptEditor?: string | null | undefined
+          setLastHumanInputIntent?: ChatMessage['intent'] | null | undefined
           smartApplyResult?: SmartApplyResult | undefined | null
           submitHumanInput?: boolean | undefined | null
       }

diff --git a/vscode/src/commands/CommandsController.ts b/vscode/src/commands/CommandsController.ts
@@ -76,7 +76,7 @@ class CommandsController implements vscode.Disposable {
 
             // Process default commands
             if (isDefaultChatCommand(commandKey) || isDefaultEditCommand(commandKey)) {
-                return executeDefaultCommand(commandKey, additionalInstruction)
+                return executeDefaultCommand(commandKey, { ...args, additionalInstruction })
             }
 
             const command = this.provider?.get(commandKey)

diff --git a/vscode/src/commands/execute/index.ts b/vscode/src/commands/execute/index.ts
@@ -8,6 +8,7 @@ import {
 } from '@sourcegraph/cody-shared'
 import type { CommandResult } from '../../CommandResult'
 import { executeEdit } from '../../edit/execute'
+import type { CodyCommandArgs } from '../types'
 import { executeDocCommand } from './doc'
 import { executeExplainCommand } from './explain'
 import { executeSmellCommand } from './smell'
@@ -47,20 +48,20 @@ export function isDefaultEditCommand(id: string): DefaultEditCommands | undefine
  */
 export async function executeDefaultCommand(
     id: DefaultCodyCommands | string,
-    additionalInstruction?: PromptString
+    args?: CodyCommandArgs
 ): Promise<CommandResult | undefined> {
     const key = id.replace(/^\//, '').trim() as DefaultCodyCommands
     switch (key) {
         case DefaultChatCommands.Explain:
-            return executeExplainCommand({ additionalInstruction })
+            return executeExplainCommand(args)
         case DefaultChatCommands.Smell:
-            return executeSmellCommand({ additionalInstruction })
+            return executeSmellCommand(args)
         case DefaultEditCommands.Test:
-            return executeTestEditCommand({ additionalInstruction })
+            return executeTestEditCommand(args)
         case DefaultEditCommands.Doc:
-            return executeDocCommand({ additionalInstruction })
+            return executeDocCommand(args)
         case DefaultEditCommands.Edit:
-            return { task: await executeEdit({}), type: 'edit' }
+            return { task: await executeEdit(args || {}), type: 'edit' }
         default:
             console.log('not a default command')
             return undefined

diff --git a/vscode/webviews/chat/cells/messageCell/human/HumanMessageCell.tsx b/vscode/webviews/chat/cells/messageCell/human/HumanMessageCell.tsx
@@ -60,13 +60,19 @@ export const HumanMessageCell: FC<HumanMessageCellProps> = ({ message, ...otherP
         [messageJSON]
     )
 
-    return <HumanMessageCellContent {...otherProps} initialEditorState={initialEditorState} />
+    return (
+        <HumanMessageCellContent
+            {...otherProps}
+            initialEditorState={initialEditorState}
+            intent={message.intent}
+        />
+    )
 }
 
-type HumanMessageCellContent = { initialEditorState: SerializedPromptEditorState } & Omit<
-    HumanMessageCellProps,
-    'message'
->
+type HumanMessageCellContent = {
+    initialEditorState: SerializedPromptEditorState
+    intent: ChatMessage['intent']
+} & Omit<HumanMessageCellProps, 'message'>
 const HumanMessageCellContent = memo<HumanMessageCellContent>(props => {
     const {
         models,
@@ -86,6 +92,7 @@ const HumanMessageCellContent = memo<HumanMessageCellContent>(props => {
         editorRef,
         __storybook__focus,
         onEditorFocusChange,
+        intent,
     } = props
 
     return (
@@ -118,6 +125,7 @@ const HumanMessageCellContent = memo<HumanMessageCellContent>(props => {
                     editorRef={editorRef}
                     __storybook__focus={__storybook__focus}
                     onEditorFocusChange={onEditorFocusChange}
+                    initialIntent={intent}
                 />
             }
             className={className}