enhance tool_calls

Author: JqRrt

cozeloop/integration/langchain/trace_callback.py

@@ -12,7 +12,7 @@
 from langchain.callbacks.base import BaseCallbackHandler
 from langchain.schema import AgentFinish, AgentAction, LLMResult
 from langchain_core.prompt_values import PromptValue, ChatPromptValue
-from langchain_core.messages import BaseMessage, AIMessageChunk
+from langchain_core.messages import BaseMessage, AIMessageChunk, AIMessage
 from langchain_core.prompts import AIMessagePromptTemplate, HumanMessagePromptTemplate, SystemMessagePromptTemplate
 from langchain_core.outputs import ChatGenerationChunk, GenerationChunk
 
@@ -416,7 +416,7 @@ def _convert_inputs(inputs: Any) -> Any:
         for each in inputs:
             format_inputs.append(_convert_inputs(each))
         return format_inputs
-    if isinstance(inputs, AIMessageChunk):
+    if isinstance(inputs, (AIMessageChunk, AIMessage)):
         """
         Must be before BaseMessage.
         """

cozeloop/integration/langchain/trace_model/llm_model.py

@@ -2,6 +2,7 @@
 # SPDX-License-Identifier: MIT
 
 import json
+import logging
 import time
 from typing import List, Optional, Union, Dict, Any
 from pydantic.dataclasses import dataclass
@@ -208,6 +209,13 @@ def convert_tool_calls_by_raw(tool_calls: list) -> List[ToolCall]:
 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', '{}')))
+        raw_args = tool_call.get('function', {}).get('arguments', '{}')
+        final_args = None
+        try:
+            final_args = json.loads(raw_args)
+        except Exception as e:
+            final_args = raw_args
+            logging.error(f"convert_tool_calls_by_additional_kwargs failed, error: {e}, tool_call.function.arguments: {raw_args}")
+        function = ToolFunction(name=tool_call.get('function', {}).get('name', ''), arguments=final_args)
         format_tool_calls.append(ToolCall(id=tool_call.get('id', ''), type=tool_call.get('type', ''), function=function))
     return format_tool_calls