Merge pull request #2841 from continuedev/dallin/current-untitled-file

Get current file untitled support

core/config/promptFile.ts

@@ -1,11 +1,12 @@
 import Handlebars from "handlebars";
 import path from "path";
 import * as YAML from "yaml";
-import type { IDE, SlashCommand } from "..";
+import type { ChatHistory, ChatHistoryItem, ChatMessage, ContextItem, ContinueSDK, IContextProvider, IDE, SlashCommand } from "..";
 import { walkDir } from "../indexing/walkDir";
 import { stripImages } from "../llm/images";
 import { renderTemplatedString } from "../promptFiles/renderTemplatedString";
 import { getBasename } from "../util/index";
+import { BaseContextProvider } from "../context";
 
 export const DEFAULT_PROMPTS_FOLDER = ".prompts";
 
@@ -147,28 +148,32 @@ function extractUserInput(input: string, commandName: string): string {
   return input;
 }
 
-async function renderPrompt(prompt: string, context: any, userInput: string) {
+async function renderPrompt(prompt: string, context: ContinueSDK, userInput: string) {
   const helpers = getContextProviderHelpers(context);
 
   // A few context providers that don't need to be in config.json to work in .prompt files
   const diff = await context.ide.getDiff(false);
-  const currentFilePath = await context.ide.getCurrentFile();
-  const currentFile = currentFilePath
-    ? await context.ide.readFile(currentFilePath)
-    : undefined;
+  const currentFile = await context.ide.getCurrentFile();
+  const inputData: Record<string, string> = {
+    diff,
+    input: userInput,
+  };
+  if (currentFile) {
+    inputData.currentFile = currentFile.path;
+  }
 
   return renderTemplatedString(
     prompt,
     context.ide.readFile.bind(context.ide),
-    { diff, currentFile, input: userInput },
+    inputData,
     helpers,
   );
 }
 
 function getContextProviderHelpers(
-  context: any,
+  context: ContinueSDK,
 ): Array<[string, Handlebars.HelperDelegate]> | undefined {
-  return context.config.contextProviders?.map((provider: any) => [
+  return context.config.contextProviders?.map((provider: IContextProvider) => [
     provider.description.title,
     async (helperContext: any) => {
       const items = await provider.getContextItems(helperContext, {
@@ -182,16 +187,16 @@ function getContextProviderHelpers(
         selectedCode: context.selectedCode,
       });
 
-      items.forEach((item: any) =>
+      items.forEach((item) =>
         context.addContextItem(createContextItem(item, provider)),
       );
 
-      return items.map((item: any) => item.content).join("\n\n");
+      return items.map((item) => item.content).join("\n\n");
     },
   ]);
 }
 
-function createContextItem(item: any, provider: any) {
+function createContextItem(item: ContextItem, provider: IContextProvider) {
   return {
     ...item,
     id: {
@@ -202,7 +207,7 @@ function createContextItem(item: any, provider: any) {
 }
 
 function updateChatHistory(
-  history: any[],
+  history: ChatMessage[],
   commandName: string,
   renderedPrompt: string,
   systemMessage?: string,

core/config/types.ts

@@ -463,7 +463,11 @@ declare global {
       stepIndex: number,
     ): Promise<void>;
     getOpenFiles(): Promise<string[]>;
-    getCurrentFile(): Promise<string | undefined>;
+    getCurrentFile(): Promise<undefined | {
+      isUntitled: boolean
+      path: string
+      contents: string
+    }>;
     getPinnedFiles(): Promise<string[]>;
     getSearchResults(query: string): Promise<string>;
     subprocess(command: string): Promise<[string, string]>;

core/context/providers/CurrentFileContextProvider.ts

@@ -24,17 +24,15 @@ class CurrentFileContextProvider extends BaseContextProvider {
     if (!currentFile) {
       return [];
     }
-    const contents = await ide.readFile(currentFile);
+    const baseName = getBasename(currentFile.path);
     return [
       {
-        description: currentFile,
-        content: `This is the currently open file:\n\n\`\`\`${getBasename(
-          currentFile,
-        )}\n${contents}\n\`\`\``,
-        name: getBasename(currentFile),
+        description: currentFile.path,
+        content: `This is the currently open file:\n\n\`\`\`${baseName}\n${currentFile.contents}\n\`\`\``,
+        name: baseName,
         uri: {
           type: "file",
-          value: currentFile,
+          value: currentFile.path,
         },
       },
     ];

core/index.d.ts

@@ -494,7 +494,11 @@ export interface IDE {
     stepIndex: number,
   ): Promise<void>;
   getOpenFiles(): Promise<string[]>;
-  getCurrentFile(): Promise<string | undefined>;
+  getCurrentFile(): Promise<undefined | {
+    isUntitled: boolean
+    path: string
+    contents: string
+  }>;
   getPinnedFiles(): Promise<string[]>;
   getSearchResults(query: string): Promise<string>;
   subprocess(command: string, cwd?: string): Promise<[string, string]>;

core/promptFiles/handlebarUtils.ts

@@ -71,7 +71,7 @@ export function registerHelpers(
 export async function prepareTemplateAndData(
   template: string,
   readFile: (filepath: string) => Promise<string>,
-  inputData: any,
+  inputData: Record<string, string>,
   ctxProviderNames: string[],
 ): Promise<[string, any]> {
   const [newTemplate, vars] = replaceFilepaths(template, ctxProviderNames);
@@ -95,7 +95,7 @@ export async function prepareTemplateAndData(
   return [newTemplate, data];
 }
 
-export function compileAndRenderTemplate(template: string, data: any): string {
+export function compileAndRenderTemplate(template: string, data: Record<string, string>): string {
   const templateFn = Handlebars.compile(template);
   return templateFn(data);
 }

core/promptFiles/renderTemplatedString.ts

@@ -8,7 +8,7 @@ import {
 export async function renderTemplatedString(
   template: string,
   readFile: (filepath: string) => Promise<string>,
-  inputData: any,
+  inputData: Record<string, string>,
   availableHelpers?: Array<[string, Handlebars.HelperDelegate]>,
 ): Promise<string> {
   const helperPromises = availableHelpers

core/promptFiles/slashCommandFromPromptFile.ts

@@ -1,4 +1,4 @@
-import { SlashCommand } from "..";
+import { ContinueSDK, SlashCommand } from "..";
 import { stripImages } from "../llm/images";
 import { getContextProviderHelpers } from "./getContextProviderHelpers";
 import { renderTemplatedString } from "./renderTemplatedString";
@@ -37,18 +37,18 @@ export function extractUserInput(input: string, commandName: string): string {
   return input;
 }
 
-export async function getDefaultVariables(context: any, userInput: string) {
-  const currentFilePath = await context.ide.getCurrentFile();
-  const currentFile = currentFilePath
-    ? await context.ide.readFile(currentFilePath)
-    : undefined;
-
-  return { currentFile, input: userInput };
+export async function getDefaultVariables(context: ContinueSDK, userInput: string): Promise<Record<string, string>> {
+  const currentFile = await context.ide.getCurrentFile();
+  const vars: Record<string, string> = { input: userInput };
+  if (currentFile) {
+    vars.currentFile = currentFile.path;
+  }
+  return vars;
 }
 
 export async function renderPrompt(
   prompt: string,
-  context: any,
+  context: ContinueSDK,
   userInput: string,
 ) {
   const helpers = getContextProviderHelpers(context);

core/protocol/ide.ts

@@ -48,7 +48,11 @@ export type ToIdeFromWebviewOrCoreProtocol = {
   ];
   getProblems: [{ filepath: string }, Problem[]];
   getOpenFiles: [undefined, string[]];
-  getCurrentFile: [undefined, string | undefined];
+  getCurrentFile: [undefined, undefined | {
+    isUntitled: boolean;
+    path: string;
+    contents: string;
+  }];
   getPinnedFiles: [undefined, string[]];
   showLines: [{ filepath: string; startLine: number; endLine: number }, void];
   readRangeInFile: [{ filepath: string; range: Range }, string];

core/util/filesystem.ts

@@ -19,7 +19,7 @@ import { GetGhTokenArgs } from "../protocol/ide.js";
 import { getContinueGlobalPath } from "./paths.js";
 
 class FileSystemIde implements IDE {
-  constructor(private readonly workspaceDir: string) {}
+  constructor(private readonly workspaceDir: string) { }
   showToast(
     type: ToastType,
     message: string,
@@ -205,6 +205,10 @@ class FileSystemIde implements IDE {
     });
   }
 
+  getCurrentFile(): Promise<undefined> {
+    return Promise.resolve(undefined);
+  }
+
   showDiff(
     filepath: string,
     newContents: string,
@@ -221,10 +225,6 @@ class FileSystemIde implements IDE {
     return Promise.resolve([]);
   }
 
-  getCurrentFile(): Promise<string | undefined> {
-    return Promise.resolve("");
-  }
-
   getPinnedFiles(): Promise<string[]> {
     return Promise.resolve([]);
   }

core/util/messageIde.ts

@@ -177,7 +177,7 @@ export class MessageIde implements IDE {
     return this.request("getOpenFiles", undefined);
   }
 
-  getCurrentFile(): Promise<string | undefined> {
+  getCurrentFile() {
     return this.request("getCurrentFile", undefined);
   }
 

extensions/intellij/src/main/kotlin/com/github/continuedev/continueintellijextension/continue/IdeProtocolClient.kt

@@ -1006,11 +1006,17 @@ class IdeProtocolClient(
 //        return openFiles.intersect(pinnedFiles).toList()
     }
 
-    private fun currentFile(): String? {
+    private fun currentFile(): Map<String, Any?>? {
         val fileEditorManager = FileEditorManager.getInstance(project)
         val editor = fileEditorManager.selectedTextEditor
         val virtualFile = editor?.document?.let { FileDocumentManager.getInstance().getFile(it) }
-        return virtualFile?.path
+        return virtualFile?.let {
+            mapOf(
+                "path" to it.path,
+                "contents" to editor.document.text,
+                "isUntitled" to false
+            )
+        }
     }
 
     suspend fun showToast(type: String, content: String, buttonTexts: Array<String> = emptyArray()): String? =

extensions/vscode/src/VsCodeIde.ts

@@ -255,6 +255,7 @@ class VsCodeIde implements IDE {
 
     return tags;
   }
+
   getIdeInfo(): Promise<IdeInfo> {
     return Promise.resolve({
       ideType: "vscode",
@@ -266,6 +267,7 @@ class VsCodeIde implements IDE {
           .version,
     });
   }
+
   readRangeInFile(filepath: string, range: Range): Promise<string> {
     return this.ideUtils.readRangeInFile(
       filepath,
@@ -342,7 +344,7 @@ class VsCodeIde implements IDE {
             type === vscode.FileType.SymbolicLink) &&
           filename === ".continuerc.json"
         ) {
-          const contents = await this.ideUtils.readFile(
+          const contents = await this.readFile(
             vscode.Uri.joinPath(workspaceDir, filename).fsPath,
           );
           configs.push(JSON.parse(contents));
@@ -421,8 +423,55 @@ class VsCodeIde implements IDE {
   async saveFile(filepath: string): Promise<void> {
     await this.ideUtils.saveFile(filepath);
   }
+
+  private static MAX_BYTES = 100000;
+
   async readFile(filepath: string): Promise<string> {
-    return await this.ideUtils.readFile(filepath);
+    try {
+      filepath = this.ideUtils.getAbsolutePath(filepath);
+      const uri = uriFromFilePath(filepath);
+
+      // First, check whether it's a notebook document
+      // Need to iterate over the cells to get full contents
+      const notebook =
+        vscode.workspace.notebookDocuments.find(
+          (doc) => doc.uri.toString() === uri.toString(),
+        ) ??
+        (uri.fsPath.endsWith("ipynb")
+          ? await vscode.workspace.openNotebookDocument(uri)
+          : undefined);
+      if (notebook) {
+        return notebook
+          .getCells()
+          .map((cell) => cell.document.getText())
+          .join("\n\n");
+      }
+
+      // Check whether it's an open document
+      const openTextDocument = vscode.workspace.textDocuments.find(
+        (doc) => doc.uri.fsPath === uri.fsPath,
+      );
+      if (openTextDocument !== undefined) {
+        return openTextDocument.getText();
+      }
+
+      const fileStats = await vscode.workspace.fs.stat(
+        uriFromFilePath(filepath),
+      );
+      if (fileStats.size > 10 * VsCodeIde.MAX_BYTES) {
+        return "";
+      }
+
+      const bytes = await vscode.workspace.fs.readFile(uri);
+
+      // Truncate the buffer to the first MAX_BYTES
+      const truncatedBytes = bytes.slice(0, VsCodeIde.MAX_BYTES);
+      const contents = new TextDecoder().decode(truncatedBytes);
+      return contents;
+    } catch (e) {
+      console.warn("Error reading file", e);
+      return "";
+    }
   }
   async showDiff(
     filepath: string,
@@ -436,8 +485,13 @@ class VsCodeIde implements IDE {
     return await this.ideUtils.getOpenFiles();
   }
 
-  async getCurrentFile(): Promise<string | undefined> {
-    return vscode.window.activeTextEditor?.document.uri.fsPath;
+  async getCurrentFile() {
+    if (!vscode.window.activeTextEditor) return undefined
+    return {
+      isUntitled: vscode.window.activeTextEditor.document.isUntitled,
+      path: vscode.window.activeTextEditor.document.uri.fsPath,
+      contents: vscode.window.activeTextEditor.document.getText()
+    }
   }
 
   async getPinnedFiles(): Promise<string[]> {
@@ -572,3 +626,4 @@ class VsCodeIde implements IDE {
 }
 
 export { VsCodeIde };
+

extensions/vscode/src/extension/VsCodeExtension.ts

@@ -329,8 +329,7 @@ export class VsCodeExtension {
 
     // Register a content provider for the readonly virtual documents
     const documentContentProvider = new (class
-      implements vscode.TextDocumentContentProvider
-    {
+      implements vscode.TextDocumentContentProvider {
       // emitter and its event
       onDidChangeEmitter = new vscode.EventEmitter<vscode.Uri>();
       onDidChange = this.onDidChangeEmitter.event;

extensions/vscode/src/test/test-suites/ideUtils.test.ts

@@ -59,13 +59,6 @@ describe("IDE Utils", () => {
     assert(regex.test(uniqueId));
   });
 
-  test("readFile", async () => {
-    const testPyContents = await utils.readFile(testPyPath);
-    assert(testPyContents === "print('Hello World!')");
-    const testJsContents = await utils.readFile(testJsPath);
-    assert(testJsContents === "console.log('Hello World!')");
-  });
-
   test.skip("getTerminalContents", async () => {
     await new Promise((resolve) => setTimeout(resolve, 1000));
     const terminal = vscode.window.createTerminal();