diff --git a/src/extension/ai/events.ts b/src/extension/ai/events.ts index 84a5dab2d..09076d13a 100644 --- a/src/extension/ai/events.ts +++ b/src/extension/ai/events.ts @@ -5,6 +5,7 @@ import { LogEventType, LogEvent, LogEvent_ExecuteStatus, + StreamGenerateRequest_Trigger, } from '@buf/jlewi_foyle.bufbuild_es/foyle/v1alpha1/agent_pb' import * as vscode from 'vscode' import { ulid } from 'ulidx' @@ -15,6 +16,7 @@ import getLogger from '../logger' import { SessionManager } from './sessions' import * as converters from './converters' +import * as stream from './stream' // Interface for the event reporter // This allows us to swap in a null op logger when AI isn't enabled @@ -29,10 +31,12 @@ export class EventReporter implements vscode.Disposable, IEventReporter { log: ReturnType queue: Subject = new Subject() subscription: Subscription + streamCreator: stream.StreamCreator - constructor(client: PromiseClient) { + constructor(client: PromiseClient, streamCreator: stream.StreamCreator) { this.client = client this.log = getLogger('AIEventReporter') + this.streamCreator = streamCreator this.subscription = this.queue .pipe( @@ -93,6 +97,13 @@ export class EventReporter implements vscode.Disposable, IEventReporter { } else { event.executeStatus = LogEvent_ExecuteStatus.FAILED } + // Fire an event to trigger the AI service + const cellChangeEvent = new stream.CellChangeEvent( + cell.notebook.uri.toString(), + cell.index, + StreamGenerateRequest_Trigger.CELL_EXECUTE, + ) + this.streamCreator.handleEvent(cellChangeEvent) return this.reportEvents([event]) } diff --git a/src/extension/ai/ghost.ts b/src/extension/ai/ghost.ts index 203a9b61f..dbfd4e179 100644 --- a/src/extension/ai/ghost.ts +++ b/src/extension/ai/ghost.ts @@ -124,6 +124,7 @@ export class GhostCellGenerator implements stream.CompletionHandlers { let notebookProto = await this.converter.notebookDataToProto(notebookData) let request = new agent_pb.StreamGenerateRequest({ contextId: SessionManager.getManager().getID(), + trigger: cellChangeEvent.trigger, request: { case: 'fullContext', value: new agent_pb.FullContext({ @@ -142,6 +143,7 @@ export class GhostCellGenerator implements stream.CompletionHandlers { let notebookProto = await this.converter.notebookDataToProto(notebookData) let request = new agent_pb.StreamGenerateRequest({ contextId: SessionManager.getManager().getID(), + trigger: cellChangeEvent.trigger, request: { case: 'update', value: new agent_pb.UpdateContext({ @@ -377,40 +379,6 @@ export class CellChangeEventGenerator { editorAsGhost(editor) } } - - handleOnDidChangeNotebookDocument = (event: vscode.NotebookDocumentChangeEvent) => { - // N.B. For non-interactive cells this will trigger each time the output is updated. - // For interactive cells this doesn't appear to trigger each time the cell output is updated. - // For example, if you have a long running command (e.g. a bash for loop with a sleep that - // echos a message on each iteration) then this won't trigger on each iteration for - // an interactive cell but will for non-interactive. - event.cellChanges.forEach(async (change) => { - log.info(`handleOnDidChangeNotebookDocument: change: ${change}`) - if (change.outputs !== undefined) { - // If outputs change then we want to trigger completions. - - // N.B. It looks like if you click the "configure" button associated with a cell then it will trigger - // an output change event. I don't think there's any easy way to filter those events out. To filter - // those events out we'd need to keep track of the output item with mime type - // application/vnd.code.notebook.stdout and then detect when the stdout changes. That would require - // keeping track of that state. If we trigger on the "configure" then we send a request to the Foyle - // server and we can rely on the Foyle server to do the debouncing. - - // It is the responsibility of the StreamCreator to decide whether the change should be processed.. - // In particular its possible that the cell that changed is not the active cell. Therefore - // we may not want to generate completions for it. For example, you can have multiple cells - // running. So in principle the active cell could be different from the cell that changed. - // - await this.streamCreator.handleEvent( - new stream.CellChangeEvent( - change.cell.notebook.uri.toString(), - change.cell.index, - StreamGenerateRequest_Trigger.CELL_OUTPUT_CHANGE, - ), - ) - } - }) - } } // editorAsGhost decorates an editor as a ghost cell. diff --git a/src/extension/ai/manager.ts b/src/extension/ai/manager.ts index 4943542fb..b37c785e9 100644 --- a/src/extension/ai/manager.ts +++ b/src/extension/ai/manager.ts @@ -33,12 +33,6 @@ export class AIManager implements vscode.Disposable { this.completionGenerator = new generate.CompletionGenerator(this.client, this.converter) if (autoComplete) { this.registerGhostCellEvents() - - const reporter = new events.EventReporter(this.client) - this.subscriptions.push(reporter) - - // Update the global event reporter to use the AI service - events.setEventReporter(reporter) } } @@ -61,6 +55,12 @@ export class AIManager implements vscode.Disposable { // and turns each window into an AsyncIterable of streaming requests. let creator = new stream.StreamCreator(cellGenerator, this.client) + const reporter = new events.EventReporter(this.client, creator) + this.subscriptions.push(reporter) + + // Update the global event reporter to use the AI service + events.setEventReporter(reporter) + let eventGenerator = new ghost.CellChangeEventGenerator(creator) // onDidChangeTextDocument fires when the contents of a cell changes. // We use this to generate completions when the contents of a cell changes. @@ -82,14 +82,6 @@ export class AIManager implements vscode.Disposable { vscode.window.onDidChangeActiveTextEditor(cellGenerator.handleOnDidChangeActiveTextEditor), ) - // We use onDidChangeNotebookDocument to listen for changes to outputs. - // We use this to trigger updates in response to a cell's output being updated. - this.subscriptions.push( - vscode.workspace.onDidChangeNotebookDocument( - eventGenerator.handleOnDidChangeNotebookDocument, - ), - ) - // Create a new status bar item aligned to the right let statusBarItem = vscode.window.createStatusBarItem(vscode.StatusBarAlignment.Right, 100) statusBarItem.text = 'Session: '