From 5d9373150da77881a34d683daa3556c84af3dfd0 Mon Sep 17 00:00:00 2001 From: Jacob Lee Date: Tue, 5 Mar 2024 16:58:17 -0800 Subject: [PATCH] Add generic JS wrapClient method (#485) CC @dqbd @hinthornw --- js/README.md | 32 +++++++++++ js/package.json | 2 +- js/src/tests/wrapped_sdk.int.test.ts | 65 ++++++++++++++++++++++ js/src/wrappers.ts | 82 ++++++++++++++++++++++++++-- 4 files changed, 174 insertions(+), 7 deletions(-) create mode 100644 js/src/tests/wrapped_sdk.int.test.ts diff --git a/js/README.md b/js/README.md index 7c9ad7732..158a0157a 100644 --- a/js/README.md +++ b/js/README.md @@ -292,6 +292,38 @@ export async function POST(req: Request) { See the [AI SDK docs](https://sdk.vercel.ai/docs) for more examples. +## Arbitrary SDKs + +You can use the generic `wrapSDK` method to add tracing for arbitrary SDKs. + +Do note that this will trace ALL methods in the SDK, not just chat completion endpoints. +If the SDK you are wrapping has other methods, we recommend using it for only LLM calls. + +Here's an example using the Anthropic SDK: + +```ts +import { wrapSDK } from "langsmith/wrappers"; +import { Anthropic } from "@anthropic-ai/sdk"; + +const originalSDK = new Anthropic(); +const sdkWithTracing = wrapSDK(originalSDK); + +const response = await sdkWithTracing.messages.create({ + messages: [ + { + role: "user", + content: `What is 1 + 1? Respond only with "2" and nothing else.`, + }, + ], + model: "claude-3-sonnet-20240229", + max_tokens: 1024, +}); +``` + +:::tip +[Click here](https://smith.langchain.com/public/0e7248af-bbed-47cf-be9f-5967fea1dec1/r) to see an example LangSmith trace of the above. +::: + #### Alternatives: **Log traces using a RunTree.** A RunTree tracks your application. Each RunTree object is required to have a name and run_type. These and other important attributes are as follows: diff --git a/js/package.json b/js/package.json index 36189d255..51156259a 100644 --- a/js/package.json +++ b/js/package.json @@ -138,4 +138,4 @@ }, "./package.json": "./package.json" } -} \ No newline at end of file +} diff --git a/js/src/tests/wrapped_sdk.int.test.ts b/js/src/tests/wrapped_sdk.int.test.ts new file mode 100644 index 000000000..70d086f5a --- /dev/null +++ b/js/src/tests/wrapped_sdk.int.test.ts @@ -0,0 +1,65 @@ +import { jest } from "@jest/globals"; +import { OpenAI } from "openai"; +import { wrapSDK } from "../wrappers.js"; +import { Client } from "../client.js"; + +test.concurrent("chat.completions", async () => { + const client = new Client(); + const callSpy = jest + // eslint-disable-next-line @typescript-eslint/no-explicit-any + .spyOn((client as any).caller, "call") + .mockResolvedValue({ ok: true, text: () => "" }); + + const originalClient = new OpenAI(); + const patchedClient = wrapSDK(new OpenAI(), { client }); + + // invoke + const original = await originalClient.chat.completions.create({ + messages: [{ role: "user", content: `Say 'foo'` }], + temperature: 0, + seed: 42, + model: "gpt-3.5-turbo", + }); + + const patched = await patchedClient.chat.completions.create({ + messages: [{ role: "user", content: `Say 'foo'` }], + temperature: 0, + seed: 42, + model: "gpt-3.5-turbo", + }); + + expect(patched.choices).toEqual(original.choices); + + // stream + const originalStream = await originalClient.chat.completions.create({ + messages: [{ role: "user", content: `Say 'foo'` }], + temperature: 0, + seed: 42, + model: "gpt-3.5-turbo", + stream: true, + }); + + const originalChoices = []; + for await (const chunk of originalStream) { + originalChoices.push(chunk.choices); + } + + const patchedStream = await patchedClient.chat.completions.create({ + messages: [{ role: "user", content: `Say 'foo'` }], + temperature: 0, + seed: 42, + model: "gpt-3.5-turbo", + stream: true, + }); + + const patchedChoices = []; + for await (const chunk of patchedStream) { + patchedChoices.push(chunk.choices); + } + + expect(patchedChoices).toEqual(originalChoices); + for (const call of callSpy.mock.calls) { + // eslint-disable-next-line @typescript-eslint/no-explicit-any + expect((call[2] as any)["method"]).toBe("POST"); + } +}); diff --git a/js/src/wrappers.ts b/js/src/wrappers.ts index 9a25c8d6b..3011fa631 100644 --- a/js/src/wrappers.ts +++ b/js/src/wrappers.ts @@ -1,18 +1,33 @@ -import type { OpenAI } from "openai"; import type { Client } from "./index.js"; import { traceable } from "./traceable.js"; -export const wrapOpenAI = ( - openai: OpenAI, +type OpenAIType = { + chat: { + completions: { + create: (...args: any[]) => any; + }; + }; + completions: { + create: (...args: any[]) => any; + }; +}; + +/** + * Wraps an OpenAI client's completion methods, enabling automatic LangSmith + * tracing. Method signatures are unchanged. + * @param openai An OpenAI client instance. + * @param options LangSmith options. + * @returns + */ +export const wrapOpenAI = ( + openai: T, options?: { client?: Client } -): OpenAI => { - // @ts-expect-error Promise> != APIPromise<...> +): T => { openai.chat.completions.create = traceable( openai.chat.completions.create.bind(openai.chat.completions), Object.assign({ name: "ChatOpenAI", run_type: "llm" }, options?.client) ); - // @ts-expect-error Promise> != APIPromise<...> openai.completions.create = traceable( openai.completions.create.bind(openai.completions), Object.assign({ name: "OpenAI", run_type: "llm" }, options?.client) @@ -20,3 +35,58 @@ export const wrapOpenAI = ( return openai; }; + +const _wrapClient = ( + sdk: T, + runName: string, + options?: { client?: Client } +): T => { + return new Proxy(sdk, { + get(target, propKey, receiver) { + const originalValue = target[propKey as keyof T]; + if (typeof originalValue === "function") { + return traceable( + originalValue.bind(target), + Object.assign( + { name: [runName, propKey.toString()].join("."), run_type: "llm" }, + options?.client + ) + ); + } else if ( + originalValue != null && + !Array.isArray(originalValue) && + // eslint-disable-next-line no-instanceof/no-instanceof + !(originalValue instanceof Date) && + typeof originalValue === "object" + ) { + return _wrapClient( + originalValue, + [runName, propKey.toString()].join("."), + options + ); + } else { + return Reflect.get(target, propKey, receiver); + } + }, + }); +}; + +/** + * Wrap an arbitrary SDK, enabling automatic LangSmith tracing. + * Method signatures are unchanged. + * + * Note that this will wrap and trace ALL SDK methods, not just + * LLM completion methods. If the passed SDK contains other methods, + * we recommend using the wrapped instance for LLM calls only. + * @param sdk An arbitrary SDK instance. + * @param options LangSmith options. + * @returns + */ +export const wrapSDK = ( + sdk: T, + options?: { client?: Client; runName?: string } +): T => { + return _wrapClient(sdk, options?.runName ?? sdk.constructor?.name, { + client: options?.client, + }); +};