[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,71 @@ 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(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 copilotEnabled = useStore(store, (state) => state.copilotEnabled);
+
   const scopedVars = useStore(store, (state) => state.scopedVars);
   const updateView = useStore(store, (state) => state.updateView);
 
@@ -492,7 +496,16 @@ export const MyMonaco = memo<MyMonacoProps>(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);

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() {
             />
           </FormGroup>
         </Tooltip>
+        <Tooltip title={"Enable Codepod Copilot"} disableInteractive>
+          <FormGroup>
+            <FormControlLabel
+              control={
+                <Switch
+                  checked={copilotEnabled}
+                  size="small"
+                  color="warning"
+                  onChange={(event: React.ChangeEvent<HTMLInputElement>) => {
+                    setCopilotEnabled(event.target.checked);
+                  }}
+                />
+              }
+              label="Enable Codepod Copilot"
+            />
+          </FormGroup>
+        </Tooltip>
         {showAnnotations && (
           <Stack spacing={0.5}>
             <Box className="myDecoration-function">Function Definition</Box>

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<monaco.languages.InlineCompletions | undefined> {
+    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<monaco.languages.InlineCompletion>,
+    item: monaco.languages.InlineCompletion
+  ): void {}
+
+  freeInlineCompletions(
+    completions: monaco.languages.InlineCompletions<monaco.languages.InlineCompletion>
+  ): void {}
+}

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<any, number>;
@@ -56,6 +58,17 @@ export const createSettingSlice: StateCreator<MyState, [], [], SettingSlice> = (
     // 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,

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<AppRouter>();
+let remoteUrl;
+
+if (import.meta.env.DEV) {
+  remoteUrl = `localhost:4000`;
+} else {
+  remoteUrl = `${window.location.hostname}:${window.location.port}`;
+}
+export const trpcProxyClient = createTRPCProxyClient<AppRouter>({
+  links: [
+    httpBatchLink({
+      url: `http://${remoteUrl}/trpc`,
+    }),
+  ],
+});