diff --git a/.github/workflows/llmobs.yml b/.github/workflows/llmobs.yml index a1e3502a8a..66ce36c238 100644 --- a/.github/workflows/llmobs.yml +++ b/.github/workflows/llmobs.yml @@ -47,3 +47,22 @@ jobs: - uses: codecov/codecov-action@v3 - if: always() uses: ./.github/actions/testagent/logs + + langchain: + runs-on: ubuntu-latest + env: + PLUGINS: langchain + steps: + - uses: actions/checkout@v4 + - uses: ./.github/actions/testagent/start + - uses: ./.github/actions/node/setup + - uses: ./.github/actions/install + - uses: ./.github/actions/node/18 + - run: yarn test:llmobs:plugins:ci + shell: bash + - uses: ./.github/actions/node/latest + - run: yarn test:llmobs:plugins:ci + shell: bash + - uses: codecov/codecov-action@v3 + - if: always() + uses: ./.github/actions/testagent/logs diff --git a/package.json b/package.json index 26fe1a5fab..f3253ebfa2 100644 --- a/package.json +++ b/package.json @@ -34,7 +34,7 @@ "test:lambda:ci": "nyc --no-clean --include \"packages/dd-trace/src/lambda/**/*.js\" -- npm run test:lambda", "test:llmobs:sdk": "mocha -r \"packages/dd-trace/test/setup/mocha.js\" --exclude \"packages/dd-trace/test/llmobs/plugins/**/*.spec.js\" \"packages/dd-trace/test/llmobs/**/*.spec.js\" ", "test:llmobs:sdk:ci": "nyc --no-clean --include \"packages/dd-trace/src/llmobs/**/*.js\" -- npm run test:llmobs:sdk", - "test:llmobs:plugins": "mocha -r \"packages/dd-trace/test/setup/mocha.js\" \"packages/dd-trace/test/llmobs/plugins/**/*.spec.js\"", + "test:llmobs:plugins": "mocha -r \"packages/dd-trace/test/setup/mocha.js\" \"packages/dd-trace/test/llmobs/plugins/@($(echo $PLUGINS))/*.spec.js\"", "test:llmobs:plugins:ci": "yarn services && nyc --no-clean --include \"packages/dd-trace/src/llmobs/**/*.js\" -- npm run test:llmobs:plugins", "test:plugins": "mocha -r \"packages/dd-trace/test/setup/mocha.js\" \"packages/datadog-instrumentations/test/@($(echo $PLUGINS)).spec.js\" \"packages/datadog-plugin-@($(echo $PLUGINS))/test/**/*.spec.js\"", "test:plugins:ci": "yarn services && nyc --no-clean --include \"packages/datadog-instrumentations/src/@($(echo $PLUGINS)).js\" --include \"packages/datadog-instrumentations/src/@($(echo $PLUGINS))/**/*.js\" --include \"packages/datadog-plugin-@($(echo $PLUGINS))/src/**/*.js\" -- npm run test:plugins", diff --git a/packages/datadog-instrumentations/src/openai.js b/packages/datadog-instrumentations/src/openai.js index 3528b1ecc1..0e921fb2b4 100644 --- a/packages/datadog-instrumentations/src/openai.js +++ b/packages/datadog-instrumentations/src/openai.js @@ -338,6 +338,8 @@ for (const shim of V4_PACKAGE_SHIMS) { }) }) + ch.end.publish(ctx) + return apiProm }) }) diff --git a/packages/datadog-plugin-langchain/src/index.js b/packages/datadog-plugin-langchain/src/index.js index 19b6e7d979..161b2791bf 100644 --- a/packages/datadog-plugin-langchain/src/index.js +++ b/packages/datadog-plugin-langchain/src/index.js @@ -1,89 +1,17 @@ 'use strict' -const { MEASURED } = require('../../../ext/tags') -const { storage } = require('../../datadog-core') -const TracingPlugin = require('../../dd-trace/src/plugins/tracing') +const LangChainTracingPlugin = require('./tracing') +const LangChainLLMObsPlugin = require('../../dd-trace/src/llmobs/plugins/langchain') +const CompositePlugin = require('../../dd-trace/src/plugins/composite') -const API_KEY = 'langchain.request.api_key' -const MODEL = 'langchain.request.model' -const PROVIDER = 'langchain.request.provider' -const TYPE = 'langchain.request.type' - -const LangChainHandler = require('./handlers/default') -const LangChainChatModelHandler = require('./handlers/language_models/chat_model') -const LangChainLLMHandler = require('./handlers/language_models/llm') -const LangChainChainHandler = require('./handlers/chain') -const LangChainEmbeddingHandler = require('./handlers/embedding') - -class LangChainPlugin extends TracingPlugin { +class LangChainPlugin extends CompositePlugin { static get id () { return 'langchain' } - static get operation () { return 'invoke' } - static get system () { return 'langchain' } - static get prefix () { - return 'tracing:apm:langchain:invoke' - } - - constructor () { - super(...arguments) - - const langchainConfig = this._tracerConfig.langchain || {} - this.handlers = { - chain: new LangChainChainHandler(langchainConfig), - chat_model: new LangChainChatModelHandler(langchainConfig), - llm: new LangChainLLMHandler(langchainConfig), - embedding: new LangChainEmbeddingHandler(langchainConfig), - default: new LangChainHandler(langchainConfig) + static get plugins () { + return { + llmobs: LangChainLLMObsPlugin, + tracing: LangChainTracingPlugin } } - - bindStart (ctx) { - const { resource, type } = ctx - const handler = this.handlers[type] - - const instance = ctx.instance - const apiKey = handler.extractApiKey(instance) - const provider = handler.extractProvider(instance) - const model = handler.extractModel(instance) - - const tags = handler.getSpanStartTags(ctx, provider) || [] - - if (apiKey) tags[API_KEY] = apiKey - if (provider) tags[PROVIDER] = provider - if (model) tags[MODEL] = model - if (type) tags[TYPE] = type - - const span = this.startSpan('langchain.request', { - service: this.config.service, - resource, - kind: 'client', - meta: { - [MEASURED]: 1, - ...tags - } - }, false) - - const store = storage.getStore() || {} - ctx.currentStore = { ...store, span } - - return ctx.currentStore - } - - asyncEnd (ctx) { - const span = ctx.currentStore.span - - const { type } = ctx - - const handler = this.handlers[type] - const tags = handler.getSpanEndTags(ctx) || {} - - span.addTags(tags) - - span.finish() - } - - getHandler (type) { - return this.handlers[type] || this.handlers.default - } } module.exports = LangChainPlugin diff --git a/packages/datadog-plugin-langchain/src/tracing.js b/packages/datadog-plugin-langchain/src/tracing.js new file mode 100644 index 0000000000..babdf88691 --- /dev/null +++ b/packages/datadog-plugin-langchain/src/tracing.js @@ -0,0 +1,89 @@ +'use strict' + +const { MEASURED } = require('../../../ext/tags') +const { storage } = require('../../datadog-core') +const TracingPlugin = require('../../dd-trace/src/plugins/tracing') + +const API_KEY = 'langchain.request.api_key' +const MODEL = 'langchain.request.model' +const PROVIDER = 'langchain.request.provider' +const TYPE = 'langchain.request.type' + +const LangChainHandler = require('./handlers/default') +const LangChainChatModelHandler = require('./handlers/language_models/chat_model') +const LangChainLLMHandler = require('./handlers/language_models/llm') +const LangChainChainHandler = require('./handlers/chain') +const LangChainEmbeddingHandler = require('./handlers/embedding') + +class LangChainTracingPlugin extends TracingPlugin { + static get id () { return 'langchain' } + static get operation () { return 'invoke' } + static get system () { return 'langchain' } + static get prefix () { + return 'tracing:apm:langchain:invoke' + } + + constructor () { + super(...arguments) + + const langchainConfig = this._tracerConfig.langchain || {} + this.handlers = { + chain: new LangChainChainHandler(langchainConfig), + chat_model: new LangChainChatModelHandler(langchainConfig), + llm: new LangChainLLMHandler(langchainConfig), + embedding: new LangChainEmbeddingHandler(langchainConfig), + default: new LangChainHandler(langchainConfig) + } + } + + bindStart (ctx) { + const { resource, type } = ctx + const handler = this.handlers[type] + + const instance = ctx.instance + const apiKey = handler.extractApiKey(instance) + const provider = handler.extractProvider(instance) + const model = handler.extractModel(instance) + + const tags = handler.getSpanStartTags(ctx, provider) || [] + + if (apiKey) tags[API_KEY] = apiKey + if (provider) tags[PROVIDER] = provider + if (model) tags[MODEL] = model + if (type) tags[TYPE] = type + + const span = this.startSpan('langchain.request', { + service: this.config.service, + resource, + kind: 'client', + meta: { + [MEASURED]: 1, + ...tags + } + }, false) + + const store = storage.getStore() || {} + ctx.currentStore = { ...store, span } + + return ctx.currentStore + } + + asyncEnd (ctx) { + const span = ctx.currentStore.span + + const { type } = ctx + + const handler = this.handlers[type] + const tags = handler.getSpanEndTags(ctx) || {} + + span.addTags(tags) + + span.finish() + } + + getHandler (type) { + return this.handlers[type] || this.handlers.default + } +} + +module.exports = LangChainTracingPlugin diff --git a/packages/dd-trace/src/llmobs/plugins/base.js b/packages/dd-trace/src/llmobs/plugins/base.js index f7f4d2b5e9..97f27d3899 100644 --- a/packages/dd-trace/src/llmobs/plugins/base.js +++ b/packages/dd-trace/src/llmobs/plugins/base.js @@ -6,7 +6,6 @@ const { storage } = require('../storage') const TracingPlugin = require('../../plugins/tracing') const LLMObsTagger = require('../tagger') -// we make this a `Plugin` so we don't have to worry about `finish` being called class LLMObsPlugin extends TracingPlugin { constructor (...args) { super(...args) @@ -14,24 +13,46 @@ class LLMObsPlugin extends TracingPlugin { this._tagger = new LLMObsTagger(this._tracerConfig, true) } - getName () {} - setLLMObsTags (ctx) { throw new Error('setLLMObsTags must be implemented by the subclass') } - getLLMObsSPanRegisterOptions (ctx) { + getLLMObsSpanRegisterOptions (ctx) { throw new Error('getLLMObsSPanRegisterOptions must be implemented by the subclass') } start (ctx) { - const oldStore = storage.getStore() - const parent = oldStore?.span + // even though llmobs span events won't be enqueued if llmobs is disabled + // we should avoid doing any computations here (these listeners aren't disabled) + const enabled = this._tracerConfig.llmobs.enabled + if (!enabled) return + + const parent = this.getLLMObsParent(ctx) const span = ctx.currentStore?.span - const registerOptions = this.getLLMObsSPanRegisterOptions(ctx) + const registerOptions = this.getLLMObsSpanRegisterOptions(ctx) + + // register options may not be set for operations we do not trace with llmobs + // ie OpenAI fine tuning jobs, file jobs, etc. + if (registerOptions) { + ctx.llmobs = {} // initialize context-based namespace + storage.enterWith({ span }) + ctx.llmobs.parent = parent - this._tagger.registerLLMObsSpan(span, { parent, ...registerOptions }) + this._tagger.registerLLMObsSpan(span, { parent, ...registerOptions }) + } + } + + end (ctx) { + const enabled = this._tracerConfig.llmobs.enabled + if (!enabled) return + + // only attempt to restore the context if the current span was an LLMObs span + const span = ctx.currentStore?.span + if (!LLMObsTagger.tagMap.has(span)) return + + const parent = ctx.llmobs.parent + storage.enterWith({ span: parent }) } asyncEnd (ctx) { @@ -60,6 +81,11 @@ class LLMObsPlugin extends TracingPlugin { } super.configure(config) } + + getLLMObsParent () { + const store = storage.getStore() + return store?.span + } } module.exports = LLMObsPlugin diff --git a/packages/dd-trace/src/llmobs/plugins/langchain/handlers/chain.js b/packages/dd-trace/src/llmobs/plugins/langchain/handlers/chain.js new file mode 100644 index 0000000000..33b3ad8488 --- /dev/null +++ b/packages/dd-trace/src/llmobs/plugins/langchain/handlers/chain.js @@ -0,0 +1,24 @@ +'use strict' + +const LangChainLLMObsHandler = require('.') +const { spanHasError } = require('../../../util') + +class LangChainLLMObsChainHandler extends LangChainLLMObsHandler { + setMetaTags ({ span, inputs, results }) { + let input, output + if (inputs) { + input = this.formatIO(inputs) + } + + if (!results || spanHasError(span)) { + output = '' + } else { + output = this.formatIO(results) + } + + // chain spans will always be workflows + this._tagger.tagTextIO(span, input, output) + } +} + +module.exports = LangChainLLMObsChainHandler diff --git a/packages/dd-trace/src/llmobs/plugins/langchain/handlers/chat_model.js b/packages/dd-trace/src/llmobs/plugins/langchain/handlers/chat_model.js new file mode 100644 index 0000000000..4e8aea269c --- /dev/null +++ b/packages/dd-trace/src/llmobs/plugins/langchain/handlers/chat_model.js @@ -0,0 +1,111 @@ +'use strict' + +const LangChainLLMObsHandler = require('.') +const LLMObsTagger = require('../../../tagger') +const { spanHasError } = require('../../../util') + +const LLM = 'llm' + +class LangChainLLMObsChatModelHandler extends LangChainLLMObsHandler { + setMetaTags ({ span, inputs, results, options, integrationName }) { + if (integrationName === 'openai' && options?.response_format) { + // langchain-openai will call a beta client if "response_format" is passed in on the options object + // we do not trace these calls, so this should be an llm span + this._tagger.changeKind(span, LLM) + } + const spanKind = LLMObsTagger.getSpanKind(span) + const isWorkflow = spanKind === 'workflow' + + const inputMessages = [] + if (!Array.isArray(inputs)) inputs = [inputs] + + for (const messageSet of inputs) { + for (const message of messageSet) { + const content = message.content || '' + const role = this.getRole(message) + inputMessages.push({ content, role }) + } + } + + if (spanHasError(span)) { + if (isWorkflow) { + this._tagger.tagTextIO(span, inputMessages, [{ content: '' }]) + } else { + this._tagger.tagLLMIO(span, inputMessages, [{ content: '' }]) + } + return + } + + const outputMessages = [] + let inputTokens = 0 + let outputTokens = 0 + let totalTokens = 0 + let tokensSetTopLevel = false + const tokensPerRunId = {} + + if (!isWorkflow) { + const tokens = this.checkTokenUsageChatOrLLMResult(results) + inputTokens = tokens.inputTokens + outputTokens = tokens.outputTokens + totalTokens = tokens.totalTokens + tokensSetTopLevel = totalTokens > 0 + } + + for (const messageSet of results.generations) { + for (const chatCompletion of messageSet) { + const chatCompletionMessage = chatCompletion.message + const role = this.getRole(chatCompletionMessage) + const content = chatCompletionMessage.text || '' + const toolCalls = this.extractToolCalls(chatCompletionMessage) + outputMessages.push({ content, role, toolCalls }) + + if (!isWorkflow && !tokensSetTopLevel) { + const { tokens, runId } = this.checkTokenUsageFromAIMessage(chatCompletionMessage) + if (!tokensPerRunId[runId]) { + tokensPerRunId[runId] = tokens + } else { + tokensPerRunId[runId].inputTokens += tokens.inputTokens + tokensPerRunId[runId].outputTokens += tokens.outputTokens + tokensPerRunId[runId].totalTokens += tokens.totalTokens + } + } + } + } + + if (!isWorkflow && !tokensSetTopLevel) { + inputTokens = Object.values(tokensPerRunId).reduce((acc, val) => acc + val.inputTokens, 0) + outputTokens = Object.values(tokensPerRunId).reduce((acc, val) => acc + val.outputTokens, 0) + totalTokens = Object.values(tokensPerRunId).reduce((acc, val) => acc + val.totalTokens, 0) + } + + if (isWorkflow) { + this._tagger.tagTextIO(span, inputMessages, outputMessages) + } else { + this._tagger.tagLLMIO(span, inputMessages, outputMessages) + this._tagger.tagMetrics(span, { + inputTokens, + outputTokens, + totalTokens + }) + } + } + + extractToolCalls (message) { + let toolCalls = message.tool_calls + if (!toolCalls) return [] + + const toolCallsInfo = [] + if (!Array.isArray(toolCalls)) toolCalls = [toolCalls] + for (const toolCall of toolCalls) { + toolCallsInfo.push({ + name: toolCall.name || '', + arguments: toolCall.args || {}, + tool_id: toolCall.id || '' + }) + } + + return toolCallsInfo + } +} + +module.exports = LangChainLLMObsChatModelHandler diff --git a/packages/dd-trace/src/llmobs/plugins/langchain/handlers/embedding.js b/packages/dd-trace/src/llmobs/plugins/langchain/handlers/embedding.js new file mode 100644 index 0000000000..285fb1f0a9 --- /dev/null +++ b/packages/dd-trace/src/llmobs/plugins/langchain/handlers/embedding.js @@ -0,0 +1,42 @@ +'use strict' + +const LangChainLLMObsHandler = require('.') +const LLMObsTagger = require('../../../tagger') +const { spanHasError } = require('../../../util') + +class LangChainLLMObsEmbeddingHandler extends LangChainLLMObsHandler { + setMetaTags ({ span, inputs, results }) { + const isWorkflow = LLMObsTagger.getSpanKind(span) === 'workflow' + let embeddingInput, embeddingOutput + + if (isWorkflow) { + embeddingInput = this.formatIO(inputs) + } else { + const input = Array.isArray(inputs) ? inputs : [inputs] + embeddingInput = input.map(doc => ({ text: doc })) + } + + if (spanHasError(span) || !results) { + embeddingOutput = '' + } else { + let embeddingDimensions, embeddingsCount + if (typeof results[0] === 'number') { + embeddingsCount = 1 + embeddingDimensions = results.length + } else { + embeddingsCount = results.length + embeddingDimensions = results[0].length + } + + embeddingOutput = `[${embeddingsCount} embedding(s) returned with size ${embeddingDimensions}]` + } + + if (isWorkflow) { + this._tagger.tagTextIO(span, embeddingInput, embeddingOutput) + } else { + this._tagger.tagEmbeddingIO(span, embeddingInput, embeddingOutput) + } + } +} + +module.exports = LangChainLLMObsEmbeddingHandler diff --git a/packages/dd-trace/src/llmobs/plugins/langchain/handlers/index.js b/packages/dd-trace/src/llmobs/plugins/langchain/handlers/index.js new file mode 100644 index 0000000000..d2a0aafdd4 --- /dev/null +++ b/packages/dd-trace/src/llmobs/plugins/langchain/handlers/index.js @@ -0,0 +1,102 @@ +'use strict' + +const ROLE_MAPPINGS = { + human: 'user', + ai: 'assistant', + system: 'system' +} + +class LangChainLLMObsHandler { + constructor (tagger) { + this._tagger = tagger + } + + setMetaTags () {} + + formatIO (messages) { + if (messages.constructor.name === 'Object') { // plain JSON + const formatted = {} + for (const [key, value] of Object.entries(messages)) { + formatted[key] = this.formatIO(value) + } + + return formatted + } else if (Array.isArray(messages)) { + return messages.map(message => this.formatIO(message)) + } else { // either a BaseMesage type or a string + return this.getContentFromMessage(messages) + } + } + + getContentFromMessage (message) { + if (typeof message === 'string') { + return message + } else { + try { + const messageContent = {} + messageContent.content = message.content || '' + + const role = this.getRole(message) + if (role) messageContent.role = role + + return messageContent + } catch { + return JSON.stringify(message) + } + } + } + + checkTokenUsageChatOrLLMResult (results) { + const llmOutput = results.llmOutput + const tokens = { + inputTokens: 0, + outputTokens: 0, + totalTokens: 0 + } + if (!llmOutput) return tokens + const tokenUsage = llmOutput.tokenUsage || llmOutput.usageMetadata || llmOutput.usage || {} + if (!tokenUsage) return tokens + + tokens.inputTokens = tokenUsage.promptTokens || tokenUsage.inputTokens || 0 + tokens.outputTokens = tokenUsage.completionTokens || tokenUsage.outputTokens || 0 + tokens.totalTokens = tokenUsage.totalTokens || tokens.inputTokens + tokens.outputTokens + + return tokens + } + + checkTokenUsageFromAIMessage (message) { + let usage = message.usage_metadata || message.additional_kwargs?.usage + const runId = message.run_id || message.id || '' + const runIdBase = runId ? runId.split('-').slice(0, -1).join('-') : '' + + const responseMetadata = message.response_metadata || {} + usage = usage || responseMetadata.usage || responseMetadata.tokenUsage || {} + + const inputTokens = usage.promptTokens || usage.inputTokens || usage.prompt_tokens || usage.input_tokens || 0 + const outputTokens = + usage.completionTokens || usage.outputTokens || usage.completion_tokens || usage.output_tokens || 0 + const totalTokens = usage.totalTokens || inputTokens + outputTokens + + return { + tokens: { + inputTokens, + outputTokens, + totalTokens + }, + runId: runIdBase + } + } + + getRole (message) { + if (message.role) return ROLE_MAPPINGS[message.role] || message.role + + const type = ( + (typeof message.getType === 'function' && message.getType()) || + (typeof message._getType === 'function' && message._getType()) + ) + + return ROLE_MAPPINGS[type] || type + } +} + +module.exports = LangChainLLMObsHandler diff --git a/packages/dd-trace/src/llmobs/plugins/langchain/handlers/llm.js b/packages/dd-trace/src/llmobs/plugins/langchain/handlers/llm.js new file mode 100644 index 0000000000..24f8db5c7c --- /dev/null +++ b/packages/dd-trace/src/llmobs/plugins/langchain/handlers/llm.js @@ -0,0 +1,32 @@ +'use strict' + +const LangChainLLMObsHandler = require('.') +const LLMObsTagger = require('../../../tagger') +const { spanHasError } = require('../../../util') + +class LangChainLLMObsLlmHandler extends LangChainLLMObsHandler { + setMetaTags ({ span, inputs, results }) { + const isWorkflow = LLMObsTagger.getSpanKind(span) === 'workflow' + const prompts = Array.isArray(inputs) ? inputs : [inputs] + + let outputs + if (spanHasError(span)) { + outputs = [{ content: '' }] + } else { + outputs = results.generations.map(completion => ({ content: completion[0].text })) + + if (!isWorkflow) { + const tokens = this.checkTokenUsageChatOrLLMResult(results) + this._tagger.tagMetrics(span, tokens) + } + } + + if (isWorkflow) { + this._tagger.tagTextIO(span, prompts, outputs) + } else { + this._tagger.tagLLMIO(span, prompts, outputs) + } + } +} + +module.exports = LangChainLLMObsLlmHandler diff --git a/packages/dd-trace/src/llmobs/plugins/langchain/index.js b/packages/dd-trace/src/llmobs/plugins/langchain/index.js new file mode 100644 index 0000000000..9dc3dd8cd9 --- /dev/null +++ b/packages/dd-trace/src/llmobs/plugins/langchain/index.js @@ -0,0 +1,138 @@ +'use strict' + +const log = require('../../../log') +const LLMObsPlugin = require('../base') + +const pluginManager = require('../../../../../..')._pluginManager + +const ANTHROPIC_PROVIDER_NAME = 'anthropic' +const BEDROCK_PROVIDER_NAME = 'amazon_bedrock' +const OPENAI_PROVIDER_NAME = 'openai' + +const SUPPORTED_INTEGRATIONS = ['openai'] +const LLM_SPAN_TYPES = ['llm', 'chat_model', 'embedding'] +const LLM = 'llm' +const WORKFLOW = 'workflow' +const EMBEDDING = 'embedding' + +const ChainHandler = require('./handlers/chain') +const ChatModelHandler = require('./handlers/chat_model') +const LlmHandler = require('./handlers/llm') +const EmbeddingHandler = require('./handlers/embedding') + +const SUPPORTED_HANLDERS = { + chain: ChainHandler, + chat_model: ChatModelHandler, + llm: LlmHandler, + embedding: EmbeddingHandler +} + +class LangChainLLMObsPlugin extends LLMObsPlugin { + static get prefix () { + return 'tracing:apm:langchain:invoke' + } + + constructor () { + super(...arguments) + + this._handlers = + Object.entries(SUPPORTED_HANLDERS) + .reduce((acc, [type, Handler]) => { + acc[type] = new Handler(this._tagger) + return acc + }, {}) + } + + getLLMObsSpanRegisterOptions (ctx) { + const span = ctx.currentStore?.span + const tags = span?.context()._tags || {} + + const modelProvider = tags['langchain.request.provider'] // could be undefined + const modelName = tags['langchain.request.model'] // could be undefined + const kind = this.getKind(ctx.type, modelProvider) + const name = tags['resource.name'] + + return { + modelProvider, + modelName, + kind, + name + } + } + + setLLMObsTags (ctx) { + const span = ctx.currentStore?.span + const type = ctx.type // langchain operation type + + if (!Object.keys(SUPPORTED_HANLDERS).includes(type)) { + log.warn(`Unsupported LangChain operation type: ${type}`) + return + } + + const provider = span?.context()._tags['langchain.request.provider'] + const integrationName = this.getIntegrationName(type, provider) + this.setMetadata(span, provider) + + const inputs = ctx.args?.[0] + const options = ctx.args?.[1] + const results = ctx.result + + this._handlers[type].setMetaTags({ span, inputs, results, options, integrationName }) + } + + setMetadata (span, provider) { + if (!provider) return + + const metadata = {} + + // these fields won't be set for non model-based operations + const temperature = + span?.context()._tags[`langchain.request.${provider}.parameters.temperature`] || + span?.context()._tags[`langchain.request.${provider}.parameters.model_kwargs.temperature`] + + const maxTokens = + span?.context()._tags[`langchain.request.${provider}.parameters.max_tokens`] || + span?.context()._tags[`langchain.request.${provider}.parameters.maxTokens`] || + span?.context()._tags[`langchain.request.${provider}.parameters.model_kwargs.max_tokens`] + + if (temperature) { + metadata.temperature = parseFloat(temperature) + } + + if (maxTokens) { + metadata.maxTokens = parseInt(maxTokens) + } + + this._tagger.tagMetadata(span, metadata) + } + + getKind (type, provider) { + if (LLM_SPAN_TYPES.includes(type)) { + const llmobsIntegration = this.getIntegrationName(type, provider) + + if (!this.isLLMIntegrationEnabled(llmobsIntegration)) { + return type === 'embedding' ? EMBEDDING : LLM + } + } + + return WORKFLOW + } + + getIntegrationName (type, provider = 'custom') { + if (provider.startsWith(BEDROCK_PROVIDER_NAME)) { + return 'bedrock' + } else if (provider.startsWith(OPENAI_PROVIDER_NAME)) { + return 'openai' + } else if (type === 'chat_model' && provider.startsWith(ANTHROPIC_PROVIDER_NAME)) { + return 'anthropic' + } + + return provider + } + + isLLMIntegrationEnabled (integration) { + return SUPPORTED_INTEGRATIONS.includes(integration) && pluginManager?._pluginsByName[integration]?.llmobs?._enabled + } +} + +module.exports = LangChainLLMObsPlugin diff --git a/packages/dd-trace/src/llmobs/plugins/openai.js b/packages/dd-trace/src/llmobs/plugins/openai.js index 431760a04f..fee41afcbe 100644 --- a/packages/dd-trace/src/llmobs/plugins/openai.js +++ b/packages/dd-trace/src/llmobs/plugins/openai.js @@ -7,7 +7,7 @@ class OpenAiLLMObsPlugin extends LLMObsPlugin { return 'tracing:apm:openai:request' } - getLLMObsSPanRegisterOptions (ctx) { + getLLMObsSpanRegisterOptions (ctx) { const resource = ctx.methodName const methodName = gateResource(normalizeOpenAIResourceName(resource)) if (!methodName) return // we will not trace all openai methods for llmobs diff --git a/packages/dd-trace/src/llmobs/tagger.js b/packages/dd-trace/src/llmobs/tagger.js index 9f1728e5d7..a7df65fbbc 100644 --- a/packages/dd-trace/src/llmobs/tagger.js +++ b/packages/dd-trace/src/llmobs/tagger.js @@ -27,6 +27,7 @@ const { // global registry of LLMObs spans // maps LLMObs spans to their annotations +// TODO: maybe move this to its own registry file? const registry = new WeakMap() class LLMObsTagger { @@ -40,6 +41,10 @@ class LLMObsTagger { return registry } + static getSpanKind (span) { + return registry.get(span)?.[SPAN_KIND] + } + registerLLMObsSpan (span, { modelName, modelProvider, @@ -136,6 +141,10 @@ class LLMObsTagger { this._setTag(span, TAGS, tags) } + changeKind (span, newKind) { + this._setTag(span, SPAN_KIND, newKind) + } + _tagText (span, data, key) { if (data) { if (typeof data === 'string') { @@ -310,7 +319,7 @@ class LLMObsTagger { _setTag (span, key, value) { if (!this._config.llmobs.enabled) return if (!registry.has(span)) { - this._handleFailure('Span must be an LLMObs generated span.') + this._handleFailure(`Span "${span._name}" must be an LLMObs generated span.`) return } diff --git a/packages/dd-trace/src/llmobs/util.js b/packages/dd-trace/src/llmobs/util.js index feba656f95..3f9127210c 100644 --- a/packages/dd-trace/src/llmobs/util.js +++ b/packages/dd-trace/src/llmobs/util.js @@ -169,8 +169,14 @@ function getFunctionArguments (fn, args = []) { } } +function spanHasError (span) { + const tags = span.context()._tags + return !!(tags.error || tags['error.type']) +} + module.exports = { encodeUnicode, validateKind, - getFunctionArguments + getFunctionArguments, + spanHasError } diff --git a/packages/dd-trace/test/llmobs/plugins/langchain/index.spec.js b/packages/dd-trace/test/llmobs/plugins/langchain/index.spec.js new file mode 100644 index 0000000000..c2c0d29495 --- /dev/null +++ b/packages/dd-trace/test/llmobs/plugins/langchain/index.spec.js @@ -0,0 +1,1107 @@ +'use strict' + +const LLMObsAgentProxySpanWriter = require('../../../../src/llmobs/writers/spans/agentProxy') +const { useEnv } = require('../../../../../../integration-tests/helpers') +const agent = require('../../../../../dd-trace/test/plugins/agent') +const { + expectedLLMObsLLMSpanEvent, + expectedLLMObsNonLLMSpanEvent, + deepEqualWithMockValues, + MOCK_ANY, + MOCK_STRING +} = require('../../util') +const chai = require('chai') + +chai.Assertion.addMethod('deepEqualWithMockValues', deepEqualWithMockValues) + +const nock = require('nock') +function stubCall ({ base = '', path = '', code = 200, response = {} }) { + const responses = Array.isArray(response) ? response : [response] + const times = responses.length + nock(base).post(path).times(times).reply(() => { + return [code, responses.shift()] + }) +} + +const openAiBaseCompletionInfo = { base: 'https://api.openai.com', path: '/v1/completions' } +const openAiBaseChatInfo = { base: 'https://api.openai.com', path: '/v1/chat/completions' } +const openAiBaseEmbeddingInfo = { base: 'https://api.openai.com', path: '/v1/embeddings' } + +describe('integrations', () => { + let langchainOpenai + let langchainAnthropic + let langchainCohere + + let langchainMessages + let langchainOutputParsers + let langchainPrompts + let langchainRunnables + + let llmobs + + // so we can verify it gets tagged properly + useEnv({ + OPENAI_API_KEY: '', + ANTHROPIC_API_KEY: '', + COHERE_API_KEY: '' + }) + + describe('langchain', () => { + before(async () => { + sinon.stub(LLMObsAgentProxySpanWriter.prototype, 'append') + + // reduce errors related to too many listeners + process.removeAllListeners('beforeExit') + + LLMObsAgentProxySpanWriter.prototype.append.reset() + + await agent.load('langchain', {}, { + llmobs: { + mlApp: 'test' + } + }) + + llmobs = require('../../../../../..').llmobs + }) + + afterEach(() => { + nock.cleanAll() + LLMObsAgentProxySpanWriter.prototype.append.reset() + }) + + after(() => { + require('../../../../../dd-trace').llmobs.disable() // unsubscribe from all events + sinon.restore() + return agent.close({ ritmReset: false, wipe: true }) + }) + + withVersions('langchain', ['@langchain/core'], version => { + describe('langchain', () => { + beforeEach(() => { + langchainOpenai = require(`../../../../../../versions/@langchain/openai@${version}`).get() + langchainAnthropic = require(`../../../../../../versions/@langchain/anthropic@${version}`).get() + langchainCohere = require(`../../../../../../versions/@langchain/cohere@${version}`).get() + + // need to specify specific import in `get(...)` + langchainMessages = require(`../../../../../../versions/@langchain/core@${version}`) + .get('@langchain/core/messages') + langchainOutputParsers = require(`../../../../../../versions/@langchain/core@${version}`) + .get('@langchain/core/output_parsers') + langchainPrompts = require(`../../../../../../versions/@langchain/core@${version}`) + .get('@langchain/core/prompts') + langchainRunnables = require(`../../../../../../versions/@langchain/core@${version}`) + .get('@langchain/core/runnables') + }) + + describe('llm', () => { + it('submits an llm span for an openai llm call', async () => { + stubCall({ + ...openAiBaseCompletionInfo, + response: { + choices: [ + { + text: 'Hello, world!' + } + ], + usage: { prompt_tokens: 8, completion_tokens: 12, otal_tokens: 20 } + } + }) + + const llm = new langchainOpenai.OpenAI({ model: 'gpt-3.5-turbo-instruct' }) + + const checkTraces = agent.use(traces => { + const span = traces[0][0] + const spanEvent = LLMObsAgentProxySpanWriter.prototype.append.getCall(0).args[0] + + const expected = expectedLLMObsLLMSpanEvent({ + span, + spanKind: 'llm', + modelName: 'gpt-3.5-turbo-instruct', + modelProvider: 'openai', + name: 'langchain.llms.openai.OpenAI', + inputMessages: [{ content: 'Hello!' }], + outputMessages: [{ content: 'Hello, world!' }], + metadata: MOCK_ANY, + tokenMetrics: { input_tokens: 8, output_tokens: 12, total_tokens: 20 }, + tags: { ml_app: 'test', language: 'javascript' } + }) + + expect(spanEvent).to.deepEqualWithMockValues(expected) + }) + + await llm.invoke('Hello!') + + await checkTraces + }) + + it('does not tag output if there is an error', async () => { + nock('https://api.openai.com').post('/v1/completions').reply(500) + + const checkTraces = agent.use(traces => { + const span = traces[0][0] + const spanEvent = LLMObsAgentProxySpanWriter.prototype.append.getCall(0).args[0] + + const expected = expectedLLMObsLLMSpanEvent({ + span, + spanKind: 'llm', + modelName: 'gpt-3.5-turbo-instruct', + modelProvider: 'openai', + name: 'langchain.llms.openai.OpenAI', + inputMessages: [{ content: 'Hello!' }], + outputMessages: [{ content: '' }], + metadata: MOCK_ANY, + tokenMetrics: { input_tokens: 0, output_tokens: 0, total_tokens: 0 }, + tags: { ml_app: 'test', language: 'javascript' }, + error: 1, + errorType: 'Error', + errorMessage: MOCK_STRING, + errorStack: MOCK_ANY + }) + + expect(spanEvent).to.deepEqualWithMockValues(expected) + }) + + const llm = new langchainOpenai.OpenAI({ model: 'gpt-3.5-turbo-instruct', maxRetries: 0 }) + + try { + await llm.invoke('Hello!') + } catch {} + + await checkTraces + }) + + it('submits an llm span for a cohere call', async function () { + if (version === '0.1.0') this.skip() // cannot patch client to mock response on lower versions + + const cohere = new langchainCohere.Cohere({ + model: 'command', + client: { + generate () { + return { + generations: [ + { + text: 'hello world!' + } + ], + meta: { + billed_units: { + input_tokens: 8, + output_tokens: 12 + } + } + } + } + } + }) + + const checkTraces = agent.use(traces => { + const span = traces[0][0] + + const spanEvent = LLMObsAgentProxySpanWriter.prototype.append.getCall(0).args[0] + + const expected = expectedLLMObsLLMSpanEvent({ + span, + spanKind: 'llm', + modelName: 'command', + modelProvider: 'cohere', + name: 'langchain.llms.cohere.Cohere', + inputMessages: [{ content: 'Hello!' }], + outputMessages: [{ content: 'hello world!' }], + metadata: MOCK_ANY, + // @langchain/cohere does not provide token usage in the response + tokenMetrics: { input_tokens: 0, output_tokens: 0, total_tokens: 0 }, + tags: { ml_app: 'test', language: 'javascript' } + }) + + expect(spanEvent).to.deepEqualWithMockValues(expected) + }) + + await cohere.invoke('Hello!') + + await checkTraces + }) + }) + + describe('chat model', () => { + it('submits an llm span for an openai chat model call', async () => { + stubCall({ + ...openAiBaseChatInfo, + response: { + choices: [ + { + message: { + content: 'Hello, world!', + role: 'assistant' + } + } + ], + usage: { prompt_tokens: 8, completion_tokens: 12, total_tokens: 20 } + } + }) + + const chat = new langchainOpenai.ChatOpenAI({ model: 'gpt-3.5-turbo' }) + + const checkTraces = agent.use(traces => { + const span = traces[0][0] + const spanEvent = LLMObsAgentProxySpanWriter.prototype.append.getCall(0).args[0] + + const expected = expectedLLMObsLLMSpanEvent({ + span, + spanKind: 'llm', + modelName: 'gpt-3.5-turbo', + modelProvider: 'openai', + name: 'langchain.chat_models.openai.ChatOpenAI', + inputMessages: [{ content: 'Hello!', role: 'user' }], + outputMessages: [{ content: 'Hello, world!', role: 'assistant' }], + metadata: MOCK_ANY, + tokenMetrics: { input_tokens: 8, output_tokens: 12, total_tokens: 20 }, + tags: { ml_app: 'test', language: 'javascript' } + }) + + expect(spanEvent).to.deepEqualWithMockValues(expected) + }) + + await chat.invoke('Hello!') + + await checkTraces + }) + + it('does not tag output if there is an error', async () => { + nock('https://api.openai.com').post('/v1/chat/completions').reply(500) + + const checkTraces = agent.use(traces => { + const span = traces[0][0] + const spanEvent = LLMObsAgentProxySpanWriter.prototype.append.getCall(0).args[0] + + const expected = expectedLLMObsLLMSpanEvent({ + span, + spanKind: 'llm', + modelName: 'gpt-3.5-turbo', + modelProvider: 'openai', + name: 'langchain.chat_models.openai.ChatOpenAI', + inputMessages: [{ content: 'Hello!', role: 'user' }], + outputMessages: [{ content: '' }], + metadata: MOCK_ANY, + tokenMetrics: { input_tokens: 0, output_tokens: 0, total_tokens: 0 }, + tags: { ml_app: 'test', language: 'javascript' }, + error: 1, + errorType: 'Error', + errorMessage: MOCK_STRING, + errorStack: MOCK_ANY + }) + + expect(spanEvent).to.deepEqualWithMockValues(expected) + }) + + const chat = new langchainOpenai.ChatOpenAI({ model: 'gpt-3.5-turbo', maxRetries: 0 }) + + try { + await chat.invoke('Hello!') + } catch {} + + await checkTraces + }) + + it('submits an llm span for an anthropic chat model call', async () => { + stubCall({ + base: 'https://api.anthropic.com', + path: '/v1/messages', + response: { + id: 'msg_01NE2EJQcjscRyLbyercys6p', + type: 'message', + role: 'assistant', + model: 'claude-2.1', + content: [ + { type: 'text', text: 'Hello!' } + ], + stop_reason: 'end_turn', + stop_sequence: null, + usage: { input_tokens: 11, output_tokens: 6 } + } + }) + + const checkTraces = agent.use(traces => { + const span = traces[0][0] + const spanEvent = LLMObsAgentProxySpanWriter.prototype.append.getCall(0).args[0] + + const expected = expectedLLMObsLLMSpanEvent({ + span, + spanKind: 'llm', + modelName: 'claude-2.1', // overriden langchain for older versions + modelProvider: 'anthropic', + name: 'langchain.chat_models.anthropic.ChatAnthropic', + inputMessages: [{ content: 'Hello!', role: 'user' }], + outputMessages: [{ content: 'Hello!', role: 'assistant' }], + metadata: MOCK_ANY, + tokenMetrics: { input_tokens: 11, output_tokens: 6, total_tokens: 17 }, + tags: { ml_app: 'test', language: 'javascript' } + }) + + expect(spanEvent).to.deepEqualWithMockValues(expected) + }) + + const chatModel = new langchainAnthropic.ChatAnthropic({ model: 'claude-2.1' }) + + await chatModel.invoke('Hello!') + + await checkTraces + }) + + it('submits an llm span with tool calls', async () => { + stubCall({ + ...openAiBaseChatInfo, + response: { + model: 'gpt-4', + choices: [{ + message: { + role: 'assistant', + content: null, + tool_calls: [ + { + id: 'tool-1', + type: 'function', + function: { + name: 'extract_fictional_info', + arguments: '{"name":"SpongeBob","origin":"Bikini Bottom"}' + } + } + ] + }, + finish_reason: 'tool_calls', + index: 0 + }] + } + }) + + const checkTraces = agent.use(traces => { + const span = traces[0][0] + const spanEvent = LLMObsAgentProxySpanWriter.prototype.append.getCall(0).args[0] + + const expected = expectedLLMObsLLMSpanEvent({ + span, + spanKind: 'llm', + modelName: 'gpt-4', + modelProvider: 'openai', + name: 'langchain.chat_models.openai.ChatOpenAI', + inputMessages: [{ content: 'My name is SpongeBob and I live in Bikini Bottom.', role: 'user' }], + outputMessages: [{ + content: '', + role: 'assistant', + tool_calls: [{ + arguments: { + name: 'SpongeBob', + origin: 'Bikini Bottom' + }, + name: 'extract_fictional_info' + }] + }], + metadata: MOCK_ANY, + // also tests tokens not sent on llm-type spans should be 0 + tokenMetrics: { input_tokens: 0, output_tokens: 0, total_tokens: 0 }, + tags: { ml_app: 'test', language: 'javascript' } + }) + + expect(spanEvent).to.deepEqualWithMockValues(expected) + }) + + const tools = [ + { + name: 'extract_fictional_info', + description: 'Get the fictional information from the body of the input text', + parameters: { + type: 'object', + properties: { + name: { type: 'string', description: 'Name of the character' }, + origin: { type: 'string', description: 'Where they live' } + } + } + } + ] + + const model = new langchainOpenai.ChatOpenAI({ model: 'gpt-4' }) + const modelWithTools = model.bindTools(tools) + + await modelWithTools.invoke('My name is SpongeBob and I live in Bikini Bottom.') + + await checkTraces + }) + }) + + describe('embedding', () => { + it('submits an embedding span for an `embedQuery` call', async () => { + stubCall({ + ...openAiBaseEmbeddingInfo, + response: { + object: 'list', + data: [{ + object: 'embedding', + index: 0, + embedding: [-0.0034387498, -0.026400521] + }] + } + }) + const embeddings = new langchainOpenai.OpenAIEmbeddings() + + const checkTraces = agent.use(traces => { + const span = traces[0][0] + const spanEvent = LLMObsAgentProxySpanWriter.prototype.append.getCall(0).args[0] + + const expected = expectedLLMObsLLMSpanEvent({ + span, + spanKind: 'embedding', + modelName: 'text-embedding-ada-002', + modelProvider: 'openai', + name: 'langchain.embeddings.openai.OpenAIEmbeddings', + inputDocuments: [{ text: 'Hello!' }], + outputValue: '[1 embedding(s) returned with size 2]', + metadata: MOCK_ANY, + tags: { ml_app: 'test', language: 'javascript' } + }) + + expect(spanEvent).to.deepEqualWithMockValues(expected) + }) + + await embeddings.embedQuery('Hello!') + + await checkTraces + }) + + it('does not tag output if there is an error', async () => { + nock('https://api.openai.com').post('/v1/embeddings').reply(500) + + const checkTraces = agent.use(traces => { + const span = traces[0][0] + const spanEvent = LLMObsAgentProxySpanWriter.prototype.append.getCall(0).args[0] + + const expected = expectedLLMObsLLMSpanEvent({ + span, + spanKind: 'embedding', + modelName: 'text-embedding-ada-002', + modelProvider: 'openai', + name: 'langchain.embeddings.openai.OpenAIEmbeddings', + inputDocuments: [{ text: 'Hello!' }], + outputValue: '', + metadata: MOCK_ANY, + tags: { ml_app: 'test', language: 'javascript' }, + error: 1, + errorType: 'Error', + errorMessage: MOCK_STRING, + errorStack: MOCK_ANY + }) + + expect(spanEvent).to.deepEqualWithMockValues(expected) + }) + + const embeddings = new langchainOpenai.OpenAIEmbeddings({ maxRetries: 0 }) + + try { + await embeddings.embedQuery('Hello!') + } catch {} + + await checkTraces + }) + + it('submits an embedding span for an `embedDocuments` call', async () => { + stubCall({ + ...openAiBaseEmbeddingInfo, + response: { + object: 'list', + data: [{ + object: 'embedding', + index: 0, + embedding: [-0.0034387498, -0.026400521] + }, { + object: 'embedding', + index: 1, + embedding: [-0.026400521, -0.0034387498] + }] + } + }) + + const embeddings = new langchainOpenai.OpenAIEmbeddings() + + const checkTraces = agent.use(traces => { + const span = traces[0][0] + const spanEvent = LLMObsAgentProxySpanWriter.prototype.append.getCall(0).args[0] + + const expected = expectedLLMObsLLMSpanEvent({ + span, + spanKind: 'embedding', + modelName: 'text-embedding-ada-002', + modelProvider: 'openai', + name: 'langchain.embeddings.openai.OpenAIEmbeddings', + inputDocuments: [{ text: 'Hello!' }, { text: 'World!' }], + outputValue: '[2 embedding(s) returned with size 2]', + metadata: MOCK_ANY, + tags: { ml_app: 'test', language: 'javascript' } + }) + + expect(spanEvent).to.deepEqualWithMockValues(expected) + }) + + await embeddings.embedDocuments(['Hello!', 'World!']) + + await checkTraces + }) + }) + + describe('chain', () => { + it('submits a workflow and llm spans for a simple chain call', async () => { + stubCall({ + ...openAiBaseCompletionInfo, + response: { + choices: [ + { + text: 'LangSmith can help with testing in several ways.' + } + ], + usage: { prompt_tokens: 8, completion_tokens: 12, otal_tokens: 20 } + } + }) + + const prompt = langchainPrompts.ChatPromptTemplate.fromMessages([ + ['system', 'You are a world class technical documentation writer'], + ['user', '{input}'] + ]) + + const llm = new langchainOpenai.OpenAI({ model: 'gpt-3.5-turbo-instruct' }) + + const chain = prompt.pipe(llm) + + const checkTraces = agent.use(traces => { + const spans = traces[0] + const workflowSpan = spans[0] + const llmSpan = spans[1] + + const workflowSpanEvent = LLMObsAgentProxySpanWriter.prototype.append.getCall(0).args[0] + const llmSpanEvent = LLMObsAgentProxySpanWriter.prototype.append.getCall(1).args[0] + + const expectedWorkflow = expectedLLMObsNonLLMSpanEvent({ + span: workflowSpan, + spanKind: 'workflow', + name: 'langchain_core.runnables.RunnableSequence', + inputValue: JSON.stringify({ input: 'Can you tell me about LangSmith?' }), + outputValue: 'LangSmith can help with testing in several ways.', + metadata: MOCK_ANY, + tags: { ml_app: 'test', language: 'javascript' } + }) + + const expectedLLM = expectedLLMObsLLMSpanEvent({ + span: llmSpan, + parentId: workflowSpan.span_id, + spanKind: 'llm', + modelName: 'gpt-3.5-turbo-instruct', + modelProvider: 'openai', + name: 'langchain.llms.openai.OpenAI', + // this is how LangChain formats these IOs for LLMs + inputMessages: [{ + content: 'System: You are a world class technical documentation writer\n' + + 'Human: Can you tell me about LangSmith?' + }], + outputMessages: [{ content: 'LangSmith can help with testing in several ways.' }], + metadata: MOCK_ANY, + tokenMetrics: { input_tokens: 8, output_tokens: 12, total_tokens: 20 }, + tags: { ml_app: 'test', language: 'javascript' } + }) + + expect(workflowSpanEvent).to.deepEqualWithMockValues(expectedWorkflow) + + expect(llmSpanEvent).to.deepEqualWithMockValues(expectedLLM) + }) + + await chain.invoke({ input: 'Can you tell me about LangSmith?' }) + + await checkTraces + }) + + it('does not tag output if there is an error', async () => { + nock('https://api.openai.com').post('/v1/completions').reply(500) + + const checkTraces = agent.use(traces => { + const spans = traces[0] + + const workflowSpan = spans[0] + + const workflowSpanEvent = LLMObsAgentProxySpanWriter.prototype.append.getCall(0).args[0] + + const expectedWorkflow = expectedLLMObsNonLLMSpanEvent({ + span: workflowSpan, + spanKind: 'workflow', + name: 'langchain_core.runnables.RunnableSequence', + inputValue: 'Hello!', + outputValue: '', + metadata: MOCK_ANY, + tags: { ml_app: 'test', language: 'javascript' }, + error: 1, + errorType: 'Error', + errorMessage: MOCK_STRING, + errorStack: MOCK_ANY + }) + + expect(workflowSpanEvent).to.deepEqualWithMockValues(expectedWorkflow) + }) + + const llm = new langchainOpenai.OpenAI({ model: 'gpt-3.5-turbo-instruct', maxRetries: 0 }) + const parser = new langchainOutputParsers.StringOutputParser() + const chain = llm.pipe(parser) + + try { + await chain.invoke('Hello!') + } catch {} + + await checkTraces + }) + + it('submits workflow and llm spans for a nested chain', async () => { + stubCall({ + ...openAiBaseChatInfo, + response: [ + { + choices: [ + { + message: { + content: 'Springfield, Illinois', + role: 'assistant' + } + } + ], + usage: { prompt_tokens: 8, completion_tokens: 12, total_tokens: 20 } + }, + { + choices: [ + { + message: { + content: 'Springfield, Illinois está en los Estados Unidos.', + role: 'assistant' + } + } + ], + usage: { prompt_tokens: 8, completion_tokens: 12, total_tokens: 20 } + } + ] + }) + + const firstPrompt = langchainPrompts.ChatPromptTemplate.fromTemplate('what is the city {person} is from?') + const secondPrompt = langchainPrompts.ChatPromptTemplate.fromTemplate( + 'what country is the city {city} in? respond in {language}' + ) + + const model = new langchainOpenai.ChatOpenAI({ model: 'gpt-3.5-turbo' }) + const parser = new langchainOutputParsers.StringOutputParser() + + const firstChain = firstPrompt.pipe(model).pipe(parser) + const secondChain = secondPrompt.pipe(model).pipe(parser) + + const completeChain = langchainRunnables.RunnableSequence.from([ + { + city: firstChain, + language: input => input.language + }, + secondChain + ]) + + const checkTraces = agent.use(traces => { + const spans = traces[0] + + const topLevelWorkflow = spans[0] + const firstSubWorkflow = spans[1] + const firstLLM = spans[2] + const secondSubWorkflow = spans[3] + const secondLLM = spans[4] + + const topLevelWorkflowSpanEvent = LLMObsAgentProxySpanWriter.prototype.append.getCall(0).args[0] + const firstSubWorkflowSpanEvent = LLMObsAgentProxySpanWriter.prototype.append.getCall(1).args[0] + const firstLLMSpanEvent = LLMObsAgentProxySpanWriter.prototype.append.getCall(2).args[0] + const secondSubWorkflowSpanEvent = LLMObsAgentProxySpanWriter.prototype.append.getCall(3).args[0] + const secondLLMSpanEvent = LLMObsAgentProxySpanWriter.prototype.append.getCall(4).args[0] + + const expectedTopLevelWorkflow = expectedLLMObsNonLLMSpanEvent({ + span: topLevelWorkflow, + spanKind: 'workflow', + name: 'langchain_core.runnables.RunnableSequence', + inputValue: JSON.stringify({ person: 'Abraham Lincoln', language: 'Spanish' }), + outputValue: 'Springfield, Illinois está en los Estados Unidos.', + tags: { ml_app: 'test', language: 'javascript' } + }) + + const expectedFirstSubWorkflow = expectedLLMObsNonLLMSpanEvent({ + span: firstSubWorkflow, + parentId: topLevelWorkflow.span_id, + spanKind: 'workflow', + name: 'langchain_core.runnables.RunnableSequence', + inputValue: JSON.stringify({ person: 'Abraham Lincoln', language: 'Spanish' }), + outputValue: 'Springfield, Illinois', + tags: { ml_app: 'test', language: 'javascript' } + }) + + const expectedFirstLLM = expectedLLMObsLLMSpanEvent({ + span: firstLLM, + parentId: firstSubWorkflow.span_id, + spanKind: 'llm', + modelName: 'gpt-3.5-turbo', + modelProvider: 'openai', + name: 'langchain.chat_models.openai.ChatOpenAI', + inputMessages: [ + { content: 'what is the city Abraham Lincoln is from?', role: 'user' } + ], + outputMessages: [{ content: 'Springfield, Illinois', role: 'assistant' }], + metadata: MOCK_ANY, + tokenMetrics: { input_tokens: 8, output_tokens: 12, total_tokens: 20 }, + tags: { ml_app: 'test', language: 'javascript' } + }) + + const expectedSecondSubWorkflow = expectedLLMObsNonLLMSpanEvent({ + span: secondSubWorkflow, + parentId: topLevelWorkflow.span_id, + spanKind: 'workflow', + name: 'langchain_core.runnables.RunnableSequence', + inputValue: JSON.stringify({ language: 'Spanish', city: 'Springfield, Illinois' }), + outputValue: 'Springfield, Illinois está en los Estados Unidos.', + tags: { ml_app: 'test', language: 'javascript' } + }) + + const expectedSecondLLM = expectedLLMObsLLMSpanEvent({ + span: secondLLM, + parentId: secondSubWorkflow.span_id, + spanKind: 'llm', + modelName: 'gpt-3.5-turbo', + modelProvider: 'openai', + name: 'langchain.chat_models.openai.ChatOpenAI', + inputMessages: [ + { content: 'what country is the city Springfield, Illinois in? respond in Spanish', role: 'user' } + ], + outputMessages: [{ content: 'Springfield, Illinois está en los Estados Unidos.', role: 'assistant' }], + metadata: MOCK_ANY, + tokenMetrics: { input_tokens: 8, output_tokens: 12, total_tokens: 20 }, + tags: { ml_app: 'test', language: 'javascript' } + }) + + expect(topLevelWorkflowSpanEvent).to.deepEqualWithMockValues(expectedTopLevelWorkflow) + expect(firstSubWorkflowSpanEvent).to.deepEqualWithMockValues(expectedFirstSubWorkflow) + expect(firstLLMSpanEvent).to.deepEqualWithMockValues(expectedFirstLLM) + expect(secondSubWorkflowSpanEvent).to.deepEqualWithMockValues(expectedSecondSubWorkflow) + expect(secondLLMSpanEvent).to.deepEqualWithMockValues(expectedSecondLLM) + }) + + const result = await completeChain.invoke({ person: 'Abraham Lincoln', language: 'Spanish' }) + expect(result).to.equal('Springfield, Illinois está en los Estados Unidos.') + + await checkTraces + }) + + it('submits workflow and llm spans for a batched chain', async () => { + stubCall({ + ...openAiBaseChatInfo, + response: [ + { + model: 'gpt-4', + usage: { + prompt_tokens: 37, + completion_tokens: 10, + total_tokens: 47 + }, + choices: [{ + message: { + role: 'assistant', + content: 'Why did the chicken cross the road? To get to the other side!' + } + }] + }, + { + model: 'gpt-4', + usage: { + prompt_tokens: 37, + completion_tokens: 10, + total_tokens: 47 + }, + choices: [{ + message: { + role: 'assistant', + content: 'Why was the dog confused? It was barking up the wrong tree!' + } + }] + } + ] + }) + + const prompt = langchainPrompts.ChatPromptTemplate.fromTemplate( + 'Tell me a joke about {topic}' + ) + const parser = new langchainOutputParsers.StringOutputParser() + const model = new langchainOpenai.ChatOpenAI({ model: 'gpt-4' }) + + const chain = langchainRunnables.RunnableSequence.from([ + { + topic: new langchainRunnables.RunnablePassthrough() + }, + prompt, + model, + parser + ]) + + const checkTraces = agent.use(traces => { + const spans = traces[0] + + const workflowSpan = spans[0] + const firstLLMSpan = spans[1] + const secondLLMSpan = spans[2] + + const workflowSpanEvent = LLMObsAgentProxySpanWriter.prototype.append.getCall(0).args[0] + const firstLLMSpanEvent = LLMObsAgentProxySpanWriter.prototype.append.getCall(1).args[0] + const secondLLMSpanEvent = LLMObsAgentProxySpanWriter.prototype.append.getCall(2).args[0] + + const expectedWorkflow = expectedLLMObsNonLLMSpanEvent({ + span: workflowSpan, + spanKind: 'workflow', + name: 'langchain_core.runnables.RunnableSequence', + inputValue: JSON.stringify(['chickens', 'dogs']), + outputValue: JSON.stringify([ + 'Why did the chicken cross the road? To get to the other side!', + 'Why was the dog confused? It was barking up the wrong tree!' + ]), + tags: { ml_app: 'test', language: 'javascript' } + }) + + const expectedFirstLLM = expectedLLMObsLLMSpanEvent({ + span: firstLLMSpan, + parentId: workflowSpan.span_id, + spanKind: 'llm', + modelName: 'gpt-4', + modelProvider: 'openai', + name: 'langchain.chat_models.openai.ChatOpenAI', + inputMessages: [{ content: 'Tell me a joke about chickens', role: 'user' }], + outputMessages: [{ + content: 'Why did the chicken cross the road? To get to the other side!', + role: 'assistant' + }], + metadata: MOCK_ANY, + tokenMetrics: { input_tokens: 37, output_tokens: 10, total_tokens: 47 }, + tags: { ml_app: 'test', language: 'javascript' } + }) + + const expectedSecondLLM = expectedLLMObsLLMSpanEvent({ + span: secondLLMSpan, + parentId: workflowSpan.span_id, + spanKind: 'llm', + modelName: 'gpt-4', + modelProvider: 'openai', + name: 'langchain.chat_models.openai.ChatOpenAI', + inputMessages: [{ content: 'Tell me a joke about dogs', role: 'user' }], + outputMessages: [{ + content: 'Why was the dog confused? It was barking up the wrong tree!', + role: 'assistant' + }], + metadata: MOCK_ANY, + tokenMetrics: { input_tokens: 37, output_tokens: 10, total_tokens: 47 }, + tags: { ml_app: 'test', language: 'javascript' } + }) + + expect(workflowSpanEvent).to.deepEqualWithMockValues(expectedWorkflow) + expect(firstLLMSpanEvent).to.deepEqualWithMockValues(expectedFirstLLM) + expect(secondLLMSpanEvent).to.deepEqualWithMockValues(expectedSecondLLM) + }) + + await chain.batch(['chickens', 'dogs']) + + await checkTraces + }) + + it('submits a workflow and llm spans for different schema IO', async () => { + stubCall({ + ...openAiBaseChatInfo, + response: { + choices: [ + { + message: { + content: 'Mitochondria', + role: 'assistant' + } + } + ], + usage: { prompt_tokens: 8, completion_tokens: 12, total_tokens: 20 } + } + }) + + const prompt = langchainPrompts.ChatPromptTemplate.fromMessages([ + ['system', 'You are an assistant who is good at {ability}. Respond in 20 words or fewer'], + new langchainPrompts.MessagesPlaceholder('history'), + ['human', '{input}'] + ]) + + const model = new langchainOpenai.ChatOpenAI({ model: 'gpt-3.5-turbo' }) + const chain = prompt.pipe(model) + + const checkTraces = agent.use(traces => { + const spans = traces[0] + + const workflowSpan = spans[0] + const llmSpan = spans[1] + + const workflowSpanEvent = LLMObsAgentProxySpanWriter.prototype.append.getCall(0).args[0] + const llmSpanEvent = LLMObsAgentProxySpanWriter.prototype.append.getCall(1).args[0] + + const expectedWorkflow = expectedLLMObsNonLLMSpanEvent({ + span: workflowSpan, + spanKind: 'workflow', + name: 'langchain_core.runnables.RunnableSequence', + inputValue: JSON.stringify({ + ability: 'world capitals', + history: [ + { + content: 'Can you be my science teacher instead?', + role: 'user' + }, + { + content: 'Yes', + role: 'assistant' + } + ], + input: 'What is the powerhouse of the cell?' + }), + // takes the form of an AIMessage struct since there is no output parser + outputValue: JSON.stringify({ + content: 'Mitochondria', + role: 'assistant' + }), + tags: { ml_app: 'test', language: 'javascript' } + }) + + const expectedLLM = expectedLLMObsLLMSpanEvent({ + span: llmSpan, + parentId: workflowSpan.span_id, + spanKind: 'llm', + modelName: 'gpt-3.5-turbo', + modelProvider: 'openai', + name: 'langchain.chat_models.openai.ChatOpenAI', + inputMessages: [ + { + content: 'You are an assistant who is good at world capitals. Respond in 20 words or fewer', + role: 'system' + }, + { + content: 'Can you be my science teacher instead?', + role: 'user' + }, + { + content: 'Yes', + role: 'assistant' + }, + { + content: 'What is the powerhouse of the cell?', + role: 'user' + } + ], + outputMessages: [{ content: 'Mitochondria', role: 'assistant' }], + metadata: MOCK_ANY, + tokenMetrics: { input_tokens: 8, output_tokens: 12, total_tokens: 20 }, + tags: { ml_app: 'test', language: 'javascript' } + }) + + expect(workflowSpanEvent).to.deepEqualWithMockValues(expectedWorkflow) + expect(llmSpanEvent).to.deepEqualWithMockValues(expectedLLM) + }) + + await chain.invoke({ + ability: 'world capitals', + history: [ + new langchainMessages.HumanMessage('Can you be my science teacher instead?'), + new langchainMessages.AIMessage('Yes') + ], + input: 'What is the powerhouse of the cell?' + }) + + await checkTraces + }) + + it('traces a manually-instrumented step', async () => { + stubCall({ + ...openAiBaseChatInfo, + response: { + choices: [ + { + message: { + content: '3 squared is 9', + role: 'assistant' + } + } + ], + usage: { prompt_tokens: 8, completion_tokens: 12, total_tokens: 20 } + } + }) + + let lengthFunction = (input = { foo: '' }) => { + llmobs.annotate({ inputData: input }) // so we don't try and tag `config` with auto-annotation + return { + length: input.foo.length.toString() + } + } + lengthFunction = llmobs.wrap({ kind: 'task' }, lengthFunction) + + const model = new langchainOpenai.ChatOpenAI({ model: 'gpt-4o' }) + + const prompt = langchainPrompts.ChatPromptTemplate.fromTemplate('What is {length} squared?') + + const chain = langchainRunnables.RunnableLambda.from(lengthFunction) + .pipe(prompt) + .pipe(model) + .pipe(new langchainOutputParsers.StringOutputParser()) + + const checkTraces = agent.use(traces => { + const spans = traces[0] + expect(spans.length).to.equal(3) + + const workflowSpan = spans[0] + const taskSpan = spans[1] + const llmSpan = spans[2] + + const workflowSpanEvent = LLMObsAgentProxySpanWriter.prototype.append.getCall(0).args[0] + const taskSpanEvent = LLMObsAgentProxySpanWriter.prototype.append.getCall(1).args[0] + const llmSpanEvent = LLMObsAgentProxySpanWriter.prototype.append.getCall(2).args[0] + + const expectedWorkflow = expectedLLMObsNonLLMSpanEvent({ + span: workflowSpan, + spanKind: 'workflow', + name: 'langchain_core.runnables.RunnableSequence', + inputValue: JSON.stringify({ foo: 'bar' }), + outputValue: '3 squared is 9', + tags: { ml_app: 'test', language: 'javascript' } + }) + + const expectedTask = expectedLLMObsNonLLMSpanEvent({ + span: taskSpan, + parentId: workflowSpan.span_id, + spanKind: 'task', + name: 'lengthFunction', + inputValue: JSON.stringify({ foo: 'bar' }), + outputValue: JSON.stringify({ length: '3' }), + tags: { ml_app: 'test', language: 'javascript' } + }) + + const expectedLLM = expectedLLMObsLLMSpanEvent({ + span: llmSpan, + parentId: workflowSpan.span_id, + spanKind: 'llm', + modelName: 'gpt-4o', + modelProvider: 'openai', + name: 'langchain.chat_models.openai.ChatOpenAI', + inputMessages: [{ content: 'What is 3 squared?', role: 'user' }], + outputMessages: [{ content: '3 squared is 9', role: 'assistant' }], + metadata: MOCK_ANY, + tokenMetrics: { input_tokens: 8, output_tokens: 12, total_tokens: 20 }, + tags: { ml_app: 'test', language: 'javascript' } + }) + + expect(workflowSpanEvent).to.deepEqualWithMockValues(expectedWorkflow) + expect(taskSpanEvent).to.deepEqualWithMockValues(expectedTask) + expect(llmSpanEvent).to.deepEqualWithMockValues(expectedLLM) + }) + + await chain.invoke({ foo: 'bar' }) + + await checkTraces + }) + }) + }) + }) + }) +}) diff --git a/packages/dd-trace/test/llmobs/sdk/typescript/index.spec.js b/packages/dd-trace/test/llmobs/sdk/typescript/index.spec.js index b792a4fbdb..111123b136 100644 --- a/packages/dd-trace/test/llmobs/sdk/typescript/index.spec.js +++ b/packages/dd-trace/test/llmobs/sdk/typescript/index.spec.js @@ -105,7 +105,9 @@ describe('typescript', () => { for (const test of testCases) { const { name, file } = test - it(name, async () => { + it(name, async function () { + this.timeout(20000) + const cwd = sandbox.folder const results = {} diff --git a/packages/dd-trace/test/llmobs/tagger.spec.js b/packages/dd-trace/test/llmobs/tagger.spec.js index 783ce91bda..2be6a9a032 100644 --- a/packages/dd-trace/test/llmobs/tagger.spec.js +++ b/packages/dd-trace/test/llmobs/tagger.spec.js @@ -472,6 +472,29 @@ describe('tagger', () => { expect(() => tagger.tagTextIO(span, data, 'output')).to.throw() }) }) + + describe('changeKind', () => { + it('changes the span kind', () => { + tagger._register(span) + tagger._setTag(span, '_ml_obs.meta.span.kind', 'old-kind') + expect(Tagger.tagMap.get(span)).to.deep.equal({ + '_ml_obs.meta.span.kind': 'old-kind' + }) + tagger.changeKind(span, 'new-kind') + expect(Tagger.tagMap.get(span)).to.deep.equal({ + '_ml_obs.meta.span.kind': 'new-kind' + }) + }) + + it('sets the kind if it is not already set', () => { + tagger._register(span) + expect(Tagger.tagMap.get(span)).to.deep.equal({}) + tagger.changeKind(span, 'new-kind') + expect(Tagger.tagMap.get(span)).to.deep.equal({ + '_ml_obs.meta.span.kind': 'new-kind' + }) + }) + }) }) describe('with softFail', () => { diff --git a/packages/dd-trace/test/llmobs/util.js b/packages/dd-trace/test/llmobs/util.js index 0106c9dd64..ba3eeb4914 100644 --- a/packages/dd-trace/test/llmobs/util.js +++ b/packages/dd-trace/test/llmobs/util.js @@ -120,7 +120,7 @@ function expectedLLMObsBaseEvent ({ const spanEvent = { trace_id: MOCK_STRING, span_id: spanId, - parent_id: parentId || 'undefined', + parent_id: parentId?.buffer ? fromBuffer(parentId) : (parentId || 'undefined'), name: spanName, tags: expectedLLMObsTags({ span, tags, error, errorType, sessionId }), start_ns: startNs, diff --git a/packages/dd-trace/test/llmobs/util.spec.js b/packages/dd-trace/test/llmobs/util.spec.js index 063e618c1e..772e4a5061 100644 --- a/packages/dd-trace/test/llmobs/util.spec.js +++ b/packages/dd-trace/test/llmobs/util.spec.js @@ -3,7 +3,8 @@ const { encodeUnicode, getFunctionArguments, - validateKind + validateKind, + spanHasError } = require('../../src/llmobs/util') describe('util', () => { @@ -139,4 +140,38 @@ describe('util', () => { }) }) }) + + describe('spanHasError', () => { + let Span + let ps + + before(() => { + Span = require('../../src/opentracing/span') + ps = { + sample () {} + } + }) + + it('returns false when there is no error', () => { + const span = new Span(null, null, ps, {}) + expect(spanHasError(span)).to.equal(false) + }) + + it('returns true if the span has an "error" tag', () => { + const span = new Span(null, null, ps, {}) + span.setTag('error', true) + expect(spanHasError(span)).to.equal(true) + }) + + it('returns true if the span has the error properties as tags', () => { + const err = new Error('boom') + const span = new Span(null, null, ps, {}) + + span.setTag('error.type', err.name) + span.setTag('error.msg', err.message) + span.setTag('error.stack', err.stack) + + expect(spanHasError(span)).to.equal(true) + }) + }) }) diff --git a/packages/dd-trace/test/plugins/externals.json b/packages/dd-trace/test/plugins/externals.json index 600df395d8..4988256986 100644 --- a/packages/dd-trace/test/plugins/externals.json +++ b/packages/dd-trace/test/plugins/externals.json @@ -275,6 +275,10 @@ { "name": "@langchain/anthropic", "versions": [">=0.1"] + }, + { + "name": "@langchain/cohere", + "versions": [">=0.1"] } ], "ldapjs": [