Skip to content

Commit c9e41f1

Browse files
xlight05xlight05
authored
Merge pull request #1009 from xlight05/edit-mode
Add edit mode to agent mode
2 parents 8bf412b + f701d25 commit c9e41f1

File tree

31 files changed

+722
-401
lines changed

31 files changed

+722
-401
lines changed

‎workspaces/ballerina/ballerina-core/src/rpc-types/ai-panel/index.ts‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -84,4 +84,5 @@ export interface AIPanelAPI {
8484
getGeneratedDocumentation: (params: DocGenerationRequest) => Promise<void>;
8585
addFilesToProject: (params: AddFilesToProjectRequest) => Promise<boolean>;
8686
isUserAuthenticated: () => Promise<boolean>;
87+
openAIPanel: (params: AIPanelPrompt) => Promise<void>;
8788
}

‎workspaces/ballerina/ballerina-core/src/rpc-types/ai-panel/interfaces.ts‎

Lines changed: 11 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -21,14 +21,14 @@ import { FunctionDefinition } from "@wso2/syntax-tree";
2121
import { AIMachineContext, AIMachineStateValue, ChatMessage } from "../../state-machine-types";
2222
import { Command, TemplateId } from "../../interfaces/ai-panel";
2323
import { AllDataMapperSourceRequest, DataMapperSourceResponse, ExtendedDataMapperMetadata } from "../../interfaces/extended-lang-client";
24-
import { ComponentInfo, DataMapperMetadata, Diagnostics, ImportStatements, Project } from "../..";
24+
import { ComponentInfo, DataMapperMetadata, Diagnostics, ImportStatements, LinePosition, Project } from "../..";
2525

2626
// ==================================
2727
// General Interfaces
2828
// ==================================
2929
export type AIPanelPrompt =
30-
| { type: 'command-template'; command: Command; templateId: TemplateId; text?: string; params?: Map<string, string>; metadata?: Record<string, any> }
31-
| { type: 'text'; text: string }
30+
| { type: 'command-template'; command: Command; templateId: TemplateId; text?: string; params?: Map<string, string>; metadata?: Record<string, any>}
31+
| { type: 'text'; text: string; planMode: boolean; codeContext?: CodeContext }
3232
| undefined;
3333

3434
export interface AIMachineSnapshot {
@@ -376,11 +376,17 @@ export interface FileAttatchment {
376376
}
377377

378378
export type OperationType = "CODE_GENERATION" | "CODE_FOR_USER_REQUIREMENT" | "TESTS_FOR_USER_REQUIREMENT";
379+
380+
export type CodeContext =
381+
| { type: 'addition'; position: LinePosition, filePath: string }
382+
| { type: 'selection'; startPosition: LinePosition; endPosition: LinePosition, filePath: string };
383+
379384
export interface GenerateCodeRequest {
380385
usecase: string;
381386
chatHistory: ChatEntry[];
382387
operationType: OperationType;
383388
fileAttachmentContents: FileAttatchment[];
389+
codeContext?: CodeContext;
384390
}
385391

386392
export interface GenerateAgentCodeRequest {
@@ -389,6 +395,8 @@ export interface GenerateAgentCodeRequest {
389395
operationType: OperationType;
390396
fileAttachmentContents: FileAttatchment[];
391397
messageId: string;
398+
isPlanMode: boolean;
399+
codeContext?: CodeContext;
392400
}
393401

394402
export interface SourceFiles {

‎workspaces/ballerina/ballerina-core/src/rpc-types/ai-panel/rpc-type.ts‎

Lines changed: 7 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -19,7 +19,7 @@
1919
*/
2020
import { DataMapperModelResponse } from "../../interfaces/extended-lang-client";
2121
import { LoginMethod } from "../../state-machine-types";
22-
import { AddToProjectRequest, GetFromFileRequest, DeleteFromProjectRequest, ProjectSource, ProjectDiagnostics, PostProcessRequest, PostProcessResponse, FetchDataRequest, FetchDataResponse, TestGenerationRequest, TestGenerationResponse, TestGenerationMentions, AIChatSummary, DeveloperDocument, RequirementSpecification, LLMDiagnostics, GetModuleDirParams, AIPanelPrompt, AIMachineSnapshot, SubmitFeedbackRequest, RelevantLibrariesAndFunctionsRequest, GenerateOpenAPIRequest, GenerateCodeRequest, TestPlanGenerationRequest, TestGeneratorIntermediaryState, RepairParams, RelevantLibrariesAndFunctionsResponse, DocGenerationRequest, AddFilesToProjectRequest, MetadataWithAttachments, DatamapperModelContext, ProcessContextTypeCreationRequest, ProcessMappingParametersRequest } from "./interfaces";
22+
import { AddToProjectRequest, GetFromFileRequest, DeleteFromProjectRequest, ProjectSource, ProjectDiagnostics, PostProcessRequest, PostProcessResponse, FetchDataRequest, FetchDataResponse, TestGenerationRequest, TestGenerationResponse, TestGenerationMentions, AIChatSummary, DeveloperDocument, RequirementSpecification, LLMDiagnostics, GetModuleDirParams, AIPanelPrompt, AIMachineSnapshot, SubmitFeedbackRequest, RelevantLibrariesAndFunctionsRequest, GenerateOpenAPIRequest, GenerateCodeRequest, GenerateAgentCodeRequest, TestPlanGenerationRequest, TestGeneratorIntermediaryState, RepairParams, RelevantLibrariesAndFunctionsResponse, DocGenerationRequest, AddFilesToProjectRequest, MetadataWithAttachments, DatamapperModelContext, ProcessContextTypeCreationRequest, ProcessMappingParametersRequest } from "./interfaces";
2323
import { RequestType, NotificationType } from "vscode-messenger-common";
2424

2525
const _preFix = "ai-panel";
@@ -39,9 +39,9 @@ export const getShadowDiagnostics: RequestType<ProjectSource, ProjectDiagnostics
3939
export const checkSyntaxError: RequestType<ProjectSource, boolean> = { method: `${_preFix}/checkSyntaxError` };
4040
export const clearInitialPrompt: NotificationType<void> = { method: `${_preFix}/clearInitialPrompt` };
4141
export const openChatWindowWithCommand: NotificationType<void> = { method: `${_preFix}/openChatWindowWithCommand` };
42-
export const generateContextTypes: RequestType<ProcessContextTypeCreationRequest, void> = { method: `${_preFix}/generateContextTypes` };
43-
export const generateMappingCode: RequestType<ProcessMappingParametersRequest, void> = { method: `${_preFix}/generateMappingCode` };
44-
export const generateInlineMappingCode: RequestType<MetadataWithAttachments, void> = { method: `${_preFix}/generateInlineMappingCode` };
42+
export const generateContextTypes: NotificationType<ProcessContextTypeCreationRequest> = { method: `${_preFix}/generateContextTypes` };
43+
export const generateMappingCode: NotificationType<ProcessMappingParametersRequest> = { method: `${_preFix}/generateMappingCode` };
44+
export const generateInlineMappingCode: NotificationType<MetadataWithAttachments> = { method: `${_preFix}/generateInlineMappingCode` };
4545
export const getGeneratedTests: RequestType<TestGenerationRequest, TestGenerationResponse> = { method: `${_preFix}/getGeneratedTests` };
4646
export const getTestDiagnostics: RequestType<TestGenerationResponse, ProjectDiagnostics> = { method: `${_preFix}/getTestDiagnostics` };
4747
export const getServiceSourceForName: RequestType<string, string> = { method: `${_preFix}/getServiceSourceForName` };
@@ -70,12 +70,13 @@ export const submitFeedback: RequestType<SubmitFeedbackRequest, boolean> = { met
7070
export const getRelevantLibrariesAndFunctions: RequestType<RelevantLibrariesAndFunctionsRequest, RelevantLibrariesAndFunctionsResponse> = { method: `${_preFix}/getRelevantLibrariesAndFunctions` };
7171
export const generateOpenAPI: NotificationType<GenerateOpenAPIRequest> = { method: `${_preFix}/generateOpenAPI` };
7272
export const generateCode: NotificationType<GenerateCodeRequest> = { method: `${_preFix}/generateCode` };
73-
export const generateDesign: RequestType<GenerateCodeRequest, boolean> = { method: `${_preFix}/generateDesign` };
73+
export const generateDesign: RequestType<GenerateAgentCodeRequest, boolean> = { method: `${_preFix}/generateDesign` };
7474
export const repairGeneratedCode: NotificationType<RepairParams> = { method: `${_preFix}/repairGeneratedCode` };
7575
export const generateTestPlan: NotificationType<TestPlanGenerationRequest> = { method: `${_preFix}/generateTestPlan` };
7676
export const generateFunctionTests: NotificationType<TestGeneratorIntermediaryState> = { method: `${_preFix}/generateFunctionTests` };
7777
export const generateHealthcareCode: NotificationType<GenerateCodeRequest> = { method: `${_preFix}/generateHealthcareCode` };
7878
export const abortAIGeneration: NotificationType<void> = { method: `${_preFix}/abortAIGeneration` };
79-
export const getGeneratedDocumentation: RequestType<DocGenerationRequest, boolean> = { method: `${_preFix}/getGeneratedDocumentation` };
79+
export const getGeneratedDocumentation: NotificationType<DocGenerationRequest> = { method: `${_preFix}/getGeneratedDocumentation` };
8080
export const addFilesToProject: RequestType<AddFilesToProjectRequest, boolean> = { method: `${_preFix}/addFilesToProject` };
8181
export const isUserAuthenticated: RequestType<void, boolean> = { method: `${_preFix}/isUserAuthenticated` };
82+
export const openAIPanel: NotificationType<AIPanelPrompt> = { method: `${_preFix}/openAIPanel` };

‎workspaces/ballerina/ballerina-core/src/rpc-types/bi-diagram/interfaces.ts‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -141,8 +141,8 @@ export interface CurrentBreakpointsResponse {
141141
}
142142

143143
export interface AIChatRequest {
144-
scafold: boolean;
145144
readme: boolean;
145+
planMode: boolean;
146146
}
147147

148148
export interface ImportStatements {

‎workspaces/ballerina/ballerina-core/src/state-machine-types.ts‎

Lines changed: 5 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -22,7 +22,7 @@ import { Command } from "./interfaces/ai-panel";
2222
import { LinePosition } from "./interfaces/common";
2323
import { Type } from "./interfaces/extended-lang-client";
2424
import { CodeData, DIRECTORY_MAP, ProjectStructureArtifactResponse, ProjectStructureResponse } from "./interfaces/bi";
25-
import { DiagnosticEntry, TestGeneratorIntermediaryState, DocumentationGeneratorIntermediaryState, SourceFile } from "./rpc-types/ai-panel/interfaces";
25+
import { DiagnosticEntry, TestGeneratorIntermediaryState, DocumentationGeneratorIntermediaryState, SourceFile, CodeContext } from "./rpc-types/ai-panel/interfaces";
2626

2727
export type MachineStateValue =
2828
| 'initialize'
@@ -327,6 +327,7 @@ export const onChatNotify: NotificationType<ChatNotify> = { method: 'onChatNotif
327327
export const onMigrationToolLogs: NotificationType<string> = { method: 'onMigrationToolLogs' };
328328
export const onMigrationToolStateChanged: NotificationType<string> = { method: 'onMigrationToolStateChanged' };
329329
export const projectContentUpdated: NotificationType<boolean> = { method: 'projectContentUpdated' };
330+
export const promptUpdated: NotificationType<void> = { method: 'promptUpdated' };
330331
export const getVisualizerLocation: RequestType<void, VisualizerLocation> = { method: 'getVisualizerLocation' };
331332
export const webviewReady: NotificationType<void> = { method: `webviewReady` };
332333

@@ -500,6 +501,8 @@ export interface AIChatMachineContext {
500501
projectId?: string;
501502
currentApproval?: UserApproval;
502503
autoApproveEnabled?: boolean;
504+
isPlanMode?: boolean;
505+
codeContext?: CodeContext;
503506
currentSpec?: {
504507
requestId: string;
505508
spec?: any;
@@ -512,7 +515,7 @@ export interface AIChatMachineContext {
512515
}
513516

514517
export type AIChatMachineSendableEvent =
515-
| { type: AIChatMachineEventType.SUBMIT_PROMPT; payload: { prompt: string } }
518+
| { type: AIChatMachineEventType.SUBMIT_PROMPT; payload: { prompt: string; isPlanMode: boolean; codeContext?: CodeContext } }
516519
| { type: AIChatMachineEventType.UPDATE_CHAT_MESSAGE; payload: { id: string; modelMessages?: any[]; uiResponse?: string } }
517520
| { type: AIChatMachineEventType.PLANNING_STARTED }
518521
| { type: AIChatMachineEventType.PLAN_GENERATED; payload: { plan: Plan } }

‎workspaces/ballerina/ballerina-extension/grammar/ballerina-grammar‎

‎workspaces/ballerina/ballerina-extension/src/RPCLayer.ts‎

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -19,7 +19,7 @@
1919
import { WebviewView, WebviewPanel, window } from 'vscode';
2020
import { Messenger } from 'vscode-messenger';
2121
import { StateMachine } from './stateMachine';
22-
import { stateChanged, getVisualizerLocation, VisualizerLocation, projectContentUpdated, aiStateChanged, sendAIStateEvent, popupStateChanged, getPopupVisualizerState, PopupVisualizerLocation, breakpointChanged, AIMachineEventType, ArtifactData, onArtifactUpdatedNotification, onArtifactUpdatedRequest, currentThemeChanged, AIMachineSendableEvent, aiChatStateChanged, sendAIChatStateEvent, getAIChatContext, getAIChatUIHistory, AIChatMachineEventType, AIChatMachineSendableEvent, checkpointCaptured, CheckpointCapturedPayload } from '@wso2/ballerina-core';
22+
import { stateChanged, getVisualizerLocation, VisualizerLocation, projectContentUpdated, aiStateChanged, sendAIStateEvent, popupStateChanged, getPopupVisualizerState, PopupVisualizerLocation, breakpointChanged, AIMachineEventType, ArtifactData, onArtifactUpdatedNotification, onArtifactUpdatedRequest, currentThemeChanged, AIMachineSendableEvent, aiChatStateChanged, sendAIChatStateEvent, getAIChatContext, getAIChatUIHistory, AIChatMachineEventType, AIChatMachineSendableEvent, checkpointCaptured, CheckpointCapturedPayload, promptUpdated } from '@wso2/ballerina-core';
2323
import { VisualizerWebview } from './views/visualizer/webview';
2424
import { registerVisualizerRpcHandlers } from './rpc-managers/visualizer/rpc-handler';
2525
import { registerLangClientRpcHandlers } from './rpc-managers/lang-client/rpc-handler';
@@ -187,6 +187,10 @@ export function notifyAiWebview() {
187187
RPCLayer._messenger.sendNotification(projectContentUpdated, { type: 'webview', webviewType: AiPanelWebview.viewType }, true);
188188
}
189189

190+
export function notifyAiPromptUpdated() {
191+
RPCLayer._messenger.sendNotification(promptUpdated, { type: 'webview', webviewType: AiPanelWebview.viewType });
192+
}
193+
190194
export function notifyBreakpointChange() {
191195
RPCLayer._messenger.sendNotification(breakpointChanged, { type: 'webview', webviewType: VisualizerWebview.viewType }, true);
192196
}

‎workspaces/ballerina/ballerina-extension/src/features/ai/constants.ts‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -18,6 +18,11 @@
1818

1919
export const CONFIG_FILE_NAME = "Config.toml";
2020
export const CONFIGURE_DEFAULT_MODEL_COMMAND = "ballerina.configureWso2DefaultModelProvider";
21+
22+
/**
23+
* @deprecated Use {@link openAIPanelWithPrompt} from `views/ai-panel/aiMachine.ts` instead for type-safe command execution.
24+
* This constant is kept for backward compatibility.
25+
*/
2126
export const OPEN_AI_PANEL_COMMAND = "ballerina.open.ai.panel";
2227
export const CLOSE_AI_PANEL_COMMAND = "ballerina.close.ai.panel";
2328
export const SIGN_IN_BI_COPILOT = "Sign in to BI Copilot";

‎workspaces/ballerina/ballerina-extension/src/features/ai/service/datamapper/datamapper.ts‎

Lines changed: 4 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -35,7 +35,8 @@ import { updateSourceCode } from "../../../../../src/utils/source-utils";
3535
import { StateMachine } from "../../../../stateMachine";
3636
import { extractVariableDefinitionSource, getHasStopped, setHasStopped } from "../../../../../src/rpc-managers/data-mapper/utils";
3737
import { commands, Uri, window } from "vscode";
38-
import { CLOSE_AI_PANEL_COMMAND, OPEN_AI_PANEL_COMMAND } from "../../constants";
38+
import { CLOSE_AI_PANEL_COMMAND } from "../../constants";
39+
import { openAIPanelWithPrompt } from "../../../../views/ai-panel/aiMachine";
3940
import path from "path";
4041
import { URI } from "vscode-uri";
4142

@@ -814,11 +815,11 @@ export async function openChatWindowWithCommand(): Promise<void> {
814815
const { identifier, dataMapperMetadata } = context;
815816

816817
commands.executeCommand(CLOSE_AI_PANEL_COMMAND);
817-
commands.executeCommand(OPEN_AI_PANEL_COMMAND, {
818+
openAIPanelWithPrompt({
818819
type: 'command-template',
819820
command: Command.DataMap,
820821
templateId: identifier ? TemplateId.MappingsForFunction : TemplateId.InlineMappings,
821-
...(identifier && { params: { functionName: identifier } }),
822+
...(identifier && { params: new Map([['functionName', identifier]]) }),
822823
metadata: {
823824
...dataMapperMetadata,
824825
mappingsModel: model.mappingsModel as DMModel

0 commit comments

Comments
 (0)