feat: add tool support

Author: k11kirky

posthog/ai/langchain/callbacks.py

@@ -516,7 +516,11 @@ def _capture_generation(
             "$ai_base_url": run.base_url,
         }
         if run.tools:
-            event_properties["$ai_tools"] = run.tools
+            event_properties["$ai_tools"] = with_privacy_mode(
+                self._client,
+                self._privacy_mode,
+                run.tools,
+            )
 
         if isinstance(output, BaseException):
             event_properties["$ai_http_status"] = _get_http_status(output)

posthog/ai/openai/openai.py

@@ -92,6 +92,7 @@ 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
@@ -100,6 +101,7 @@ def _create_streaming(
         def generator():
             nonlocal usage_stats
             nonlocal accumulated_content
+            nonlocal accumulated_tools
 
             try:
                 for chunk in response:
@@ -122,12 +124,25 @@ 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,
@@ -138,6 +153,7 @@ def generator():
                     usage_stats,
                     latency,
                     output,
+                    tools,
                 )
 
         return generator()
@@ -153,6 +169,7 @@ def _capture_streaming_event(
         usage_stats: Dict[str, int],
         latency: float,
         output: str,
+        tool_calls=None,
     ):
         if posthog_trace_id is None:
             posthog_trace_id = uuid.uuid4()
@@ -177,6 +194,13 @@ def _capture_streaming_event(
             **posthog_properties,
         }
 
+        if tool_calls:
+            event_properties["$ai_tools"] = with_privacy_mode(
+                self._client._ph_client,
+                posthog_privacy_mode,
+                tool_calls,
+            )
+
         if posthog_distinct_id is None:
             event_properties["$process_person_profile"] = False
 

posthog/ai/openai/openai_async.py

@@ -93,13 +93,14 @@ 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"] = {}
         kwargs["stream_options"]["include_usage"] = True
         response = await super().create(**kwargs)
 
         async def async_generator():
-            nonlocal usage_stats, accumulated_content
+            nonlocal usage_stats, accumulated_content, accumulated_tools
             try:
                 async for chunk in response:
                     if hasattr(chunk, "usage") and chunk.usage:
@@ -121,12 +122,25 @@ 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,
@@ -137,6 +151,7 @@ async def async_generator():
                     usage_stats,
                     latency,
                     output,
+                    tools,
                 )
 
         return async_generator()
@@ -152,6 +167,7 @@ async def _capture_streaming_event(
         usage_stats: Dict[str, int],
         latency: float,
         output: str,
+        tool_calls=None,
     ):
         if posthog_trace_id is None:
             posthog_trace_id = uuid.uuid4()
@@ -176,6 +192,13 @@ async def _capture_streaming_event(
             **posthog_properties,
         }
 
+        if tool_calls:
+            event_properties["$ai_tools"] = with_privacy_mode(
+                self._client._ph_client,
+                posthog_privacy_mode,
+                tool_calls,
+            )
+
         if posthog_distinct_id is None:
             event_properties["$process_person_profile"] = False
 

posthog/ai/utils.py

@@ -94,6 +94,16 @@ def format_response_openai(response):
     return output
 
 
+def format_tool_calls(response, provider: str):
+    if provider == "anthropic":
+        if hasattr(response, "tools") and response.tools and len(response.tools) > 0:
+            return response.tools
+    elif provider == "openai":
+        if (hasattr(response, "choices") and response.choices and hasattr(response.choices[0].message, "tool_calls") and response.choices[0].message.tool_calls):
+            return response.choices[0].message.tool_calls
+    return None
+
+
 def merge_system_prompt(kwargs: Dict[str, Any], provider: str):
     if provider != "anthropic":
         return kwargs.get("messages")
@@ -165,6 +175,10 @@ def call_llm_and_track_usage(
             **(error_params or {}),
         }
 
+        tool_calls = format_tool_calls(response, provider)
+        if tool_calls:
+            event_properties["$ai_tools"] = with_privacy_mode(ph_client, posthog_privacy_mode, tool_calls)
+
         if usage.get("cache_read_input_tokens", 0) > 0:
             event_properties["$ai_cache_read_input_tokens"] = usage.get("cache_read_input_tokens", 0)
 
@@ -247,6 +261,10 @@ async def call_llm_and_track_usage_async(
             **(error_params or {}),
         }
 
+        tool_calls = format_tool_calls(response, provider)
+        if tool_calls:
+            event_properties["$ai_tools"] = with_privacy_mode(ph_client, posthog_privacy_mode, tool_calls)
+
         if usage.get("cache_read_input_tokens", 0) > 0:
             event_properties["$ai_cache_read_input_tokens"] = usage.get("cache_read_input_tokens", 0)