feat: include other cells as context with ai autocomplete

srcbookdev · Oct 10, 2024 · b3cd5d8 · b3cd5d8
1 parent 9859780
commit b3cd5d8
Show file tree

Hide file tree

Showing 2 changed files with 15 additions and 1 deletion.
diff --git a/packages/web/src/components/cells/code.tsx b/packages/web/src/components/cells/code.tsx
@@ -167,6 +167,7 @@ export default function ControlledCodeCell(props: Props) {
   );
 
   const {
+    cells,
     updateCell: updateCellOnClient,
     clearOutput,
     getTsServerDiagnostics,
@@ -390,6 +391,7 @@ export default function ControlledCodeCell(props: Props) {
           prefix + suffix,
           cell.language,
           prefix.length,
+          cells.filter((c): c is CodeCellType => c.type === 'code' && c.id !== cell.id),
         );
       } catch (err) {
         console.error('Error fetching ai autocomplete suggestion:', err);

diff --git a/packages/web/src/lib/ai-autocomplete/index.ts b/packages/web/src/lib/ai-autocomplete/index.ts
@@ -1,6 +1,7 @@
 import protobuf from 'protobufjs';
 import Long from 'long';
 
+import { CodeCellType } from '@srcbook/shared';
 import languageServerProto from './language-server-proto';
 
 // NOTE: this EDITOR_API_KEY value was just included as a raw string in
@@ -15,6 +16,7 @@ export async function runCodeiumAiAutocomplete(
   source: string,
   sourceLanguage: 'javascript' | 'typescript',
   cursorOffset: number,
+  otherCodeCells: Array<CodeCellType> = [],
 ): Promise<CodiumCompletionResult> {
   const protos = protobuf.Root.fromJSON(languageServerProto as protobuf.INamespace);
   const GetCompletionsRequest = protos.lookupType('exa.language_server_pb.GetCompletionsRequest');
@@ -28,7 +30,17 @@ export async function runCodeiumAiAutocomplete(
   const apiKey = optionalApiKey ?? EDITOR_API_KEY;
 
   const payload = {
-    otherDocuments: [],
+    otherDocuments: otherCodeCells.map((otherCodeCell) =>
+      DocumentInfo.create({
+        absolutePath: otherCodeCell.filename,
+        relativePath: otherCodeCell.filename,
+        text: otherCodeCell.source,
+        editorLanguage: sourceLanguage,
+        language: Language.getOption(sourceLanguage === 'javascript' ? 'JAVASCRIPT' : 'TYPESCRIPT'),
+        cursorOffset: Long.fromValue(0), // NOTE: how do I represent the cursor not being in here?
+        lineEnding: '\n',
+      }),
+    ),
     metadata: Metadata.create({
       ideName: 'web',
       extensionVersion: '1.0.12',