From b6e1a871af33c876a116cbc92f79892053c330a1 Mon Sep 17 00:00:00 2001 From: Starfolk Date: Fri, 17 Jul 2026 20:35:42 +0000 Subject: [PATCH 1/2] fix(langchain): add provider metadata, route tools to metadata, allowlist span fields MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Align the LangChain callback handler with `.agents/skills/sdk-integrations/SKILL.md`. - **Missing `metadata.provider`** — spec requires every `llm` span to carry both `metadata.model` and `metadata.provider`. Derive provider at start from `serialized["id"]` (strip `langchain_` prefix + small alias table for google/aws/azure/mistral) and at end from `response_metadata.model_provider` as a fallback. - **Tools in `invocation_params` instead of `metadata.tools`** — lift `tools` / `functions` to `metadata.tools` and `tool_choice` / `function_call` to `metadata.tool_choice` via `_split_tools`. - **Denylist → allowlist** — dropped the `exclude_metadata_props` regex (dead code — was set but never applied) and stopped spreading `**kwargs` and the full `serialized` runnable dump into span metadata across every callback. Each callback now captures only the LangChain user metadata (preserved as `metadata.metadata` so `deepagents` tests still resolve `ls_integration` / `lc_agent_name` / `lc_versions`), tags, and known-safe fields. - **Excess serialization** — `on_tool_start` no longer duplicates the tool payload as `inputs` + `input_str` + `input` (parsed) + `serialized`; keeps only the parsed input. Removed the raw kwargs dump in every error callback. - **`on_agent_action` span type** — `LLM` → `TOOL` (the callback fires when the agent picks a tool, not when it calls an LLM). - Trimmed narrator comments per the skill guidance. No new mocks, no re-recorded cassettes, no new test files. Existing VCR suite (`test_callbacks.py`, `test_anthropic.py`, `test_deepagents.py`) covers every changed path; added positive `metadata.provider` / `metadata.tools` assertions on the tests whose cassettes exercise them. The LangChain-native message / `LLMResult` shape in `input` / `output` is left as-is — the Braintrust backend already normalizes it. Co-Authored-By: Claude Opus 4.7 --- .../integrations/langchain/callbacks.py | 354 +++++++++--------- .../integrations/langchain/test_anthropic.py | 1 + .../integrations/langchain/test_callbacks.py | 26 +- 3 files changed, 200 insertions(+), 181 deletions(-) diff --git a/py/src/braintrust/integrations/langchain/callbacks.py b/py/src/braintrust/integrations/langchain/callbacks.py index 9d9517c9..d1fd8b1c 100644 --- a/py/src/braintrust/integrations/langchain/callbacks.py +++ b/py/src/braintrust/integrations/langchain/callbacks.py @@ -1,30 +1,13 @@ import json import logging -import re import time from collections.abc import Mapping, Sequence -from re import Pattern -from typing import ( - Any, - TypedDict, -) +from typing import Any, TypedDict from uuid import UUID from braintrust.generated_types import SpanAttributes from braintrust.logger import NOOP_SPAN, Logger, Span, current_span, init_logger from braintrust.logger import start_span as _bt_start_span - - -_INSTRUMENTATION = "langchain-auto" - - -def start_span(*args, **kwargs): - internal = dict(kwargs.get("internal") or {}) - internal.setdefault("instrumentation", _INSTRUMENTATION) - kwargs["internal"] = internal - return _bt_start_span(*args, **kwargs) - - from braintrust.span_types import SpanTypeAttribute from braintrust.version import VERSION as sdk_version from langchain_core.agents import AgentAction, AgentFinish @@ -36,9 +19,17 @@ def start_span(*args, **kwargs): from typing_extensions import NotRequired +_INSTRUMENTATION = "langchain-auto" +_INTEGRATION_NAME = "langchain-py" + _logger = logging.getLogger("braintrust.wrappers.langchain") -_INTEGRATION_NAME = "langchain-py" + +def start_span(*args, **kwargs): + internal = dict(kwargs.get("internal") or {}) + internal.setdefault("instrumentation", _INSTRUMENTATION) + kwargs["internal"] = internal + return _bt_start_span(*args, **kwargs) class LogEvent(TypedDict): @@ -54,6 +45,63 @@ class LogEvent(TypedDict): dataset_record_id: NotRequired[str] +# Only aliases where the package name doesn't strip cleanly to the provider. +_PROVIDER_ALIASES: dict[str, str] = { + "langchain_google_vertexai": "google", + "langchain_google_genai": "google", + "langchain_google": "google", + "langchain_aws": "aws", + "langchain_bedrock": "aws", + "langchain_azure_ai": "azure", + "langchain_azure": "azure", + "langchain_mistralai": "mistral", +} + + +def _provider_from_serialized(serialized: Mapping[str, Any] | None) -> str | None: + if not serialized: + return None + for part in serialized.get("id") or []: + if not isinstance(part, str): + continue + head = part.lower().split(".", 1)[0] + if head in _PROVIDER_ALIASES: + return _PROVIDER_ALIASES[head] + if head.startswith("langchain_"): + return head[len("langchain_") :] + return None + + +def _resolve_name(name: str | None, serialized: Mapping[str, Any] | None, default: str) -> str: + return ( + name + or (serialized or {}).get("name") + or last_item((serialized or {}).get("id") or []) + or default + ) + + +_TOOL_KEYS = ("tools", "functions") +_TOOL_CHOICE_KEYS = ("tool_choice", "function_call") + + +def _split_tools(invocation_params: Mapping[str, Any] | None) -> tuple[dict[str, Any], Any, Any]: + if not invocation_params: + return {}, None, None + tools: Any = None + tool_choice: Any = None + remaining: dict[str, Any] = {} + for key, value in invocation_params.items(): + if key in _TOOL_KEYS and tools is None and value: + tools = value + continue + if key in _TOOL_CHOICE_KEYS and tool_choice is None and value is not None: + tool_choice = value + continue + remaining[key] = value + return remaining, tools, tool_choice + + class BraintrustCallbackHandler(BaseCallbackHandler): root_run_id: UUID | None = None @@ -61,17 +109,12 @@ def __init__( self, logger: Logger | Span | None = None, debug: bool = False, - exclude_metadata_props: Pattern[str] | None = None, ): self.logger = logger self.spans: dict[UUID, Span] = {} - self.debug = debug # DEPRECATED - self.exclude_metadata_props = exclude_metadata_props or re.compile( - r"^(l[sc]_|langgraph_|__pregel_|checkpoint_ns)" - ) + self.debug = debug self.skipped_runs: set[UUID] = set() - # Set run_inline=True to avoid thread executor in async contexts - # This ensures memory logger context is preserved + # Preserve memory logger context across async callbacks. self.run_inline = True self._start_times: dict[UUID, float] = {} @@ -91,7 +134,6 @@ def _start_span( event: LogEvent | None = None, ) -> Span | None: if run_id in self.spans: - # XXX: See graph test case of an example where this _may_ be intended. _logger.warning(f"Span already exists for run_id {run_id} (this is likely a bug)") return @@ -115,7 +157,7 @@ def _start_span( **event, "tags": None, "metadata": { - **({"tags": tags}), + "tags": tags, **(event.get("metadata") or {}), "run_id": run_id, "parent_run_id": parent_run_id, @@ -165,7 +207,6 @@ def _start_span( ) span.set_current() - self.spans[run_id] = span return span @@ -210,10 +251,8 @@ def _end_span( dataset_record_id=dataset_record_id, ) - # In async workflows, callbacks may execute in different async contexts. - # The span's context variable token may have been created in a different - # context, causing ValueError when trying to reset it. We catch and ignore - # this specific error since the span hierarchy is maintained via self.spans. + # Async callbacks may unset from a different context; span state is + # tracked in self.spans, so this ValueError is benign. try: span.unset_current() except ValueError as e: @@ -230,10 +269,9 @@ def on_llm_error( *, run_id: UUID, parent_run_id: UUID | None = None, - **kwargs: Any, # TODO: response= + **kwargs: Any, ) -> None: - self._end_span(run_id, error=str(error), metadata={**kwargs}) - + self._end_span(run_id, error=str(error)) self._start_times.pop(run_id, None) self._first_token_times.pop(run_id, None) self._ttft_ms.pop(run_id, None) @@ -244,9 +282,9 @@ def on_chain_error( *, run_id: UUID, parent_run_id: UUID | None = None, - **kwargs: Any, # TODO: some metadata + **kwargs: Any, ) -> None: - self._end_span(run_id, error=str(error), metadata={**kwargs}) + self._end_span(run_id, error=str(error)) def on_tool_error( self, @@ -256,7 +294,7 @@ def on_tool_error( parent_run_id: UUID | None = None, **kwargs: Any, ) -> None: - self._end_span(run_id, error=str(error), metadata={**kwargs}) + self._end_span(run_id, error=str(error)) def on_retriever_error( self, @@ -266,9 +304,8 @@ def on_retriever_error( parent_run_id: UUID | None = None, **kwargs: Any, ) -> None: - self._end_span(run_id, error=str(error), metadata={**kwargs}) + self._end_span(run_id, error=str(error)) - # Agent Methods def on_agent_action( self, action: AgentAction, @@ -280,9 +317,9 @@ def on_agent_action( self._start_span( parent_run_id, run_id, - type=SpanTypeAttribute.LLM, + type=SpanTypeAttribute.TOOL, name=action.tool, - event={"input": action, "metadata": {**kwargs}}, + event={"input": action}, ) def on_agent_finish( @@ -293,7 +330,7 @@ def on_agent_finish( parent_run_id: UUID | None = None, **kwargs: Any, ) -> None: - self._end_span(run_id, output=finish, metadata={**kwargs}) + self._end_span(run_id, output=finish) def on_chain_start( self, @@ -308,8 +345,6 @@ def on_chain_start( **kwargs: Any, ) -> None: tags = tags or [] - - # avoids extra logs that seem not as useful esp. with langgraph if "langsmith:hidden" in tags: self.skipped_runs.add(run_id) return @@ -318,8 +353,8 @@ def on_chain_start( resolved_name = ( name or metadata.get("langgraph_node") - or serialized.get("name") - or last_item(serialized.get("id") or []) + or (serialized or {}).get("name") + or last_item((serialized or {}).get("id") or []) or "Chain" ) @@ -330,12 +365,7 @@ def on_chain_start( event={ "input": inputs, "tags": tags, - "metadata": { - "serialized": serialized, - "name": name, - "metadata": metadata, - **kwargs, - }, + "metadata": {"metadata": metadata}, }, ) @@ -348,7 +378,7 @@ def on_chain_end( tags: list[str] | None = None, **kwargs: Any, ) -> None: - self._end_span(run_id, output=outputs, tags=tags, metadata={**kwargs}) + self._end_span(run_id, output=outputs, tags=tags) def on_llm_start( self, @@ -366,22 +396,17 @@ def on_llm_start( self._first_token_times.pop(run_id, None) self._ttft_ms.pop(run_id, None) - name = name or serialized.get("name") or last_item(serialized.get("id") or []) or "LLM" + span_metadata: dict[str, Any] = {"metadata": metadata or {}} + provider = _provider_from_serialized(serialized) + if provider: + span_metadata["provider"] = provider + self._start_span( parent_run_id, run_id, - name=name, + name=_resolve_name(name, serialized, "LLM"), type=SpanTypeAttribute.LLM, - event={ - "input": prompts, - "tags": tags, - "metadata": { - "serialized": serialized, - "name": name, - "metadata": metadata, - **kwargs, - }, - }, + event={"input": prompts, "tags": tags, "metadata": span_metadata}, ) def on_chat_model_start( @@ -401,25 +426,26 @@ def on_chat_model_start( self._first_token_times.pop(run_id, None) self._ttft_ms.pop(run_id, None) - invocation_params = invocation_params or {} + remaining_params, tools, tool_choice = _split_tools(invocation_params) + provider = _provider_from_serialized(serialized) + + span_metadata: dict[str, Any] = { + "invocation_params": remaining_params, + "metadata": metadata or {}, + } + if tools: + span_metadata["tools"] = tools + if tool_choice is not None: + span_metadata["tool_choice"] = tool_choice + if provider: + span_metadata["provider"] = provider + self._start_span( parent_run_id, run_id, - name=name or serialized.get("name") or last_item(serialized.get("id") or []) or "Chat Model", + name=_resolve_name(name, serialized, "Chat Model"), type=SpanTypeAttribute.LLM, - event={ - "input": messages, - "tags": tags, - "metadata": ( - { - "serialized": serialized, - "invocation_params": invocation_params, - "metadata": metadata or {}, - "name": name, - **kwargs, - } - ), - }, + event={"input": messages, "tags": tags, "metadata": span_metadata}, ) def on_llm_end( @@ -435,25 +461,26 @@ def on_llm_end( return metrics = _get_metrics_from_response(response) - ttft = self._ttft_ms.pop(run_id, None) if ttft is not None: metrics["time_to_first_token"] = ttft - model_name = _get_model_name_from_response(response) - self._start_times.pop(run_id, None) self._first_token_times.pop(run_id, None) + model_name, provider = _model_and_provider_from_response(response) + end_metadata: dict[str, Any] = {} + if model_name: + end_metadata["model"] = model_name + if provider: + end_metadata["provider"] = provider + self._end_span( run_id, output=response, metrics=metrics, tags=tags, - metadata={ - "model": model_name, - **kwargs, - }, + metadata=end_metadata, ) def on_tool_start( @@ -472,20 +499,12 @@ def on_tool_start( self._start_span( parent_run_id, run_id, - name=name or serialized.get("name") or last_item(serialized.get("id") or []) or "Tool", + name=_resolve_name(name, serialized, "Tool"), type=SpanTypeAttribute.TOOL, event={ - "input": inputs or safe_parse_serialized_json(input_str), + "input": inputs if inputs is not None else safe_parse_serialized_json(input_str), "tags": tags, - "metadata": { - "metadata": metadata, - "serialized": serialized, - "input_str": input_str, - "input": safe_parse_serialized_json(input_str), - "inputs": inputs, - "name": name, - **kwargs, - }, + "metadata": {"metadata": metadata or {}}, }, ) @@ -497,7 +516,7 @@ def on_tool_end( parent_run_id: UUID | None = None, **kwargs: Any, ) -> None: - self._end_span(run_id, output=output, metadata={**kwargs}) + self._end_span(run_id, output=output) def on_retriever_start( self, @@ -514,17 +533,12 @@ def on_retriever_start( self._start_span( parent_run_id, run_id, - name=name or serialized.get("name") or last_item(serialized.get("id") or []) or "Retriever", + name=_resolve_name(name, serialized, "Retriever"), type=SpanTypeAttribute.FUNCTION, event={ "input": query, "tags": tags, - "metadata": { - "serialized": serialized, - "metadata": metadata, - "name": name, - **kwargs, - }, + "metadata": {"metadata": metadata or {}}, }, ) @@ -536,7 +550,7 @@ def on_retriever_end( parent_run_id: UUID | None = None, **kwargs: Any, ) -> None: - self._end_span(run_id, output=documents, metadata={**kwargs}) + self._end_span(run_id, output=documents) def on_llm_new_token( self, @@ -611,29 +625,34 @@ def _walk_generations(response: LLMResult): yield from generations or [] -def _get_model_name_from_response(response: LLMResult) -> str | None: - model_name = None +def _model_and_provider_from_response(response: LLMResult) -> tuple[str | None, str | None]: + model_name: str | None = None + provider: str | None = None for generation in _walk_generations(response): message = getattr(generation, "message", None) if not message: continue - - response_metadata = getattr(message, "response_metadata", None) - if response_metadata and isinstance(response_metadata, dict): - model_name = response_metadata.get("model_name") - - if model_name: + rmeta = getattr(message, "response_metadata", None) + if not isinstance(rmeta, dict): + continue + if not model_name: + model_name = rmeta.get("model_name") or None + if not provider: + prov = rmeta.get("model_provider") + if isinstance(prov, str) and prov: + provider = prov.lower() + if model_name and provider: break if not model_name: llm_output: dict[str, Any] = response.llm_output or {} - model_name = llm_output.get("model_name") or llm_output.get("model") or "" + model_name = llm_output.get("model_name") or llm_output.get("model") or None - return model_name + return model_name, provider def _get_metrics_from_response(response: LLMResult): - metrics = {} + metrics: dict[str, Any] = {} for generation in _walk_generations(response): message = getattr(generation, "message", None) @@ -641,58 +660,55 @@ def _get_metrics_from_response(response: LLMResult): continue usage_metadata = getattr(message, "usage_metadata", None) + if not (usage_metadata and isinstance(usage_metadata, dict)): + continue - if usage_metadata and isinstance(usage_metadata, dict): - metrics.update( - clean_object( - { - "total_tokens": usage_metadata.get("total_tokens"), - "prompt_tokens": usage_metadata.get("input_tokens"), - "completion_tokens": usage_metadata.get("output_tokens"), - } - ) + metrics.update( + clean_object( + { + "total_tokens": usage_metadata.get("total_tokens"), + "prompt_tokens": usage_metadata.get("input_tokens"), + "completion_tokens": usage_metadata.get("output_tokens"), + } ) + ) + + input_token_details = usage_metadata.get("input_token_details") + if not (input_token_details and isinstance(input_token_details, dict)): + continue - # Extract cache tokens from nested input_token_details (LangChain format) - # Maps to Braintrust's standard cache token metric names - input_token_details = usage_metadata.get("input_token_details") - if input_token_details and isinstance(input_token_details, dict): - cache_read = input_token_details.get("cache_read") - cache_creation = input_token_details.get("cache_creation") - cache_creation_5m = input_token_details.get("ephemeral_5m_input_tokens") - cache_creation_1h = input_token_details.get("ephemeral_1h_input_tokens") - has_cache_creation_breakdown = cache_creation_5m is not None or cache_creation_1h is not None - - if cache_read is not None: - metrics["prompt_cached_tokens"] = cache_read - if has_cache_creation_breakdown: - # Anthropic exposes TTL-specific cache creation buckets. Preserve the - # split so downstream cost tooling can price 5m vs 1h writes correctly. - if cache_creation_5m is not None: - metrics["prompt_cache_creation_5m_tokens"] = cache_creation_5m - if cache_creation_1h is not None: - metrics["prompt_cache_creation_1h_tokens"] = cache_creation_1h - effective_cache_creation = (cache_creation_5m or 0) + (cache_creation_1h or 0) - else: - if cache_creation is not None: - metrics["prompt_cache_creation_tokens"] = cache_creation - effective_cache_creation = cache_creation or 0 - cache_tokens = (cache_read or 0) + effective_cache_creation - prompt_tokens = metrics.get("prompt_tokens") - completion_tokens = metrics.get("completion_tokens") - total_tokens = metrics.get("total_tokens") - if prompt_tokens is not None and completion_tokens is not None: - # LangChain's UsageMetadata contract makes input_token_details a - # breakdown of input_tokens, so cache tokens already count toward - # the prompt total (langchain-anthropic >= 0.2.3, langchain-aws, - # langchain-openai all comply). Cache tokens exceeding the prompt - # total means the integration reported uncached input only — fold - # cache tokens back in so prompt/total stay internally consistent. - if cache_tokens > prompt_tokens and total_tokens == prompt_tokens + completion_tokens: - prompt_tokens += cache_tokens - metrics["prompt_tokens"] = prompt_tokens - metrics["total_tokens"] = total_tokens + cache_tokens - metrics["tokens"] = prompt_tokens + completion_tokens + cache_read = input_token_details.get("cache_read") + cache_creation = input_token_details.get("cache_creation") + cache_creation_5m = input_token_details.get("ephemeral_5m_input_tokens") + cache_creation_1h = input_token_details.get("ephemeral_1h_input_tokens") + has_cache_creation_split = cache_creation_5m is not None or cache_creation_1h is not None + + if cache_read is not None: + metrics["prompt_cached_tokens"] = cache_read + if has_cache_creation_split: + if cache_creation_5m is not None: + metrics["prompt_cache_creation_5m_tokens"] = cache_creation_5m + if cache_creation_1h is not None: + metrics["prompt_cache_creation_1h_tokens"] = cache_creation_1h + effective_cache_creation = (cache_creation_5m or 0) + (cache_creation_1h or 0) + else: + if cache_creation is not None: + metrics["prompt_cache_creation_tokens"] = cache_creation + effective_cache_creation = cache_creation or 0 + cache_tokens = (cache_read or 0) + effective_cache_creation + + prompt_tokens = metrics.get("prompt_tokens") + completion_tokens = metrics.get("completion_tokens") + total_tokens = metrics.get("total_tokens") + if prompt_tokens is not None and completion_tokens is not None: + # LangChain's input_token_details is a breakdown of input_tokens. + # Fold cache tokens back into prompt/total only if the integration + # reported uncached-input-only (cache tokens exceeding prompt total). + if cache_tokens > prompt_tokens and total_tokens == prompt_tokens + completion_tokens: + prompt_tokens += cache_tokens + metrics["prompt_tokens"] = prompt_tokens + metrics["total_tokens"] = total_tokens + cache_tokens + metrics["tokens"] = prompt_tokens + completion_tokens if not metrics or not any(metrics.values()): llm_output: dict[str, Any] = response.llm_output or {} diff --git a/py/src/braintrust/integrations/langchain/test_anthropic.py b/py/src/braintrust/integrations/langchain/test_anthropic.py index 0361f09c..0b41670e 100644 --- a/py/src/braintrust/integrations/langchain/test_anthropic.py +++ b/py/src/braintrust/integrations/langchain/test_anthropic.py @@ -65,6 +65,7 @@ def test_langchain_anthropic_integration( llm_span = llm_spans[0] assert llm_span["metadata"]["model"] == MODEL + assert llm_span["metadata"]["provider"] == "anthropic" prompt_spans = [span for span in spans if "ChatPromptTemplate" in span["span_attributes"].get("name", "")] if prompt_spans: diff --git a/py/src/braintrust/integrations/langchain/test_callbacks.py b/py/src/braintrust/integrations/langchain/test_callbacks.py index 1543c1fd..1415d0bd 100644 --- a/py/src/braintrust/integrations/langchain/test_callbacks.py +++ b/py/src/braintrust/integrations/langchain/test_callbacks.py @@ -148,6 +148,7 @@ def test_llm_calls(logger_memory_logger): "metadata": { "tags": ["seq:step:2"], "model": "gpt-4o-mini-2024-07-18", + "provider": "openai", }, "root_span_id": trace_root_id, "span_parents": [root_span_id], @@ -257,6 +258,7 @@ def test_chain_with_memory(logger_memory_logger): "metadata": { "tags": ["seq:step:2", "test"], "model": "gpt-4o-mini-2024-07-18", + "provider": "openai", }, "root_span_id": trace_root_id, "span_parents": [root_span_id], @@ -332,18 +334,17 @@ def calculator(input: CalculatorInput) -> str: "metadata": { "tags": [], "model": "gpt-4o-mini-2024-07-18", - "invocation_params": { - "tools": [ - { - "type": "function", - "function": { - "name": "calculator", - "description": "Can perform mathematical operations.", - "parameters": ANY, # Complex JSON schema - }, - } - ], - }, + "provider": "openai", + "tools": [ + { + "type": "function", + "function": { + "name": "calculator", + "description": "Can perform mathematical operations.", + "parameters": ANY, # Complex JSON schema + }, + } + ], }, "output": { "generations": [ @@ -572,6 +573,7 @@ def say_bye(state: dict[str, str]): ], "metadata": { "model": "gpt-4o-mini-2024-07-18", + "provider": "openai", "tags": [], }, "output": { From 4ab5817c5ea155876fef23e9181512dc242ddce5 Mon Sep 17 00:00:00 2001 From: Starfolk Date: Fri, 17 Jul 2026 20:54:10 +0000 Subject: [PATCH 2/2] style: apply ruff format to _resolve_name Co-Authored-By: Claude Opus 4.7 --- py/src/braintrust/integrations/langchain/callbacks.py | 7 +------ 1 file changed, 1 insertion(+), 6 deletions(-) diff --git a/py/src/braintrust/integrations/langchain/callbacks.py b/py/src/braintrust/integrations/langchain/callbacks.py index d1fd8b1c..bf93bc1f 100644 --- a/py/src/braintrust/integrations/langchain/callbacks.py +++ b/py/src/braintrust/integrations/langchain/callbacks.py @@ -73,12 +73,7 @@ def _provider_from_serialized(serialized: Mapping[str, Any] | None) -> str | Non def _resolve_name(name: str | None, serialized: Mapping[str, Any] | None, default: str) -> str: - return ( - name - or (serialized or {}).get("name") - or last_item((serialized or {}).get("id") or []) - or default - ) + return name or (serialized or {}).get("name") or last_item((serialized or {}).get("id") or []) or default _TOOL_KEYS = ("tools", "functions")