feat: proxy (#271)

* initial version of the pezzo proxy

* add proxy to ci

* fix formatting

.github/workflows/ci.yaml

@@ -82,7 +82,7 @@ jobs:
     runs-on: ubuntu-20.04
     strategy:
       matrix:
-        project: ["server", "console"]
+        project: ["server", "console", "proxy"]
     steps:
       - uses: actions/checkout@v3
         with:

.github/workflows/release.yaml

@@ -62,7 +62,7 @@ jobs:
     runs-on: ubuntu-20.04
     strategy:
       matrix:
-        project: ["server", "console"]
+        project: ["server", "console", "proxy"]
     steps:
       - uses: actions/checkout@v2
         with:

.gitignore

@@ -48,5 +48,6 @@ schema.graphql
 
 # temp
 temp
+tmp
 
 kms/data

apps/proxy/.eslintrc.json

@@ -0,0 +1,18 @@
+{
+  "extends": ["../../.eslintrc.json"],
+  "ignorePatterns": ["!**/*"],
+  "overrides": [
+    {
+      "files": ["*.ts", "*.tsx", "*.js", "*.jsx"],
+      "rules": {}
+    },
+    {
+      "files": ["*.ts", "*.tsx"],
+      "rules": {}
+    },
+    {
+      "files": ["*.js", "*.jsx"],
+      "rules": {}
+    }
+  ]
+}

apps/proxy/Dockerfile

@@ -0,0 +1,18 @@
+FROM node:20-alpine
+LABEL org.opencontainers.image.source https://github.com/pezzolabs/pezzo/proxy
+
+RUN apk update
+
+WORKDIR /app
+
+COPY ./dist/apps/proxy/package*.json ./
+
+RUN npm i --omit=dev
+
+COPY ./dist/apps/proxy .
+
+ENV HOST=0.0.0.0
+ENV PORT=3000
+EXPOSE $PORT
+
+ENTRYPOINT ["node", "main.js"]

apps/proxy/jest.config.ts

@@ -0,0 +1,11 @@
+/* eslint-disable */
+export default {
+  displayName: "proxy",
+  preset: "../../jest.preset.js",
+  testEnvironment: "node",
+  transform: {
+    "^.+\\.[tj]s$": ["ts-jest", { tsconfig: "<rootDir>/tsconfig.spec.json" }],
+  },
+  moduleFileExtensions: ["ts", "js", "html"],
+  coverageDirectory: "../../coverage/apps/proxy",
+};

apps/proxy/project.json

@@ -0,0 +1,91 @@
+{
+  "name": "proxy",
+  "$schema": "../../node_modules/nx/schemas/project-schema.json",
+  "sourceRoot": "apps/proxy/src",
+  "projectType": "application",
+  "targets": {
+    "build": {
+      "dependsOn": ["^build"],
+      "executor": "@nx/esbuild:esbuild",
+      "outputs": ["{options.outputPath}"],
+      "defaultConfiguration": "production",
+      "options": {
+        "platform": "node",
+        "outputPath": "dist/apps/proxy",
+        "format": ["cjs"],
+        "bundle": false,
+        "main": "apps/proxy/src/main.ts",
+        "tsConfig": "apps/proxy/tsconfig.app.json",
+        "assets": ["apps/proxy/src/assets"],
+        "generatePackageJson": true,
+        "esbuildOptions": {
+          "sourcemap": true,
+          "outExtension": {
+            ".js": ".js"
+          }
+        }
+      },
+      "configurations": {
+        "development": {},
+        "production": {
+          "esbuildOptions": {
+            "sourcemap": false,
+            "outExtension": {
+              ".js": ".js"
+            }
+          }
+        }
+      }
+    },
+    "serve": {
+      "executor": "@nx/js:node",
+      "defaultConfiguration": "development",
+      "options": {
+        "buildTarget": "proxy:build"
+      },
+      "configurations": {
+        "development": {
+          "buildTarget": "proxy:build:development"
+        },
+        "production": {
+          "buildTarget": "proxy:build:production"
+        }
+      }
+    },
+    "lint": {
+      "executor": "@nx/linter:eslint",
+      "outputs": ["{options.outputFile}"],
+      "options": {
+        "lintFilePatterns": ["apps/proxy/**/*.ts"]
+      }
+    },
+    "test": {
+      "executor": "@nx/jest:jest",
+      "outputs": ["{workspaceRoot}/coverage/{projectRoot}"],
+      "options": {
+        "jestConfig": "apps/proxy/jest.config.ts",
+        "passWithNoTests": true
+      },
+      "configurations": {
+        "ci": {
+          "ci": true,
+          "codeCoverage": true
+        }
+      }
+    },
+    "docker:build": {
+      "dependsOn": ["build"],
+      "executor": "@nx-tools/nx-container:build",
+      "inputs": ["{projectRoot}/../../dist/apps/proxy"],
+      "defaultConfiguration": "local",
+      "options": {},
+      "configurations": {
+        "local": {
+          "tags": ["pezzolabs/pezzo/proxy"],
+          "push": false
+        }
+      }
+    }
+  },
+  "tags": []
+}

apps/proxy/src/lib/OpenAIHandler.ts

@@ -0,0 +1,167 @@
+import {
+  ObservabilityReportMetadata,
+  PromptExecutionType,
+  Provider,
+} from "@pezzo/types";
+import { RequestWithPezzoClient } from "../types/common.types";
+import { Response } from "express";
+import axios from "axios";
+
+export class OpenAIV1Handler {
+  constructor(private req: RequestWithPezzoClient, private res: Response) {}
+
+  async handleRequest() {
+    const method = this.req.method;
+    const { headers, originalUrl } = this.req;
+    const url = originalUrl.replace("/openai/v1", "");
+    console.log(`[openai] ${method} ${url}`);
+
+    const execFn = async () => {
+      try {
+        const result = await axios({
+          method,
+          url: `https://api.openai.com/v1/${url}`,
+          data: this.req.body,
+          headers: {
+            Authorization: headers.authorization,
+          },
+        });
+
+        const status = result.status;
+        const data = result.data;
+        this.res.status(result.status).send(result.data);
+        return { status, data };
+      } catch (err) {
+        this.res.status(err.response.status).send(err.response.data);
+        return { status: err.response.status, data: err.response.data };
+      }
+    };
+
+    if (url.startsWith("/chat/completions")) {
+      await this.handleCreateChatCompletion(this.req, this.res, execFn);
+    } else {
+      await execFn();
+    }
+  }
+
+  async handleCreateChatCompletion(
+    originalRequest: RequestWithPezzoClient,
+    originalResponse: Response,
+    execFn: any
+  ) {
+    const pezzo = originalRequest.pezzo;
+
+    let properties = {};
+    const isCacheEnabled =
+      originalRequest.headers["x-pezzo-cache-enabled"] === "true";
+    const hasProperties =
+      originalRequest.headers["x-pezzo-properties"] !== undefined;
+
+    if (hasProperties) {
+      properties = JSON.parse(
+        originalRequest.headers["x-pezzo-properties"] as string
+      );
+    }
+
+    const baseMetadata: ObservabilityReportMetadata = {
+      environment: pezzo.options.environment,
+      provider: Provider.OpenAI,
+      type: PromptExecutionType.ChatCompletion,
+      client: "pezzo-proxy",
+      clientVersion: "0.0.1",
+    };
+
+    const requestTimestamp = new Date().toISOString();
+
+    // Report Execution
+    const baseReport = {
+      cacheEnabled: isCacheEnabled,
+      cacheHit: false,
+      metadata: baseMetadata,
+      properties,
+      request: {
+        timestamp: requestTimestamp,
+        body: originalRequest.body,
+      },
+    };
+
+    let response;
+    let reportPayload;
+
+    if (isCacheEnabled) {
+      const cachedRequest = await pezzo.fetchCachedRequest(
+        originalRequest.body
+      );
+
+      if (cachedRequest.hit === true) {
+        baseReport.cacheHit = true;
+
+        console.log("cachedRequest", cachedRequest);
+
+        response = {
+          ...cachedRequest.data,
+          usage: {
+            prompt_tokens: 0,
+            completion_tokens: 0,
+            total_tokens: 0,
+          },
+        };
+
+        reportPayload = {
+          ...baseReport,
+          response: {
+            timestamp: requestTimestamp,
+            body: response,
+            status: 200,
+          },
+        };
+
+        originalResponse.status(200).json(response);
+      } else {
+        baseReport.cacheHit = false;
+      }
+    }
+
+    if (!isCacheEnabled || (isCacheEnabled && !baseReport.cacheHit)) {
+      const { status, data } = await execFn();
+
+      if (status === 200) {
+        reportPayload = {
+          ...baseReport,
+          response: {
+            timestamp: new Date().toISOString(),
+            body: data,
+            status: 200,
+          },
+        };
+      } else {
+        reportPayload = {
+          ...baseReport,
+          response: {
+            timestamp: new Date().toISOString(),
+            body: data,
+            status: status,
+          },
+        };
+      }
+    }
+
+    const shouldWriteToCache =
+      isCacheEnabled &&
+      reportPayload.cacheHit === false &&
+      reportPayload.response.status === 200;
+
+    try {
+      if (shouldWriteToCache) {
+        await Promise.all([
+          pezzo.reportPromptExecution(reportPayload),
+          pezzo.cacheRequest(originalRequest.body, reportPayload.response.body),
+        ]);
+      } else {
+        await pezzo.reportPromptExecution(reportPayload);
+      }
+    } catch (error) {
+      console.error("Error reporting prompt execution", error);
+    }
+  }
+}

apps/proxy/src/lib/middleware/create-openai-client-from-request.ts

@@ -0,0 +1,39 @@
+import { NextFunction, Response } from "express";
+import { RequestWithPezzoClient } from "../../types/common.types";
+import { Pezzo } from "@pezzo/client";
+
+export function createPezzoClientFromRequest(
+  req: RequestWithPezzoClient,
+  res: Response,
+  next: NextFunction
+) {
+  if (!req.headers["x-pezzo-api-key"]) {
+    return res.status(400).send("Missing x-pezzo-api-key header");
+  }
+
+  if (!req.headers["x-pezzo-project-id"]) {
+    return res.status(400).send("Missing x-pezzo-project-id header");
+  }
+
+  if (!req.headers["x-pezzo-environment"]) {
+    return res.status(400).send("Missing x-pezzo-environment header");
+  }
+
+  const options: {
+    apiKey: string;
+    projectId: string;
+    environment: string;
+    serverUrl?: string;
+  } = {
+    apiKey: req.headers["x-pezzo-api-key"] as string,
+    projectId: req.headers["x-pezzo-project-id"] as string,
+    environment: req.headers["x-pezzo-environment"] as string,
+  };
+
+  if (req.headers["x-pezzo-server-url"]) {
+    options.serverUrl = req.headers["x-pezzo-server-url"] as string;
+  }
+
+  req.pezzo = new Pezzo(options);
+  next();
+}

apps/proxy/src/main.ts

@@ -0,0 +1,21 @@
+import "reflect-metadata";
+import express from "express";
+import { openaiRouter } from "./routers/openai.router";
+
+const host = process.env.HOST ?? "localhost";
+const port = process.env.PORT ? Number(process.env.PORT) : 3001;
+
+const app = express();
+
+app.use(express.json());
+app.use(express.urlencoded({ extended: true }));
+
+app.use("/openai/v1", openaiRouter);
+
+app.get("/healthz", (_, res) => {
+  res.status(200).send("OK");
+});
+
+app.listen(port, host, () => {
+  console.log(`[ ready ] http://${host}:${port}`);
+});