@ -1,7 +1,7 @@
from __future__ import annotations
from copy import deepcopy
from typing import Any , Dict , List , Optional
from typing import Any , Dict , List , Literal, Optional
from uuid import UUID
from langchain . load . serializable import Serializable
@ -19,6 +19,8 @@ class Generation(Serializable):
""" Raw response from the provider. May include things like the
reason for finishing or token log probabilities .
"""
type : Literal [ " Generation " ] = " Generation "
""" Type is used exclusively for serialization purposes. """
# TODO: add log probs as separate attribute
@classmethod
@ -54,11 +56,17 @@ class ChatGeneration(Generation):
""" *SHOULD NOT BE SET DIRECTLY* The text contents of the output message. """
message : BaseMessage
""" 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]
""" Type is used exclusively for serialization purposes. """
@root_validator
def set_text ( cls , values : Dict [ str , Any ] ) - > Dict [ str , Any ] :
""" Set the text attribute to be the contents of the message. """
values [ " text " ] = values [ " message " ] . content
try :
values [ " text " ] = values [ " message " ] . content
except ( KeyError , AttributeError ) as e :
raise ValueError ( " Error while initializing ChatGeneration " ) from e
return values
@ -71,6 +79,9 @@ class ChatGenerationChunk(ChatGeneration):
"""
message : BaseMessageChunk
# Override type to be ChatGeneration, ignore mypy error as this is intentional
type : Literal [ " ChatGenerationChunk " ] = " ChatGenerationChunk " # type: ignore[assignment] # noqa: E501
""" Type is used exclusively for serialization purposes. """
def __add__ ( self , other : ChatGenerationChunk ) - > ChatGenerationChunk :
if isinstance ( other , ChatGenerationChunk ) :