fix(langchain): record tool span input as structured data

LangChain's BaseTool.run/arun call on_tool_start with the tool input twice:
positionally as str(tool_input), and as the original dict under the `inputs`
keyword. The handler stored the positional value, so dict inputs landed in
$ai_input_state as a Python repr with single quotes, which no JSON parser can
read.

Prefer the `inputs` dict when LangChain supplies one, matching on_chain_start.
Tools invoked with a plain string keep recording that string.

Generated-By: PostHog Desktop
Task-Id: 75cf9ec9-fd9e-49fa-af56-ac505bc09cbe

Author: GeneralistDev

.sampo/changesets/langchain-tool-span-structured-input.md

@@ -0,0 +1,5 @@
+---
+pypi/posthog: patch
+---
+
+Capture LangChain tool inputs on `$ai_span` as structured data rather than a Python `repr` string. `BaseTool.run`/`arun` pass the tool input to `on_tool_start` twice — positionally as `str(tool_input)`, and as the original dict under the `inputs` keyword. The handler was storing the positional value, so a dict input landed in `$ai_input_state` as `{'query': 'SELECT 1'}` (single quotes), which no JSON parser can read: `JSONExtract*` in ClickHouse returns empty, and any downstream consumer has to fall back to substring matching. Tool spans now record the `inputs` dict when LangChain supplies one, matching what `on_chain_start` already does; tools invoked with a plain string are unchanged.

posthog/ai/langchain/callbacks.py

@@ -283,8 +283,15 @@ def on_tool_start(
             "on_tool_start", run_id, parent_run_id, input_str=input_str
         )
         self._set_parent_of_run(run_id, parent_run_id)
+        # LangChain hands us the tool input twice: `input_str` is `str(tool_input)`, which
+        # renders a dict as a Python repr (single quotes) rather than JSON, and the `inputs`
+        # kwarg carries the original dict. Prefer the structured value so `$ai_input_state`
+        # stays machine-readable and consistent with `on_chain_start`; tools invoked with a
+        # plain string have no `inputs` and keep using `input_str`.
+        inputs = kwargs.get("inputs")
+        tool_input = inputs if isinstance(inputs, dict) else input_str
         self._set_trace_or_span_metadata(
-            serialized, input_str, run_id, parent_run_id, **kwargs
+            serialized, tool_input, run_id, parent_run_id, **kwargs
         )
 
     def on_tool_end(

posthog/test/ai/langchain/test_callbacks.py

@@ -2442,6 +2442,48 @@ def test_agent_action_and_finish_imports():
     assert call_args["event"] == "$ai_span"
 
 
+def test_tool_span_input_state_prefers_structured_inputs(mock_client):
+    """A dict tool input is captured as a dict, not LangChain's `str(tool_input)` repr."""
+    callbacks = CallbackHandler(mock_client)
+    run_id = uuid.uuid4()
+    parent_run_id = uuid.uuid4()
+    tool_input = {"query": "SELECT 1", "truncate": True}
+
+    # Mirrors how langchain_core's BaseTool.run/arun calls on_tool_start: the positional
+    # argument is str(tool_input), while the original dict comes through `inputs`.
+    callbacks.on_tool_start(
+        {"name": "execute_sql"},
+        str(tool_input),
+        run_id=run_id,
+        parent_run_id=parent_run_id,
+        inputs=tool_input,
+    )
+    callbacks.on_tool_end("1", run_id=run_id, parent_run_id=parent_run_id)
+
+    props = mock_client.capture.call_args[1]["properties"]
+    assert props["$ai_input_state"] == tool_input
+
+
+def test_tool_span_input_state_falls_back_to_string(mock_client):
+    """A tool invoked with a plain string still records that string."""
+    callbacks = CallbackHandler(mock_client)
+    run_id = uuid.uuid4()
+    parent_run_id = uuid.uuid4()
+
+    # langchain_core passes inputs=None when tool_input isn't a dict.
+    callbacks.on_tool_start(
+        {"name": "get_weather"},
+        "sf",
+        run_id=run_id,
+        parent_run_id=parent_run_id,
+        inputs=None,
+    )
+    callbacks.on_tool_end("sunny", run_id=run_id, parent_run_id=parent_run_id)
+
+    props = mock_client.capture.call_args[1]["properties"]
+    assert props["$ai_input_state"] == "sf"
+
+
 def test_posthog_properties_field_in_generation_metadata(mock_client):
     """Test that posthog_properties is properly stored in GenerationMetadata."""
     callbacks = CallbackHandler(mock_client)