update

Author: JqRrt

cozeloop/integration/langchain/trace_model/llm_model.py

@@ -129,15 +129,15 @@ def __init__(self, messages: List[Union[BaseMessage, List[BaseMessage]]], invoca
                 for tool_call in message.additional_kwargs.get('tool_calls', []):
                     if tool_call.get('id', ''):
                         tool_call_id_name_map[tool_call.get('id', '')] = tool_call.get('function', {}).get('name', '')
-                for tool_call in message.get('tool_calls', []):
+                for tool_call in message.tool_calls:
                     if tool_call.get('id', ''):
                         tool_call_id_name_map[tool_call.get('id', '')] = tool_call.get('name', '')
 
         for message in process_messages:
             if isinstance(message, (AIMessageChunk, AIMessage)):
                 tool_calls = convert_tool_calls_by_additional_kwargs(message.additional_kwargs.get('tool_calls', []))
                 if len(tool_calls) == 0:
-                    tool_calls = convert_tool_calls_by_raw(message.get('tool_calls', []))
+                    tool_calls = convert_tool_calls_by_raw(message.tool_calls)
                 self._messages.append(Message(role=message.type, content=message.content, tool_calls=tool_calls))
             elif isinstance(message, ToolMessage):
                 name = ''
@@ -200,14 +200,14 @@ def to_json(self):
 def convert_tool_calls_by_raw(tool_calls: list) -> List[ToolCall]:
     format_tool_calls: List[ToolCall] = []
     for tool_call in tool_calls:
-        function = ToolFunction(name=tool_call.get('name', ''), arguments=json.loads(tool_call.get('args', {})))
+        function = ToolFunction(name=tool_call.get('name', ''), arguments=tool_call.get('args', {}))
         format_tool_calls.append(ToolCall(id=tool_call.get('id', ''), type=tool_call.get('type', ''), function=function))
     return format_tool_calls
 
 
 def convert_tool_calls_by_additional_kwargs(tool_calls: list) -> List[ToolCall]:
     format_tool_calls: List[ToolCall] = []
     for tool_call in tool_calls:
-        function = ToolFunction(name=tool_call.get('function', {}).get('name', ''), arguments=json.loads(tool_call.get('function', {}).get('arguments', {})))
+        function = ToolFunction(name=tool_call.get('function', {}).get('name', ''), arguments=json.loads(tool_call.get('function', {}).get('arguments', '{}')))
         format_tool_calls.append(ToolCall(id=tool_call.get('id', ''), type=tool_call.get('type', ''), function=function))
     return format_tool_calls