From 79b2e5b6fda320532e565979189618fc80e6258c Mon Sep 17 00:00:00 2001
From: nai-degen <khoners@protonmail.com>
Date: Wed, 24 Jan 2024 16:42:26 -0600
Subject: [PATCH] adds very basic support for OpenAI function calling

---
 src/config.ts                                      |  9 +++++++++
 src/proxy/middleware/common.ts                     |  4 +++-
 .../preprocessors/transform-outbound-payload.ts    | 14 ++++++++++++++
 3 files changed, 26 insertions(+), 1 deletion(-)

diff --git a/src/config.ts b/src/config.ts
index 45e65ea..00c07a6 100644
--- a/src/config.ts
+++ b/src/config.ts
@@ -236,6 +236,14 @@ type Config = {
    * Defaults to 1, as most deployments are on HuggingFace or Cloudflare Tunnel.
    */
   trustedProxies?: number;
+  /**
+   * Whether to allow OpenAI tool usage.  The proxy doesn't impelment any
+   * support for tools/function calling but can pass requests and responses as
+   * is. Note that the proxy also cannot accurately track quota usage for
+   * requests involving tools, so you must opt in to this feature at your own
+   * risk.
+   */
+  allowOpenAIToolUsage?: boolean;
 };
 
 // To change configs, create a file called .env in the root directory.
@@ -326,6 +334,7 @@ export const config: Config = {
   useInsecureCookies: getEnvWithDefault("USE_INSECURE_COOKIES", isDev),
   staticServiceInfo: getEnvWithDefault("STATIC_SERVICE_INFO", false),
   trustedProxies: getEnvWithDefault("TRUSTED_PROXIES", 1),
+  allowOpenAIToolUsage: getEnvWithDefault("ALLOW_OPENAI_TOOL_USAGE", false),
 } as const;
 
 function generateCookieSecret() {
diff --git a/src/proxy/middleware/common.ts b/src/proxy/middleware/common.ts
index c0928e2..cf5ed39 100644
--- a/src/proxy/middleware/common.ts
+++ b/src/proxy/middleware/common.ts
@@ -194,7 +194,9 @@ export function getCompletionFromBody(req: Request, body: Record<string, any>) {
   switch (format) {
     case "openai":
     case "mistral-ai":
-      return body.choices[0].message.content;
+      // Can be null if the model wants to invoke tools rather than return a
+      // completion.
+      return body.choices[0].message.content || "";
     case "openai-text":
       return body.choices[0].text;
     case "anthropic":
diff --git a/src/proxy/middleware/request/preprocessors/transform-outbound-payload.ts b/src/proxy/middleware/request/preprocessors/transform-outbound-payload.ts
index ddad3f4..7a6a487 100644
--- a/src/proxy/middleware/request/preprocessors/transform-outbound-payload.ts
+++ b/src/proxy/middleware/request/preprocessors/transform-outbound-payload.ts
@@ -92,7 +92,21 @@ export const OpenAIV1ChatCompletionSchema = z
     // special cased it in `addAzureKey` rather than expecting clients to do it.
     logprobs: z.boolean().optional().default(false),
     top_logprobs: z.number().int().optional(),
+    // Quickly adding some newer tool usage params, not tested. They will be
+    // passed through to the API as-is.
+    tools: z.array(z.any()).optional(),
+    functions: z.array(z.any()).optional(),
+    tool_choice: z.any().optional(),
+    function_choice: z.any().optional(),
+    response_format: z.any(),
   })
+  // Tool usage must be enabled via config because we currently have no way to
+  // track quota usage for them or enforce limits.
+  .omit(
+    Boolean(config.allowOpenAIToolUsage)
+      ? {}
+      : { tools: true, functions: true }
+  )
   .strip();
 
 export type OpenAIChatMessage = z.infer<