[Prototype] Self-hosted CodeLlama LLM for code autocompletion

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 });

api/src/server.ts

@@ -10,7 +10,12 @@ import { bindState, writeState } from "./yjs/yjs-blob";
 import cors from "cors";
 import { createSpawnerRouter, router } from "./spawner/trpc";
 
-export async function startServer({ port, repoDir }) {
+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" }));
@@ -27,7 +32,11 @@ export async function startServer({ port, repoDir }) {
     "/trpc",
     trpcExpress.createExpressMiddleware({
       router: router({
-        spawner: createSpawnerRouter(yjsServerUrl),
+        spawner: createSpawnerRouter(
+          yjsServerUrl,
+          copilotIpAddress,
+          copilotPort
+        ),
       }),
     })
   );
@@ -59,6 +68,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}`
+    );
   });
 }

api/src/spawner/trpc.ts

@@ -3,6 +3,7 @@ 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";
@@ -17,6 +18,12 @@ import { connectSocket, runtime2socket, RuntimeInfo } from "./yjs_runtime";
 // FIXME need to have a TTL to clear the ydoc.
 const docs: Map<string, Y.Doc> = 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<Y.Doc> {
   return new Promise((resolve, reject) => {
     const oldydoc = docs.get(repoId);
@@ -52,7 +59,11 @@ async function getMyYDoc({ repoId, yjsServerUrl }): Promise<Y.Doc> {
 
 const routingTable: Map<string, string> = 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() }))
@@ -227,11 +238,100 @@ export function createSpawnerRouter(yjsServerUrl) {
         );
         return true;
       }),
+    codeAutoComplete: publicProcedure
+      .input(
+        z.object({
+          inputPrefix: z.string(),
+          inputSuffix: z.string(),
+          podId: z.string(),
+        })
+      )
+      .mutation(async ({ input: { inputPrefix, inputSuffix, podId } }) => {
+        console.log(
+          `======= codeAutoComplete of pod ${podId} ========\n`,
+          inputPrefix,
+          inputSuffix
+        );
+        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,
+            },
+          };
+        }
+
+        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(res.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();
+        });
+      }),
   });
 }
 
 // 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;

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,
@@ -404,6 +406,8 @@ export const MyMonaco = memo<MyMonacoProps>(function MyMonaco({
     (state) => state.parseResult[id]?.annotations
   );
   const showAnnotations = useStore(store, (state) => state.showAnnotations);
+  const copilotManualMode = useStore(store, (state) => state.copilotManualMode);
+
   const scopedVars = useStore(store, (state) => state.scopedVars);
   const updateView = useStore(store, (state) => state.updateView);
 
@@ -437,6 +441,7 @@ export const MyMonaco = memo<MyMonacoProps>(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(
@@ -488,11 +493,33 @@ export const MyMonaco = memo<MyMonacoProps>(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);
     // });
 
+    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);

ui/src/components/Sidebar.tsx

@@ -64,6 +64,12 @@ function SidebarSettings() {
   );
   const devMode = useStore(store, (state) => state.devMode);
   const setDevMode = useStore(store, (state) => state.setDevMode);
+  const copilotManualMode = useStore(store, (state) => state.copilotManualMode);
+  const setCopilotManualMode = useStore(
+    store,
+    (state) => state.setCopilotManualMode
+  );
+
   const showLineNumbers = useStore(store, (state) => state.showLineNumbers);
   const setShowLineNumbers = useStore(
     store,
@@ -488,6 +494,26 @@ function SidebarSettings() {
             />
           </FormGroup>
         </Tooltip>
+        <Tooltip
+          title={"Ctrl+Shift+Space to trigger Copilot manually"}
+          disableInteractive
+        >
+          <FormGroup>
+            <FormControlLabel
+              control={
+                <Switch
+                  checked={copilotManualMode}
+                  size="small"
+                  color="warning"
+                  onChange={(event: React.ChangeEvent<HTMLInputElement>) => {
+                    setCopilotManualMode(event.target.checked);
+                  }}
+                />
+              }
+              label="Trigger Copilot Manually"
+            />
+          </FormGroup>
+        </Tooltip>
         {showAnnotations && (
           <Stack spacing={0.5}>
             <Box className="myDecoration-function">Function Definition</Box>