Skip to content

Commit 5a52af6

Browse files
authored
fix(ai): capture tool calls in reasoning models (#292)
* fix: capture tool calls in reasoning models * fix: check for empty tool calls
1 parent 722c887 commit 5a52af6

File tree

3 files changed

+67
-2
lines changed

3 files changed

+67
-2
lines changed

posthog/ai/langchain/callbacks.py

Lines changed: 29 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -5,6 +5,7 @@
55
"Please install LangChain to use this feature: 'pip install langchain'"
66
)
77

8+
import json
89
import logging
910
import time
1011
from dataclasses import dataclass
@@ -29,6 +30,7 @@
2930
HumanMessage,
3031
SystemMessage,
3132
ToolMessage,
33+
ToolCall,
3234
)
3335
from langchain_core.outputs import ChatGeneration, LLMResult
3436
from pydantic import BaseModel
@@ -629,12 +631,35 @@ def _extract_raw_esponse(last_response):
629631
return ""
630632

631633

632-
def _convert_message_to_dict(message: BaseMessage) -> Dict[str, Any]:
634+
def _convert_lc_tool_calls_to_oai(
635+
tool_calls: list[ToolCall],
636+
) -> list[dict[str, Any]]:
637+
try:
638+
return [
639+
{
640+
"type": "function",
641+
"id": tool_call["id"],
642+
"function": {
643+
"name": tool_call["name"],
644+
"arguments": json.dumps(tool_call["args"]),
645+
},
646+
}
647+
for tool_call in tool_calls
648+
]
649+
except KeyError:
650+
return tool_calls
651+
652+
653+
def _convert_message_to_dict(message: BaseMessage) -> dict[str, Any]:
633654
# assistant message
634655
if isinstance(message, HumanMessage):
635656
message_dict = {"role": "user", "content": message.content}
636657
elif isinstance(message, AIMessage):
637658
message_dict = {"role": "assistant", "content": message.content}
659+
if message.tool_calls:
660+
message_dict["tool_calls"] = _convert_lc_tool_calls_to_oai(
661+
message.tool_calls
662+
)
638663
elif isinstance(message, SystemMessage):
639664
message_dict = {"role": "system", "content": message.content}
640665
elif isinstance(message, ToolMessage):
@@ -647,6 +672,9 @@ def _convert_message_to_dict(message: BaseMessage) -> Dict[str, Any]:
647672
if message.additional_kwargs:
648673
message_dict.update(message.additional_kwargs)
649674

675+
if "content" in message_dict and not message_dict["content"]:
676+
message_dict["content"] = ""
677+
650678
return message_dict
651679

652680

posthog/test/ai/langchain/test_callbacks.py

Lines changed: 37 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1753,3 +1753,40 @@ def test_callback_handler_without_client():
17531753

17541754
# Verify that the mock client was used for capturing events
17551755
assert mock_client.capture.call_count == 3
1756+
1757+
1758+
def test_convert_message_to_dict_tool_calls():
1759+
"""Test that _convert_message_to_dict properly converts tool calls in AIMessage."""
1760+
from posthog.ai.langchain.callbacks import _convert_message_to_dict
1761+
from langchain_core.messages import AIMessage
1762+
from langchain_core.messages.tool import ToolCall
1763+
1764+
# Create an AIMessage with tool calls
1765+
tool_calls = [
1766+
ToolCall(
1767+
id="call_123",
1768+
name="get_weather",
1769+
args={"city": "San Francisco", "units": "celsius"},
1770+
)
1771+
]
1772+
1773+
ai_message = AIMessage(
1774+
content="I'll check the weather for you.", tool_calls=tool_calls
1775+
)
1776+
1777+
# Convert to dict
1778+
result = _convert_message_to_dict(ai_message)
1779+
1780+
# Verify the conversion
1781+
assert result["role"] == "assistant"
1782+
assert result["content"] == "I'll check the weather for you."
1783+
assert result["tool_calls"] == [
1784+
{
1785+
"type": "function",
1786+
"id": "call_123",
1787+
"function": {
1788+
"name": "get_weather",
1789+
"arguments": '{"city": "San Francisco", "units": "celsius"}',
1790+
},
1791+
}
1792+
]

posthog/version.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,4 @@
1-
VERSION = "6.3.0"
1+
VERSION = "6.3.1"
22

33
if __name__ == "__main__":
44
print(VERSION, end="") # noqa: T201

0 commit comments

Comments
 (0)