@cloudbase/agent-observability 0.0.16

package/src/langchain/CallbackHandler.ts

@@ -0,0 +1,896 @@
+/**
+ * LangChain Callback Handler for AG-Kit Observability
+ *
+ * Converts LangChain callback events into AG-Kit observations with OpenInference semantics.
+ */
+
+import type { AgentAction, AgentFinish } from "@langchain/core/agents";
+import { BaseCallbackHandler } from "@langchain/core/callbacks/base";
+import type { Document } from "@langchain/core/documents";
+import type { Serialized } from "@langchain/core/load/serializable";
+import {
+  AIMessage,
+  AIMessageChunk,
+  BaseMessage,
+  type UsageMetadata,
+  type BaseMessageFields,
+  type MessageContent,
+} from "@langchain/core/messages";
+import type { Generation, LLMResult } from "@langchain/core/outputs";
+import type { ChainValues } from "@langchain/core/utils/types";
+
+import {
+  startObservation,
+  type ObservationLLM,
+  type ObservationSpan,
+  type ObservationTool,
+  type Observation,
+  type ObservationAttributes,
+} from "../index.js";
+import type { SpanContext } from "@opentelemetry/api";
+import { type Logger, noopLogger } from "@cloudbase/agent-shared";
+
+/**
+ * Constructor parameters for CallbackHandler.
+ *
+ * @public
+ */
+type ConstructorParams = {
+  userId?: string;
+  sessionId?: string;
+  tags?: string[];
+  version?: string;
+  traceMetadata?: Record<string, unknown>;
+  adapterName?: string;  // e.g., "LangGraph" or "LangChain"
+  /** Logger for debug output. Defaults to noopLogger (silent). */
+  logger?: Logger;
+};
+
+/**
+ * Message format for LLM input/output.
+ *
+ * @public
+ */
+export type LlmMessage = {
+  role: string;
+  content: BaseMessageFields["content"];
+  additional_kwargs?: BaseMessageFields["additional_kwargs"];
+};
+
+/**
+ * Anonymous message format (without role).
+ *
+ * @public
+ */
+export type AnonymousLlmMessage = {
+  content: BaseMessageFields["content"];
+  additional_kwargs?: BaseMessageFields["additional_kwargs"];
+};
+
+/**
+ * Prompt information for linking to generations.
+ *
+ * @public
+ */
+type PromptInfo = {
+  name: string;
+  version: number;
+  isFallback: boolean;
+};
+
+/**
+ * LangChain Callback Handler for AG-Kit Observability.
+ *
+ * This handler intercepts LangChain callbacks and converts them into
+ * AG-Kit observations following OpenInference semantic conventions.
+ *
+ * @public
+ */
+export class CallbackHandler extends BaseCallbackHandler {
+  name = "ObservabilityCallbackHandler";
+
+  private userId?: string;
+  private version?: string;
+  private sessionId?: string;
+  private tags: string[];
+  private traceMetadata?: Record<string, unknown>;
+
+  private completionStartTimes: Record<string, Date> = {};
+  private promptToParentRunMap;
+  private runMap: Map<string, Observation> = new Map();
+
+  public last_trace_id: string | null = null;
+
+  // External parent context from AG-UI.Server span
+  private externalParentSpanContext?: SpanContext;
+
+  // Adapter name for ROOT span prefix
+  private adapterName?: string;
+
+  // Logger for debug output (defaults to noopLogger for silent operation)
+  private logger: Logger;
+
+  constructor(params?: ConstructorParams) {
+    super();
+
+    this.sessionId = params?.sessionId;
+    this.userId = params?.userId;
+    this.tags = params?.tags ?? [];
+    this.traceMetadata = params?.traceMetadata;
+    this.version = params?.version;
+    this.adapterName = params?.adapterName;
+    this.logger = params?.logger ?? noopLogger;
+
+    this.promptToParentRunMap = new Map<string, PromptInfo>();
+  }
+
+  /**
+   * Set external parent SpanContext from AG-UI.Server span.
+   * This allows the CallbackHandler to link LangChain/LangGraph spans
+   * to the server-level span, creating a unified trace hierarchy.
+   *
+   * @param spanContext - SpanContext from the AG-UI.Server span
+   * @public
+   */
+  setExternalParentContext(spanContext: SpanContext): void {
+    this.externalParentSpanContext = spanContext;
+  }
+
+  async handleLLMNewToken(
+    token: string,
+    _idx: any,
+    runId: string,
+    _parentRunId?: string,
+    _tags?: string[],
+    _fields?: any
+  ): Promise<void> {
+    if (runId && !(runId in this.completionStartTimes)) {
+      this.logger.debug?.(`LLM first streaming token: ${runId}`);
+      this.completionStartTimes[runId] = new Date();
+    }
+  }
+
+  async handleChainStart(
+    chain: Serialized,
+    inputs: ChainValues,
+    runId: string,
+    parentRunId?: string | undefined,
+    tags?: string[] | undefined,
+    metadata?: Record<string, unknown> | undefined,
+    runType?: string,
+    name?: string
+  ): Promise<void> {
+    try {
+      this.logger.debug?.(`Chain start with Id: ${runId}`);
+
+      const runName = name ?? chain.id.at(-1)?.toString() ?? "Langchain Run";
+
+      this.registerPromptInfo(parentRunId, metadata);
+
+      let finalInput: string | ChainValues = inputs;
+      if (
+        typeof inputs === "object" &&
+        "input" in inputs &&
+        Array.isArray(inputs["input"]) &&
+        inputs["input"].every((m: unknown) => m instanceof BaseMessage)
+      ) {
+        finalInput = inputs["input"].map((m: BaseMessage) =>
+          this.extractChatMessageContent(m)
+        );
+      } else if (
+        typeof inputs === "object" &&
+        "messages" in inputs &&
+        Array.isArray(inputs["messages"]) &&
+        inputs["messages"].every((m: unknown) => m instanceof BaseMessage)
+      ) {
+        finalInput = inputs["messages"].map((m: BaseMessage) =>
+          this.extractChatMessageContent(m)
+        );
+      } else if (
+        typeof inputs === "object" &&
+        "content" in inputs &&
+        typeof inputs["content"] === "string"
+      ) {
+        finalInput = inputs["content"];
+      }
+
+      const observation = this.startAndRegisterObservation({
+        runName,
+        parentRunId,
+        runId,
+        tags,
+        metadata,
+        attributes: {
+          input: finalInput,
+        },
+        asType: "span",
+      });
+
+      const traceTags = [...new Set([...(tags ?? []), ...this.tags])];
+
+      if (!parentRunId) {
+        observation.updateTrace({
+          tags: traceTags,
+          userId:
+            metadata &&
+            "userId" in metadata &&
+            typeof metadata["userId"] === "string"
+              ? metadata["userId"]
+              : this.userId,
+          sessionId:
+            metadata &&
+            "sessionId" in metadata &&
+            typeof metadata["sessionId"] === "string"
+              ? metadata["sessionId"]
+              : this.sessionId,
+          metadata: this.traceMetadata,
+          version: this.version,
+        });
+      }
+    } catch (e) {
+      this.logger.debug?.(e instanceof Error ? e.message : String(e));
+    }
+  }
+
+  async handleAgentAction(
+    action: AgentAction,
+    runId: string,
+    parentRunId?: string
+  ): Promise<void> {
+    try {
+      this.logger.debug?.(`Agent action ${action.tool} with ID: ${runId}`);
+      this.startAndRegisterObservation({
+        runId,
+        parentRunId,
+        runName: action.tool,
+        attributes: {
+          input: action,
+        },
+        asType: "tool",
+      });
+    } catch (e) {
+      this.logger.debug?.(e instanceof Error ? e.message : String(e));
+    }
+  }
+
+  async handleAgentEnd?(
+    action: AgentFinish,
+    runId: string,
+    _parentRunId?: string
+  ): Promise<void> {
+    try {
+      this.logger.debug?.(`Agent finish with ID: ${runId}`);
+      this.handleObservationEnd({
+        runId,
+        attributes: { output: action },
+      });
+    } catch (e) {
+      this.logger.debug?.(e instanceof Error ? e.message : String(e));
+    }
+  }
+
+  async handleChainError(
+    err: any,
+    runId: string,
+    _parentRunId?: string | undefined
+  ): Promise<void> {
+    try {
+      this.logger.debug?.(`Chain error: ${err} with ID: ${runId}`);
+      this.handleObservationEnd({
+        runId,
+        attributes: {
+          level: "ERROR",
+          statusMessage: err.toString(),
+        },
+      });
+    } catch (e) {
+      this.logger.debug?.(e instanceof Error ? e.message : String(e));
+    }
+  }
+
+  async handleGenerationStart(
+    llm: Serialized,
+    messages: (LlmMessage | MessageContent | AnonymousLlmMessage)[],
+    runId: string,
+    parentRunId?: string | undefined,
+    extraParams?: Record<string, unknown> | undefined,
+    tags?: string[] | undefined,
+    metadata?: Record<string, unknown> | undefined,
+    name?: string
+  ): Promise<void> {
+    this.logger.debug?.(
+      `Generation start with ID: ${runId} and parentRunId ${parentRunId}`
+    );
+
+    const runName = name ?? llm.id.at(-1)?.toString() ?? "Langchain Generation";
+
+    const modelParameters: Record<string, any> = {};
+    const invocationParams = extraParams?.["invocation_params"];
+
+    for (const [key, value] of Object.entries({
+      temperature: (invocationParams as any)?.temperature,
+      max_tokens: (invocationParams as any)?.max_tokens,
+      top_p: (invocationParams as any)?.top_p,
+      frequency_penalty: (invocationParams as any)?.frequency_penalty,
+      presence_penalty: (invocationParams as any)?.presence_penalty,
+      request_timeout: (invocationParams as any)?.request_timeout,
+    })) {
+      if (value !== undefined && value !== null) {
+        modelParameters[key] = value;
+      }
+    }
+
+    interface InvocationParams {
+      _type?: string;
+      model?: string;
+      model_name?: string;
+      repo_id?: string;
+    }
+
+    let extractedModelName: string | undefined;
+    if (extraParams) {
+      const invocationParamsModelName = (
+        extraParams.invocation_params as InvocationParams
+      ).model;
+      const metadataModelName =
+        metadata && "ls_model_name" in metadata
+          ? (metadata["ls_model_name"] as string)
+          : undefined;
+
+      extractedModelName = invocationParamsModelName ?? metadataModelName;
+    }
+
+    const registeredPrompt = this.promptToParentRunMap.get(
+      parentRunId ?? "root"
+    );
+    if (registeredPrompt && parentRunId) {
+      this.deregisterPromptInfo(parentRunId);
+    }
+
+    this.startAndRegisterObservation({
+      runId,
+      parentRunId,
+      metadata,
+      tags,
+      runName,
+      attributes: {
+        input: messages,
+        model: extractedModelName,
+        modelParameters: modelParameters,
+      },
+      asType: "llm",
+    });
+  }
+
+  async handleChatModelStart(
+    llm: Serialized,
+    messages: BaseMessage[][],
+    runId: string,
+    parentRunId?: string | undefined,
+    extraParams?: Record<string, unknown> | undefined,
+    tags?: string[] | undefined,
+    metadata?: Record<string, unknown> | undefined,
+    name?: string
+  ): Promise<void> {
+    try {
+      this.logger.debug?.(`Chat model start with ID: ${runId}`);
+
+      const prompts = messages.flatMap((message) =>
+        message.map((m) => this.extractChatMessageContent(m))
+      );
+
+      this.handleGenerationStart(
+        llm,
+        prompts,
+        runId,
+        parentRunId,
+        extraParams,
+        tags,
+        metadata,
+        name
+      );
+    } catch (e) {
+      this.logger.debug?.(e instanceof Error ? e.message : String(e));
+    }
+  }
+
+  async handleChainEnd(
+    outputs: ChainValues,
+    runId: string,
+    _parentRunId?: string | undefined
+  ): Promise<void> {
+    try {
+      this.logger.debug?.(`Chain end with ID: ${runId}`);
+
+      let finalOutput: ChainValues | string = outputs;
+      if (
+        typeof outputs === "object" &&
+        "output" in outputs &&
+        typeof outputs["output"] === "string"
+      ) {
+        finalOutput = outputs["output"];
+      } else if (
+        typeof outputs === "object" &&
+        "messages" in outputs &&
+        Array.isArray(outputs["messages"]) &&
+        outputs["messages"].every((m: unknown) => m instanceof BaseMessage)
+      ) {
+        finalOutput = {
+          messages: outputs.messages.map((message: BaseMessage) =>
+            this.extractChatMessageContent(message)
+          ),
+        };
+      }
+
+      this.handleObservationEnd({
+        runId,
+        attributes: {
+          output: finalOutput,
+        },
+      });
+      this.deregisterPromptInfo(runId);
+    } catch (e) {
+      this.logger.debug?.(e instanceof Error ? e.message : String(e));
+    }
+  }
+
+  async handleLLMStart(
+    llm: Serialized,
+    prompts: string[],
+    runId: string,
+    parentRunId?: string | undefined,
+    extraParams?: Record<string, unknown> | undefined,
+    tags?: string[] | undefined,
+    metadata?: Record<string, unknown> | undefined,
+    name?: string
+  ): Promise<void> {
+    try {
+      this.logger.debug?.(`LLM start with ID: ${runId}`);
+      this.handleGenerationStart(
+        llm,
+        prompts,
+        runId,
+        parentRunId,
+        extraParams,
+        tags,
+        metadata,
+        name
+      );
+    } catch (e) {
+      this.logger.debug?.(e instanceof Error ? e.message : String(e));
+    }
+  }
+
+  async handleToolStart(
+    tool: Serialized,
+    input: string,
+    runId: string,
+    parentRunId?: string | undefined,
+    tags?: string[] | undefined,
+    metadata?: Record<string, unknown> | undefined,
+    name?: string
+  ): Promise<void> {
+    try {
+      this.logger.debug?.(`Tool start with ID: ${runId}`);
+      this.startAndRegisterObservation({
+        runId,
+        parentRunId,
+        runName: name ?? tool.id.at(-1)?.toString() ?? "Tool execution",
+        attributes: {
+          input,
+        },
+        metadata,
+        tags,
+        asType: "tool",
+      });
+    } catch (e) {
+      this.logger.debug?.(e instanceof Error ? e.message : String(e));
+    }
+  }
+
+  async handleRetrieverStart(
+    retriever: Serialized,
+    query: string,
+    runId: string,
+    parentRunId?: string | undefined,
+    tags?: string[] | undefined,
+    metadata?: Record<string, unknown> | undefined,
+    name?: string
+  ): Promise<void> {
+    try {
+      this.logger.debug?.(`Retriever start with ID: ${runId}`);
+      this.startAndRegisterObservation({
+        runId,
+        parentRunId,
+        runName: name ?? retriever.id.at(-1)?.toString() ?? "Retriever",
+        attributes: {
+          input: query,
+        },
+        tags,
+        metadata,
+        asType: "span",
+      });
+    } catch (e) {
+      this.logger.debug?.(e instanceof Error ? e.message : String(e));
+    }
+  }
+
+  async handleRetrieverEnd(
+    documents: Document<Record<string, any>>[],
+    runId: string,
+    _parentRunId?: string | undefined
+  ): Promise<void> {
+    try {
+      this.logger.debug?.(`Retriever end with ID: ${runId}`);
+      this.handleObservationEnd({
+        runId,
+        attributes: {
+          output: documents,
+        },
+      });
+    } catch (e) {
+      this.logger.debug?.(e instanceof Error ? e.message : String(e));
+    }
+  }
+
+  async handleRetrieverError(
+    err: any,
+    runId: string,
+    _parentRunId?: string | undefined
+  ): Promise<void> {
+    try {
+      this.logger.debug?.(`Retriever error: ${err} with ID: ${runId}`);
+      this.handleObservationEnd({
+        runId,
+        attributes: {
+          level: "ERROR",
+          statusMessage: err.toString(),
+        },
+      });
+    } catch (e) {
+      this.logger.debug?.(e instanceof Error ? e.message : String(e));
+    }
+  }
+
+  async handleToolEnd(
+    output: string,
+    runId: string,
+    _parentRunId?: string | undefined
+  ): Promise<void> {
+    try {
+      this.logger.debug?.(`Tool end with ID: ${runId}`);
+      this.handleObservationEnd({
+        runId,
+        attributes: { output },
+      });
+    } catch (e) {
+      this.logger.debug?.(e instanceof Error ? e.message : String(e));
+    }
+  }
+
+  async handleToolError(
+    err: any,
+    runId: string,
+    _parentRunId?: string | undefined
+  ): Promise<void> {
+    try {
+      this.logger.debug?.(`Tool error ${err} with ID: ${runId}`);
+      this.handleObservationEnd({
+        runId,
+        attributes: {
+          level: "ERROR",
+          statusMessage: err.toString(),
+        },
+      });
+    } catch (e) {
+      this.logger.debug?.(e instanceof Error ? e.message : String(e));
+    }
+  }
+
+  async handleLLMEnd(
+    output: LLMResult,
+    runId: string,
+    _parentRunId?: string | undefined
+  ): Promise<void> {
+    try {
+      this.logger.debug?.(`LLM end with ID: ${runId}`);
+
+      const lastResponse =
+        output.generations[output.generations.length - 1][
+          output.generations[output.generations.length - 1].length - 1
+        ];
+      const llmUsage =
+        this.extractUsageMetadata(lastResponse) ??
+        output.llmOutput?.["tokenUsage"];
+      const modelName = this.extractModelNameFromMetadata(lastResponse);
+
+      const usageDetails: Record<string, any> = {
+        input:
+          llmUsage?.input_tokens ??
+          ("promptTokens" in llmUsage ? llmUsage?.promptTokens : undefined),
+        output:
+          llmUsage?.output_tokens ??
+          ("completionTokens" in llmUsage
+            ? llmUsage?.completionTokens
+            : undefined),
+        total:
+          llmUsage?.total_tokens ??
+          ("totalTokens" in llmUsage ? llmUsage?.totalTokens : undefined),
+      };
+
+      if (llmUsage && "input_token_details" in llmUsage) {
+        for (const [key, val] of Object.entries(
+          llmUsage["input_token_details"] ?? {}
+        )) {
+          usageDetails[`input_${key}`] = val;
+          if ("input" in usageDetails && typeof val === "number") {
+            usageDetails["input"] = Math.max(0, usageDetails["input"] - val);
+          }
+        }
+      }
+
+      if (llmUsage && "output_token_details" in llmUsage) {
+        for (const [key, val] of Object.entries(
+          llmUsage["output_token_details"] ?? {}
+        )) {
+          usageDetails[`output_${key}`] = val;
+          if ("output" in usageDetails && typeof val === "number") {
+            usageDetails["output"] = Math.max(0, usageDetails["output"] - val);
+          }
+        }
+      }
+
+      const extractedOutput =
+        "message" in lastResponse
+          ? this.extractChatMessageContent(
+              lastResponse["message"] as BaseMessage
+            )
+          : lastResponse.text;
+
+      this.handleObservationEnd({
+        runId,
+        attributes: {
+          model: modelName,
+          output: extractedOutput,
+          completionStartTime:
+            runId in this.completionStartTimes
+              ? this.completionStartTimes[runId]
+              : undefined,
+          usageDetails: usageDetails,
+        },
+      });
+
+      if (runId in this.completionStartTimes) {
+        delete this.completionStartTimes[runId];
+      }
+    } catch (e) {
+      this.logger.debug?.(e instanceof Error ? e.message : String(e));
+    }
+  }
+
+  async handleLLMError(
+    err: any,
+    runId: string,
+    _parentRunId?: string | undefined
+  ): Promise<void> {
+    try {
+      this.logger.debug?.(`LLM error ${err} with ID: ${runId}`);
+      this.handleObservationEnd({
+        runId,
+        attributes: {
+          level: "ERROR",
+          statusMessage: err.toString(),
+        },
+      });
+    } catch (e) {
+      this.logger.debug?.(e instanceof Error ? e.message : String(e));
+    }
+  }
+
+  private registerPromptInfo(
+    parentRunId?: string,
+    metadata?: Record<string, unknown>
+  ): void {
+    if (metadata && "promptInfo" in metadata && parentRunId) {
+      this.promptToParentRunMap.set(
+        parentRunId,
+        metadata.promptInfo as PromptInfo
+      );
+    }
+  }
+
+  private deregisterPromptInfo(runId: string): void {
+    this.promptToParentRunMap.delete(runId);
+  }
+
+  private startAndRegisterObservation(params: {
+    runName: string;
+    runId: string;
+    parentRunId?: string;
+    attributes: Record<string, unknown>;
+    metadata?: Record<string, unknown>;
+    tags?: string[];
+    asType?: "span" | "llm" | "tool";
+  }): Observation {
+    const { runName, runId, parentRunId, attributes, metadata, tags, asType } =
+      params;
+
+    // Determine parent context:
+    // 1. If parentRunId exists, use the parent span from runMap (internal LangChain/LangGraph hierarchy)
+    // 2. If no parentRunId (ROOT span) but externalParentSpanContext exists, use it (link to AG-UI.Server)
+    // 3. Otherwise, create a new root span
+    let parentSpanContext: SpanContext | undefined;
+
+    if (parentRunId) {
+      // Internal parent from LangChain/LangGraph
+      parentSpanContext = this.runMap.get(parentRunId)?.otelSpan.spanContext();
+    } else if (this.externalParentSpanContext) {
+      // External parent from AG-UI.Server
+      parentSpanContext = this.externalParentSpanContext;
+    }
+
+    // Add adapter name prefix to ROOT span
+    let finalRunName = runName;
+    if (!parentRunId && this.adapterName) {
+      // ROOT span: add Adapter.LangGraph or Adapter.LangChain prefix
+      finalRunName = `Adapter.${this.adapterName}`;
+    }
+
+    const observation = startObservation(
+      finalRunName,
+      {
+        version: this.version,
+        metadata: this.joinTagsAndMetaData(tags, metadata),
+        ...attributes,
+      },
+      {
+        asType: asType ?? "span",
+        parentSpanContext,
+      }
+    );
+    this.runMap.set(runId, observation);
+
+    return observation;
+  }
+
+  private handleObservationEnd(params: {
+    runId: string;
+    attributes?: Record<string, unknown>;
+  }) {
+    const { runId, attributes = {} } = params;
+
+    const observation = this.runMap.get(runId);
+    if (!observation) {
+      this.logger.warn?.("Observation not found in runMap. Skipping operation.");
+      return;
+    }
+
+    // Type-safe update: cast to ObservationAttributes which is the union of all observation attribute types
+    observation.update(attributes as ObservationAttributes).end();
+
+    this.last_trace_id = observation.traceId;
+    this.runMap.delete(runId);
+  }
+
+  private joinTagsAndMetaData(
+    tags?: string[] | undefined,
+    metadata1?: Record<string, unknown> | undefined,
+    metadata2?: Record<string, unknown> | undefined
+  ): Record<string, unknown> | undefined {
+    const finalDict: Record<string, unknown> = {};
+    if (tags && tags.length > 0) {
+      finalDict.tags = tags;
+    }
+    if (metadata1) {
+      Object.assign(finalDict, metadata1);
+    }
+    if (metadata2) {
+      Object.assign(finalDict, metadata2);
+    }
+    return this.stripObservabilityKeysFromMetadata(finalDict);
+  }
+
+  private stripObservabilityKeysFromMetadata(
+    metadata?: Record<string, unknown>
+  ): Record<string, unknown> | undefined {
+    if (!metadata) {
+      return;
+    }
+
+    const reservedKeys = ["promptInfo", "userId", "sessionId"];
+
+    return Object.fromEntries(
+      Object.entries(metadata).filter(([key, _]) => !reservedKeys.includes(key))
+    );
+  }
+
+  private extractUsageMetadata(
+    generation: Generation
+  ): UsageMetadata | undefined {
+    try {
+      const usageMetadata =
+        "message" in generation &&
+        (AIMessage.isInstance(generation["message"]) ||
+          AIMessageChunk.isInstance(generation["message"]))
+          ? generation["message"].usage_metadata
+          : undefined;
+      return usageMetadata;
+    } catch (err) {
+      this.logger.debug?.(`Error extracting usage metadata: ${err}`);
+      return;
+    }
+  }
+
+  private extractModelNameFromMetadata(generation: any): string | undefined {
+    try {
+      return "message" in generation &&
+        (AIMessage.isInstance(generation["message"]) ||
+          AIMessageChunk.isInstance(generation["message"]))
+        ? generation["message"].response_metadata.model_name
+        : undefined;
+    } catch {}
+  }
+
+  private extractChatMessageContent(
+    message: BaseMessage
+  ): LlmMessage | AnonymousLlmMessage | MessageContent {
+    let response = undefined;
+
+    if (message.getType() === "human") {
+      response = { content: message.content, role: "user" };
+    } else if (message.getType() === "generic") {
+      response = {
+        content: message.content,
+        role: "human",
+      };
+    } else if (message.getType() === "ai") {
+      response = { content: message.content, role: "assistant" };
+
+      if (
+        "tool_calls" in message &&
+        Array.isArray(message.tool_calls) &&
+        (message.tool_calls?.length ?? 0) > 0
+      ) {
+        (response as any)["tool_calls"] = message["tool_calls"];
+      }
+      if (
+        "additional_kwargs" in message &&
+        "tool_calls" in message["additional_kwargs"]
+      ) {
+        (response as any)["tool_calls"] =
+          message["additional_kwargs"]["tool_calls"];
+      }
+    } else if (message.getType() === "system") {
+      response = { content: message.content, role: "system" };
+    } else if (message.getType() === "function") {
+      response = {
+        content: message.content,
+        additional_kwargs: message.additional_kwargs,
+        role: message.name,
+      };
+    } else if (message.getType() === "tool") {
+      response = {
+        content: message.content,
+        additional_kwargs: message.additional_kwargs,
+        role: message.name,
+      };
+    } else if (!message.name) {
+      response = { content: message.content };
+    } else {
+      response = {
+        role: message.name,
+        content: message.content,
+      };
+    }
+
+    if (
+      (message.additional_kwargs.function_call ||
+        message.additional_kwargs.tool_calls) &&
+      (response as any)["tool_calls"] === undefined
+    ) {
+      return { ...response, additional_kwargs: message.additional_kwargs };
+    }
+
+    return response;
+  }
+}