Add OpenAI traceable wrapper for JS

js/.gitignore

@@ -46,6 +46,9 @@ Chinook_Sqlite.sql
 /schemas.cjs
 /schemas.js
 /schemas.d.ts
+/wrappers.cjs
+/wrappers.js
+/wrappers.d.ts
 /index.cjs
 /index.js
 /index.d.ts

js/package.json

@@ -20,6 +20,9 @@
     "schemas.cjs",
     "schemas.js",
     "schemas.d.ts",
+    "wrappers.cjs",
+    "wrappers.js",
+    "wrappers.d.ts",
     "index.cjs",
     "index.js",
     "index.d.ts"
@@ -80,6 +83,7 @@
     "eslint-plugin-no-instanceof": "^1.0.1",
     "eslint-plugin-prettier": "^4.2.1",
     "jest": "^29.5.0",
+    "openai": "^4.28.0",
     "prettier": "^2.8.8",
     "ts-jest": "^29.1.0",
     "ts-node": "^10.9.1",
@@ -129,6 +133,11 @@
       "import": "./schemas.js",
       "require": "./schemas.cjs"
     },
+    "./wrappers": {
+      "types": "./wrappers.d.ts",
+      "import": "./wrappers.js",
+      "require": "./wrappers.cjs"
+    },
     "./package.json": "./package.json"
   }
-}
+}

js/scripts/create-entrypoints.js

@@ -12,6 +12,7 @@ const entrypoints = {
   traceable: "traceable",
   evaluation: "evaluation/index",
   schemas: "schemas",
+  wrappers: "wrappers",
 };
 const updateJsonFile = (relativePath, updateFunction) => {
   const contents = fs.readFileSync(relativePath).toString();

js/src/wrappers.ts

@@ -0,0 +1,18 @@
+import type { OpenAI } from "openai";
+import { traceable } from "./traceable.js";
+
+export const wrapOpenAI = (client: OpenAI): OpenAI => {
+  // @ts-expect-error Promise<APIPromise<...>> != APIPromise<...>
+  client.chat.completions.create = traceable(
+    client.chat.completions.create.bind(client.chat.completions),
+    { name: "ChatOpenAI", run_type: "llm" }
+  );
+
+  // @ts-expect-error Promise<APIPromise<...>> != APIPromise<...>
+  client.completions.create = traceable(
+    client.completions.create.bind(client.completions),
+    { name: "OpenAI", run_type: "llm" }
+  );
+
+  return client;
+};

js/tsconfig.json

@@ -36,7 +36,8 @@
       "src/run_trees.ts",
       "src/traceable.ts",
       "src/evaluation/index.ts",
-      "src/schemas.ts"
+      "src/schemas.ts",
+      "src/wrappers.ts"
     ]
   }
 }