Skip to content

Commit 63e9802

Browse files
committed
fix: capture tool calls in reasoning models
1 parent 7a8b091 commit 63e9802

File tree

3 files changed

+69
-3
lines changed

3 files changed

+69
-3
lines changed

posthog/ai/langchain/callbacks.py

Lines changed: 30 additions & 2 deletions
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
@@ -630,12 +632,35 @@ def _extract_raw_esponse(last_response):
630632
return ""
631633

632634

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

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

653681

posthog/test/ai/langchain/test_callbacks.py

Lines changed: 38 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1727,3 +1727,41 @@ def test_openai_reasoning_tokens(mock_client):
17271727
assert call["properties"]["$ai_reasoning_tokens"] is not None
17281728
assert call["properties"]["$ai_input_tokens"] is not None
17291729
assert call["properties"]["$ai_output_tokens"] is not None
1730+
1731+
1732+
def test_convert_message_to_dict_tool_calls():
1733+
"""Test that _convert_message_to_dict properly converts tool calls in AIMessage."""
1734+
from posthog.ai.langchain.callbacks import _convert_message_to_dict
1735+
from langchain_core.messages import AIMessage
1736+
from langchain_core.messages.tool import ToolCall
1737+
1738+
# Create an AIMessage with tool calls
1739+
tool_calls = [
1740+
ToolCall(
1741+
id="call_123",
1742+
name="get_weather",
1743+
args={"city": "San Francisco", "units": "celsius"}
1744+
)
1745+
]
1746+
1747+
ai_message = AIMessage(
1748+
content="I'll check the weather for you.",
1749+
tool_calls=tool_calls
1750+
)
1751+
1752+
# Convert to dict
1753+
result = _convert_message_to_dict(ai_message)
1754+
1755+
# Verify the conversion
1756+
assert result["role"] == "assistant"
1757+
assert result["content"] == "I'll check the weather for you."
1758+
assert result["tool_calls"] == [
1759+
{
1760+
"type": "function",
1761+
"id": "call_123",
1762+
"function": {
1763+
"name": "get_weather",
1764+
"arguments": '{"city": "San Francisco", "units": "celsius"}',
1765+
},
1766+
}
1767+
]

posthog/version.py

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

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

0 commit comments

Comments
 (0)