Run oai-compatible api via separate koa app on a separate port

Author: zenoverflow

server.ts

@@ -27,39 +27,30 @@ const DIR_FRONTEND = path.join(__dirname, "dist");
 const DIR_DATA = path.join(__dirname, "data");
 const DIR_CUSTOM_NODES = path.join(__dirname, "custom_nodes");
 
-// API setup
-
-const app = new Koa();
-const router = new Router();
+ensureDirExists(DIR_DATA);
+ensureDirExists(DIR_CUSTOM_NODES);
 
-setupResourcesApi(router, DIR_DATA, DIR_CUSTOM_NODES);
-setupExecutorApi(router, DIR_DATA, DIR_CUSTOM_NODES);
-// setupExecutorWs(router);
+// Server setup
 
-// Setup folders
+const appMain = new Koa();
+const routerMain = new Router();
 
-ensureDirExists(DIR_DATA);
-ensureDirExists(DIR_CUSTOM_NODES);
+const portOpenAi: number = argv.port_openai || 5002;
 
-// Config: middleware
+setupResourcesApi(routerMain, DIR_DATA, DIR_CUSTOM_NODES);
+setupExecutorApi(routerMain, DIR_DATA, DIR_CUSTOM_NODES, portOpenAi);
 
-app
+appMain
     // body parsing
-    .use(
-        koaBody({
-            jsonLimit: "1gb",
-        })
-    )
+    .use(koaBody({ jsonLimit: "10240gb" }))
     // routing
-    .use(router.routes())
-    .use(router.allowedMethods())
+    .use(routerMain.routes())
+    .use(routerMain.allowedMethods())
     // frontend
     .use(serve(DIR_FRONTEND));
 
-// Config: server
-
 const port: number = argv.port || 12538;
 
-app.listen(port, () => {
-    console.log(`Server started on http://localhost:${port.toString()}`);
+appMain.listen(port, () => {
+    console.log(`App started on http://localhost:${port}`);
 });

server/executor.ts

@@ -344,7 +344,8 @@ const runGraph = async (
 export const setupExecutorApi = (
     router: Router,
     dirData: string,
-    dirCustomNodes: string
+    dirCustomNodes: string,
+    portOpenAi: number
 ) => {
     // Executor connection for ping api
     executorStorage.subscribe((instance) => {
@@ -401,36 +402,40 @@ export const setupExecutorApi = (
     });
 
     // Endpoints to receive messages from an OpenAI-compatible client
-    setupOpenAiCompatibleAPI(router, async (message, checkRequestActive) => {
-        messageQueue.push(message);
-        addMessageToSession(message);
-        console.log("Received message from OAI API:", message.content);
-
-        // Ensure the correct chain is running
-        await runGraph(message.chainId, dirData, dirCustomNodes);
-
-        // If the graph did not run, something went wrong. Return null.
-        if (!executorStorage.get()) return null;
-
-        // Wait for message with the assistant role to appear in the executor session
-        let result: ChatMessage | null = null;
-        while (checkRequestActive()) {
-            const session = executorStorage.get();
-            if (!session) break;
-
-            const messages = session.sessionMessages;
-            if (messages.length) {
-                const lastMessage = messages[messages.length - 1];
-                if (lastMessage.role === "assistant") {
-                    result = lastMessage;
-                    break;
+    // Runs as a separate app on a separate port
+    setupOpenAiCompatibleAPI(
+        portOpenAi,
+        async (message, checkRequestActive) => {
+            messageQueue.push(message);
+            addMessageToSession(message);
+            console.log("Received message from OAI API:", message.content);
+
+            // Ensure the correct chain is running
+            await runGraph(message.chainId, dirData, dirCustomNodes);
+
+            // If the graph did not run, something went wrong. Return null.
+            if (!executorStorage.get()) return null;
+
+            // Wait for message with the assistant role to appear in the executor session
+            let result: ChatMessage | null = null;
+            while (checkRequestActive()) {
+                const session = executorStorage.get();
+                if (!session) break;
+
+                const messages = session.sessionMessages;
+                if (messages.length) {
+                    const lastMessage = messages[messages.length - 1];
+                    if (lastMessage.role === "assistant") {
+                        result = lastMessage;
+                        break;
+                    }
                 }
-            }
 
-            await new Promise((resolve) => setTimeout(resolve, 100));
+                await new Promise((resolve) => setTimeout(resolve, 100));
+            }
+            return result;
         }
-        return result;
-    });
+    );
 
     // Endpoint to run graph
     router.post("/api/executor/run/:id", async (ctx) => {

server/openai.ts

@@ -1,18 +1,32 @@
-import type Router from "koa-router";
+import Koa from "koa";
+import Router from "koa-router";
+import { koaBody } from "koa-body";
 import { v4 as uuid } from "uuid";
 import mime from "mime-types";
 
 import { ChatMessage } from "../src/data/types.ts";
 import { MsgUtils } from "../src/util/MsgUtils.ts";
 
+const appOpenAi = new Koa();
+const routerOpenAi = new Router();
+
+/**
+ * Set up the OpenAI-compatible API.
+ *
+ * Runs on a separate port, to allow easy forwarding
+ * separately from the main app for special use-cases.
+ *
+ * @param port
+ * @param onMessage
+ */
 export const setupOpenAiCompatibleAPI = (
-    router: Router,
+    port: number,
     onMessage: (
         message: ChatMessage,
         checkRequestActive: () => boolean
     ) => Promise<ChatMessage | null>
 ) => {
-    router.post("/v1/completions", async (ctx) => {
+    routerOpenAi.post("/v1/completions", async (ctx) => {
         try {
             const { model, prompt } = ctx.request.body;
 
@@ -50,7 +64,7 @@ export const setupOpenAiCompatibleAPI = (
         }
     });
 
-    router.post("/v1/chat/completions", async (ctx) => {
+    routerOpenAi.post("/v1/chat/completions", async (ctx) => {
         try {
             const { model, messages } = ctx.request.body;
 
@@ -127,4 +141,18 @@ export const setupOpenAiCompatibleAPI = (
             ctx.status = 400;
         }
     });
+
+    // Set up the server itself
+    appOpenAi
+        // body parsing
+        .use(koaBody({ jsonLimit: "10240gb" }))
+        // routing
+        .use(routerOpenAi.routes())
+        .use(routerOpenAi.allowedMethods());
+
+    appOpenAi.listen(port, () => {
+        console.log(
+            `OpenAI-compatible API started on http://localhost:${port}`
+        );
+    });
 };