diff --git a/.sampo/changesets/langchain-web-search-count.md b/.sampo/changesets/langchain-web-search-count.md new file mode 100644 index 000000000..465d4d912 --- /dev/null +++ b/.sampo/changesets/langchain-web-search-count.md @@ -0,0 +1,5 @@ +--- +pypi/posthog: patch +--- + +The LangChain callback handler now records `$ai_web_search_count` for Anthropic responses that report `server_tool_use.web_search_requests`, as the OpenAI, Anthropic and Gemini wrappers and `posthog-js` already do. diff --git a/posthog/ai/langchain/callbacks.py b/posthog/ai/langchain/callbacks.py index 5f996fb6c..8f51fd6b3 100644 --- a/posthog/ai/langchain/callbacks.py +++ b/posthog/ai/langchain/callbacks.py @@ -706,6 +706,8 @@ def _capture_generation( ) if usage.reasoning_tokens is not None: event_properties["$ai_reasoning_tokens"] = usage.reasoning_tokens + if usage.web_search_count is not None and usage.web_search_count > 0: + event_properties["$ai_web_search_count"] = usage.web_search_count # Generation results generation_result = output.generations[-1] @@ -838,6 +840,23 @@ class ModelUsage: reasoning_tokens: Optional[int] cache_write_5m_tokens: Optional[int] = None cache_write_1h_tokens: Optional[int] = None + web_search_count: Optional[int] = None + + +def _extract_web_search_requests(usage: Union[BaseModel, dict]) -> Optional[int]: + if isinstance(usage, BaseModel): + usage = usage.__dict__ + server_tool_use = usage.get("server_tool_use") + web_search_requests = ( + server_tool_use.get("web_search_requests") + if isinstance(server_tool_use, dict) + else getattr(server_tool_use, "web_search_requests", None) + ) + if isinstance(web_search_requests, int) and not isinstance( + web_search_requests, bool + ): + return web_search_requests + return None def _parse_usage_model( @@ -906,6 +925,11 @@ def _parse_usage_model( ): parsed_usage["reasoning"] = usage["output_token_details"].get("reasoning") + # Anthropic reports the exact number of web searches of a response + web_search_requests = _extract_web_search_requests(usage) + if web_search_requests is not None: + parsed_usage["web_search"] = web_search_requests + field_mapping = { "input": "input_tokens", "output": "output_tokens", @@ -920,6 +944,7 @@ def _parse_usage_model( }, cache_write_5m_tokens=parsed_usage.get("cache_write_5m"), cache_write_1h_tokens=parsed_usage.get("cache_write_1h"), + web_search_count=parsed_usage.get("web_search"), ) # For Anthropic providers, LangChain reports input_tokens as the sum of all input tokens. # Our cost calculation expects them to be separate for Anthropic, so we subtract cache tokens. @@ -957,12 +982,19 @@ def _parse_usage( reasoning_tokens=None, ) + # The generations loop below can replace the llm_output usage, which is the + # only place langchain-anthropic reports server_tool_use, so read it here. + web_search_count: Optional[int] = None + if response.llm_output is not None: for key in llm_usage_keys: if response.llm_output.get(key): llm_usage = _parse_usage_model( response.llm_output[key], provider, model ) + web_search_count = _extract_web_search_requests( + response.llm_output[key] + ) break if hasattr(response, "generations"): @@ -1008,6 +1040,9 @@ def _parse_usage( llm_usage = _parse_usage_model(chunk_usage, provider, model) break + if llm_usage.web_search_count is None: + llm_usage.web_search_count = web_search_count + return llm_usage diff --git a/posthog/test/ai/langchain/test_callbacks.py b/posthog/test/ai/langchain/test_callbacks.py index b33e154fc..1980ed425 100644 --- a/posthog/test/ai/langchain/test_callbacks.py +++ b/posthog/test/ai/langchain/test_callbacks.py @@ -2955,6 +2955,57 @@ def test_stop_reason_resolution( assert props.get("$ai_stop_reason") == expected +@pytest.mark.parametrize( + "server_tool_use,expected", + [ + # Anthropic reports the exact number of searches + ({"web_search_requests": 2}, 2), + # no search is not worth an event property + ({"web_search_requests": 0}, None), + # the provider reports nothing + (None, None), + ], +) +def test_web_search_count(mock_client, server_tool_use, expected): + from langchain_core.outputs import ChatGeneration, LLMResult + + cb = CallbackHandler(mock_client) + run_id = uuid.uuid4() + cb._set_llm_metadata( + serialized={}, + run_id=run_id, + messages=[{"role": "user", "content": "test"}], + metadata={"ls_provider": "anthropic", "ls_model_name": "claude-sonnet"}, + ) + usage = {"input_tokens": 10, "output_tokens": 5} + if server_tool_use is not None: + usage["server_tool_use"] = server_tool_use + # langchain-anthropic always sets usage_metadata on the message, which the + # generations loop prefers over llm_output["usage"] + response = LLMResult( + generations=[ + [ + ChatGeneration( + message=AIMessage( + content="r", + usage_metadata={ + "input_tokens": 10, + "output_tokens": 5, + "total_tokens": 15, + }, + ) + ) + ] + ], + llm_output={"usage": usage}, + ) + + cb._pop_run_and_capture_generation(run_id, None, response) + + props = mock_client.capture.call_args.kwargs["properties"] + assert props.get("$ai_web_search_count") == expected + + class TestGetHttpStatus: def test_uses_the_status_code_of_an_http_error(self): from posthog.ai.langchain.callbacks import _get_http_status diff --git a/references/public_api_snapshot.txt b/references/public_api_snapshot.txt index a7d09a3ff..f37ae758e 100644 --- a/references/public_api_snapshot.txt +++ b/references/public_api_snapshot.txt @@ -593,6 +593,7 @@ attribute posthog.ai.langchain.callbacks.ModelUsage.cache_write_tokens: Optional attribute posthog.ai.langchain.callbacks.ModelUsage.input_tokens: Optional[int] attribute posthog.ai.langchain.callbacks.ModelUsage.output_tokens: Optional[int] attribute posthog.ai.langchain.callbacks.ModelUsage.reasoning_tokens: Optional[int] +attribute posthog.ai.langchain.callbacks.ModelUsage.web_search_count: Optional[int] = None attribute posthog.ai.langchain.callbacks.RunMetadata = Union[SpanMetadata, GenerationMetadata] attribute posthog.ai.langchain.callbacks.RunMetadataStorage = Dict[UUID, RunMetadata] attribute posthog.ai.langchain.callbacks.SpanMetadata.end_time: Optional[float] @@ -1231,7 +1232,7 @@ class posthog.ai.gemini.gemini_converter.GeminiMessage class posthog.ai.gemini.gemini_converter.GeminiPart class posthog.ai.langchain.callbacks.CallbackHandler(client: Optional[Client] = None, *, distinct_id: Optional[Union[str, int, UUID]] = None, trace_id: Optional[Union[str, int, float, UUID]] = None, properties: Optional[Dict[str, Any]] = None, privacy_mode: bool = False, groups: Optional[Dict[str, Any]] = None) class posthog.ai.langchain.callbacks.GenerationMetadata(name: str, start_time: float, end_time: Optional[float], input: Optional[Any], provider: Optional[str] = None, model: Optional[str] = None, model_params: Optional[Dict[str, Any]] = None, base_url: Optional[str] = None, tools: Optional[List[Dict[str, Any]]] = None, posthog_properties: Optional[Dict[str, Any]] = None) -class posthog.ai.langchain.callbacks.ModelUsage(input_tokens: Optional[int], output_tokens: Optional[int], cache_write_tokens: Optional[int], cache_read_tokens: Optional[int], reasoning_tokens: Optional[int], cache_write_5m_tokens: Optional[int] = None, cache_write_1h_tokens: Optional[int] = None) +class posthog.ai.langchain.callbacks.ModelUsage(input_tokens: Optional[int], output_tokens: Optional[int], cache_write_tokens: Optional[int], cache_read_tokens: Optional[int], reasoning_tokens: Optional[int], cache_write_5m_tokens: Optional[int] = None, cache_write_1h_tokens: Optional[int] = None, web_search_count: Optional[int] = None) class posthog.ai.langchain.callbacks.SpanMetadata(name: str, start_time: float, end_time: Optional[float], input: Optional[Any]) class posthog.ai.langchain.middleware.PostHogMiddleware(client: Optional[Client] = None, *, distinct_id: Optional[Union[str, int, UUID]] = None, trace_id: Optional[Union[str, int, float, UUID]] = None, properties: Optional[dict[str, Any]] = None, privacy_mode: bool = False, groups: Optional[dict[str, Any]] = None) class posthog.ai.openai.openai.OpenAI(posthog_client: Optional[PostHogClient] = None, **kwargs)