fix(llmo): set the $ai_tools properly for all providers

Author: Radu-Raicea

posthog/ai/langchain/callbacks.py

@@ -556,6 +556,7 @@ def _capture_generation(
             "$ai_latency": run.latency,
             "$ai_base_url": run.base_url,
         }
+
         if run.tools:
             event_properties["$ai_tools"] = with_privacy_mode(
                 self._ph_client,

posthog/ai/openai/openai.py

@@ -11,6 +11,7 @@
 
 from posthog.ai.utils import (
     call_llm_and_track_usage,
+    extract_available_tool_calls,
     get_model_params,
     with_privacy_mode,
 )
@@ -167,6 +168,7 @@ def generator():
                     usage_stats,
                     latency,
                     output,
+                    extract_available_tool_calls("openai", kwargs),
                 )
 
         return generator()
@@ -341,7 +343,6 @@ def _create_streaming(
         start_time = time.time()
         usage_stats: Dict[str, int] = {}
         accumulated_content = []
-        accumulated_tools = {}
         if "stream_options" not in kwargs:
             kwargs["stream_options"] = {}
         kwargs["stream_options"]["include_usage"] = True
@@ -350,7 +351,6 @@ def _create_streaming(
         def generator():
             nonlocal usage_stats
             nonlocal accumulated_content  # noqa: F824
-            nonlocal accumulated_tools  # noqa: F824
 
             try:
                 for chunk in response:
@@ -389,31 +389,12 @@ def generator():
                             if content:
                                 accumulated_content.append(content)
 
-                        # Process tool calls
-                        tool_calls = getattr(chunk.choices[0].delta, "tool_calls", None)
-                        if tool_calls:
-                            for tool_call in tool_calls:
-                                index = tool_call.index
-                                if index not in accumulated_tools:
-                                    accumulated_tools[index] = tool_call
-                                else:
-                                    # Append arguments for existing tool calls
-                                    if hasattr(tool_call, "function") and hasattr(
-                                        tool_call.function, "arguments"
-                                    ):
-                                        accumulated_tools[
-                                            index
-                                        ].function.arguments += (
-                                            tool_call.function.arguments
-                                        )
-
                     yield chunk
 
             finally:
                 end_time = time.time()
                 latency = end_time - start_time
                 output = "".join(accumulated_content)
-                tools = list(accumulated_tools.values()) if accumulated_tools else None
                 self._capture_streaming_event(
                     posthog_distinct_id,
                     posthog_trace_id,
@@ -424,7 +405,7 @@ def generator():
                     usage_stats,
                     latency,
                     output,
-                    tools,
+                    extract_available_tool_calls("openai", kwargs),
                 )
 
         return generator()

posthog/ai/openai/openai_async.py

@@ -12,6 +12,7 @@
 from posthog import setup
 from posthog.ai.utils import (
     call_llm_and_track_usage_async,
+    extract_available_tool_calls,
     get_model_params,
     with_privacy_mode,
 )
@@ -168,6 +169,7 @@ async def async_generator():
                     usage_stats,
                     latency,
                     output,
+                    extract_available_tool_calls("openai", kwargs),
                 )
 
         return async_generator()
@@ -344,7 +346,6 @@ async def _create_streaming(
         start_time = time.time()
         usage_stats: Dict[str, int] = {}
         accumulated_content = []
-        accumulated_tools = {}
 
         if "stream_options" not in kwargs:
             kwargs["stream_options"] = {}
@@ -354,7 +355,6 @@ async def _create_streaming(
         async def async_generator():
             nonlocal usage_stats
             nonlocal accumulated_content  # noqa: F824
-            nonlocal accumulated_tools  # noqa: F824
 
             try:
                 async for chunk in response:
@@ -393,31 +393,12 @@ async def async_generator():
                             if content:
                                 accumulated_content.append(content)
 
-                        # Process tool calls
-                        tool_calls = getattr(chunk.choices[0].delta, "tool_calls", None)
-                        if tool_calls:
-                            for tool_call in tool_calls:
-                                index = tool_call.index
-                                if index not in accumulated_tools:
-                                    accumulated_tools[index] = tool_call
-                                else:
-                                    # Append arguments for existing tool calls
-                                    if hasattr(tool_call, "function") and hasattr(
-                                        tool_call.function, "arguments"
-                                    ):
-                                        accumulated_tools[
-                                            index
-                                        ].function.arguments += (
-                                            tool_call.function.arguments
-                                        )
-
                     yield chunk
 
             finally:
                 end_time = time.time()
                 latency = end_time - start_time
                 output = "".join(accumulated_content)
-                tools = list(accumulated_tools.values()) if accumulated_tools else None
                 await self._capture_streaming_event(
                     posthog_distinct_id,
                     posthog_trace_id,
@@ -428,7 +409,7 @@ async def async_generator():
                     usage_stats,
                     latency,
                     output,
-                    tools,
+                    extract_available_tool_calls("openai", kwargs),
                 )
 
         return async_generator()

posthog/ai/utils.py

@@ -228,41 +228,22 @@ def format_response_gemini(response):
     return output
 
 
-def format_tool_calls(response, provider: str):
+def extract_available_tool_calls(provider: str, kwargs: Dict[str, Any]):
     if provider == "anthropic":
-        if hasattr(response, "content") and response.content:
-            tool_calls = []
+        if "tools" in kwargs:
+            return kwargs["tools"]
 
-            for content_item in response.content:
-                if hasattr(content_item, "type") and content_item.type == "tool_use":
-                    tool_calls.append(
-                        {
-                            "type": content_item.type,
-                            "id": content_item.id,
-                            "name": content_item.name,
-                            "input": content_item.input,
-                        }
-                    )
+        return None
+    elif provider == "gemini":
+        if "config" in kwargs and hasattr(kwargs["config"], "tools"):
+            return kwargs["config"].tools
 
-            return tool_calls if tool_calls else None
+        return None
     elif provider == "openai":
-        # Handle both Chat Completions and Responses API
-        if hasattr(response, "choices") and response.choices:
-            # Check for tool_calls in message (Chat Completions format)
-            if (
-                hasattr(response.choices[0], "message")
-                and hasattr(response.choices[0].message, "tool_calls")
-                and response.choices[0].message.tool_calls
-            ):
-                return response.choices[0].message.tool_calls
-
-            # Check for tool_calls directly in response (Responses API format)
-            if (
-                hasattr(response.choices[0], "tool_calls")
-                and response.choices[0].tool_calls
-            ):
-                return response.choices[0].tool_calls
-    return None
+        if "tools" in kwargs:
+            return kwargs["tools"]
+
+        return None
 
 
 def merge_system_prompt(kwargs: Dict[str, Any], provider: str):
@@ -395,11 +376,11 @@ def call_llm_and_track_usage(
             **(error_params or {}),
         }
 
-        tool_calls = format_tool_calls(response, provider)
+        available_tool_calls = extract_available_tool_calls(provider, kwargs)
 
-        if tool_calls:
+        if available_tool_calls:
             event_properties["$ai_tools"] = with_privacy_mode(
-                ph_client, posthog_privacy_mode, tool_calls
+                ph_client, posthog_privacy_mode, available_tool_calls
             )
 
         if (
@@ -511,10 +492,11 @@ async def call_llm_and_track_usage_async(
             **(error_params or {}),
         }
 
-        tool_calls = format_tool_calls(response, provider)
-        if tool_calls:
+        available_tool_calls = extract_available_tool_calls(provider, kwargs)
+
+        if available_tool_calls:
             event_properties["$ai_tools"] = with_privacy_mode(
-                ph_client, posthog_privacy_mode, tool_calls
+                ph_client, posthog_privacy_mode, available_tool_calls
             )
 
         if (

posthog/test/ai/anthropic/test_anthropic.py

@@ -88,31 +88,6 @@ def mock_anthropic_response_with_cached_tokens():
     )
 
 
-@pytest.fixture
-def mock_anthropic_response_with_tool_use():
-    return Message(
-        id="msg_123",
-        type="message",
-        role="assistant",
-        content=[
-            {"type": "text", "text": "I'll help you with that."},
-            {
-                "type": "tool_use",
-                "id": "tool_1",
-                "name": "get_weather",
-                "input": {"location": "New York"},
-            },
-        ],
-        model="claude-3-opus-20240229",
-        usage=Usage(
-            input_tokens=20,
-            output_tokens=10,
-        ),
-        stop_reason="end_turn",
-        stop_sequence=None,
-    )
-
-
 def test_basic_completion(mock_client, mock_anthropic_response):
     with patch(
         "anthropic.resources.Messages.create", return_value=mock_anthropic_response
@@ -461,20 +436,41 @@ def test_cached_tokens(mock_client, mock_anthropic_response_with_cached_tokens):
         assert isinstance(props["$ai_latency"], float)
 
 
-def test_tool_use_response(mock_client, mock_anthropic_response_with_tool_use):
+def test_tool_definition(mock_client, mock_anthropic_response):
     with patch(
         "anthropic.resources.Messages.create",
-        return_value=mock_anthropic_response_with_tool_use,
+        return_value=mock_anthropic_response,
     ):
         client = Anthropic(api_key="test-key", posthog_client=mock_client)
+        
+        tools = [
+            {
+                "name": "get_weather",
+                "description": "Get the current weather for a specific location",
+                "input_schema": {
+                    "type": "object",
+                    "properties": {
+                        "location": {
+                            "type": "string",
+                            "description": "The city or location name to get weather for"
+                        }
+                    },
+                    "required": ["location"]
+                }
+            }
+        ]
+        
         response = client.messages.create(
-            model="claude-3-opus-20240229",
-            messages=[{"role": "user", "content": "What's the weather like?"}],
+            model="claude-3-5-sonnet-20241022",
+            max_tokens=200,
+            temperature=0.7,
+            tools=tools,
+            messages=[{"role": "user", "content": "hey"}],
             posthog_distinct_id="test-id",
             posthog_properties={"foo": "bar"},
         )
 
-        assert response == mock_anthropic_response_with_tool_use
+        assert response == mock_anthropic_response
         assert mock_client.capture.call_count == 1
 
         call_args = mock_client.capture.call_args[1]
@@ -483,25 +479,15 @@ def test_tool_use_response(mock_client, mock_anthropic_response_with_tool_use):
         assert call_args["distinct_id"] == "test-id"
         assert call_args["event"] == "$ai_generation"
         assert props["$ai_provider"] == "anthropic"
-        assert props["$ai_model"] == "claude-3-opus-20240229"
-        assert props["$ai_input"] == [
-            {"role": "user", "content": "What's the weather like?"}
-        ]
-        # Should only include text content, not tool_use content
+        assert props["$ai_model"] == "claude-3-5-sonnet-20241022"
+        assert props["$ai_input"] == [{"role": "user", "content": "hey"}]
         assert props["$ai_output_choices"] == [
-            {"role": "assistant", "content": "I'll help you with that."}
+            {"role": "assistant", "content": "Test response"}
         ]
         assert props["$ai_input_tokens"] == 20
         assert props["$ai_output_tokens"] == 10
         assert props["$ai_http_status"] == 200
         assert props["foo"] == "bar"
         assert isinstance(props["$ai_latency"], float)
-        # Verify that tools are captured separately
-        assert props["$ai_tools"] == [
-            {
-                "type": "tool_use",
-                "id": "tool_1",
-                "name": "get_weather",
-                "input": {"location": "New York"},
-            }
-        ]
+        # Verify that tools are captured in the $ai_tools property
+        assert props["$ai_tools"] == tools