Skip to content
Closed
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
5 changes: 5 additions & 0 deletions .sampo/changesets/langchain-tool-span-structured-input.md
Original file line number Diff line number Diff line change
@@ -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.
9 changes: 8 additions & 1 deletion posthog/ai/langchain/callbacks.py
Original file line number Diff line number Diff line change
Expand Up @@ -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(
Expand Down
42 changes: 42 additions & 0 deletions posthog/test/ai/langchain/test_callbacks.py
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand Down