feat: support tool calls langchain

Chainlit · Aug 28, 2024 · 8195797 · 8195797
1 parent 493d937
commit 8195797
Showing 1 changed file with 11 additions and 0 deletions.
diff --git a/literalai/callback/langchain_callback.py b/literalai/callback/langchain_callback.py
@@ -101,6 +101,7 @@ def _convert_message_dict(
             class_name = message["id"][-1]
             kwargs = message.get("kwargs", {})
             function_call = kwargs.get("additional_kwargs", {}).get("function_call")
+            tool_calls = kwargs.get("additional_kwargs", {}).get("tool_calls")
 
             msg = GenerationMessage(
                 name=kwargs.get("name"),
@@ -113,6 +114,9 @@ def _convert_message_dict(
             else:
                 msg["content"] = kwargs.get("content", "")
 
+            if tool_calls:
+                msg["tool_calls"] = tool_calls
+
             return msg
 
         def _convert_message(
@@ -124,6 +128,7 @@ def _convert_message(
                     message,
                 )
             function_call = message.additional_kwargs.get("function_call")
+            tool_calls = message.additional_kwargs.get("tool_calls")
             msg = GenerationMessage(
                 name=getattr(message, "name", None),
                 role=_convert_message_role(message.type),
@@ -139,6 +144,9 @@ def _convert_message(
             else:
                 msg["content"] = message.content  # type: ignore
 
+            if tool_calls:
+                msg["tool_calls"] = tool_calls
+
             return msg
 
         def _build_llm_settings(
@@ -360,6 +368,7 @@ def _start_trace(self, run: Run) -> None:
             if ignore:
                 return
 
+            print(f"Starting trace for {run.id} and type {run.run_type}")
             step_type: "TrueStepType" = "undefined"
             if run.run_type == "agent":
                 step_type = "run"
@@ -417,7 +426,9 @@ def _on_run_update(self, run: Run) -> None:
                         throughput = chat_start["token_count"] / duration
                     else:
                         throughput = None
+                    print("message_completion before conversion", message)
                     message_completion = self._convert_message(message)
+                    print("message_completion", message_completion)
                     current_step.generation = ChatGeneration(
                         provider=provider,
                         model=model,