feat(ai): add LangChain v1 agent middleware #4556
+634
−15
There are no files selected for viewing
This file contains hidden or 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,5 @@ | ||
| --- | ||
| '@posthog/ai': minor | ||
| --- | ||
|
|
||
| Add LangChain v1 agent middleware for AI observability. |
This file contains hidden or 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 hidden or 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 hidden or 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 hidden or 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,197 @@ | ||
| import { AIMessage, BaseMessage, ToolMessage } from '@langchain/core/messages' | ||
| import type { ChatGeneration, LLMResult } from '@langchain/core/outputs' | ||
| import type { Serialized } from '@langchain/core/load/serializable' | ||
| import { createMiddleware } from 'langchain' | ||
| import { v7 as uuidv7 } from 'uuid' | ||
| import { z } from 'zod' | ||
| import { toContentString } from '../../utils' | ||
| import { LangChainCallbackHandler, LangChainCallbackHandlerOptions } from '../callbacks' | ||
|
|
||
| const postHogStateSchema = z.object({ | ||
| _posthogRunId: z.string().optional(), | ||
| _posthogStartTime: z.number().optional(), | ||
| _posthogInput: z.record(z.string(), z.unknown()).optional(), | ||
| }) | ||
|
|
||
| type PostHogState = z.infer<typeof postHogStateSchema> | ||
|
|
||
| const withoutPostHogState = <T extends Record<string, unknown>>(state: T): Omit<T, keyof PostHogState> => { | ||
| const { _posthogRunId: _, _posthogStartTime: __, _posthogInput: ___, ...rest } = state | ||
| return rest | ||
| } | ||
|
|
||
| const getRunId = (state: PostHogState): string => state._posthogRunId ?? uuidv7() | ||
|
|
||
| const stringify = (value: unknown): string => { | ||
| try { | ||
| return JSON.stringify(value) ?? String(value) | ||
| } catch { | ||
| try { | ||
| return String(value) | ||
| } catch { | ||
| return '' | ||
| } | ||
| } | ||
| } | ||
|
|
||
| const toError = (error: unknown): Error => (error instanceof Error ? error : new Error(stringify(error))) | ||
|
|
||
| const safely = (callback: () => void): void => { | ||
| try { | ||
| callback() | ||
| } catch { | ||
| // Telemetry must never affect the LangChain middleware lifecycle. | ||
| } | ||
| } | ||
|
|
||
| const serializeModel = (model: unknown): Serialized => { | ||
| if (model && typeof model === 'object' && 'toJSON' in model && typeof model.toJSON === 'function') { | ||
| try { | ||
| return model.toJSON() as Serialized | ||
| } catch { | ||
| // Fall back to a minimal LangChain serialization below. | ||
| } | ||
| } | ||
|
|
||
| return { lc: 1, type: 'constructor', id: ['langchain', 'chat_models', 'unknown'], kwargs: {} } | ||
| } | ||
|
|
||
| const getModelMetadata = (model: unknown, modelSettings: unknown): Record<string, unknown> | undefined => { | ||
| if (model && typeof model === 'object' && 'getLsParams' in model && typeof model.getLsParams === 'function') { | ||
| try { | ||
| return model.getLsParams(modelSettings) as Record<string, unknown> | ||
| } catch { | ||
| return undefined | ||
| } | ||
| } | ||
| return undefined | ||
| } | ||
|
|
||
| const toLLMResult = (response: unknown): LLMResult => { | ||
| if (!AIMessage.isInstance(response)) { | ||
| return { generations: [] } | ||
| } | ||
|
|
||
| const generation: ChatGeneration = { | ||
| text: toContentString(response.content), | ||
| message: response, | ||
| } | ||
| return { generations: [[generation]] } | ||
| } | ||
|
|
||
| /** Options shared with the LangChain callback integration. */ | ||
| export type PostHogLangChainMiddlewareOptions = LangChainCallbackHandlerOptions | ||
|
|
||
| /** | ||
| * Creates PostHog AI observability middleware for LangChain v1 agents. | ||
| * | ||
| * Use either this middleware or `LangChainCallbackHandler`, not both, to avoid | ||
| * capturing the same model and tool calls twice. | ||
| * | ||
| * LangChain only invokes `afterAgent` for completed runs. A terminal agent | ||
| * failure still captures the failed model or tool call, but not a root trace. | ||
| */ | ||
| export const createPostHogMiddleware = (options: PostHogLangChainMiddlewareOptions) => { | ||
| const callback = new LangChainCallbackHandler(options) | ||
|
|
||
| return createMiddleware({ | ||
| name: 'PostHogMiddleware', | ||
| stateSchema: postHogStateSchema, | ||
|
|
||
| beforeAgent: (state) => { | ||
| return { | ||
| _posthogRunId: uuidv7(), | ||
| _posthogStartTime: Date.now(), | ||
| _posthogInput: withoutPostHogState(state), | ||
| } | ||
| }, | ||
|
|
||
| afterAgent: (state) => { | ||
| safely(() => { | ||
| const runId = getRunId(state) | ||
| callback.handleChainStart( | ||
| { lc: 1, type: 'constructor', id: ['langchain', 'agents', 'PostHogMiddleware'], kwargs: {} }, | ||
| state._posthogInput ?? withoutPostHogState(state), | ||
| runId, | ||
| undefined, | ||
| undefined, | ||
| undefined, | ||
| undefined, | ||
| 'LangChain Agent', | ||
| { posthogStartTime: state._posthogStartTime } | ||
| ) | ||
| callback.handleChainEnd(withoutPostHogState(state), runId) | ||
| }) | ||
| }, | ||
|
|
||
| wrapModelCall: async (request, handler) => { | ||
| const runId = uuidv7() | ||
| const parentRunId = getRunId(request.state) | ||
|
gouveags marked this conversation as resolved.
|
||
| const messages = [request.systemMessage, ...request.messages].filter( | ||
| (message): message is BaseMessage => message !== undefined | ||
| ) | ||
| const invocationParams = { | ||
| ...request.modelSettings, | ||
| tools: request.tools, | ||
|
Member
There was a problem hiding this comment. blocking: Normalize tools before capture — |
||
| } | ||
|
|
||
| safely(() => | ||
| callback.handleChatModelStart( | ||
| serializeModel(request.model), | ||
| [messages], | ||
| runId, | ||
| parentRunId, | ||
| { invocation_params: invocationParams }, | ||
| undefined, | ||
| getModelMetadata(request.model, request.modelSettings) | ||
| ) | ||
| ) | ||
|
|
||
| try { | ||
| const response = await handler(request) | ||
| safely(() => callback.handleLLMEnd(toLLMResult(response), runId, parentRunId)) | ||
| return response | ||
| } catch (error) { | ||
| safely(() => callback.handleLLMError(toError(error), runId, parentRunId)) | ||
| throw error | ||
| } | ||
| }, | ||
|
|
||
| wrapToolCall: async (request, handler) => { | ||
| const runId = uuidv7() | ||
| const parentRunId = getRunId(request.state) | ||
| const toolName = String(request.tool?.name ?? request.toolCall.name) | ||
| const serializedTool: Serialized = { | ||
| lc: 1, | ||
| type: 'constructor', | ||
| id: ['langchain', 'tools', toolName], | ||
| kwargs: {}, | ||
| } | ||
|
|
||
| safely(() => | ||
| callback.handleToolStart( | ||
| serializedTool, | ||
| stringify(request.toolCall.args), | ||
| runId, | ||
| parentRunId, | ||
| undefined, | ||
| undefined, | ||
| toolName | ||
| ) | ||
| ) | ||
|
|
||
| try { | ||
| const result = await handler(request) | ||
| if (ToolMessage.isInstance(result) && result.status === 'error') { | ||
| safely(() => callback.handleToolError(new Error(toContentString(result.content)), runId, parentRunId)) | ||
| } else { | ||
| safely(() => callback.handleToolEnd(result, runId, parentRunId)) | ||
| } | ||
| return result | ||
| } catch (error) { | ||
| safely(() => callback.handleToolError(toError(error), runId, parentRunId)) | ||
| throw error | ||
| } | ||
| }, | ||
| }) | ||
| } | ||
Oops, something went wrong.
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
blocking: Preserve custom agent state in traces — LangChain projects lifecycle-hook state through this middleware schema, which only declares the
_posthog*fields plus built-ins. With a customcreateAgent({ stateSchema }), fields such as tenant or workflow context remain in the agent result but are silently absent from$ai_input_stateand$ai_output_state; a custom-state probe reproduced the omission.