-
Notifications
You must be signed in to change notification settings - Fork 85
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
Add
wrapAISDKModel method for Vercel's AI SDK (#896)
Improve support for tracing Vercel AI SDK calls with a new
`wrapAISDKModel` method
Because AI SDK calls nest the text/object stream within the response,
this PR adds an optional parameter to `traceable` that specifies a key
in the response object to wrap. It also adds support for specially
tapping `ReadableStream` so that `traceable` does not change them to
plain async iterators.
Here's an example:
```ts
import { anthropic } from "@ai-sdk/anthropic";
import { streamObject } from "ai";
import { wrapAISDKModel } from "langsmith/wrappers/vercel";
const modelWithTracing = wrapAISDKModel(anthropic("claude-3-haiku-20240307"));
const { partialObjectStream } = await streamObject({
model: modelWithTracing,
prompt: "Write a vegetarian lasagna recipe for 4 people.",
schema: z.object({
ingredients: z.array(z.string()),
}),
});
for await (const chunk of partialObjectStream) {
console.log(chunk);
}
```
TODO: improve support in the UI for properly rendering token counts and
aggregating chunks. The data is returned but is not in one of our
currently supported formats.
Example traces:
Generate text:
https://smith.langchain.com/public/fbd12847-9485-43cf-a0a3-82c0b3318594/r
Generate object:
https://smith.langchain.com/public/c25944c7-2428-44d0-991b-ef3b5e47cab5/r
Stream text:
https://smith.langchain.com/public/2fb34c85-fec5-4361-a487-ffcf67718750/r
Stream object:
https://smith.langchain.com/public/27f5c9ba-006d-4a2f-b569-2e8da09dae23/r
CC @dqbd- Loading branch information
Showing
11 changed files
with
568 additions
and
73 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,77 @@ | ||
| import { openai } from "@ai-sdk/openai"; | ||
| import { | ||
| generateObject, | ||
| generateText, | ||
| streamObject, | ||
| streamText, | ||
| tool, | ||
| } from "ai"; | ||
| import { z } from "zod"; | ||
| import { wrapAISDKModel } from "../wrappers/vercel.js"; | ||
|
|
||
| test("AI SDK generateText", async () => { | ||
| const modelWithTracing = wrapAISDKModel(openai("gpt-4o-mini")); | ||
| const { text } = await generateText({ | ||
| model: modelWithTracing, | ||
| prompt: "Write a vegetarian lasagna recipe for 4 people.", | ||
| }); | ||
| console.log(text); | ||
| }); | ||
|
|
||
| test("AI SDK generateText with a tool", async () => { | ||
| const modelWithTracing = wrapAISDKModel(openai("gpt-4o-mini")); | ||
| const { text } = await generateText({ | ||
| model: modelWithTracing, | ||
| prompt: | ||
| "Write a vegetarian lasagna recipe for 4 people. Get ingredients first.", | ||
| tools: { | ||
| getIngredients: tool({ | ||
| description: "get a list of ingredients", | ||
| parameters: z.object({ | ||
| ingredients: z.array(z.string()), | ||
| }), | ||
| execute: async () => | ||
| JSON.stringify(["pasta", "tomato", "cheese", "onions"]), | ||
| }), | ||
| }, | ||
| maxToolRoundtrips: 2, | ||
| }); | ||
| console.log(text); | ||
| }); | ||
|
|
||
| test("AI SDK generateObject", async () => { | ||
| const modelWithTracing = wrapAISDKModel(openai("gpt-4o-mini")); | ||
| const { object } = await generateObject({ | ||
| model: modelWithTracing, | ||
| prompt: "Write a vegetarian lasagna recipe for 4 people.", | ||
| schema: z.object({ | ||
| ingredients: z.array(z.string()), | ||
| }), | ||
| }); | ||
| console.log(object); | ||
| }); | ||
|
|
||
| test("AI SDK streamText", async () => { | ||
| const modelWithTracing = wrapAISDKModel(openai("gpt-4o-mini")); | ||
| const { textStream } = await streamText({ | ||
| model: modelWithTracing, | ||
| prompt: "Write a vegetarian lasagna recipe for 4 people.", | ||
| }); | ||
| for await (const chunk of textStream) { | ||
| console.log(chunk); | ||
| } | ||
| }); | ||
|
|
||
| test("AI SDK streamObject", async () => { | ||
| const modelWithTracing = wrapAISDKModel(openai("gpt-4o-mini")); | ||
| const { partialObjectStream } = await streamObject({ | ||
| model: modelWithTracing, | ||
| prompt: "Write a vegetarian lasagna recipe for 4 people.", | ||
| schema: z.object({ | ||
| ingredients: z.array(z.string()), | ||
| }), | ||
| }); | ||
| for await (const chunk of partialObjectStream) { | ||
| console.log(chunk); | ||
| } | ||
| }); |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,72 @@ | ||
| import type { RunTreeConfig } from "../index.js"; | ||
| import { traceable } from "../traceable.js"; | ||
|
|
||
| export const _wrapClient = <T extends object>( | ||
| sdk: T, | ||
| runName: string, | ||
| options?: Omit<RunTreeConfig, "name"> | ||
| ): 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), { | ||
| run_type: "llm", | ||
| ...options, | ||
| name: [runName, propKey.toString()].join("."), | ||
| }); | ||
| } 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); | ||
| } | ||
| }, | ||
| }); | ||
| }; | ||
|
|
||
| type WrapSDKOptions = Partial< | ||
| RunTreeConfig & { | ||
| /** | ||
| * @deprecated Use `name` instead. | ||
| */ | ||
| runName: string; | ||
| } | ||
| >; | ||
|
|
||
| /** | ||
| * 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 = <T extends object>( | ||
| sdk: T, | ||
| options?: WrapSDKOptions | ||
| ): T => { | ||
| const traceableOptions = options ? { ...options } : undefined; | ||
| if (traceableOptions != null) { | ||
| delete traceableOptions.runName; | ||
| delete traceableOptions.name; | ||
| } | ||
|
|
||
| return _wrapClient( | ||
| sdk, | ||
| options?.name ?? options?.runName ?? sdk.constructor?.name, | ||
| traceableOptions | ||
| ); | ||
| }; |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -1 +1,2 @@ | ||
| export * from "./openai.js"; | ||
| export { wrapSDK } from "./generic.js"; |
Oops, something went wrong.