From 493d93790f18fc93d03a44e77409a66068e4cd9c Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Hugues=20de=20Saxc=C3=A9?= <hugues@chainlit.io>
Date: Wed, 28 Aug 2024 15:00:21 +0200
Subject: [PATCH 1/3] fix: async invokes langchain for agent executors

---
 examples/langchain_toolcall.py         | 2 ++
 literalai/prompt_engineering/prompt.py | 3 +++
 2 files changed, 5 insertions(+)

diff --git a/examples/langchain_toolcall.py b/examples/langchain_toolcall.py
index c7b6415..6fc3939 100644
--- a/examples/langchain_toolcall.py
+++ b/examples/langchain_toolcall.py
@@ -11,6 +11,7 @@
 
 from dotenv import load_dotenv
 
+# Add OPENAI_API_KEY and TAVILY_API_KEY for this example.
 load_dotenv()
 
 model = ChatOpenAI(model="gpt-4o")
@@ -43,6 +44,7 @@
 
 cb = lai_client.langchain_callback()
 
+# Replace with ainvoke for asynchronous execution.
 agent_executor.invoke(
     {
         "chat_history": [
diff --git a/literalai/prompt_engineering/prompt.py b/literalai/prompt_engineering/prompt.py
index fdaf609..19e3933 100644
--- a/literalai/prompt_engineering/prompt.py
+++ b/literalai/prompt_engineering/prompt.py
@@ -231,6 +231,9 @@ def format_messages(self, **kwargs: Any) -> List[BaseMessage]:
 
                 return rendered_messages
 
+            async def aformat_messages(self, **kwargs: Any) -> List[BaseMessage]:
+                return self.format_messages(**kwargs)
+
         lc_messages = [(m["role"], m["content"]) for m in self.template_messages]
 
         chat_template = CustomChatPromptTemplate.from_messages(lc_messages)

From 81957975f138ba890c91e97f4f1aeb96f3494118 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Hugues=20de=20Saxc=C3=A9?= <hugues@chainlit.io>
Date: Wed, 28 Aug 2024 16:56:49 +0200
Subject: [PATCH 2/3] feat: support tool calls langchain

---
 literalai/callback/langchain_callback.py | 11 +++++++++++
 1 file changed, 11 insertions(+)

diff --git a/literalai/callback/langchain_callback.py b/literalai/callback/langchain_callback.py
index 703b506..806e66d 100644
--- 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,

From c64e602700fa916f7ea79a3f85edd3a394221d52 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Hugues=20de=20Saxc=C3=A9?= <hugues@chainlit.io>
Date: Wed, 28 Aug 2024 17:01:55 +0200
Subject: [PATCH 3/3] fix: remove prints

---
 literalai/callback/langchain_callback.py | 3 ---
 1 file changed, 3 deletions(-)

diff --git a/literalai/callback/langchain_callback.py b/literalai/callback/langchain_callback.py
index 806e66d..57db15a 100644
--- a/literalai/callback/langchain_callback.py
+++ b/literalai/callback/langchain_callback.py
@@ -368,7 +368,6 @@ 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"
@@ -426,9 +425,7 @@ 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,