@@ -109,10 +109,10 @@ def on_chat_model_start(
109109 input_messages : list [InputMessage ] = []
110110 for sub_messages in messages : # type: ignore[reportUnknownVariableType]
111111 for message in sub_messages : # type: ignore[reportUnknownVariableType]
112- content = get_property_value ( message , " content" ) # type: ignore[reportUnknownVariableType]
113- role = get_property_value ( message , " type" ) # type: ignore[reportUnknownArgumentType, reportUnknownVariableType]
114- parts = [Text (content = content , type = "text" )]
115- input_messages .append (InputMessage (parts = parts , role = role ))
112+ content = message . content # type: ignore[reportUnknownVariableType]
113+ role = message . type # type: ignore[reportUnknownVariableType]
114+ parts = [Text (content = content , type = "text" )] # type: ignore[reportUnknownVariableType]
115+ input_messages .append (InputMessage (parts = parts , role = role )) # type: ignore[reportUnknownVariableType]
116116
117117 llm_invocation = LLMInvocation (
118118 request_model = request_model ,
@@ -177,13 +177,11 @@ def on_llm_end(
177177 # Get message content
178178 parts = [
179179 Text (
180- content = get_property_value (
181- chat_generation .message , "content"
182- ),
180+ content = chat_generation .message .content ,
183181 type = "text" ,
184182 )
185183 ]
186- role = get_property_value ( chat_generation .message , " type" )
184+ role = chat_generation .message . type
187185 output_message = OutputMessage (
188186 role = role ,
189187 parts = parts ,
@@ -249,12 +247,3 @@ def on_llm_error(
249247 )
250248 if not llm_invocation .span .is_recording (): # type: ignore[reportOptionalMemberAccess]
251249 self ._invocation_manager .delete_invocation_state (run_id = run_id )
252-
253-
254- def get_property_value (
255- obj : dict [str , Any ] | object , property_name : str
256- ) -> Any :
257- if isinstance (obj , dict ):
258- return obj .get (property_name , None ) # type: ignore[reportUnknownParameterType]
259-
260- return getattr (obj , property_name , None )
0 commit comments