-
Notifications
You must be signed in to change notification settings - Fork 16
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
Use the connect protocol to generate unary AI completions (#1559)
* Update Foyle protos to use the new version which has the completion method. * Fix code. * Remove the Foyle Address option. * Print out more information about the error. * Format. * Try adding a log statement to debug the tests. * Try using a console message because log has problems with the mocks. * Try setting runme.baseURL * Add a comment. * Use 8877 as the default port to avoid conflicts with 8080.
- Loading branch information
Showing
10 changed files
with
119 additions
and
123 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.
Oops, something went wrong.
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -1,84 +1,82 @@ | ||
| import * as vscode from 'vscode' | ||
| import { PromiseClient } from '@connectrpc/connect' | ||
| import { AIService } from '@buf/jlewi_foyle.connectrpc_es/foyle/v1alpha1/agent_connect' | ||
| import { | ||
| GenerateCellsRequest, | ||
| GenerateCellsResponse, | ||
| } from '@buf/jlewi_foyle.bufbuild_es/foyle/v1alpha1/agent_pb' | ||
|
|
||
| import { GenerateCellsRequest, GenerateCellsResponse } from '../grpc/aiTypes' | ||
| import * as serializer from '../serializer' | ||
| import getLogger from '../logger' | ||
| import { initAIServiceClient } from '../grpc/aiClient' | ||
| import { AIServiceClient } from '../grpc/aiTypes' | ||
|
|
||
| import * as protos from './protos' | ||
| import * as converters from './converters' | ||
| const log = getLogger('AIGenerate') | ||
|
|
||
| const clients = new Map<string, AIServiceClient>() | ||
| // CompletionGenerator is a class that generates completions for a notebook. | ||
| // It generates a single completion | ||
| export class CompletionGenerator { | ||
| client: PromiseClient<typeof AIService> | ||
|
|
||
| const extName = 'runme' | ||
|
|
||
| export async function generateCompletion() { | ||
| const editor = vscode.window.activeNotebookEditor | ||
|
|
||
| if (!editor) { | ||
| return | ||
| } | ||
|
|
||
| if (editor?.selection.isEmpty) { | ||
| return | ||
| } | ||
|
|
||
| // We subtract 1 because end is non-inclusive | ||
| const lastSelectedCell = editor?.selection.end - 1 | ||
| log.trace(`generateCompletion: lastSelectedCell: ${lastSelectedCell}`) | ||
|
|
||
| // Notebook uses the vscode interface types NotebookDocument and NotebookCell. We | ||
| // need to convert this to NotebookCellData which is the concrete type used by the serializer. | ||
| // This allows us to reuse the existing serializer code. | ||
| let cellData = editor?.notebook.getCells().map((cell) => converters.cellToCellData(cell)) | ||
| let notebookData = new vscode.NotebookData(cellData) | ||
|
|
||
| let notebookProto = serializer.GrpcSerializer.marshalNotebook(notebookData) | ||
|
|
||
| const req = GenerateCellsRequest.create() | ||
| req.notebook = notebookProto | ||
|
|
||
| const config = vscode.workspace.getConfiguration(extName) | ||
| // Include a default so that address is always well defined | ||
| const address = config.get<string>('foyleAddress', 'localhost:9080') | ||
|
|
||
| let client = clients.get(address) | ||
| if (!client) { | ||
| log.info(`Creating new client for address: ${address}`) | ||
| client = initAIServiceClient(address) | ||
| clients.set(address, client) | ||
| constructor(client: PromiseClient<typeof AIService>) { | ||
| this.client = client | ||
| } | ||
|
|
||
| client | ||
| .generateCells(req) | ||
| .then((finished) => { | ||
| let response = finished.response | ||
| // TODO(jeremy): We should have the server add the traceId to the response and then we should | ||
| // log it here. This is for debugging purposes as it allows to easily link to the logs | ||
| log.info('Generate request succeeded traceId') | ||
| public generateCompletion = async () => { | ||
| const editor = vscode.window.activeNotebookEditor | ||
|
|
||
| const insertCells = addAIGeneratedCells(lastSelectedCell + 1, response) | ||
| if (!editor) { | ||
| return | ||
| } | ||
|
|
||
| const edit = new vscode.WorkspaceEdit() | ||
| const notebookUri = editor?.notebook.uri | ||
| edit.set(notebookUri, [insertCells]) | ||
| vscode.workspace.applyEdit(edit).then((result: boolean) => { | ||
| log.trace(`applyedit resolved with ${result}`) | ||
| }) | ||
| }) | ||
| .catch((error) => { | ||
| log.error(`AI Generate request failed ${error}`) | ||
| if (editor?.selection.isEmpty) { | ||
| return | ||
| }) | ||
| } | ||
|
|
||
| // We subtract 1 because end is non-inclusive | ||
| const lastSelectedCell = editor?.selection.end - 1 | ||
| log.trace(`generateCompletion: lastSelectedCell: ${lastSelectedCell}`) | ||
|
|
||
| // Notebook uses the vscode interface types NotebookDocument and NotebookCell. We | ||
| // need to convert this to NotebookCellData which is the concrete type used by the serializer. | ||
| // This allows us to reuse the existing serializer code. | ||
| let cellData = editor?.notebook.getCells().map((cell) => converters.cellToCellData(cell)) | ||
| let notebookData = new vscode.NotebookData(cellData) | ||
|
|
||
| let notebookProto = serializer.GrpcSerializer.marshalNotebook(notebookData) | ||
|
|
||
| const req = new GenerateCellsRequest() | ||
| req.notebook = protos.notebookTSToES(notebookProto) | ||
|
|
||
| this.client | ||
| .generateCells(req) | ||
| .then((response) => { | ||
| // TODO(jeremy): We should have the server add the traceId to the response and then we should | ||
| // log it here. This is for debugging purposes as it allows to easily link to the logs | ||
| log.info('Generate request succeeded traceId') | ||
|
|
||
| const insertCells = addAIGeneratedCells(lastSelectedCell + 1, response) | ||
|
|
||
| const edit = new vscode.WorkspaceEdit() | ||
| const notebookUri = editor?.notebook.uri | ||
| edit.set(notebookUri, [insertCells]) | ||
| vscode.workspace.applyEdit(edit).then((result: boolean) => { | ||
| log.trace(`applyedit resolved with ${result}`) | ||
| }) | ||
| }) | ||
| .catch((error) => { | ||
| log.error(`AI Generate request failed ${error}`) | ||
| return | ||
| }) | ||
| } | ||
| } | ||
|
|
||
| // addAIGeneratedCells turns the response from the AI model into a set of cells that can be inserted into the notebook. | ||
| // This is done by returning a mutation to add the new cells to the notebook. | ||
| // index is the position in the notebook at which the new the new cells should be inserted. | ||
| // | ||
| function addAIGeneratedCells(index: number, response: GenerateCellsResponse): vscode.NotebookEdit { | ||
| let newCellData = converters.cellProtosToCellData(response.cells) | ||
| let newCellData = converters.cellProtosToCellData(protos.cellsESToTS(response.cells)) | ||
| // Now insert the new cells at the end of the notebook | ||
| return vscode.NotebookEdit.insertCells(index, newCellData) | ||
| } |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file was deleted.
Oops, something went wrong.
This file was deleted.
Oops, something went wrong.
Oops, something went wrong.