Add tracing, factor out some helper functions, fix a few bugs with Nodejs mode

apis/cloudflare/src/proxy.ts

@@ -19,6 +19,10 @@ function apiCacheKey(key: string) {
   return `http://apikey.cache/${encodeURIComponent(key)}.jpg`;
 }
 
+export function braintrustAppUrl(env: Env) {
+  return new URL(env.BRAINTRUST_APP_URL || "https://www.braintrust.dev");
+}
+
 export function originWhitelist(env: Env) {
   return env.WHITELISTED_ORIGINS && env.WHITELISTED_ORIGINS.length > 0
     ? env.WHITELISTED_ORIGINS.split(",")
@@ -113,8 +117,7 @@ export async function handleProxyV1(
         cacheSetLatency.record(end - start);
       },
     },
-    braintrustApiUrl:
-      env.BRAINTRUST_APP_URL || "https://www.braintrustdata.com",
+    braintrustApiUrl: braintrustAppUrl(env).toString(),
     meterProvider,
     whitelist,
   })(request, ctx);

apis/node/package.json

@@ -30,7 +30,7 @@
     "esbuild": "^0.19.9",
     "eventsource-parser": "^1.1.1",
     "express": "^4.19.2",
-    "openai": "^4.23.0",
+    "openai": "^4.42.0",
     "redis": "^4.6.8"
   },
   "devDependencies": {

packages/proxy/package.json

@@ -75,9 +75,9 @@
     "@opentelemetry/core": "^1.19.0",
     "@opentelemetry/resources": "^1.19.0",
     "@opentelemetry/sdk-metrics": "^1.19.0",
-    "ai": "2.2.22",
+    "ai": "2.2.37",
     "eventsource-parser": "^1.1.1",
-    "openai": "4.23.0",
+    "openai": "4.47.1",
     "uuid": "^9.0.1",
     "zod": "^3.22.4"
   }

packages/proxy/src/encrypt.ts

@@ -1,3 +1,5 @@
+import { z } from "zod";
+
 // This is copied from duckdb.tsx in the app directory
 function base64ToArrayBuffer(base64: string) {
   var binaryString = atob(base64);
@@ -34,10 +36,11 @@ export async function decryptMessage(
   return new TextDecoder().decode(decoded);
 }
 
-export interface EncryptedMessage {
-  iv: string;
-  data: string;
-}
+export const encryptedMessageSchema = z.strictObject({
+  iv: z.string(),
+  data: z.string(),
+});
+export type EncryptedMessage = z.infer<typeof encryptedMessageSchema>;
 
 export async function encryptMessage(
   keyString: string,

packages/proxy/src/providers/openai.ts

@@ -0,0 +1,249 @@
+// This is copied from the Vercel AI SDK commit bfa1182c7f5379d7a3d81878ea00ec84682cb046
+// We just need the OpenAI parser, but not the streaming code.
+
+import { CompletionUsage, FunctionCall, trimStartOfStreamHelper } from "ai";
+
+// https://github.com/openai/openai-node/blob/07b3504e1c40fd929f4aae1651b83afc19e3baf8/src/resources/chat/completions.ts#L28-L40
+interface ChatCompletionChunk {
+  id: string;
+  choices: Array<ChatCompletionChunkChoice>;
+  created: number;
+  model: string;
+  object: string;
+}
+
+// https://github.com/openai/openai-node/blob/07b3504e1c40fd929f4aae1651b83afc19e3baf8/src/resources/chat/completions.ts#L43-L49
+// Updated for https://github.com/openai/openai-node/commit/f10c757d831d90407ba47b4659d9cd34b1a35b1d
+// Updated to https://github.com/openai/openai-node/commit/84b43280089eacdf18f171723591856811beddce
+interface ChatCompletionChunkChoice {
+  delta: ChoiceDelta;
+  finish_reason:
+    | "stop"
+    | "length"
+    | "tool_calls"
+    | "content_filter"
+    | "function_call"
+    | null;
+  index: number;
+}
+
+// https://github.com/openai/openai-node/blob/07b3504e1c40fd929f4aae1651b83afc19e3baf8/src/resources/chat/completions.ts#L123-L139
+// Updated to https://github.com/openai/openai-node/commit/84b43280089eacdf18f171723591856811beddce
+interface ChoiceDelta {
+  /**
+   * The contents of the chunk message.
+   */
+  content?: string | null;
+
+  /**
+   * The name and arguments of a function that should be called, as generated by the
+   * model.
+   */
+  function_call?: FunctionCall;
+
+  /**
+   * The role of the author of this message.
+   */
+  role?: "system" | "user" | "assistant" | "tool";
+
+  tool_calls?: Array<DeltaToolCall>;
+}
+
+// From https://github.com/openai/openai-node/blob/master/src/resources/chat/completions.ts
+// Updated to https://github.com/openai/openai-node/commit/84b43280089eacdf18f171723591856811beddce
+interface DeltaToolCall {
+  index: number;
+
+  /**
+   * The ID of the tool call.
+   */
+  id?: string;
+
+  /**
+   * The function that the model called.
+   */
+  function?: ToolCallFunction;
+
+  /**
+   * The type of the tool. Currently, only `function` is supported.
+   */
+  type?: "function";
+}
+
+// From https://github.com/openai/openai-node/blob/master/src/resources/chat/completions.ts
+// Updated to https://github.com/openai/openai-node/commit/84b43280089eacdf18f171723591856811beddce
+interface ToolCallFunction {
+  /**
+   * The arguments to call the function with, as generated by the model in JSON
+   * format. Note that the model does not always generate valid JSON, and may
+   * hallucinate parameters not defined by your function schema. Validate the
+   * arguments in your code before calling your function.
+   */
+  arguments?: string;
+
+  /**
+   * The name of the function to call.
+   */
+  name?: string;
+}
+
+/**
+ * https://github.com/openai/openai-node/blob/3ec43ee790a2eb6a0ccdd5f25faa23251b0f9b8e/src/resources/completions.ts#L28C1-L64C1
+ * Completions API. Streamed and non-streamed responses are the same.
+ */
+interface Completion {
+  /**
+   * A unique identifier for the completion.
+   */
+  id: string;
+
+  /**
+   * The list of completion choices the model generated for the input prompt.
+   */
+  choices: Array<CompletionChoice>;
+
+  /**
+   * The Unix timestamp of when the completion was created.
+   */
+  created: number;
+
+  /**
+   * The model used for completion.
+   */
+  model: string;
+
+  /**
+   * The object type, which is always "text_completion"
+   */
+  object: string;
+
+  /**
+   * Usage statistics for the completion request.
+   */
+  usage?: CompletionUsage;
+}
+
+interface CompletionChoice {
+  /**
+   * The reason the model stopped generating tokens. This will be `stop` if the model
+   * hit a natural stop point or a provided stop sequence, or `length` if the maximum
+   * number of tokens specified in the request was reached.
+   */
+  finish_reason: "stop" | "length" | "content_filter";
+
+  index: number;
+
+  // edited: Removed CompletionChoice.logProbs and replaced with any
+  logprobs: any | null;
+
+  text: string;
+}
+
+/**
+ * Creates a parser function for processing the OpenAI stream data.
+ * The parser extracts and trims text content from the JSON data. This parser
+ * can handle data for chat or completion models.
+ *
+ * @return {(data: string) => string | void} A parser function that takes a JSON string as input and returns the extracted text content or nothing.
+ */
+export function parseOpenAIStream(): (data: string) => string | void {
+  const extract = chunkToText();
+  return (data) => extract(JSON.parse(data) as OpenAIStreamReturnTypes);
+}
+
+function chunkToText(): (chunk: OpenAIStreamReturnTypes) => string | void {
+  const trimStartOfStream = trimStartOfStreamHelper();
+  let isFunctionStreamingIn: boolean;
+  return (json) => {
+    if (isChatCompletionChunk(json)) {
+      const delta = json.choices[0]?.delta;
+      if (delta.function_call?.name) {
+        isFunctionStreamingIn = true;
+        return `{"function_call": {"name": "${delta.function_call.name}", "arguments": "`;
+      } else if (delta.tool_calls?.[0]?.function?.name) {
+        isFunctionStreamingIn = true;
+        const toolCall = delta.tool_calls[0];
+        if (toolCall.index === 0) {
+          return `{"tool_calls":[ {"id": "${toolCall.id}", "type": "function", "function": {"name": "${toolCall.function?.name}", "arguments": "`;
+        } else {
+          return `"}}, {"id": "${toolCall.id}", "type": "function", "function": {"name": "${toolCall.function?.name}", "arguments": "`;
+        }
+      } else if (delta.function_call?.arguments) {
+        return cleanupArguments(delta.function_call?.arguments);
+      } else if (delta.tool_calls?.[0]?.function?.arguments) {
+        return cleanupArguments(delta.tool_calls?.[0]?.function?.arguments);
+      } else if (
+        isFunctionStreamingIn &&
+        (json.choices[0]?.finish_reason === "function_call" ||
+          json.choices[0]?.finish_reason === "stop")
+      ) {
+        isFunctionStreamingIn = false; // Reset the flag
+        return '"}}';
+      } else if (
+        isFunctionStreamingIn &&
+        json.choices[0]?.finish_reason === "tool_calls"
+      ) {
+        isFunctionStreamingIn = false; // Reset the flag
+        return '"}}]}';
+      }
+    }
+
+    const text = trimStartOfStream(
+      isChatCompletionChunk(json) && json.choices[0].delta.content
+        ? json.choices[0].delta.content
+        : isCompletion(json)
+          ? json.choices[0].text
+          : "",
+    );
+    return text;
+  };
+
+  function cleanupArguments(argumentChunk: string) {
+    let escapedPartialJson = argumentChunk
+      .replace(/\\/g, "\\\\") // Replace backslashes first to prevent double escaping
+      .replace(/\//g, "\\/") // Escape slashes
+      .replace(/"/g, '\\"') // Escape double quotes
+      .replace(/\n/g, "\\n") // Escape new lines
+      .replace(/\r/g, "\\r") // Escape carriage returns
+      .replace(/\t/g, "\\t") // Escape tabs
+      .replace(/\f/g, "\\f"); // Escape form feeds
+
+    return `${escapedPartialJson}`;
+  }
+}
+
+const __internal__OpenAIFnMessagesSymbol = Symbol(
+  "internal_openai_fn_messages",
+);
+
+type AzureChatCompletions = any;
+
+type AsyncIterableOpenAIStreamReturnTypes =
+  | AsyncIterable<ChatCompletionChunk>
+  | AsyncIterable<Completion>
+  | AsyncIterable<AzureChatCompletions>;
+
+type ExtractType<T> = T extends AsyncIterable<infer U> ? U : never;
+
+type OpenAIStreamReturnTypes =
+  ExtractType<AsyncIterableOpenAIStreamReturnTypes>;
+
+function isChatCompletionChunk(
+  data: OpenAIStreamReturnTypes,
+): data is ChatCompletionChunk {
+  return (
+    "choices" in data &&
+    data.choices &&
+    data.choices[0] &&
+    "delta" in data.choices[0]
+  );
+}
+
+function isCompletion(data: OpenAIStreamReturnTypes): data is Completion {
+  return (
+    "choices" in data &&
+    data.choices &&
+    data.choices[0] &&
+    "text" in data.choices[0]
+  );
+}