From 5ec346c94ea7a1c79995e5d02e8bf7fcacaf45c2 Mon Sep 17 00:00:00 2001 From: Tat Dat Duong Date: Thu, 17 Oct 2024 19:32:36 +0200 Subject: [PATCH] feat(vercel): add OTEL based LangSmith trace exporter --- js/package.json | 1 + js/src/wrappers/vercel.ts | 2 + js/src/wrappers/vercel/exporter.ts | 414 +++++++++++++++++++++++ js/src/wrappers/vercel/exporter.types.ts | 231 +++++++++++++ js/yarn.lock | 29 ++ 5 files changed, 677 insertions(+) create mode 100644 js/src/wrappers/vercel/exporter.ts create mode 100644 js/src/wrappers/vercel/exporter.types.ts diff --git a/js/package.json b/js/package.json index fde21f5b..4b67df72 100644 --- a/js/package.json +++ b/js/package.json @@ -112,6 +112,7 @@ "@langchain/core": "^0.3.1", "@langchain/langgraph": "^0.2.3", "@langchain/openai": "^0.3.0", + "@opentelemetry/sdk-trace-base": "^1.26.0", "@tsconfig/recommended": "^1.0.2", "@types/jest": "^29.5.1", "@typescript-eslint/eslint-plugin": "^5.59.8", diff --git a/js/src/wrappers/vercel.ts b/js/src/wrappers/vercel.ts index dc022d7c..e143647a 100644 --- a/js/src/wrappers/vercel.ts +++ b/js/src/wrappers/vercel.ts @@ -107,3 +107,5 @@ export const wrapAISDKModel = ( }, }); }; + +export { LangSmithAISDKExporter } from "./vercel/exporter.js"; diff --git a/js/src/wrappers/vercel/exporter.ts b/js/src/wrappers/vercel/exporter.ts new file mode 100644 index 00000000..9e414b24 --- /dev/null +++ b/js/src/wrappers/vercel/exporter.ts @@ -0,0 +1,414 @@ +import type { ReadableSpan, SpanExporter } from "@opentelemetry/sdk-trace-base"; +import type { ExportResult } from "@opentelemetry/core"; +import type { CoreAssistantMessage, CoreMessage } from "ai"; +import type { AISDKSpan } from "./exporter.types.js"; +import type { + BaseMessageFields, + MessageContentText, +} from "@langchain/core/messages"; +import { Client, ClientConfig, RunTree, RunTreeConfig } from "../../index.js"; +import { KVMap } from "../../schemas.js"; +import { AsyncLocalStorageProviderSingleton } from "../../singletons/traceable.js"; + +function assertNever(x: never): never { + throw new Error("Unreachable state: " + x); +} + +// TODO: remove dependency on @langchain/core +type $SmithMessage = { type: string; data: BaseMessageFields } | CoreMessage; + +function convertCoreToSmith( + message: CoreMessage +): $SmithMessage | $SmithMessage[] { + if (message.role === "assistant") { + const data: BaseMessageFields = { content: message.content }; + + if (Array.isArray(message.content)) { + data.content = message.content.map((part) => { + if (part.type === "text") { + return { + type: "text", + text: part.text, + } satisfies MessageContentText; + } + + if (part.type === "tool-call") { + return { + type: "tool_use", + name: part.toolName, + id: part.toolCallId, + input: part.args, + }; + } + + return part; + }); + + const toolCalls = message.content.filter( + (part) => part.type === "tool-call" + ); + + if (toolCalls.length > 0) { + data.additional_kwargs ??= {}; + data.additional_kwargs.tool_calls = toolCalls.map((part) => { + return { + id: part.toolCallId, + type: "function", + function: { + name: part.toolName, + id: part.toolCallId, + arguments: JSON.stringify(part.args), + }, + }; + }); + } + } + + return { type: "ai", data }; + } + + if (message.role === "user") { + // TODO: verify user content + return { type: "human", data: { content: message.content } }; + } + + if (message.role === "system") { + return { type: "system", data: { content: message.content } }; + } + + if (message.role === "tool") { + const res = message.content.map((toolCall) => { + return { + type: "tool", + data: { + content: JSON.stringify(toolCall.result), + name: toolCall.toolName, + tool_call_id: toolCall.toolCallId, + }, + }; + }); + if (res.length === 1) return res[0]; + return res; + } + + return message as any; +} + +const tryJson = ( + str: + | string + | number + | boolean + | Array + | Array + | Array + | undefined +) => { + try { + if (!str) return str; + if (typeof str !== "string") return str; + return JSON.parse(str); + } catch { + return str; + } +}; + +const sortByHrTime = (a: ReadableSpan, b: ReadableSpan) => { + return ( + Math.sign(a.startTime[0] - b.startTime[0]) || + Math.sign(a.startTime[1] - b.startTime[1]) + ); +}; + +export class LangSmithAISDKExporter implements SpanExporter { + private client: Client; + + constructor(config?: ClientConfig) { + this.client = new Client(config); + } + + export( + spans: ReadableSpan[], + resultCallback: (result: ExportResult) => void + ): void { + const runTreeMap: Record = {}; + const sortedSpans = [...spans].sort(sortByHrTime) as AISDKSpan[]; + + for (const span of sortedSpans) { + const spanId = span.spanContext().spanId; + const parentSpanId = span.parentSpanId; + let parentRunTree = parentSpanId ? runTreeMap[parentSpanId] : null; + + if (parentRunTree == null) { + try { + parentRunTree = + AsyncLocalStorageProviderSingleton.getInstance().getStore() ?? null; + } catch { + // pass + } + } + + const toRunTree = (rawConfig: RunTreeConfig) => { + const aiMetadata = Object.keys(span.attributes) + .filter((key) => key.startsWith("ai.telemetry.metadata.")) + .reduce((acc, key) => { + acc[key.slice("ai.telemetry.metadata.".length)] = + span.attributes[key as keyof typeof span.attributes]; + + return acc; + }, {} as Record); + + if ( + ("ai.telemetry.functionId" in span.attributes && + span.attributes["ai.telemetry.functionId"]) || + ("resource.name" in span.attributes && + span.attributes["resource.name"]) + ) { + aiMetadata["functionId"] = + span.attributes["ai.telemetry.functionId"] || + span.attributes["resource.name"]; + } + + const config: RunTreeConfig = { + ...rawConfig, + metadata: { + ...rawConfig.metadata, + ...aiMetadata, + "ai.operationId": span.attributes["ai.operationId"], + }, + start_time: +( + String(span.startTime[0]) + String(span.startTime[1]).slice(0, 3) + ), + end_time: +( + String(span.endTime[0]) + String(span.endTime[1]).slice(0, 3) + ), + client: this.client, + }; + return parentRunTree?.createChild(config) ?? new RunTree(config); + }; + + switch (span.name) { + case "ai.generateText.doGenerate": + case "ai.generateText": + case "ai.streamText.doStream": + case "ai.streamText": { + const inputs = ((): KVMap | undefined => { + if ("ai.prompt.messages" in span.attributes) { + return { + messages: tryJson( + span.attributes["ai.prompt.messages"] + ).flatMap((i: CoreMessage) => convertCoreToSmith(i)), + }; + } + + if ("ai.prompt" in span.attributes) { + const input = tryJson(span.attributes["ai.prompt"]); + + if ( + typeof input === "object" && + input != null && + "messages" in input && + Array.isArray(input.messages) + ) { + return { + messages: input.messages.flatMap((i: CoreMessage) => + convertCoreToSmith(i) + ), + }; + } + + return { input }; + } + + return undefined; + })(); + + const outputs = ((): KVMap | undefined => { + let result: KVMap | undefined = undefined; + if (span.attributes["ai.response.toolCalls"]) { + result = { + llm_output: convertCoreToSmith({ + role: "assistant", + content: tryJson(span.attributes["ai.response.toolCalls"]), + } satisfies CoreAssistantMessage), + }; + } else if (span.attributes["ai.response.text"]) { + result = { + llm_output: convertCoreToSmith({ + role: "assistant", + content: span.attributes["ai.response.text"], + }), + }; + } + + if (span.attributes["ai.usage.completionTokens"]) { + result ??= {}; + result.llm_output ??= {}; + result.llm_output.token_usage ??= {}; + result.llm_output.token_usage["completion_tokens"] = + span.attributes["ai.usage.completionTokens"]; + } + + if (span.attributes["ai.usage.promptTokens"]) { + result ??= {}; + result.llm_output ??= {}; + result.llm_output.token_usage ??= {}; + result.llm_output.token_usage["prompt_tokens"] = + span.attributes["ai.usage.promptTokens"]; + } + + return result; + })(); + + // TODO: add first_token_time + runTreeMap[spanId] = toRunTree({ + run_type: "llm", + name: span.attributes["ai.model.provider"], + inputs, + outputs, + metadata: { + ls_provider: span.attributes["ai.model.provider"] + .split(".") + .at(0), + ls_model_type: span.attributes["ai.model.provider"] + .split(".") + .at(1), + ls_model_name: span.attributes["ai.model.id"], + }, + extra: { batch_size: 1 }, + }); + break; + } + + case "ai.toolCall": { + const args = tryJson(span.attributes["ai.toolCall.args"]); + let inputs: KVMap | undefined = { args }; + + if (typeof args === "object" && args != null) { + inputs = args; + } + + const output = tryJson(span.attributes["ai.toolCall.result"]); + let outputs: KVMap | undefined = { output }; + + if (typeof output === "object" && output != null) { + outputs = output; + } + + runTreeMap[spanId] = toRunTree({ + run_type: "tool", + name: span.attributes["ai.toolCall.name"], + inputs, + outputs, + }); + break; + } + + case "ai.streamObject": + case "ai.streamObject.doStream": + case "ai.generateObject": + case "ai.generateObject.doGenerate": { + const inputs = ((): KVMap | undefined => { + if ("ai.prompt.messages" in span.attributes) { + return { + messages: tryJson( + span.attributes["ai.prompt.messages"] + ).flatMap((i: CoreMessage) => convertCoreToSmith(i)), + }; + } + + if ("ai.prompt" in span.attributes) { + return { input: span.attributes["ai.prompt"] }; + } + + return undefined; + })(); + + const outputs = ((): KVMap | undefined => { + let result: KVMap | undefined = undefined; + + if (span.attributes["ai.response.object"]) { + result = { + output: tryJson(span.attributes["ai.response.object"]), + }; + } + + if (span.attributes["ai.usage.completionTokens"]) { + result ??= {}; + result.llm_output ??= {}; + result.llm_output.token_usage ??= {}; + result.llm_output.token_usage["completion_tokens"] = + span.attributes["ai.usage.completionTokens"]; + } + + if (span.attributes["ai.usage.promptTokens"]) { + result ??= {}; + result.llm_output ??= {}; + result.llm_output.token_usage ??= {}; + result.llm_output.token_usage["prompt_tokens"] = + +span.attributes["ai.usage.promptTokens"]; + } + + return result; + })(); + + runTreeMap[spanId] = toRunTree({ + run_type: "llm", + name: span.attributes["ai.model.provider"], + inputs, + outputs, + metadata: { + ls_provider: span.attributes["ai.model.provider"] + .split(".") + .at(0), + ls_model_type: span.attributes["ai.model.provider"] + .split(".") + .at(1), + ls_model_name: span.attributes["ai.model.id"], + }, + extra: { batch_size: 1 }, + }); + break; + } + + case "ai.embed": { + runTreeMap[spanId] = toRunTree({ + run_type: "chain", + name: span.attributes["ai.model.provider"], + inputs: { value: span.attributes["ai.value"] }, + outputs: { embedding: span.attributes["ai.embedding"] }, + }); + break; + } + case "ai.embed.doEmbed": + case "ai.embedMany": + case "ai.embedMany.doEmbed": { + runTreeMap[spanId] = toRunTree({ + run_type: "chain", + name: span.attributes["ai.model.provider"], + inputs: { values: span.attributes["ai.values"] }, + outputs: { embeddings: span.attributes["ai.embeddings"] }, + }); + break; + } + + default: + assertNever(span); + } + } + + Promise.all( + Object.values(runTreeMap).map((runTree) => runTree.postRun()) + ).then( + () => resultCallback({ code: 0 }), + (error) => resultCallback({ code: 1, error }) + ); + } + + async shutdown(): Promise { + // pass + } + async forceFlush?(): Promise { + // pass + } +} diff --git a/js/src/wrappers/vercel/exporter.types.ts b/js/src/wrappers/vercel/exporter.types.ts new file mode 100644 index 00000000..8e3d2a32 --- /dev/null +++ b/js/src/wrappers/vercel/exporter.types.ts @@ -0,0 +1,231 @@ +import type { ReadableSpan } from "@opentelemetry/sdk-trace-base"; + +// eslint-disable-next-line @typescript-eslint/ban-types +type AnyString = string & {}; + +interface TypedReadableSpan + extends Omit { + name: Name; + attributes: Attributes; +} + +interface BaseLLMSpanAttributes { + "ai.model.id": string; + "ai.model.provider": string; + + "ai.usage.promptTokens": number; + "ai.usage.completionTokens": number; + + "ai.telemetry.functionId"?: string; + "resource.name"?: string; +} + +interface CallLLMSpanAttributes extends BaseLLMSpanAttributes { + "ai.response.model": string; + "ai.response.id": string; + "ai.response.timestamp": number; +} + +interface BaseEmbedSpanAttributes { + "ai.model.id": string; + "ai.model.provider": string; + "ai.usage.tokens": number; + + "ai.telemetry.functionId"?: string; + "resource.name"?: string; +} + +export type ToolCallSpan = TypedReadableSpan< + "ai.toolCall", + { + "operation.name": "ai.toolCall"; + "ai.operationId": "ai.toolCall"; + "ai.toolCall.name": string; + "ai.toolCall.id": string; + "ai.toolCall.args": string; + "ai.toolCall.result"?: string; + } +>; + +export type GenerateTextSpan = TypedReadableSpan< + "ai.generateText", + BaseLLMSpanAttributes & { + "operation.name": "ai.generateText"; + "ai.operationId": "ai.generateText"; + "ai.prompt": string; + "ai.response.text": string; + "ai.response.toolCalls": string; + "ai.response.finishReason": string; + "ai.settings.maxSteps": number; + } +>; + +export type DoGenerateTextSpan = TypedReadableSpan< + "ai.generateText.doGenerate", + CallLLMSpanAttributes & { + "operation.name": "ai.generateText.doGenerate"; + "ai.operationId": "ai.generateText.doGenerate"; + "ai.prompt.format": string; + "ai.prompt.messages": string; + "ai.response.text": string; + "ai.response.toolCalls": string; + "ai.response.finishReason": string; + } +>; + +export type StreamTextSpan = TypedReadableSpan< + "ai.streamText", + BaseLLMSpanAttributes & { + "operation.name": "ai.streamText"; + "ai.operationId": "ai.streamText"; + "ai.prompt": string; + "ai.response.text": string; + "ai.response.toolCalls": string; + "ai.response.finishReason": string; + "ai.settings.maxSteps": number; + } +>; + +export type DoStreamTextSpan = TypedReadableSpan< + "ai.streamText.doStream", + CallLLMSpanAttributes & { + "operation.name": "ai.streamText.doStream"; + "ai.operationId": "ai.streamText.doStream"; + "ai.prompt.format": string; + "ai.prompt.messages": string; + "ai.response.text": string; + "ai.response.toolCalls": string; + "ai.response.msToFirstChunk": number; + "ai.response.msToFinish": number; + "ai.response.avgCompletionTokensPerSecond": number; + "ai.response.finishReason": string; + } +>; + +export type GenerateObjectSpan = TypedReadableSpan< + "ai.generateObject", + BaseLLMSpanAttributes & { + "operation.name": "ai.generateObject"; + "ai.operationId": "ai.generateObject"; + "ai.prompt": string; + + "ai.schema": string; + "ai.schema.name": string; + "ai.schema.description": string; + + "ai.response.object": string; + + "ai.settings.mode": "json" | AnyString; + "ai.settings.output": "object" | "no-schema" | AnyString; + } +>; +export type DoGenerateObjectSpan = TypedReadableSpan< + "ai.generateObject.doGenerate", + CallLLMSpanAttributes & { + "operation.name": "ai.generateObject.doGenerate"; + "ai.operationId": "ai.generateObject.doGenerate"; + + "ai.prompt.format": string; + "ai.prompt.messages": string; + + "ai.response.object": string; + "ai.response.finishReason": string; + + "ai.settings.mode": "json" | AnyString; + "ai.settings.output": "object" | "no-schema" | AnyString; + } +>; + +export type StreamObjectSpan = TypedReadableSpan< + "ai.streamObject", + BaseLLMSpanAttributes & { + "operation.name": "ai.streamObject"; + "ai.operationId": "ai.streamObject"; + "ai.prompt": string; + + "ai.schema": string; + "ai.schema.name": string; + "ai.schema.description": string; + + "ai.response.object": string; + + "ai.settings.mode": "json" | AnyString; + "ai.settings.output": "object" | "no-schema" | AnyString; + } +>; +export type DoStreamObjectSpan = TypedReadableSpan< + "ai.streamObject.doStream", + CallLLMSpanAttributes & { + "operation.name": "ai.streamObject.doStream"; + "ai.operationId": "ai.streamObject.doStream"; + + "ai.prompt.format": string; + "ai.prompt.messages": string; + + "ai.response.object": string; + "ai.response.finishReason": string; + "ai.response.msToFirstChunk": number; + + "ai.settings.mode": "json" | AnyString; + } +>; + +export type EmbedSpan = TypedReadableSpan< + "ai.embed", + BaseEmbedSpanAttributes & { + "operation.name": "ai.embed"; + "ai.operationId": "ai.embed"; + + // TODO: is this correct? + "ai.value": string; + "ai.embedding": string; + } +>; + +export type DoEmbedSpan = TypedReadableSpan< + "ai.embed.doEmbed", + BaseEmbedSpanAttributes & { + "operation.name": "ai.embed.doEmbed"; + "ai.operationId": "ai.embed.doEmbed"; + + "ai.values": string[]; + "ai.embeddings": string[]; + } +>; + +export type EmbedManySpan = TypedReadableSpan< + "ai.embedMany", + BaseEmbedSpanAttributes & { + "operation.name": "ai.embedMany"; + "ai.operationId": "ai.embedMany"; + + "ai.values": string[]; + "ai.embeddings": string[]; + } +>; + +export type DoEmbedManySpan = TypedReadableSpan< + "ai.embedMany.doEmbed", + BaseEmbedSpanAttributes & { + "operation.name": "ai.embedMany.doEmbed"; + "ai.operationId": "ai.embedMany.doEmbed"; + + "ai.values": string[]; + "ai.embeddings": string[]; + } +>; + +export type AISDKSpan = + | ToolCallSpan + | GenerateTextSpan + | DoGenerateTextSpan + | StreamTextSpan + | DoStreamTextSpan + | GenerateObjectSpan + | DoGenerateObjectSpan + | StreamObjectSpan + | DoStreamObjectSpan + | EmbedSpan + | DoEmbedSpan + | EmbedManySpan + | DoEmbedManySpan; diff --git a/js/yarn.lock b/js/yarn.lock index 2a3feae3..531ebf49 100644 --- a/js/yarn.lock +++ b/js/yarn.lock @@ -1445,6 +1445,35 @@ resolved "https://registry.yarnpkg.com/@opentelemetry/api/-/api-1.9.0.tgz#d03eba68273dc0f7509e2a3d5cba21eae10379fe" integrity sha512-3giAOQvZiH5F9bMlMiv8+GSPMeqg0dbaeo58/0SlA9sxSqZhnUtxzX9/2FzyhS9sWQf5S0GJE0AKBrFqjpeYcg== +"@opentelemetry/core@1.26.0": + version "1.26.0" + resolved "https://registry.yarnpkg.com/@opentelemetry/core/-/core-1.26.0.tgz#7d84265aaa850ed0ca5813f97d831155be42b328" + integrity sha512-1iKxXXE8415Cdv0yjG3G6hQnB5eVEsJce3QaawX8SjDn0mAS0ZM8fAbZZJD4ajvhC15cePvosSCut404KrIIvQ== + dependencies: + "@opentelemetry/semantic-conventions" "1.27.0" + +"@opentelemetry/resources@1.26.0": + version "1.26.0" + resolved "https://registry.yarnpkg.com/@opentelemetry/resources/-/resources-1.26.0.tgz#da4c7366018bd8add1f3aa9c91c6ac59fd503cef" + integrity sha512-CPNYchBE7MBecCSVy0HKpUISEeJOniWqcHaAHpmasZ3j9o6V3AyBzhRc90jdmemq0HOxDr6ylhUbDhBqqPpeNw== + dependencies: + "@opentelemetry/core" "1.26.0" + "@opentelemetry/semantic-conventions" "1.27.0" + +"@opentelemetry/sdk-trace-base@^1.26.0": + version "1.26.0" + resolved "https://registry.yarnpkg.com/@opentelemetry/sdk-trace-base/-/sdk-trace-base-1.26.0.tgz#0c913bc6d2cfafd901de330e4540952269ae579c" + integrity sha512-olWQldtvbK4v22ymrKLbIcBi9L2SpMO84sCPY54IVsJhP9fRsxJT194C/AVaAuJzLE30EdhhM1VmvVYR7az+cw== + dependencies: + "@opentelemetry/core" "1.26.0" + "@opentelemetry/resources" "1.26.0" + "@opentelemetry/semantic-conventions" "1.27.0" + +"@opentelemetry/semantic-conventions@1.27.0": + version "1.27.0" + resolved "https://registry.yarnpkg.com/@opentelemetry/semantic-conventions/-/semantic-conventions-1.27.0.tgz#1a857dcc95a5ab30122e04417148211e6f945e6c" + integrity sha512-sAay1RrB+ONOem0OZanAR1ZI/k7yDpnOQSQmTMuGImUQb2y8EbSaCJ94FQluM74xoU03vlb2d2U90hZluL6nQg== + "@sinclair/typebox@^0.25.16": version "0.25.24" resolved "https://registry.npmjs.org/@sinclair/typebox/-/typebox-0.25.24.tgz"