langsmith
Version:
Client library to connect to the LangSmith LLM Tracing and Evaluation Platform.
302 lines (301 loc) • 11.9 kB
JavaScript
/* eslint-disable import/no-extraneous-dependencies */
import { LangSmithMiddleware, populateToolCallsForTracing, } from "./middleware.js";
import { traceable } from "../../traceable.js";
const _wrapTools = (tools) => {
// eslint-disable-next-line @typescript-eslint/no-explicit-any
const wrappedTools = {};
if (tools) {
for (const [key, tool] of Object.entries(tools)) {
wrappedTools[key] = tool;
if (tool != null &&
typeof tool === "object" &&
"execute" in tool &&
typeof tool.execute === "function") {
wrappedTools[key].execute = traceable(tool.execute.bind(tool), {
name: key,
run_type: "tool",
});
}
}
}
return wrappedTools;
};
const _getModelDisplayName = (model) => {
if (typeof model === "string") {
return model;
}
if (model.config != null &&
typeof model.config === "object" &&
"provider" in model.config &&
typeof model.config.provider === "string") {
return model.config.provider;
}
if (model.modelId != null && typeof model.modelId === "string") {
return model.modelId;
}
return "unknown";
};
const _getModelId = (model) => {
if (typeof model === "string") {
return model;
}
return typeof model.modelId === "string" ? model.modelId : undefined;
};
const _formatTracedInputs = (params) => {
const { prompt, messages, model, tools, ...rest } = params;
if (Array.isArray(prompt)) {
return { ...rest, messages: prompt.map(populateToolCallsForTracing) };
}
else if (Array.isArray(messages)) {
return { ...rest, messages: messages.map(populateToolCallsForTracing) };
}
else {
return { ...rest, prompt, messages };
}
};
/**
* Wraps Vercel AI SDK 5 functions with LangSmith tracing capabilities.
*
* @param methods - Object containing AI SDK methods to wrap
* @param methods.wrapLanguageModel - AI SDK's wrapLanguageModel function
* @param methods.generateText - AI SDK's generateText function
* @param methods.streamText - AI SDK's streamText function
* @param methods.streamObject - AI SDK's streamObject function
* @param methods.generateObject - AI SDK's generateObject function
*
* @returns Object containing wrapped versions of the AI SDK functions with LangSmith tracing
* @returns returns.generateText - Wrapped generateText function that traces calls to LangSmith
* @returns returns.generateObject - Wrapped generateObject function that traces calls to LangSmith
* @returns returns.streamText - Wrapped streamText function that traces calls to LangSmith
* @returns returns.streamObject - Wrapped streamObject function that traces calls to LangSmith
*/
const wrapAISDK = ({ wrapLanguageModel, generateText, streamText, streamObject, generateObject, }, lsConfig) => {
/**
* Wrapped version of AI SDK 5's generateText with LangSmith tracing.
*
* This function has the same signature and behavior as the original generateText,
* but adds automatic tracing to LangSmith for observability.
*
* ```ts
* import * as ai from "ai";
* import { wrapAISDK } from "langsmith/experimental/vercel";
*
* const { generateText } = wrapAISDK(ai);
* const { text } = await generateText(...);
* ```
*
* @see {@link https://sdk.vercel.ai/docs/ai-sdk-core/generating-text} Original generateText documentation
* @param params - Same parameters as the original generateText function
* @returns Promise resolving to the same result as generateText, with tracing applied
*/
const wrappedGenerateText = async (params) => {
const traceableFunc = traceable(async (params) => {
const wrappedModel = wrapLanguageModel({
model: params.model,
middleware: LangSmithMiddleware({
name: _getModelDisplayName(params.model),
modelId: params.model.modelId,
}),
});
return generateText({
...params,
tools: _wrapTools(params.tools),
model: wrappedModel,
});
}, {
name: _getModelDisplayName(params.model),
...lsConfig,
metadata: {
ai_sdk_method: "ai.generateText",
...lsConfig?.metadata,
},
processInputs: (inputs) => _formatTracedInputs(inputs),
processOutputs: (outputs) => {
if (outputs.outputs == null || typeof outputs.outputs !== "object") {
return outputs;
}
const { steps } = outputs.outputs;
if (Array.isArray(steps)) {
const lastStep = steps.at(-1);
if (lastStep == null || typeof lastStep !== "object") {
return outputs;
}
const { content } = lastStep;
return populateToolCallsForTracing({
content,
role: "assistant",
});
}
else {
return outputs;
}
},
});
return traceableFunc(params);
};
/**
* Wrapped version of AI SDK 5's generateObject with LangSmith tracing.
*
* This function has the same signature and behavior as the original generateObject,
* but adds automatic tracing to LangSmith for observability.
*
* ```ts
* import * as ai from "ai";
* import { wrapAISDK } from "langsmith/experimental/vercel";
*
* const { generateObject } = wrapAISDK(ai);
* const { object } = await generateObject(...);
* ```
*
* @see {@link https://sdk.vercel.ai/docs/ai-sdk-core/generating-structured-data} Original generateObject documentation
* @param params - Same parameters as the original generateObject function
* @returns Promise resolving to the same result as generateObject, with tracing applied
*/
const wrappedGenerateObject = async (params) => {
const traceableFunc = traceable(async (params) => {
const wrappedModel = wrapLanguageModel({
model: params.model,
middleware: LangSmithMiddleware({
name: _getModelDisplayName(params.model),
modelId: _getModelId(params.model),
}),
});
return generateObject({
...params,
model: wrappedModel,
});
}, {
name: _getModelDisplayName(params.model),
...lsConfig,
metadata: {
ai_sdk_method: "ai.generateObject",
...lsConfig?.metadata,
},
processInputs: (inputs) => _formatTracedInputs(inputs),
processOutputs: (outputs) => {
if (outputs.outputs == null || typeof outputs.outputs !== "object") {
return outputs;
}
return outputs.outputs.object ?? outputs;
},
});
return traceableFunc(params);
};
/**
* Wrapped version of AI SDK 5's streamText with LangSmith tracing.
*
* Must be called with `await`, but otherwise behaves the same as the
* original streamText and adds adds automatic tracing to LangSmith
* for observability.
*
* ```ts
* import * as ai from "ai";
* import { wrapAISDK } from "langsmith/experimental/vercel";
*
* const { streamText } = wrapAISDK(ai);
* const { textStream } = await streamText(...);
* ```
*
* @see {@link https://sdk.vercel.ai/docs/ai-sdk-core/generating-text} Original streamText documentation
* @param params - Same parameters as the original streamText function
* @returns Promise resolving to the same result as streamText, with tracing applied
*/
const wrappedStreamText = async (params) => {
const traceableFunc = traceable(async (params) => {
const wrappedModel = wrapLanguageModel({
model: params.model,
middleware: LangSmithMiddleware({
name: _getModelDisplayName(params.model),
modelId: _getModelId(params.model),
}),
});
return streamText({
...params,
tools: _wrapTools(params.tools),
model: wrappedModel,
});
}, {
name: _getModelDisplayName(params.model),
...lsConfig,
metadata: {
ai_sdk_method: "ai.streamText",
...lsConfig?.metadata,
},
processInputs: (inputs) => _formatTracedInputs(inputs),
processOutputs: async (outputs) => {
if (outputs.outputs == null || typeof outputs.outputs !== "object") {
return outputs;
}
const content = await outputs.outputs.content;
if (content == null || typeof content !== "object") {
return outputs;
}
return populateToolCallsForTracing({
content,
role: "assistant",
});
},
});
return traceableFunc(params);
};
/**
* Wrapped version of AI SDK 5's streamObject with LangSmith tracing.
*
* Must be called with `await`, but otherwise behaves the same as the
* original streamObject and adds adds automatic tracing to LangSmith
* for observability.
*
* ```ts
* import * as ai from "ai";
* import { wrapAISDK } from "langsmith/experimental/vercel";
*
* const { streamObject } = wrapAISDK(ai);
* const { partialObjectStream } = await streamObject(...);
* ```
*
* @see {@link https://sdk.vercel.ai/docs/ai-sdk-core/generating-structured-data} Original streamObject documentation
* @param params - Same parameters as the original streamObject function
* @returns Promise resolving to the same result as streamObject, with tracing applied
*/
const wrappedStreamObject = async (params) => {
const traceableFunc = traceable(async (params) => {
const wrappedModel = wrapLanguageModel({
model: params.model,
middleware: LangSmithMiddleware({
name: _getModelDisplayName(params.model),
modelId: _getModelId(params.model),
}),
});
return streamObject({
...params,
model: wrappedModel,
});
}, {
name: _getModelDisplayName(params.model),
...lsConfig,
metadata: {
ai_sdk_method: "ai.streamObject",
...lsConfig?.metadata,
},
processInputs: (inputs) => _formatTracedInputs(inputs),
processOutputs: async (outputs) => {
if (outputs.outputs == null || typeof outputs.outputs !== "object") {
return outputs;
}
const object = await outputs.outputs.object;
if (object == null || typeof object !== "object") {
return outputs;
}
return object;
},
});
return traceableFunc(params);
};
return {
generateText: wrappedGenerateText,
generateObject: wrappedGenerateObject,
streamText: wrappedStreamText,
streamObject: wrappedStreamObject,
};
};
export { wrapAISDK };