Skip to content

Commit 936682e

Browse files
OlegIvanivOlegIvaniv
authored
feat(core): Improve Langsmith traces for AI executions (#9081)
Signed-off-by: Oleg Ivaniv <[email protected]>
1 parent 3bcfef9 commit 936682e

File tree

18 files changed

+99
-26
lines changed

18 files changed

+99
-26
lines changed

‎packages/@n8n/nodes-langchain/nodes/agents/Agent/agents/ConversationalAgent/execute.ts

Lines changed: 6 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -1,9 +1,5 @@
1-
import {
2-
type IExecuteFunctions,
3-
type INodeExecutionData,
4-
NodeConnectionType,
5-
NodeOperationError,
6-
} from 'n8n-workflow';
1+
import { NodeConnectionType, NodeOperationError } from 'n8n-workflow';
2+
import type { IExecuteFunctions, INodeExecutionData } from 'n8n-workflow';
73

84
import { initializeAgentExecutorWithOptions } from 'langchain/agents';
95
import type { BaseChatMemory } from '@langchain/community/memory/chat_memory';
@@ -16,13 +12,13 @@ import {
1612
getOptionalOutputParsers,
1713
getConnectedTools,
1814
} from '../../../../../utils/helpers';
15+
import { getTracingConfig } from '../../../../../utils/tracing';
1916

2017
export async function conversationalAgentExecute(
2118
this: IExecuteFunctions,
2219
nodeVersion: number,
2320
): Promise<INodeExecutionData[][]> {
2421
this.logger.verbose('Executing Conversational Agent');
25-
2622
const model = await this.getInputConnectionData(NodeConnectionType.AiLanguageModel, 0);
2723

2824
if (!isChatInstance(model)) {
@@ -104,7 +100,9 @@ export async function conversationalAgentExecute(
104100
input = (await prompt.invoke({ input })).value;
105101
}
106102

107-
let response = await agentExecutor.call({ input, outputParsers });
103+
let response = await agentExecutor
104+
.withConfig(getTracingConfig(this))
105+
.invoke({ input, outputParsers });
108106

109107
if (outputParser) {
110108
response = { output: await outputParser.parse(response.output as string) };

‎packages/@n8n/nodes-langchain/nodes/agents/Agent/agents/OpenAiFunctionsAgent/execute.ts

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,7 @@ import {
1717
getOptionalOutputParsers,
1818
getPromptInputByType,
1919
} from '../../../../../utils/helpers';
20+
import { getTracingConfig } from '../../../../../utils/tracing';
2021

2122
export async function openAiFunctionsAgentExecute(
2223
this: IExecuteFunctions,
@@ -104,7 +105,9 @@ export async function openAiFunctionsAgentExecute(
104105
input = (await prompt.invoke({ input })).value;
105106
}
106107

107-
let response = await agentExecutor.call({ input, outputParsers });
108+
let response = await agentExecutor
109+
.withConfig(getTracingConfig(this))
110+
.invoke({ input, outputParsers });
108111

109112
if (outputParser) {
110113
response = { output: await outputParser.parse(response.output as string) };

‎packages/@n8n/nodes-langchain/nodes/agents/Agent/agents/PlanAndExecuteAgent/execute.ts

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -15,6 +15,7 @@ import {
1515
getOptionalOutputParsers,
1616
getPromptInputByType,
1717
} from '../../../../../utils/helpers';
18+
import { getTracingConfig } from '../../../../../utils/tracing';
1819

1920
export async function planAndExecuteAgentExecute(
2021
this: IExecuteFunctions,
@@ -79,7 +80,9 @@ export async function planAndExecuteAgentExecute(
7980
input = (await prompt.invoke({ input })).value;
8081
}
8182

82-
let response = await agentExecutor.call({ input, outputParsers });
83+
let response = await agentExecutor
84+
.withConfig(getTracingConfig(this))
85+
.invoke({ input, outputParsers });
8386

8487
if (outputParser) {
8588
response = { output: await outputParser.parse(response.output as string) };

‎packages/@n8n/nodes-langchain/nodes/agents/Agent/agents/ReActAgent/execute.ts

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,7 @@ import {
1717
getPromptInputByType,
1818
isChatInstance,
1919
} from '../../../../../utils/helpers';
20+
import { getTracingConfig } from '../../../../../utils/tracing';
2021

2122
export async function reActAgentAgentExecute(
2223
this: IExecuteFunctions,
@@ -100,7 +101,10 @@ export async function reActAgentAgentExecute(
100101
input = (await prompt.invoke({ input })).value;
101102
}
102103

103-
let response = await agentExecutor.call({ input, outputParsers });
104+
let response = await agentExecutor
105+
.withConfig(getTracingConfig(this))
106+
.invoke({ input, outputParsers });
107+
104108
if (outputParser) {
105109
response = { output: await outputParser.parse(response.output as string) };
106110
}

‎packages/@n8n/nodes-langchain/nodes/agents/Agent/agents/SqlAgent/execute.ts

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -14,6 +14,7 @@ import type { BaseChatMemory } from '@langchain/community/memory/chat_memory';
1414
import type { DataSource } from '@n8n/typeorm';
1515

1616
import { getPromptInputByType, serializeChatHistory } from '../../../../../utils/helpers';
17+
import { getTracingConfig } from '../../../../../utils/tracing';
1718
import { getSqliteDataSource } from './other/handlers/sqlite';
1819
import { getPostgresDataSource } from './other/handlers/postgres';
1920
import { SQL_PREFIX, SQL_SUFFIX } from './other/prompts';
@@ -126,7 +127,7 @@ export async function sqlAgentAgentExecute(
126127

127128
let response: IDataObject;
128129
try {
129-
response = await agentExecutor.call({
130+
response = await agentExecutor.withConfig(getTracingConfig(this)).invoke({
130131
input,
131132
signal: this.getExecutionCancelSignal(),
132133
chatHistory,

‎packages/@n8n/nodes-langchain/nodes/agents/OpenAiAssistant/OpenAiAssistant.node.ts

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -10,6 +10,7 @@ import type {
1010
} from 'n8n-workflow';
1111
import type { OpenAIToolType } from 'langchain/dist/experimental/openai_assistant/schema';
1212
import { getConnectedTools } from '../../../utils/helpers';
13+
import { getTracingConfig } from '../../../utils/tracing';
1314
import { formatToOpenAIAssistantTool } from './utils';
1415

1516
export class OpenAiAssistant implements INodeType {
@@ -373,7 +374,7 @@ export class OpenAiAssistant implements INodeType {
373374
tools,
374375
});
375376

376-
const response = await agentExecutor.call({
377+
const response = await agentExecutor.withConfig(getTracingConfig(this)).invoke({
377378
content: input,
378379
signal: this.getExecutionCancelSignal(),
379380
timeout: options.timeout ?? 10000,

‎packages/@n8n/nodes-langchain/nodes/chains/ChainLLM/ChainLlm.node.ts

Lines changed: 6 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -27,6 +27,7 @@ import {
2727
getPromptInputByType,
2828
isChatInstance,
2929
} from '../../../utils/helpers';
30+
import { getTracingConfig } from '../../../utils/tracing';
3031

3132
interface MessagesTemplate {
3233
type: string;
@@ -154,9 +155,9 @@ async function createSimpleLLMChain(
154155
const chain = new LLMChain({
155156
llm,
156157
prompt,
157-
});
158+
}).withConfig(getTracingConfig(context));
158159

159-
const response = (await chain.call({
160+
const response = (await chain.invoke({
160161
query,
161162
signal: context.getExecutionCancelSignal(),
162163
})) as string[];
@@ -203,8 +204,9 @@ async function getChain(
203204
);
204205

205206
const chain = prompt.pipe(llm).pipe(combinedOutputParser);
206-
207-
const response = (await chain.invoke({ query })) as string | string[];
207+
const response = (await chain.withConfig(getTracingConfig(context)).invoke({ query })) as
208+
| string
209+
| string[];
208210

209211
return Array.isArray(response) ? response : [response];
210212
}

‎packages/@n8n/nodes-langchain/nodes/chains/ChainRetrievalQA/ChainRetrievalQa.node.ts

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -12,6 +12,7 @@ import type { BaseLanguageModel } from '@langchain/core/language_models/base';
1212
import type { BaseRetriever } from '@langchain/core/retrievers';
1313
import { getTemplateNoticeField } from '../../../utils/sharedFields';
1414
import { getPromptInputByType } from '../../../utils/helpers';
15+
import { getTracingConfig } from '../../../utils/tracing';
1516

1617
export class ChainRetrievalQa implements INodeType {
1718
description: INodeTypeDescription = {
@@ -176,7 +177,7 @@ export class ChainRetrievalQa implements INodeType {
176177
throw new NodeOperationError(this.getNode(), 'The ‘query‘ parameter is empty.');
177178
}
178179

179-
const response = await chain.call({ query });
180+
const response = await chain.withConfig(getTracingConfig(this)).invoke({ query });
180181
returnData.push({ json: { response } });
181182
}
182183
return await this.prepareOutputData(returnData);

‎packages/@n8n/nodes-langchain/nodes/chains/ChainSummarization/V2/ChainSummarizationV2.node.ts

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -18,6 +18,7 @@ import { N8nBinaryLoader } from '../../../../utils/N8nBinaryLoader';
1818
import { getTemplateNoticeField } from '../../../../utils/sharedFields';
1919
import { REFINE_PROMPT_TEMPLATE, DEFAULT_PROMPT_TEMPLATE } from '../prompt';
2020
import { getChainPromptsArgs } from '../helpers';
21+
import { getTracingConfig } from '../../../../utils/tracing';
2122

2223
function getInputs(parameters: IDataObject) {
2324
const chunkingMode = parameters?.chunkingMode;
@@ -364,7 +365,7 @@ export class ChainSummarizationV2 implements INodeType {
364365
? await documentInput.processItem(item, itemIndex)
365366
: documentInput;
366367

367-
const response = await chain.call({
368+
const response = await chain.withConfig(getTracingConfig(this)).invoke({
368369
input_documents: processedDocuments,
369370
});
370371

‎packages/@n8n/nodes-langchain/nodes/retrievers/RetrieverWorkflow/RetrieverWorkflow.node.ts

Lines changed: 6 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -16,6 +16,7 @@ import { Document } from '@langchain/core/documents';
1616

1717
import type { SetField, SetNodeOptions } from 'n8n-nodes-base/dist/nodes/Set/v2/helpers/interfaces';
1818
import * as manual from 'n8n-nodes-base/dist/nodes/Set/v2/manual.mode';
19+
import type { CallbackManagerForRetrieverRun } from '@langchain/core/callbacks/manager';
1920
import { logWrapper } from '../../../utils/logWrapper';
2021

2122
function objectToString(obj: Record<string, string> | IDataObject, level = 0) {
@@ -287,7 +288,10 @@ export class RetrieverWorkflow implements INodeType {
287288
this.executeFunctions = executeFunctions;
288289
}
289290

290-
async getRelevantDocuments(query: string): Promise<Document[]> {
291+
async _getRelevantDocuments(
292+
query: string,
293+
config?: CallbackManagerForRetrieverRun,
294+
): Promise<Document[]> {
291295
const source = this.executeFunctions.getNodeParameter('source', itemIndex) as string;
292296

293297
const baseMetadata: IDataObject = {
@@ -360,6 +364,7 @@ export class RetrieverWorkflow implements INodeType {
360364
receivedItems = (await this.executeFunctions.executeWorkflow(
361365
workflowInfo,
362366
items,
367+
config?.getChild(),
363368
)) as INodeExecutionData[][];
364369
} catch (error) {
365370
// Make sure a valid error gets returned that can by json-serialized else it will

‎packages/@n8n/nodes-langchain/nodes/tools/ToolWorkflow/ToolWorkflow.node.ts

Lines changed: 12 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -16,6 +16,7 @@ import * as manual from 'n8n-nodes-base/dist/nodes/Set/v2/manual.mode';
1616
import { DynamicTool } from '@langchain/core/tools';
1717
import get from 'lodash/get';
1818
import isObject from 'lodash/isObject';
19+
import type { CallbackManagerForToolRun } from '@langchain/core/callbacks/manager';
1920
import { getConnectionHintNoticeField } from '../../../utils/sharedFields';
2021

2122
export class ToolWorkflow implements INodeType {
@@ -320,7 +321,10 @@ export class ToolWorkflow implements INodeType {
320321
const name = this.getNodeParameter('name', itemIndex) as string;
321322
const description = this.getNodeParameter('description', itemIndex) as string;
322323

323-
const runFunction = async (query: string): Promise<string> => {
324+
const runFunction = async (
325+
query: string,
326+
runManager?: CallbackManagerForToolRun,
327+
): Promise<string> => {
324328
const source = this.getNodeParameter('source', itemIndex) as string;
325329
const responsePropertyName = this.getNodeParameter(
326330
'responsePropertyName',
@@ -385,7 +389,11 @@ export class ToolWorkflow implements INodeType {
385389

386390
let receivedData: INodeExecutionData;
387391
try {
388-
receivedData = (await this.executeWorkflow(workflowInfo, items)) as INodeExecutionData;
392+
receivedData = (await this.executeWorkflow(
393+
workflowInfo,
394+
items,
395+
runManager?.getChild(),
396+
)) as INodeExecutionData;
389397
} catch (error) {
390398
// Make sure a valid error gets returned that can by json-serialized else it will
391399
// not show up in the frontend
@@ -413,13 +421,13 @@ export class ToolWorkflow implements INodeType {
413421
name,
414422
description,
415423

416-
func: async (query: string): Promise<string> => {
424+
func: async (query: string, runManager?: CallbackManagerForToolRun): Promise<string> => {
417425
const { index } = this.addInputData(NodeConnectionType.AiTool, [[{ json: { query } }]]);
418426

419427
let response: string = '';
420428
let executionError: ExecutionError | undefined;
421429
try {
422-
response = await runFunction(query);
430+
response = await runFunction(query, runManager);
423431
} catch (error) {
424432
// TODO: Do some more testing. Issues here should actually fail the workflow
425433
// eslint-disable-next-line @typescript-eslint/no-unsafe-assignment

‎packages/@n8n/nodes-langchain/nodes/vendors/OpenAi/actions/assistant/message.operation.ts

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -11,6 +11,7 @@ import { formatToOpenAIAssistantTool } from '../../helpers/utils';
1111
import { assistantRLC } from '../descriptions';
1212

1313
import { getConnectedTools } from '../../../../../utils/helpers';
14+
import { getTracingConfig } from '../../../../../utils/tracing';
1415

1516
const properties: INodeProperties[] = [
1617
assistantRLC,
@@ -181,7 +182,7 @@ export async function execute(this: IExecuteFunctions, i: number): Promise<INode
181182
tools: tools ?? [],
182183
});
183184

184-
const response = await agentExecutor.invoke({
185+
const response = await agentExecutor.withConfig(getTracingConfig(this)).invoke({
185186
content: input,
186187
signal: this.getExecutionCancelSignal(),
187188
timeout: options.timeout ?? 10000,

‎packages/@n8n/nodes-langchain/utils/tracing.ts

Lines changed: 26 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,26 @@
1+
import type { BaseCallbackConfig } from '@langchain/core/callbacks/manager';
2+
import type { IExecuteFunctions } from 'n8n-workflow';
3+
4+
interface TracingConfig {
5+
additionalMetadata?: Record<string, unknown>;
6+
}
7+
8+
export function getTracingConfig(
9+
context: IExecuteFunctions,
10+
config: TracingConfig = {},
11+
): BaseCallbackConfig {
12+
const parentRunManager = context.getParentCallbackManager
13+
? context.getParentCallbackManager()
14+
: undefined;
15+
16+
return {
17+
runName: `[${context.getWorkflow().name}] ${context.getNode().name}`,
18+
metadata: {
19+
execution_id: context.getExecutionId(),
20+
workflow: context.getWorkflow(),
21+
node: context.getNode().name,
22+
...(config.additionalMetadata ?? {}),
23+
},
24+
callbacks: parentRunManager,
25+
};
26+
}

‎packages/cli/src/WorkflowExecuteAdditionalData.ts

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -24,6 +24,7 @@ import type {
2424
ExecutionStatus,
2525
ExecutionError,
2626
EventNamesAiNodesType,
27+
CallbackManager,
2728
} from 'n8n-workflow';
2829
import {
2930
ApplicationError,
@@ -754,6 +755,7 @@ async function executeWorkflow(
754755
loadedWorkflowData?: IWorkflowBase;
755756
loadedRunData?: IWorkflowExecutionDataProcess;
756757
parentWorkflowSettings?: IWorkflowSettings;
758+
parentCallbackManager?: CallbackManager;
757759
},
758760
): Promise<Array<INodeExecutionData[] | null> | IWorkflowExecuteProcess> {
759761
const internalHooks = Container.get(InternalHooks);
@@ -815,6 +817,7 @@ async function executeWorkflow(
815817
workflowData,
816818
);
817819
additionalDataIntegrated.executionId = executionId;
820+
additionalDataIntegrated.parentCallbackManager = options.parentCallbackManager;
818821

819822
// Make sure we pass on the original executeWorkflow function we received
820823
// This one already contains changes to talk to parent process

‎packages/core/src/NodeExecuteFunctions.ts

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -98,6 +98,7 @@ import type {
9898
Workflow,
9999
WorkflowActivateMode,
100100
WorkflowExecuteMode,
101+
CallbackManager,
101102
} from 'n8n-workflow';
102103
import {
103104
ExpressionError,
@@ -3487,13 +3488,15 @@ export function getExecuteFunctions(
34873488
async executeWorkflow(
34883489
workflowInfo: IExecuteWorkflowInfo,
34893490
inputData?: INodeExecutionData[],
3491+
parentCallbackManager?: CallbackManager,
34903492
): Promise<any> {
34913493
return await additionalData
34923494
.executeWorkflow(workflowInfo, additionalData, {
34933495
parentWorkflowId: workflow.id?.toString(),
34943496
inputData,
34953497
parentWorkflowSettings: workflow.settings,
34963498
node,
3499+
parentCallbackManager,
34973500
})
34983501
.then(
34993502
async (result) =>
@@ -3719,6 +3722,7 @@ export function getExecuteFunctions(
37193722
msg,
37203723
});
37213724
},
3725+
getParentCallbackManager: () => additionalData.parentCallbackManager,
37223726
};
37233727
})(workflow, runExecutionData, connectionInputData, inputData, node) as IExecuteFunctions;
37243728
}

‎packages/workflow/package.json

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -65,6 +65,7 @@
6565
"recast": "0.21.5",
6666
"title-case": "3.0.3",
6767
"transliteration": "2.3.5",
68-
"xml2js": "0.6.2"
68+
"xml2js": "0.6.2",
69+
"@langchain/core": "0.1.41"
6970
}
7071
}

0 commit comments

Comments
 (0)