Skip to content

Commit 6523d94

Browse files
sabrennersabrenner
authored
[MLOB-1858] feat(llmobs): langchain submits llmobs span events (#4923)
* llmobs langchain plugin * starting test changes * refactor llmobs langchain plugin to use handlers * wip * finish adding most tests * higher timeout ts tests * add missing tests * change lowerbound node version for langchain test to 18 * add cohere test for newer versions * add externals * add error tests * more consistent parentage behavior * review comments
1 parent 95f82a9 commit 6523d94

21 files changed

Lines changed: 1796 additions & 98 deletions

File tree

‎.github/workflows/llmobs.yml‎

Lines changed: 19 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -47,3 +47,22 @@ jobs:
4747
- uses: codecov/codecov-action@v3
4848
- if: always()
4949
uses: ./.github/actions/testagent/logs
50+
51+
langchain:
52+
runs-on: ubuntu-latest
53+
env:
54+
PLUGINS: langchain
55+
steps:
56+
- uses: actions/checkout@v4
57+
- uses: ./.github/actions/testagent/start
58+
- uses: ./.github/actions/node/setup
59+
- uses: ./.github/actions/install
60+
- uses: ./.github/actions/node/18
61+
- run: yarn test:llmobs:plugins:ci
62+
shell: bash
63+
- uses: ./.github/actions/node/latest
64+
- run: yarn test:llmobs:plugins:ci
65+
shell: bash
66+
- uses: codecov/codecov-action@v3
67+
- if: always()
68+
uses: ./.github/actions/testagent/logs

‎package.json‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -33,7 +33,7 @@
3333
"test:lambda:ci": "nyc --no-clean --include \"packages/dd-trace/src/lambda/**/*.js\" -- npm run test:lambda",
3434
"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\" ",
3535
"test:llmobs:sdk:ci": "nyc --no-clean --include \"packages/dd-trace/src/llmobs/**/*.js\" -- npm run test:llmobs:sdk",
36-
"test:llmobs:plugins": "mocha -r \"packages/dd-trace/test/setup/mocha.js\" \"packages/dd-trace/test/llmobs/plugins/**/*.spec.js\"",
36+
"test:llmobs:plugins": "mocha -r \"packages/dd-trace/test/setup/mocha.js\" \"packages/dd-trace/test/llmobs/plugins/@($(echo $PLUGINS))/*.spec.js\"",
3737
"test:llmobs:plugins:ci": "yarn services && nyc --no-clean --include \"packages/dd-trace/src/llmobs/**/*.js\" -- npm run test:llmobs:plugins",
3838
"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\"",
3939
"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",

‎packages/datadog-instrumentations/src/openai.js‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -338,6 +338,8 @@ for (const shim of V4_PACKAGE_SHIMS) {
338338
})
339339
})
340340

341+
ch.end.publish(ctx)
342+
341343
return apiProm
342344
})
343345
})

‎packages/datadog-plugin-langchain/src/index.js‎

Lines changed: 12 additions & 80 deletions
Original file line numberDiff line numberDiff line change
@@ -1,89 +1,21 @@
11
'use strict'
22

3-
const { MEASURED } = require('../../../ext/tags')
4-
const { storage } = require('../../datadog-core')
5-
const TracingPlugin = require('../../dd-trace/src/plugins/tracing')
3+
const LangChainTracingPlugin = require('./tracing')
4+
const LangChainLLMObsPlugin = require('../../dd-trace/src/llmobs/plugins/langchain')
5+
const CompositePlugin = require('../../dd-trace/src/plugins/composite')
66

7-
const API_KEY = 'langchain.request.api_key'
8-
const MODEL = 'langchain.request.model'
9-
const PROVIDER = 'langchain.request.provider'
10-
const TYPE = 'langchain.request.type'
11-
12-
const LangChainHandler = require('./handlers/default')
13-
const LangChainChatModelHandler = require('./handlers/language_models/chat_model')
14-
const LangChainLLMHandler = require('./handlers/language_models/llm')
15-
const LangChainChainHandler = require('./handlers/chain')
16-
const LangChainEmbeddingHandler = require('./handlers/embedding')
17-
18-
class LangChainPlugin extends TracingPlugin {
7+
class LangChainPlugin extends CompositePlugin {
198
static get id () { return 'langchain' }
20-
static get operation () { return 'invoke' }
21-
static get system () { return 'langchain' }
22-
static get prefix () {
23-
return 'tracing:apm:langchain:invoke'
24-
}
25-
26-
constructor () {
27-
super(...arguments)
28-
29-
const langchainConfig = this._tracerConfig.langchain || {}
30-
this.handlers = {
31-
chain: new LangChainChainHandler(langchainConfig),
32-
chat_model: new LangChainChatModelHandler(langchainConfig),
33-
llm: new LangChainLLMHandler(langchainConfig),
34-
embedding: new LangChainEmbeddingHandler(langchainConfig),
35-
default: new LangChainHandler(langchainConfig)
9+
static get plugins () {
10+
return {
11+
// ordering here is important - the llm observability plugin must come first
12+
// so that we can add annotations associated with the span before it finishes.
13+
// however, because the tracing plugin uses `bindStart` vs the llmobs' `start`,
14+
// the span is guaranteed to be created in the tracing plugin before the llmobs one is called
15+
llmobs: LangChainLLMObsPlugin,
16+
tracing: LangChainTracingPlugin
3617
}
3718
}
38-
39-
bindStart (ctx) {
40-
const { resource, type } = ctx
41-
const handler = this.handlers[type]
42-
43-
const instance = ctx.instance
44-
const apiKey = handler.extractApiKey(instance)
45-
const provider = handler.extractProvider(instance)
46-
const model = handler.extractModel(instance)
47-
48-
const tags = handler.getSpanStartTags(ctx, provider) || []
49-
50-
if (apiKey) tags[API_KEY] = apiKey
51-
if (provider) tags[PROVIDER] = provider
52-
if (model) tags[MODEL] = model
53-
if (type) tags[TYPE] = type
54-
55-
const span = this.startSpan('langchain.request', {
56-
service: this.config.service,
57-
resource,
58-
kind: 'client',
59-
meta: {
60-
[MEASURED]: 1,
61-
...tags
62-
}
63-
}, false)
64-
65-
const store = storage.getStore() || {}
66-
ctx.currentStore = { ...store, span }
67-
68-
return ctx.currentStore
69-
}
70-
71-
asyncEnd (ctx) {
72-
const span = ctx.currentStore.span
73-
74-
const { type } = ctx
75-
76-
const handler = this.handlers[type]
77-
const tags = handler.getSpanEndTags(ctx) || {}
78-
79-
span.addTags(tags)
80-
81-
span.finish()
82-
}
83-
84-
getHandler (type) {
85-
return this.handlers[type] || this.handlers.default
86-
}
8719
}
8820

8921
module.exports = LangChainPlugin

‎packages/datadog-plugin-langchain/src/tracing.js‎

Lines changed: 89 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,89 @@
1+
'use strict'
2+
3+
const { MEASURED } = require('../../../ext/tags')
4+
const { storage } = require('../../datadog-core')
5+
const TracingPlugin = require('../../dd-trace/src/plugins/tracing')
6+
7+
const API_KEY = 'langchain.request.api_key'
8+
const MODEL = 'langchain.request.model'
9+
const PROVIDER = 'langchain.request.provider'
10+
const TYPE = 'langchain.request.type'
11+
12+
const LangChainHandler = require('./handlers/default')
13+
const LangChainChatModelHandler = require('./handlers/language_models/chat_model')
14+
const LangChainLLMHandler = require('./handlers/language_models/llm')
15+
const LangChainChainHandler = require('./handlers/chain')
16+
const LangChainEmbeddingHandler = require('./handlers/embedding')
17+
18+
class LangChainTracingPlugin extends TracingPlugin {
19+
static get id () { return 'langchain' }
20+
static get operation () { return 'invoke' }
21+
static get system () { return 'langchain' }
22+
static get prefix () {
23+
return 'tracing:apm:langchain:invoke'
24+
}
25+
26+
constructor () {
27+
super(...arguments)
28+
29+
const langchainConfig = this._tracerConfig.langchain || {}
30+
this.handlers = {
31+
chain: new LangChainChainHandler(langchainConfig),
32+
chat_model: new LangChainChatModelHandler(langchainConfig),
33+
llm: new LangChainLLMHandler(langchainConfig),
34+
embedding: new LangChainEmbeddingHandler(langchainConfig),
35+
default: new LangChainHandler(langchainConfig)
36+
}
37+
}
38+
39+
bindStart (ctx) {
40+
const { resource, type } = ctx
41+
const handler = this.handlers[type]
42+
43+
const instance = ctx.instance
44+
const apiKey = handler.extractApiKey(instance)
45+
const provider = handler.extractProvider(instance)
46+
const model = handler.extractModel(instance)
47+
48+
const tags = handler.getSpanStartTags(ctx, provider) || []
49+
50+
if (apiKey) tags[API_KEY] = apiKey
51+
if (provider) tags[PROVIDER] = provider
52+
if (model) tags[MODEL] = model
53+
if (type) tags[TYPE] = type
54+
55+
const span = this.startSpan('langchain.request', {
56+
service: this.config.service,
57+
resource,
58+
kind: 'client',
59+
meta: {
60+
[MEASURED]: 1,
61+
...tags
62+
}
63+
}, false)
64+
65+
const store = storage.getStore() || {}
66+
ctx.currentStore = { ...store, span }
67+
68+
return ctx.currentStore
69+
}
70+
71+
asyncEnd (ctx) {
72+
const span = ctx.currentStore.span
73+
74+
const { type } = ctx
75+
76+
const handler = this.handlers[type]
77+
const tags = handler.getSpanEndTags(ctx) || {}
78+
79+
span.addTags(tags)
80+
81+
span.finish()
82+
}
83+
84+
getHandler (type) {
85+
return this.handlers[type] || this.handlers.default
86+
}
87+
}
88+
89+
module.exports = LangChainTracingPlugin

‎packages/dd-trace/src/llmobs/plugins/base.js‎

Lines changed: 40 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -1,37 +1,60 @@
11
'use strict'
22

33
const log = require('../../log')
4-
const { storage } = require('../storage')
4+
const { storage: llmobsStorage } = require('../storage')
55

66
const TracingPlugin = require('../../plugins/tracing')
77
const LLMObsTagger = require('../tagger')
88

9-
// we make this a `Plugin` so we don't have to worry about `finish` being called
109
class LLMObsPlugin extends TracingPlugin {
1110
constructor (...args) {
1211
super(...args)
1312

1413
this._tagger = new LLMObsTagger(this._tracerConfig, true)
1514
}
1615

17-
getName () {}
18-
1916
setLLMObsTags (ctx) {
2017
throw new Error('setLLMObsTags must be implemented by the subclass')
2118
}
2219

23-
getLLMObsSPanRegisterOptions (ctx) {
20+
getLLMObsSpanRegisterOptions (ctx) {
2421
throw new Error('getLLMObsSPanRegisterOptions must be implemented by the subclass')
2522
}
2623

2724
start (ctx) {
28-
const oldStore = storage.getStore()
29-
const parent = oldStore?.span
30-
const span = ctx.currentStore?.span
25+
// even though llmobs span events won't be enqueued if llmobs is disabled
26+
// we should avoid doing any computations here (these listeners aren't disabled)
27+
const enabled = this._tracerConfig.llmobs.enabled
28+
if (!enabled) return
29+
30+
const parent = this.getLLMObsParent(ctx)
31+
const apmStore = ctx.currentStore
32+
const span = apmStore?.span
33+
34+
const registerOptions = this.getLLMObsSpanRegisterOptions(ctx)
35+
36+
// register options may not be set for operations we do not trace with llmobs
37+
// ie OpenAI fine tuning jobs, file jobs, etc.
38+
if (registerOptions) {
39+
ctx.llmobs = {} // initialize context-based namespace
40+
llmobsStorage.enterWith({ span })
41+
ctx.llmobs.parent = parent
3142

32-
const registerOptions = this.getLLMObsSPanRegisterOptions(ctx)
43+
this._tagger.registerLLMObsSpan(span, { parent, ...registerOptions })
44+
}
45+
}
46+
47+
end (ctx) {
48+
const enabled = this._tracerConfig.llmobs.enabled
49+
if (!enabled) return
50+
51+
// only attempt to restore the context if the current span was an LLMObs span
52+
const apmStore = ctx.currentStore
53+
const span = apmStore?.span
54+
if (!LLMObsTagger.tagMap.has(span)) return
3355

34-
this._tagger.registerLLMObsSpan(span, { parent, ...registerOptions })
56+
const parent = ctx.llmobs.parent
57+
llmobsStorage.enterWith({ span: parent })
3558
}
3659

3760
asyncEnd (ctx) {
@@ -40,7 +63,8 @@ class LLMObsPlugin extends TracingPlugin {
4063
const enabled = this._tracerConfig.llmobs.enabled
4164
if (!enabled) return
4265

43-
const span = ctx.currentStore?.span
66+
const apmStore = ctx.currentStore
67+
const span = apmStore?.span
4468
if (!span) {
4569
log.debug(
4670
`Tried to start an LLMObs span for ${this.constructor.name} without an active APM span.
@@ -60,6 +84,11 @@ class LLMObsPlugin extends TracingPlugin {
6084
}
6185
super.configure(config)
6286
}
87+
88+
getLLMObsParent () {
89+
const store = llmobsStorage.getStore()
90+
return store?.span
91+
}
6392
}
6493

6594
module.exports = LLMObsPlugin

‎packages/dd-trace/src/llmobs/plugins/langchain/handlers/chain.js‎

Lines changed: 24 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,24 @@
1+
'use strict'
2+
3+
const LangChainLLMObsHandler = require('.')
4+
const { spanHasError } = require('../../../util')
5+
6+
class LangChainLLMObsChainHandler extends LangChainLLMObsHandler {
7+
setMetaTags ({ span, inputs, results }) {
8+
let input, output
9+
if (inputs) {
10+
input = this.formatIO(inputs)
11+
}
12+
13+
if (!results || spanHasError(span)) {
14+
output = ''
15+
} else {
16+
output = this.formatIO(results)
17+
}
18+
19+
// chain spans will always be workflows
20+
this._tagger.tagTextIO(span, input, output)
21+
}
22+
}
23+
24+
module.exports = LangChainLLMObsChainHandler

0 commit comments

Comments
 (0)