diff --git a/libs/core/langchain_core/outputs/chat_generation.py b/libs/core/langchain_core/outputs/chat_generation.py index d64015ed6a2..5fe542cecd2 100644 --- a/libs/core/langchain_core/outputs/chat_generation.py +++ b/libs/core/langchain_core/outputs/chat_generation.py @@ -5,7 +5,7 @@ from typing import List, Literal, Union from pydantic import model_validator from typing_extensions import Self -from langchain_core.messages import BaseMessage, BaseMessageChunk +from langchain_core.messages import AnyMessage, BaseMessageChunk from langchain_core.outputs.generation import Generation from langchain_core.utils._merge import merge_dicts @@ -26,7 +26,7 @@ class ChatGeneration(Generation): text: str = "" """*SHOULD NOT BE SET DIRECTLY* The text contents of the output message.""" - message: BaseMessage + message: AnyMessage """The message output by the chat model.""" # Override type to be ChatGeneration, ignore mypy error as this is intentional type: Literal["ChatGeneration"] = "ChatGeneration" # type: ignore[assignment]