feat(llm-observability): $ai_tools capture in Langchain (#199)

Author: Twixes

CHANGELOG.md

@@ -1,3 +1,7 @@
+## 3.17.0 - 2025-02-27
+
+1. The LangChain handler now captures tools in `$ai_generation` events, in property `$ai_tools`. This allows for displaying tools provided to the LLM call in PostHog UI. Note that support for `$ai_tools` in OpenAI and Anthropic SDKs is coming soon.
+
 ## 3.16.0 - 2025-02-26
 
 1. feat: add some platform info to events (#198)

posthog/ai/langchain/callbacks.py

@@ -60,6 +60,8 @@ class GenerationMetadata(SpanMetadata):
     """Model parameters of the run: temperature, max_tokens, etc."""
     base_url: Optional[str] = None
     """Base URL of the provider's API used in the run."""
+    tools: Optional[List[Dict[str, Any]]] = None
+    """Tools provided to the model."""
 
 
 RunMetadata = Union[SpanMetadata, GenerationMetadata]
@@ -377,6 +379,8 @@ def _set_llm_metadata(
         generation = GenerationMetadata(name=run_name, input=messages, start_time=time.time(), end_time=None)
         if isinstance(invocation_params, dict):
             generation.model_params = get_model_params(invocation_params)
+            if tools := invocation_params.get("tools"):
+                generation.tools = tools
         if isinstance(metadata, dict):
             if model := metadata.get("ls_model_name"):
                 generation.model = model
@@ -424,7 +428,11 @@ def _pop_run_and_capture_trace_or_span(self, run_id: UUID, parent_run_id: Option
             log.warning(f"Run {run_id} is a generation, but attempted to be captured as a trace or span.")
             return
         self._capture_trace_or_span(
-            trace_id, run_id, run, outputs, self._get_parent_run_id(trace_id, run_id, parent_run_id)
+            trace_id,
+            run_id,
+            run,
+            outputs,
+            self._get_parent_run_id(trace_id, run_id, parent_run_id),
         )
 
     def _capture_trace_or_span(
@@ -465,7 +473,10 @@ def _capture_trace_or_span(
         )
 
     def _pop_run_and_capture_generation(
-        self, run_id: UUID, parent_run_id: Optional[UUID], response: Union[LLMResult, BaseException]
+        self,
+        run_id: UUID,
+        parent_run_id: Optional[UUID],
+        response: Union[LLMResult, BaseException],
     ):
         trace_id = self._get_trace_id(run_id)
         self._pop_parent_of_run(run_id)
@@ -476,7 +487,11 @@ def _pop_run_and_capture_generation(
             log.warning(f"Run {run_id} is not a generation, but attempted to be captured as a generation.")
             return
         self._capture_generation(
-            trace_id, run_id, run, response, self._get_parent_run_id(trace_id, run_id, parent_run_id)
+            trace_id,
+            run_id,
+            run,
+            response,
+            self._get_parent_run_id(trace_id, run_id, parent_run_id),
         )
 
     def _capture_generation(
@@ -500,6 +515,8 @@ def _capture_generation(
             "$ai_latency": run.latency,
             "$ai_base_url": run.base_url,
         }
+        if run.tools:
+            event_properties["$ai_tools"] = run.tools
 
         if isinstance(output, BaseException):
             event_properties["$ai_http_status"] = _get_http_status(output)

posthog/test/ai/langchain/test_callbacks.py

@@ -1168,6 +1168,61 @@ async def test_async_anthropic_streaming(mock_client):
     assert isinstance(trace_props["$ai_output_state"], AIMessage)
 
 
+def test_metadata_tools(mock_client):
+    callbacks = CallbackHandler(mock_client)
+    run_id = uuid.uuid4()
+    tools = [
+        [
+            {
+                "type": "function",
+                "function": {
+                    "name": "foo",
+                    "description": "The foo.",
+                    "parameters": {
+                        "properties": {
+                            "bar": {
+                                "description": "The bar of foo.",
+                                "type": "string",
+                            },
+                        },
+                        "required": ["query_description", "query_kind"],
+                        "type": "object",
+                        "additionalProperties": False,
+                    },
+                    "strict": True,
+                },
+            }
+        ]
+    ]
+
+    with patch("time.time", return_value=1234567890):
+        callbacks._set_llm_metadata(
+            {"kwargs": {"openai_api_base": "https://us.posthog.com"}},
+            run_id,
+            messages=[{"role": "user", "content": "What's the weather like in SF?"}],
+            invocation_params={"temperature": 0.5, "tools": tools},
+            metadata={"ls_model_name": "hog-mini", "ls_provider": "posthog"},
+            name="test",
+        )
+    expected = GenerationMetadata(
+        model="hog-mini",
+        input=[{"role": "user", "content": "What's the weather like in SF?"}],
+        start_time=1234567890,
+        model_params={"temperature": 0.5},
+        provider="posthog",
+        base_url="https://us.posthog.com",
+        name="test",
+        tools=tools,
+        end_time=None,
+    )
+    assert callbacks._runs[run_id] == expected
+    with patch("time.time", return_value=1234567891):
+        run = callbacks._pop_run_metadata(run_id)
+    expected.end_time = 1234567891
+    assert run == expected
+    assert callbacks._runs == {}
+
+
 def test_tool_calls(mock_client):
     prompt = ChatPromptTemplate.from_messages([("user", "Foo")])
     model = FakeMessagesListChatModel(

posthog/version.py

@@ -1,4 +1,4 @@
-VERSION = "3.16.0"
+VERSION = "3.17.0"
 
 if __name__ == "__main__":
     print(VERSION, end="")  # noqa: T201