From e1548fe250a169b58e6406a59632581ed992bdab Mon Sep 17 00:00:00 2001 From: Sen Wang Date: Fri, 3 Nov 2023 17:35:56 -0700 Subject: [PATCH 1/6] Codepod self-hosted code copilot --- api/src/run.ts | 19 ++++++- api/src/server.ts | 16 +++++- api/src/spawner/trpc.ts | 69 ++++++++++++++++++++++++ ui/src/components/MyMonaco.tsx | 11 +++- ui/src/lib/llamaCompletionProvider.ts | 75 +++++++++++++++++++++++++++ ui/src/lib/trpc.ts | 20 ++++++- 6 files changed, 205 insertions(+), 5 deletions(-) create mode 100644 ui/src/lib/llamaCompletionProvider.ts diff --git a/api/src/run.ts b/api/src/run.ts index 1624808b..455084f0 100644 --- a/api/src/run.ts +++ b/api/src/run.ts @@ -3,4 +3,21 @@ import { startServer } from "./server"; const repoDir = `${process.cwd()}/example-repo`; console.log("repoDir", repoDir); -startServer({ port: 4000, repoDir }); +const args = process.argv.slice(2); + +const options = { + copilotIpAddress: "127.0.0.1", + copilotPort: 9090, +}; + +for (let i = 0; i < args.length; i++) { + if (args[i] === "--copilotIP" && i + 1 < args.length) { + options.copilotIpAddress = args[i + 1]; + i++; + } else if (args[i] === "--copilotPort" && i + 1 < args.length) { + options.copilotPort = Number(args[i + 1]); + i++; + } +} + +startServer({ port: 4000, repoDir, ...options }); diff --git a/api/src/server.ts b/api/src/server.ts index c8ad55f9..e7a37757 100644 --- a/api/src/server.ts +++ b/api/src/server.ts @@ -10,7 +10,15 @@ import { bindState, writeState } from "./yjs/yjs-blob"; import cors from "cors"; import { createSpawnerRouter, router } from "./spawner/trpc"; -export async function startServer({ port, repoDir }) { +export const copilotIpAddress = "127.0.0.1"; +export const copilotPort = 9090; + +export async function startServer({ + port, + repoDir, + copilotIpAddress = "127.0.0.1", + copilotPort = 9090, +}) { console.log("starting server .."); const app = express(); app.use(express.json({ limit: "20mb" })); @@ -33,6 +41,8 @@ export async function startServer({ port, repoDir }) { ); const http_server = http.createServer(app); + copilotIpAddress = copilotIpAddress; + copilotPort = copilotPort; // Yjs websocket const wss = new WebSocketServer({ noServer: true }); @@ -59,6 +69,8 @@ export async function startServer({ port, repoDir }) { }); http_server.listen({ port }, () => { - console.log(`🚀 Server ready at http://localhost:${port}`); + console.log( + `🚀 Server ready at http://localhost:${port}, LLM Copilot is hosted at ${copilotIpAddress}:${copilotPort}` + ); }); } diff --git a/api/src/spawner/trpc.ts b/api/src/spawner/trpc.ts index 87d1631e..35ac351e 100644 --- a/api/src/spawner/trpc.ts +++ b/api/src/spawner/trpc.ts @@ -3,10 +3,13 @@ const t = initTRPC.create(); export const router = t.router; export const publicProcedure = t.procedure; +import express from "express"; import Y from "yjs"; import WebSocket from "ws"; import { z } from "zod"; +import { copilotIpAddress, copilotPort } from "../server"; + // import { WebsocketProvider } from "../../ui/src/lib/y-websocket"; import { WebsocketProvider } from "../yjs/y-websocket"; @@ -17,6 +20,12 @@ import { connectSocket, runtime2socket, RuntimeInfo } from "./yjs_runtime"; // FIXME need to have a TTL to clear the ydoc. const docs: Map = new Map(); +// FIXME hard-coded yjs server url +const yjsServerUrl = `ws://localhost:4000/socket`; + +const app = express(); +const http = require("http"); + async function getMyYDoc({ repoId, yjsServerUrl }): Promise { return new Promise((resolve, reject) => { const oldydoc = docs.get(repoId); @@ -227,6 +236,66 @@ export function createSpawnerRouter(yjsServerUrl) { ); return true; }), + codeAutoComplete: publicProcedure + .input( + z.object({ + code: z.string(), + podId: z.string(), + }) + ) + .mutation(async ({ input: { code, podId } }) => { + console.log( + `======= codeAutoComplete of pod ${podId} ========\n`, + code + ); + const data = JSON.stringify({ + prompt: code, + temperature: 0.1, + top_k: 40, + top_p: 0.9, + repeat_penalty: 1.05, + // large n_predict significantly slows down the server, a small value is good enough for testing purposes + n_predict: 128, + stream: false, + }); + + const options = { + hostname: copilotIpAddress, + port: copilotPort, + path: "/completion", + method: "POST", + headers: { + "Content-Type": "application/json", + "Content-Length": data.length, + }, + }; + return new Promise((resolve, reject) => { + const req = http.request(options, (res) => { + let responseData = ""; + + res.on("data", (chunk) => { + responseData += chunk; + }); + + res.on("end", () => { + if (responseData.toString() === "") { + resolve(""); // Resolve with an empty string if no data + } + const resData = JSON.parse(responseData.toString()); + console.log(req.statusCode, resData["content"]); + resolve(resData["content"]); // Resolve the Promise with the response data + }); + }); + + req.on("error", (error) => { + console.error(error); + reject(error); // Reject the Promise if an error occurs + }); + + req.write(data); + req.end(); + }); + }), }); } diff --git a/ui/src/components/MyMonaco.tsx b/ui/src/components/MyMonaco.tsx index c1ba2570..778775db 100644 --- a/ui/src/components/MyMonaco.tsx +++ b/ui/src/components/MyMonaco.tsx @@ -13,6 +13,8 @@ import { Annotation } from "../lib/parser"; import { useApolloClient } from "@apollo/client"; import { trpc } from "../lib/trpc"; +import { llamaInlineCompletionProvider } from "../lib/llamaCompletionProvider"; + const theme: monaco.editor.IStandaloneThemeData = { base: "vs", inherit: true, @@ -492,7 +494,14 @@ export const MyMonaco = memo(function MyMonaco({ // // content is value? // updateGitGutter(editor); // }); - + const llamaCompletionProvider = new llamaInlineCompletionProvider( + id, + editor + ); + monaco.languages.registerInlineCompletionsProvider( + "python", + llamaCompletionProvider + ); // bind it to the ytext with pod id if (!codeMap.has(id)) { throw new Error("codeMap doesn't have pod " + id); diff --git a/ui/src/lib/llamaCompletionProvider.ts b/ui/src/lib/llamaCompletionProvider.ts new file mode 100644 index 00000000..8e83cfee --- /dev/null +++ b/ui/src/lib/llamaCompletionProvider.ts @@ -0,0 +1,75 @@ +import { monaco } from "react-monaco-editor"; +import { trpcProxyClient } from "./trpc"; + +export class llamaInlineCompletionProvider + implements monaco.languages.InlineCompletionsProvider +{ + private readonly podId: string; + private readonly editor: monaco.editor.IStandaloneCodeEditor; + private isFetchingSuggestions: boolean; // Flag to track if a fetch operation is in progress + + constructor(podId: string, editor: monaco.editor.IStandaloneCodeEditor) { + this.podId = podId; + this.editor = editor; + this.isFetchingSuggestions = false; // Initialize the flag + } + + private async provideSuggestions(input: string) { + if (/^\s*$/.test(input || " ")) { + return ""; + } + + const suggestion = await trpcProxyClient.spawner.codeAutoComplete.mutate({ + code: input, + podId: this.podId, + }); + return suggestion; + } + public async provideInlineCompletions( + model: monaco.editor.ITextModel, + position: monaco.IPosition, + context: monaco.languages.InlineCompletionContext, + token: monaco.CancellationToken + ): Promise { + if (!this.editor.hasTextFocus()) { + return; + } + if (token.isCancellationRequested) { + return; + } + + if (!this.isFetchingSuggestions) { + this.isFetchingSuggestions = true; + try { + const suggestion = await this.provideSuggestions( + model.getValue() || " " + ); + + return { + items: [ + { + insertText: suggestion, + range: new monaco.Range( + position.lineNumber, + position.column, + position.lineNumber, + position.column + ), + }, + ], + }; + } finally { + this.isFetchingSuggestions = false; + } + } + } + + handleItemDidShow?( + completions: monaco.languages.InlineCompletions, + item: monaco.languages.InlineCompletion + ): void {} + + freeInlineCompletions( + completions: monaco.languages.InlineCompletions + ): void {} +} diff --git a/ui/src/lib/trpc.ts b/ui/src/lib/trpc.ts index 2308e687..54b5638b 100644 --- a/ui/src/lib/trpc.ts +++ b/ui/src/lib/trpc.ts @@ -1,3 +1,21 @@ -import { createTRPCReact } from "@trpc/react-query"; +import { + createTRPCReact, + createTRPCProxyClient, + httpBatchLink, +} from "@trpc/react-query"; import type { AppRouter } from "../../../api/src/spawner/trpc"; export const trpc = createTRPCReact(); +let remoteUrl; + +if (import.meta.env.DEV) { + remoteUrl = `localhost:4000`; +} else { + remoteUrl = `${window.location.hostname}:${window.location.port}`; +} +export const trpcProxyClient = createTRPCProxyClient({ + links: [ + httpBatchLink({ + url: `http://${remoteUrl}/trpc`, + }), + ], +}); From 5e9276c0dcd350655cce6da381ac5e9d500c0f90 Mon Sep 17 00:00:00 2001 From: Sen Wang Date: Fri, 3 Nov 2023 18:30:45 -0700 Subject: [PATCH 2/6] Add a Sidebar setting to turn on/off Codepod copilot --- ui/src/components/MyMonaco.tsx | 28 +++++++++++++++++++--------- ui/src/components/Sidebar.tsx | 20 ++++++++++++++++++++ ui/src/lib/store/settingSlice.ts | 13 +++++++++++++ 3 files changed, 52 insertions(+), 9 deletions(-) diff --git a/ui/src/components/MyMonaco.tsx b/ui/src/components/MyMonaco.tsx index 778775db..ddcce3b3 100644 --- a/ui/src/components/MyMonaco.tsx +++ b/ui/src/components/MyMonaco.tsx @@ -406,6 +406,8 @@ export const MyMonaco = memo(function MyMonaco({ (state) => state.parseResult[id]?.annotations ); const showAnnotations = useStore(store, (state) => state.showAnnotations); + const copilotEnabled = useStore(store, (state) => state.copilotEnabled); + const scopedVars = useStore(store, (state) => state.scopedVars); const updateView = useStore(store, (state) => state.updateView); @@ -427,7 +429,23 @@ export const MyMonaco = memo(function MyMonaco({ } else { highlightAnnotations(editor, []); } - }, [annotations, editor, showAnnotations, scopedVars]); + + let llamaCompletionProvider; // Define the provider reference + let decompose; + if (copilotEnabled) { + llamaCompletionProvider = new llamaInlineCompletionProvider(id, editor); + decompose = monaco.languages.registerInlineCompletionsProvider( + "python", + llamaCompletionProvider + ); + } + return () => { + if (llamaCompletionProvider && !copilotEnabled) { + // Unregister the provider if it exists + decompose(); + } + }; + }, [annotations, editor, showAnnotations, scopedVars, copilotEnabled]); if (lang === "racket") { lang = "scheme"; @@ -494,14 +512,6 @@ export const MyMonaco = memo(function MyMonaco({ // // content is value? // updateGitGutter(editor); // }); - const llamaCompletionProvider = new llamaInlineCompletionProvider( - id, - editor - ); - monaco.languages.registerInlineCompletionsProvider( - "python", - llamaCompletionProvider - ); // bind it to the ytext with pod id if (!codeMap.has(id)) { throw new Error("codeMap doesn't have pod " + id); diff --git a/ui/src/components/Sidebar.tsx b/ui/src/components/Sidebar.tsx index 22f4bc03..3f51699a 100644 --- a/ui/src/components/Sidebar.tsx +++ b/ui/src/components/Sidebar.tsx @@ -64,6 +64,9 @@ function SidebarSettings() { ); const devMode = useStore(store, (state) => state.devMode); const setDevMode = useStore(store, (state) => state.setDevMode); + const copilotEnabled = useStore(store, (state) => state.copilotEnabled); + const setCopilotEnabled = useStore(store, (state) => state.setCopilotEnabled); + const showLineNumbers = useStore(store, (state) => state.showLineNumbers); const setShowLineNumbers = useStore( store, @@ -488,6 +491,23 @@ function SidebarSettings() { /> + + + ) => { + setCopilotEnabled(event.target.checked); + }} + /> + } + label="Enable Codepod Copilot" + /> + + {showAnnotations && ( Function Definition diff --git a/ui/src/lib/store/settingSlice.ts b/ui/src/lib/store/settingSlice.ts index a942402a..9c57c4db 100644 --- a/ui/src/lib/store/settingSlice.ts +++ b/ui/src/lib/store/settingSlice.ts @@ -8,6 +8,8 @@ export interface SettingSlice { setShowAnnotations: (b: boolean) => void; devMode?: boolean; setDevMode: (b: boolean) => void; + copilotEnabled?: boolean; + setCopilotEnabled: (b: boolean) => void; autoRunLayout?: boolean; setAutoRunLayout: (b: boolean) => void; contextualZoomParams: Record; @@ -56,6 +58,17 @@ export const createSettingSlice: StateCreator = ( // also write to local storage localStorage.setItem("devMode", JSON.stringify(b)); }, + + copilotEnabled: localStorage.getItem("copilotEnabled") + ? JSON.parse(localStorage.getItem("copilotEnabled")!) + : false, + setCopilotEnabled: (b: boolean) => { + // set it + set({ copilotEnabled: b }); + // also write to local storage + localStorage.setItem("copilotEnabled", JSON.stringify(b)); + }, + autoRunLayout: localStorage.getItem("autoRunLayout") ? JSON.parse(localStorage.getItem("autoRunLayout")!) : true, From 1974f45e73286d243e3cb807968969a14129dd75 Mon Sep 17 00:00:00 2001 From: Sen Wang Date: Fri, 3 Nov 2023 21:08:02 -0700 Subject: [PATCH 3/6] deactive useEffect() on copilotEnabled --- api/src/server.ts | 11 +++++------ api/src/spawner/trpc.ts | 12 +++++++----- ui/src/components/MyMonaco.tsx | 28 +++++++++++----------------- 3 files changed, 23 insertions(+), 28 deletions(-) diff --git a/api/src/server.ts b/api/src/server.ts index e7a37757..1a523a40 100644 --- a/api/src/server.ts +++ b/api/src/server.ts @@ -10,9 +10,6 @@ import { bindState, writeState } from "./yjs/yjs-blob"; import cors from "cors"; import { createSpawnerRouter, router } from "./spawner/trpc"; -export const copilotIpAddress = "127.0.0.1"; -export const copilotPort = 9090; - export async function startServer({ port, repoDir, @@ -35,14 +32,16 @@ export async function startServer({ "/trpc", trpcExpress.createExpressMiddleware({ router: router({ - spawner: createSpawnerRouter(yjsServerUrl), + spawner: createSpawnerRouter( + yjsServerUrl, + copilotIpAddress, + copilotPort + ), }), }) ); const http_server = http.createServer(app); - copilotIpAddress = copilotIpAddress; - copilotPort = copilotPort; // Yjs websocket const wss = new WebSocketServer({ noServer: true }); diff --git a/api/src/spawner/trpc.ts b/api/src/spawner/trpc.ts index 35ac351e..c6246a9f 100644 --- a/api/src/spawner/trpc.ts +++ b/api/src/spawner/trpc.ts @@ -8,8 +8,6 @@ import Y from "yjs"; import WebSocket from "ws"; import { z } from "zod"; -import { copilotIpAddress, copilotPort } from "../server"; - // import { WebsocketProvider } from "../../ui/src/lib/y-websocket"; import { WebsocketProvider } from "../yjs/y-websocket"; @@ -61,7 +59,11 @@ async function getMyYDoc({ repoId, yjsServerUrl }): Promise { const routingTable: Map = new Map(); -export function createSpawnerRouter(yjsServerUrl) { +export function createSpawnerRouter( + yjsServerUrl, + copilotIpAddress, + copilotPort +) { return router({ spawnRuntime: publicProcedure .input(z.object({ runtimeId: z.string(), repoId: z.string() })) @@ -282,7 +284,7 @@ export function createSpawnerRouter(yjsServerUrl) { resolve(""); // Resolve with an empty string if no data } const resData = JSON.parse(responseData.toString()); - console.log(req.statusCode, resData["content"]); + console.log(res.statusCode, resData["content"]); resolve(resData["content"]); // Resolve the Promise with the response data }); }); @@ -301,6 +303,6 @@ export function createSpawnerRouter(yjsServerUrl) { // This is only used for frontend to get the type of router. const _appRouter_for_type = router({ - spawner: createSpawnerRouter(null), // put procedures under "post" namespace + spawner: createSpawnerRouter(null, null, null), // put procedures under "post" namespace }); export type AppRouter = typeof _appRouter_for_type; diff --git a/ui/src/components/MyMonaco.tsx b/ui/src/components/MyMonaco.tsx index ddcce3b3..6252ad4a 100644 --- a/ui/src/components/MyMonaco.tsx +++ b/ui/src/components/MyMonaco.tsx @@ -429,23 +429,7 @@ export const MyMonaco = memo(function MyMonaco({ } else { highlightAnnotations(editor, []); } - - let llamaCompletionProvider; // Define the provider reference - let decompose; - if (copilotEnabled) { - llamaCompletionProvider = new llamaInlineCompletionProvider(id, editor); - decompose = monaco.languages.registerInlineCompletionsProvider( - "python", - llamaCompletionProvider - ); - } - return () => { - if (llamaCompletionProvider && !copilotEnabled) { - // Unregister the provider if it exists - decompose(); - } - }; - }, [annotations, editor, showAnnotations, scopedVars, copilotEnabled]); + }, [annotations, editor, showAnnotations, scopedVars]); if (lang === "racket") { lang = "scheme"; @@ -512,6 +496,16 @@ export const MyMonaco = memo(function MyMonaco({ // // content is value? // updateGitGutter(editor); // }); + if (copilotEnabled) { + const llamaCompletionProvider = new llamaInlineCompletionProvider( + id, + editor + ); + monaco.languages.registerInlineCompletionsProvider( + "python", + llamaCompletionProvider + ); + } // bind it to the ytext with pod id if (!codeMap.has(id)) { throw new Error("codeMap doesn't have pod " + id); From 33951f45bbceef77b3dd93051b2eed45c0d461c7 Mon Sep 17 00:00:00 2001 From: Sen Wang Date: Thu, 9 Nov 2023 18:31:24 -0800 Subject: [PATCH 4/6] Using trpc.useUtils() hook to get the global trpcClient --- ui/src/components/MyMonaco.tsx | 4 +++- ui/src/lib/llamaCompletionProvider.ts | 11 ++++++++--- ui/src/lib/trpc.ts | 20 +------------------- 3 files changed, 12 insertions(+), 23 deletions(-) diff --git a/ui/src/components/MyMonaco.tsx b/ui/src/components/MyMonaco.tsx index 6252ad4a..a8aed8d7 100644 --- a/ui/src/components/MyMonaco.tsx +++ b/ui/src/components/MyMonaco.tsx @@ -441,6 +441,7 @@ export const MyMonaco = memo(function MyMonaco({ const selectPod = useStore(store, (state) => state.selectPod); const resetSelection = useStore(store, (state) => state.resetSelection); const editMode = useStore(store, (state) => state.editMode); + const { client } = trpc.useUtils(); // FIXME useCallback? function onEditorDidMount( @@ -499,7 +500,8 @@ export const MyMonaco = memo(function MyMonaco({ if (copilotEnabled) { const llamaCompletionProvider = new llamaInlineCompletionProvider( id, - editor + editor, + client ); monaco.languages.registerInlineCompletionsProvider( "python", diff --git a/ui/src/lib/llamaCompletionProvider.ts b/ui/src/lib/llamaCompletionProvider.ts index 8e83cfee..e67bb187 100644 --- a/ui/src/lib/llamaCompletionProvider.ts +++ b/ui/src/lib/llamaCompletionProvider.ts @@ -1,16 +1,21 @@ import { monaco } from "react-monaco-editor"; -import { trpcProxyClient } from "./trpc"; export class llamaInlineCompletionProvider implements monaco.languages.InlineCompletionsProvider { private readonly podId: string; private readonly editor: monaco.editor.IStandaloneCodeEditor; + private readonly trpc: any; private isFetchingSuggestions: boolean; // Flag to track if a fetch operation is in progress - constructor(podId: string, editor: monaco.editor.IStandaloneCodeEditor) { + constructor( + podId: string, + editor: monaco.editor.IStandaloneCodeEditor, + trpc: any + ) { this.podId = podId; this.editor = editor; + this.trpc = trpc; this.isFetchingSuggestions = false; // Initialize the flag } @@ -19,7 +24,7 @@ export class llamaInlineCompletionProvider return ""; } - const suggestion = await trpcProxyClient.spawner.codeAutoComplete.mutate({ + const suggestion = await this.trpc.spawner.codeAutoComplete.mutate({ code: input, podId: this.podId, }); diff --git a/ui/src/lib/trpc.ts b/ui/src/lib/trpc.ts index 54b5638b..2308e687 100644 --- a/ui/src/lib/trpc.ts +++ b/ui/src/lib/trpc.ts @@ -1,21 +1,3 @@ -import { - createTRPCReact, - createTRPCProxyClient, - httpBatchLink, -} from "@trpc/react-query"; +import { createTRPCReact } from "@trpc/react-query"; import type { AppRouter } from "../../../api/src/spawner/trpc"; export const trpc = createTRPCReact(); -let remoteUrl; - -if (import.meta.env.DEV) { - remoteUrl = `localhost:4000`; -} else { - remoteUrl = `${window.location.hostname}:${window.location.port}`; -} -export const trpcProxyClient = createTRPCProxyClient({ - links: [ - httpBatchLink({ - url: `http://${remoteUrl}/trpc`, - }), - ], -}); From 6fb6b316e9c62c9ce8b533c60d48607fa35d2ca7 Mon Sep 17 00:00:00 2001 From: Sen Wang Date: Fri, 10 Nov 2023 15:55:24 -0800 Subject: [PATCH 5/6] Allow manually triggering Copilot code-autocompletion --- ui/src/components/MyMonaco.tsx | 36 ++++++++++++++++++--------- ui/src/components/Sidebar.tsx | 18 +++++++++----- ui/src/lib/llamaCompletionProvider.ts | 13 +++++++--- ui/src/lib/store/settingSlice.ts | 14 +++++------ 4 files changed, 53 insertions(+), 28 deletions(-) diff --git a/ui/src/components/MyMonaco.tsx b/ui/src/components/MyMonaco.tsx index a8aed8d7..49d1188f 100644 --- a/ui/src/components/MyMonaco.tsx +++ b/ui/src/components/MyMonaco.tsx @@ -406,7 +406,7 @@ export const MyMonaco = memo(function MyMonaco({ (state) => state.parseResult[id]?.annotations ); const showAnnotations = useStore(store, (state) => state.showAnnotations); - const copilotEnabled = useStore(store, (state) => state.copilotEnabled); + const copilotManualMode = useStore(store, (state) => state.copilotManualMode); const scopedVars = useStore(store, (state) => state.scopedVars); const updateView = useStore(store, (state) => state.updateView); @@ -493,21 +493,33 @@ export const MyMonaco = memo(function MyMonaco({ }, }); + editor.addAction({ + id: "trigger-inline-suggest", + label: "Trigger Suggest", + keybindings: [ + monaco.KeyMod.WinCtrl | monaco.KeyMod.Shift | monaco.KeyCode.Space, + ], + run: () => { + editor.trigger(null, "editor.action.inlineSuggest.trigger", null); + }, + }); + // editor.onDidChangeModelContent(async (e) => { // // content is value? // updateGitGutter(editor); // }); - if (copilotEnabled) { - const llamaCompletionProvider = new llamaInlineCompletionProvider( - id, - editor, - client - ); - monaco.languages.registerInlineCompletionsProvider( - "python", - llamaCompletionProvider - ); - } + + const llamaCompletionProvider = new llamaInlineCompletionProvider( + id, + editor, + client, + copilotManualMode || false + ); + monaco.languages.registerInlineCompletionsProvider( + "python", + llamaCompletionProvider + ); + // bind it to the ytext with pod id if (!codeMap.has(id)) { throw new Error("codeMap doesn't have pod " + id); diff --git a/ui/src/components/Sidebar.tsx b/ui/src/components/Sidebar.tsx index 3f51699a..3aec8edf 100644 --- a/ui/src/components/Sidebar.tsx +++ b/ui/src/components/Sidebar.tsx @@ -64,8 +64,11 @@ function SidebarSettings() { ); const devMode = useStore(store, (state) => state.devMode); const setDevMode = useStore(store, (state) => state.setDevMode); - const copilotEnabled = useStore(store, (state) => state.copilotEnabled); - const setCopilotEnabled = useStore(store, (state) => state.setCopilotEnabled); + const copilotManualMode = useStore(store, (state) => state.copilotManualMode); + const setCopilotManualMode = useStore( + store, + (state) => state.setCopilotManualMode + ); const showLineNumbers = useStore(store, (state) => state.showLineNumbers); const setShowLineNumbers = useStore( @@ -491,20 +494,23 @@ function SidebarSettings() { /> - + ) => { - setCopilotEnabled(event.target.checked); + setCopilotManualMode(event.target.checked); }} /> } - label="Enable Codepod Copilot" + label="Trigger Copilot Manually" /> diff --git a/ui/src/lib/llamaCompletionProvider.ts b/ui/src/lib/llamaCompletionProvider.ts index e67bb187..c3969a3e 100644 --- a/ui/src/lib/llamaCompletionProvider.ts +++ b/ui/src/lib/llamaCompletionProvider.ts @@ -6,16 +6,19 @@ export class llamaInlineCompletionProvider private readonly podId: string; private readonly editor: monaco.editor.IStandaloneCodeEditor; private readonly trpc: any; + private readonly manualMode: boolean; private isFetchingSuggestions: boolean; // Flag to track if a fetch operation is in progress constructor( podId: string, editor: monaco.editor.IStandaloneCodeEditor, - trpc: any + trpc: any, + manualMode: boolean ) { this.podId = podId; this.editor = editor; this.trpc = trpc; + this.manualMode = manualMode; this.isFetchingSuggestions = false; // Initialize the flag } @@ -36,10 +39,14 @@ export class llamaInlineCompletionProvider context: monaco.languages.InlineCompletionContext, token: monaco.CancellationToken ): Promise { - if (!this.editor.hasTextFocus()) { + if (!this.editor.hasTextFocus() || token.isCancellationRequested) { return; } - if (token.isCancellationRequested) { + if ( + context.triggerKind === + monaco.languages.InlineCompletionTriggerKind.Automatic && + this.manualMode + ) { return; } diff --git a/ui/src/lib/store/settingSlice.ts b/ui/src/lib/store/settingSlice.ts index 9c57c4db..e7c0b689 100644 --- a/ui/src/lib/store/settingSlice.ts +++ b/ui/src/lib/store/settingSlice.ts @@ -8,8 +8,8 @@ export interface SettingSlice { setShowAnnotations: (b: boolean) => void; devMode?: boolean; setDevMode: (b: boolean) => void; - copilotEnabled?: boolean; - setCopilotEnabled: (b: boolean) => void; + copilotManualMode?: boolean; + setCopilotManualMode: (b: boolean) => void; autoRunLayout?: boolean; setAutoRunLayout: (b: boolean) => void; contextualZoomParams: Record; @@ -59,14 +59,14 @@ export const createSettingSlice: StateCreator = ( localStorage.setItem("devMode", JSON.stringify(b)); }, - copilotEnabled: localStorage.getItem("copilotEnabled") - ? JSON.parse(localStorage.getItem("copilotEnabled")!) + copilotManualMode: localStorage.getItem("copilotManualMode") + ? JSON.parse(localStorage.getItem("copilotManualMode")!) : false, - setCopilotEnabled: (b: boolean) => { + setCopilotManualMode: (b: boolean) => { // set it - set({ copilotEnabled: b }); + set({ copilotManualMode: b }); // also write to local storage - localStorage.setItem("copilotEnabled", JSON.stringify(b)); + localStorage.setItem("copilotManualMode", JSON.stringify(b)); }, autoRunLayout: localStorage.getItem("autoRunLayout") From e2f24be78745999e18d4c00d31e716b041a5f98f Mon Sep 17 00:00:00 2001 From: Sen Wang Date: Tue, 14 Nov 2023 21:26:57 -0800 Subject: [PATCH 6/6] Add infilling mode for code auto-completion --- api/src/spawner/trpc.ts | 75 +++++++++++++++++++-------- ui/src/lib/llamaCompletionProvider.ts | 41 ++++++++++++--- 2 files changed, 86 insertions(+), 30 deletions(-) diff --git a/api/src/spawner/trpc.ts b/api/src/spawner/trpc.ts index c6246a9f..49c0b801 100644 --- a/api/src/spawner/trpc.ts +++ b/api/src/spawner/trpc.ts @@ -241,36 +241,65 @@ export function createSpawnerRouter( codeAutoComplete: publicProcedure .input( z.object({ - code: z.string(), + inputPrefix: z.string(), + inputSuffix: z.string(), podId: z.string(), }) ) - .mutation(async ({ input: { code, podId } }) => { + .mutation(async ({ input: { inputPrefix, inputSuffix, podId } }) => { console.log( `======= codeAutoComplete of pod ${podId} ========\n`, - code + inputPrefix, + inputSuffix ); - const data = JSON.stringify({ - prompt: code, - temperature: 0.1, - top_k: 40, - top_p: 0.9, - repeat_penalty: 1.05, - // large n_predict significantly slows down the server, a small value is good enough for testing purposes - n_predict: 128, - stream: false, - }); + let data = ""; + let options = {}; + if (inputSuffix.length == 0) { + data = JSON.stringify({ + prompt: inputPrefix, + temperature: 0.1, + top_k: 40, + top_p: 0.9, + repeat_penalty: 1.05, + // large n_predict significantly slows down the server, a small value is good enough for testing purposes + n_predict: 128, + stream: false, + }); + + options = { + hostname: copilotIpAddress, + port: copilotPort, + path: "/completion", + method: "POST", + headers: { + "Content-Type": "application/json", + "Content-Length": data.length, + }, + }; + } else { + data = JSON.stringify({ + input_prefix: inputPrefix, + input_suffix: inputSuffix, + temperature: 0.1, + top_k: 40, + top_p: 0.9, + repeat_penalty: 1.05, + // large n_predict significantly slows down the server, a small value is good enough for testing purposes + n_predict: 128, + }); + + options = { + hostname: copilotIpAddress, + port: copilotPort, + path: "/infill", + method: "POST", + headers: { + "Content-Type": "application/json", + "Content-Length": data.length, + }, + }; + } - const options = { - hostname: copilotIpAddress, - port: copilotPort, - path: "/completion", - method: "POST", - headers: { - "Content-Type": "application/json", - "Content-Length": data.length, - }, - }; return new Promise((resolve, reject) => { const req = http.request(options, (res) => { let responseData = ""; diff --git a/ui/src/lib/llamaCompletionProvider.ts b/ui/src/lib/llamaCompletionProvider.ts index c3969a3e..58e08586 100644 --- a/ui/src/lib/llamaCompletionProvider.ts +++ b/ui/src/lib/llamaCompletionProvider.ts @@ -22,13 +22,10 @@ export class llamaInlineCompletionProvider this.isFetchingSuggestions = false; // Initialize the flag } - private async provideSuggestions(input: string) { - if (/^\s*$/.test(input || " ")) { - return ""; - } - + private async provideSuggestions(prefix: string, suffix: string) { const suggestion = await this.trpc.spawner.codeAutoComplete.mutate({ - code: input, + inputPrefix: prefix, + inputSuffix: suffix, podId: this.podId, }); return suggestion; @@ -53,8 +50,38 @@ export class llamaInlineCompletionProvider if (!this.isFetchingSuggestions) { this.isFetchingSuggestions = true; try { + // Get text before the position + let inputPrefix = model.getValueInRange({ + startLineNumber: 1, + startColumn: 1, + endLineNumber: position.lineNumber, + endColumn: position.column, + }); + + // Get text after the position + let inputSuffix = model.getValueInRange({ + startLineNumber: position.lineNumber, + startColumn: position.column, + endLineNumber: model.getLineCount(), + endColumn: model.getLineMaxColumn(model.getLineCount()), + }); + + console.log(inputPrefix); + console.log(inputSuffix); + + if (/^\s*$/.test(inputPrefix || " ")) { + inputPrefix = inputPrefix.trim(); + } + if (/^\s*$/.test(inputSuffix || " ")) { + inputSuffix = inputSuffix.trim(); + } + + if (inputPrefix === "" && inputSuffix === "") { + return; + } const suggestion = await this.provideSuggestions( - model.getValue() || " " + inputPrefix, + inputSuffix ); return {