Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
feat: Add instrumentation for RunnableSequence (#1996)
Co-authored-by: Bob Evans <robert.evans25@gmail.com> Co-authored-by: Bob Evans <revans@newrelic.com> Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com> Co-authored-by: James Sumners <jsumners@newrelic.com>
- Loading branch information
5 people
committed
Feb 22, 2024
1 parent
ead8c25
commit 71ffa37
Showing
10 changed files
with
664 additions
and
11 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,98 @@ | ||
| /* | ||
| * Copyright 2024 New Relic Corporation. All rights reserved. | ||
| * SPDX-License-Identifier: Apache-2.0 | ||
| */ | ||
|
|
||
| 'use strict' | ||
|
|
||
| const common = require('./common') | ||
| const { | ||
| AI: { LANGCHAIN } | ||
| } = require('../../metrics/names') | ||
| const { | ||
| LangChainCompletionMessage, | ||
| LangChainCompletionSummary | ||
| } = require('../../llm-events/langchain/') | ||
| const { DESTINATIONS } = require('../../config/attribute-filter') | ||
| const { langchainRunId } = require('../../symbols') | ||
| const { RecorderSpec } = require('../../shim/specs') | ||
|
|
||
| module.exports = function initialize(shim, langchain) { | ||
| const { agent, pkgVersion } = shim | ||
|
|
||
| if (common.shouldSkipInstrumentation(agent.config)) { | ||
| shim.logger.debug( | ||
| 'langchain instrumentation is disabled. To enable set `config.ai_monitoring.enabled` to true' | ||
| ) | ||
| return | ||
| } | ||
|
|
||
| shim.record( | ||
| langchain.RunnableSequence.prototype, | ||
| 'invoke', | ||
| function wrapCall(shim, invoke, fnName, args) { | ||
| const [request, params] = args | ||
| const metadata = params?.metadata ?? {} | ||
| const tags = params?.tags ?? [] | ||
|
|
||
| return new RecorderSpec({ | ||
| name: `${LANGCHAIN.CHAIN}/${fnName}`, | ||
| promise: true, | ||
| // eslint-disable-next-line max-params | ||
| after(_shim, _fn, _name, _err, output, segment) { | ||
| segment.end() | ||
| const completionSummary = new LangChainCompletionSummary({ | ||
| agent, | ||
| messages: [{ output }], | ||
| metadata, | ||
| tags, | ||
| segment, | ||
| runId: segment[langchainRunId] | ||
| }) | ||
|
|
||
| common.recordEvent({ | ||
| agent, | ||
| type: 'LlmChatCompletionSummary', | ||
| pkgVersion, | ||
| msg: completionSummary | ||
| }) | ||
|
|
||
| const data = [request, output] | ||
|
|
||
| // output can be BaseMessage with a content property https://js.langchain.com/docs/modules/model_io/concepts#messages | ||
| // or an output parser https://js.langchain.com/docs/modules/model_io/concepts#output-parsers | ||
| data.forEach((msg, sequence) => { | ||
| if (msg?.content) { | ||
| msg = msg.content | ||
| } | ||
|
|
||
| let msgString | ||
| try { | ||
| msgString = typeof msg === 'string' ? msg : JSON.stringify(msg) | ||
| } catch (err) { | ||
| shim.logger.error(err, 'Failed to stringify message') | ||
| msgString = '' | ||
| } | ||
|
|
||
| const completionMsg = new LangChainCompletionMessage({ | ||
| sequence, | ||
| agent, | ||
| content: msgString, | ||
| completionId: completionSummary.id, | ||
| segment, | ||
| runId: segment[langchainRunId] | ||
| }) | ||
|
|
||
| common.recordEvent({ | ||
| agent, | ||
| type: 'LlmChatCompletionMessage', | ||
| pkgVersion, | ||
| msg: completionMsg | ||
| }) | ||
| }) | ||
| segment.transaction.trace.attributes.addAttribute(DESTINATIONS.TRANS_EVENT, 'llm', true) | ||
| } | ||
| }) | ||
| } | ||
| ) | ||
| } |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,68 @@ | ||
| /* | ||
| * Copyright 2023 New Relic Corporation. All rights reserved. | ||
| * SPDX-License-Identifier: Apache-2.0 | ||
| */ | ||
|
|
||
| 'use strict' | ||
|
|
||
| const { test } = require('tap') | ||
| const helper = require('../../../lib/agent_helper') | ||
| const GenericShim = require('../../../../lib/shim/shim') | ||
| const sinon = require('sinon') | ||
|
|
||
| test('langchain/core/runnables unit tests', (t) => { | ||
| t.beforeEach(function (t) { | ||
| const sandbox = sinon.createSandbox() | ||
| const agent = helper.loadMockedAgent() | ||
| agent.config.ai_monitoring = { enabled: true } | ||
| agent.config.feature_flag = { langchain_instrumentation: true } | ||
| const shim = new GenericShim(agent, 'langchain') | ||
| shim.pkgVersion = '0.1.26' | ||
| sandbox.stub(shim.logger, 'debug') | ||
| sandbox.stub(shim.logger, 'warn') | ||
|
|
||
| t.context.agent = agent | ||
| t.context.shim = shim | ||
| t.context.sandbox = sandbox | ||
| t.context.initialize = require('../../../../lib/instrumentation/langchain/runnable') | ||
| }) | ||
|
|
||
| t.afterEach(function (t) { | ||
| helper.unloadAgent(t.context.agent) | ||
| t.context.sandbox.restore() | ||
| }) | ||
|
|
||
| function getMockModule() { | ||
| function RunnableSequence() {} | ||
| RunnableSequence.prototype.invoke = async function call() {} | ||
| return { RunnableSequence } | ||
| } | ||
|
|
||
| ;[ | ||
| { aiMonitoring: false, langChain: true }, | ||
| { aiMonitoring: true, langChain: false }, | ||
| { aiMonitoring: false, langChain: false } | ||
| ].forEach(({ aiMonitoring, langChain }) => { | ||
| t.test( | ||
| `should not register instrumentation if ai_monitoring is ${aiMonitoring} and langchain_instrumentation is ${langChain}`, | ||
| (t) => { | ||
| const { shim, agent, initialize } = t.context | ||
| const MockRunnable = getMockModule() | ||
| agent.config.ai_monitoring.enabled = aiMonitoring | ||
| agent.config.feature_flag.langchain_instrumentation = langChain | ||
|
|
||
| initialize(shim, MockRunnable) | ||
| t.equal(shim.logger.debug.callCount, 1, 'should log 1 debug messages') | ||
| t.equal( | ||
| shim.logger.debug.args[0][0], | ||
| 'langchain instrumentation is disabled. To enable set `config.ai_monitoring.enabled` to true' | ||
| ) | ||
| const isWrapped = shim.isWrapped(MockRunnable.RunnableSequence.prototype.invoke) | ||
| t.equal(isWrapped, false, 'should not wrap runnable invoke') | ||
| t.end() | ||
| } | ||
| ) | ||
| }) | ||
|
|
||
| t.end() | ||
| }) |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,91 @@ | ||
| /* | ||
| * Copyright 2024 New Relic Corporation. All rights reserved. | ||
| * SPDX-License-Identifier: Apache-2.0 | ||
| */ | ||
|
|
||
| 'use strict' | ||
|
|
||
| const tap = require('tap') | ||
|
|
||
| function filterLangchainEvents(events) { | ||
| return events.filter((event) => { | ||
| const [, chainEvent] = event | ||
| return chainEvent.vendor === 'langchain' | ||
| }) | ||
| } | ||
|
|
||
| function filterLangchainMessages(events, msgType) { | ||
| return events.filter((event) => { | ||
| const [{ type }] = event | ||
| return type === msgType | ||
| }) | ||
| } | ||
|
|
||
| function assertLangChainChatCompletionSummary({ tx, chatSummary, withCallback }) { | ||
| const expectedSummary = { | ||
| 'id': /[a-f0-9]{36}/, | ||
| 'appName': 'New Relic for Node.js tests', | ||
| 'span_id': tx.trace.root.children[0].id, | ||
| 'trace_id': tx.traceId, | ||
| 'transaction_id': tx.id, | ||
| 'request_id': undefined, | ||
| 'ingest_source': 'Node', | ||
| 'vendor': 'langchain', | ||
| 'metadata.key': 'value', | ||
| 'metadata.hello': 'world', | ||
| 'tags': 'tag1,tag2', | ||
| 'virtual_llm': true, | ||
| ['response.number_of_messages']: 1, | ||
| 'duration': tx.trace.root.children[0].getDurationInMillis() | ||
| } | ||
|
|
||
| if (withCallback) { | ||
| expectedSummary.request_id = /[a-f0-9\-]{36}/ | ||
| expectedSummary.id = /[a-f0-9\-]{36}/ | ||
| } | ||
|
|
||
| this.equal(chatSummary[0].type, 'LlmChatCompletionSummary') | ||
| this.match(chatSummary[1], expectedSummary, 'should match chat summary message') | ||
| } | ||
|
|
||
| function assertLangChainChatCompletionMessages({ tx, chatMsgs, chatSummary, withCallback }) { | ||
| const baseMsg = { | ||
| id: /[a-f0-9]{36}/, | ||
| appName: 'New Relic for Node.js tests', | ||
| span_id: tx.trace.root.children[0].id, | ||
| trace_id: tx.traceId, | ||
| transaction_id: tx.id, | ||
| ingest_source: 'Node', | ||
| vendor: 'langchain', | ||
| completion_id: chatSummary.id, | ||
| virtual_llm: true, | ||
| request_id: undefined | ||
| } | ||
|
|
||
| if (withCallback) { | ||
| baseMsg.request_id = /[a-f0-9\-]{36}/ | ||
| baseMsg.id = /[a-f0-9\-]{36}/ | ||
| } | ||
|
|
||
| chatMsgs.forEach((msg) => { | ||
| const expectedChatMsg = { ...baseMsg } | ||
| if (msg[1].sequence === 0) { | ||
| expectedChatMsg.sequence = 0 | ||
| expectedChatMsg.content = '{"topic":"scientist"}' | ||
| } else if (msg[1].sequence === 1) { | ||
| expectedChatMsg.sequence = 1 | ||
| expectedChatMsg.content = '212 degrees Fahrenheit is equal to 100 degrees Celsius.' | ||
| } | ||
|
|
||
| this.equal(msg[0].type, 'LlmChatCompletionMessage') | ||
| this.match(msg[1], expectedChatMsg, 'should match chat completion message') | ||
| }) | ||
| } | ||
|
|
||
| tap.Test.prototype.addAssert('langchainMessages', 1, assertLangChainChatCompletionMessages) | ||
| tap.Test.prototype.addAssert('langchainSummary', 1, assertLangChainChatCompletionSummary) | ||
|
|
||
| module.exports = { | ||
| filterLangchainEvents, | ||
| filterLangchainMessages | ||
| } |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Oops, something went wrong.