chore(core): reduce streaming metadata / perf (#36588)

- looking into reducing streaming metadata / perfm

---------

Co-authored-by: William Fu-Hinthorn <13333726+hinthornw@users.noreply.github.com>
This commit is contained in:
Eugene Yurtsev
2026-04-10 10:47:54 -04:00
committed by GitHub
parent f0c5a28fa0
commit af4d711a2f
11 changed files with 1195 additions and 30 deletions

View File

@@ -7,7 +7,7 @@ import atexit
import functools
import logging
from abc import ABC, abstractmethod
from collections.abc import Callable
from collections.abc import Callable, Mapping
from concurrent.futures import ThreadPoolExecutor
from contextlib import asynccontextmanager, contextmanager
from contextvars import copy_context
@@ -1614,6 +1614,9 @@ class CallbackManager(BaseCallbackManager):
local_tags: list[str] | None = None,
inheritable_metadata: dict[str, Any] | None = None,
local_metadata: dict[str, Any] | None = None,
*,
langsmith_inheritable_metadata: Mapping[str, Any] | None = None,
langsmith_inheritable_tags: list[str] | None = None,
) -> CallbackManager:
"""Configure the callback manager.
@@ -1625,6 +1628,10 @@ class CallbackManager(BaseCallbackManager):
local_tags: The local tags.
inheritable_metadata: The inheritable metadata.
local_metadata: The local metadata.
langsmith_inheritable_metadata: Default inheritable metadata applied
to any `LangChainTracer` handlers via `set_defaults`.
langsmith_inheritable_tags: Default inheritable tags applied to any
`LangChainTracer` handlers via `set_defaults`.
Returns:
The configured callback manager.
@@ -1638,6 +1645,8 @@ class CallbackManager(BaseCallbackManager):
inheritable_metadata,
local_metadata,
verbose=verbose,
langsmith_inheritable_metadata=langsmith_inheritable_metadata,
langsmith_inheritable_tags=langsmith_inheritable_tags,
)
@@ -2134,6 +2143,9 @@ class AsyncCallbackManager(BaseCallbackManager):
local_tags: list[str] | None = None,
inheritable_metadata: dict[str, Any] | None = None,
local_metadata: dict[str, Any] | None = None,
*,
langsmith_inheritable_metadata: Mapping[str, Any] | None = None,
langsmith_inheritable_tags: list[str] | None = None,
) -> AsyncCallbackManager:
"""Configure the async callback manager.
@@ -2145,6 +2157,10 @@ class AsyncCallbackManager(BaseCallbackManager):
local_tags: The local tags.
inheritable_metadata: The inheritable metadata.
local_metadata: The local metadata.
langsmith_inheritable_metadata: Default inheritable metadata applied
to any `LangChainTracer` handlers via `set_defaults`.
langsmith_inheritable_tags: Default inheritable tags applied to any
`LangChainTracer` handlers via `set_defaults`.
Returns:
The configured async callback manager.
@@ -2158,6 +2174,8 @@ class AsyncCallbackManager(BaseCallbackManager):
inheritable_metadata,
local_metadata,
verbose=verbose,
langsmith_inheritable_metadata=langsmith_inheritable_metadata,
langsmith_inheritable_tags=langsmith_inheritable_tags,
)
@@ -2304,6 +2322,8 @@ def _configure(
local_metadata: dict[str, Any] | None = None,
*,
verbose: bool = False,
langsmith_inheritable_metadata: Mapping[str, Any] | None = None,
langsmith_inheritable_tags: list[str] | None = None,
) -> T:
"""Configure the callback manager.
@@ -2316,6 +2336,10 @@ def _configure(
inheritable_metadata: The inheritable metadata.
local_metadata: The local metadata.
verbose: Whether to enable verbose mode.
langsmith_inheritable_metadata: Default inheritable metadata applied to
any `LangChainTracer` handlers via `set_defaults`.
langsmith_inheritable_tags: Default inheritable tags applied to any
`LangChainTracer` handlers via `set_defaults`.
Raises:
RuntimeError: If `LANGCHAIN_TRACING` is set but `LANGCHAIN_TRACING_V2` is not.
@@ -2387,8 +2411,6 @@ def _configure(
if inheritable_metadata or local_metadata:
callback_manager.add_metadata(inheritable_metadata or {})
callback_manager.add_metadata(local_metadata or {}, inherit=False)
if tracing_metadata:
callback_manager.add_metadata(tracing_metadata.copy())
if tracing_tags:
callback_manager.add_tags(tracing_tags.copy())
@@ -2440,6 +2462,7 @@ def _configure(
else tracing_context["client"]
),
tags=tracing_tags,
metadata=tracing_metadata,
)
callback_manager.add_handler(handler)
except Exception as e:
@@ -2479,6 +2502,32 @@ def _configure(
for handler in callback_manager.handlers
):
callback_manager.add_handler(var_handler, inheritable)
if tracing_metadata:
langsmith_inheritable_metadata = {
**tracing_metadata,
**(langsmith_inheritable_metadata or {}),
}
if langsmith_inheritable_metadata or langsmith_inheritable_tags:
callback_manager.handlers = [
handler.copy_with_metadata_defaults(
metadata=langsmith_inheritable_metadata,
tags=langsmith_inheritable_tags,
)
if isinstance(handler, LangChainTracer)
else handler
for handler in callback_manager.handlers
]
callback_manager.inheritable_handlers = [
handler.copy_with_metadata_defaults(
metadata=langsmith_inheritable_metadata,
tags=langsmith_inheritable_tags,
)
if isinstance(handler, LangChainTracer)
else handler
for handler in callback_manager.inheritable_handlers
]
return callback_manager

View File

@@ -138,6 +138,28 @@ COPIABLE_KEYS = [
"configurable",
]
# Users are expected to use the `context` API with a context object
# (which does not get traced)
CONFIGURABLE_TO_TRACING_METADATA_EXCLUDED_KEYS = frozenset(("api_key",))
def _get_langsmith_inheritable_metadata_from_config(
config: RunnableConfig,
) -> dict[str, Any] | None:
"""Get LangSmith-only inheritable metadata defaults derived from config."""
configurable = config.get("configurable") or {}
metadata = {
key: value
for key, value in configurable.items()
if not key.startswith("__")
and isinstance(value, (str, int, float, bool))
and key not in config.get("metadata", {})
and key not in CONFIGURABLE_TO_TRACING_METADATA_EXCLUDED_KEYS
}
return metadata or None
DEFAULT_RECURSION_LIMIT = 25
@@ -264,14 +286,11 @@ def ensure_config(config: RunnableConfig | None = None) -> RunnableConfig:
for k, v in config.items():
if k not in CONFIG_KEYS and v is not None:
empty["configurable"][k] = v
for key, value in empty.get("configurable", {}).items():
if (
not key.startswith("__")
and isinstance(value, (str, int, float, bool))
and key not in empty["metadata"]
and key != "api_key"
):
empty["metadata"][key] = value
if (
isinstance(model := empty.get("configurable", {}).get("model"), str)
and "model" not in empty["metadata"]
):
empty["metadata"]["model"] = model
return empty
@@ -508,6 +527,9 @@ def get_callback_manager_for_config(config: RunnableConfig) -> CallbackManager:
inheritable_callbacks=config.get("callbacks"),
inheritable_tags=config.get("tags"),
inheritable_metadata=config.get("metadata"),
langsmith_inheritable_metadata=_get_langsmith_inheritable_metadata_from_config(
config
),
)
@@ -526,6 +548,9 @@ def get_async_callback_manager_for_config(
inheritable_callbacks=config.get("callbacks"),
inheritable_tags=config.get("tags"),
inheritable_metadata=config.get("metadata"),
langsmith_inheritable_metadata=_get_langsmith_inheritable_metadata_from_config(
config
),
)

View File

@@ -51,6 +51,8 @@ class _TracerCore(ABC):
_schema_format: Literal[
"original", "streaming_events", "original+chat"
] = "original",
run_map: dict[str, Run] | None = None,
order_map: dict[UUID, tuple[UUID, str]] | None = None,
**kwargs: Any,
) -> None:
"""Initialize the tracer.
@@ -70,6 +72,8 @@ class _TracerCore(ABC):
streaming events.
- `'original+chat'` is a format that is the same as `'original'` except
it does NOT raise an attribute error `on_chat_model_start`
run_map: Optional shared map of run ID to run.
order_map: Optional shared map of run ID to trace ordering data.
**kwargs: Additional keyword arguments that will be passed to the
superclass.
"""
@@ -77,10 +81,10 @@ class _TracerCore(ABC):
self._schema_format = _schema_format # For internal use only API will change.
self.run_map: dict[str, Run] = {}
self.run_map = run_map if run_map is not None else {}
"""Map of run ID to run. Cleared on run end."""
self.order_map: dict[UUID, tuple[UUID, str]] = {}
self.order_map = order_map if order_map is not None else {}
"""Map of run ID to (trace_id, dotted_order). Cleared when tracer GCed."""
@abstractmethod

View File

@@ -27,6 +27,8 @@ from langchain_core.tracers.base import BaseTracer
from langchain_core.tracers.schemas import Run
if TYPE_CHECKING:
from collections.abc import Mapping
from langchain_core.messages import BaseMessage
from langchain_core.outputs import ChatGenerationChunk, GenerationChunk
@@ -124,6 +126,8 @@ class LangChainTracer(BaseTracer):
project_name: str | None = None,
client: Client | None = None,
tags: list[str] | None = None,
*,
metadata: Mapping[str, str] | None = None,
**kwargs: Any,
) -> None:
"""Initialize the LangChain tracer.
@@ -139,6 +143,9 @@ class LangChainTracer(BaseTracer):
tags: The tags.
Defaults to an empty list.
metadata: Additional metadata to include if it isn't already in the run.
Defaults to None.
**kwargs: Additional keyword arguments.
"""
super().__init__(**kwargs)
@@ -150,6 +157,39 @@ class LangChainTracer(BaseTracer):
self.tags = tags or []
self.latest_run: Run | None = None
self.run_has_token_event_map: dict[str, bool] = {}
self.tracing_metadata: dict[str, str] | None = (
dict(metadata) if metadata is not None else None
)
def copy_with_metadata_defaults(
self,
*,
metadata: Mapping[str, str] | None = None,
tags: list[str] | None = None,
) -> LangChainTracer:
"""Return a new tracer with merged tracer-only defaults."""
base_metadata = self.tracing_metadata
if metadata is None:
merged_metadata = dict(base_metadata) if base_metadata is not None else None
elif base_metadata is None:
merged_metadata = dict(metadata)
else:
merged_metadata = dict(base_metadata)
for key, value in metadata.items():
if key not in merged_metadata:
merged_metadata[key] = value
merged_tags = sorted(set(self.tags + tags)) if tags else self.tags
return self.__class__(
example_id=self.example_id,
project_name=self.project_name,
client=self.client,
tags=merged_tags,
metadata=merged_metadata,
run_map=self.run_map,
order_map=self.order_map,
)
def _start_trace(self, run: Run) -> None:
if self.project_name:
@@ -263,6 +303,7 @@ class LangChainTracer(BaseTracer):
try:
run.extra["runtime"] = get_runtime_environment()
run.tags = self._get_tags(run)
_patch_missing_metadata(self, run)
if run.ls_client is not self.client:
run.ls_client = self.client
run.post()
@@ -398,3 +439,17 @@ class LangChainTracer(BaseTracer):
"""Wait for the given futures to complete."""
if self.client is not None:
self.client.flush()
def _patch_missing_metadata(self: LangChainTracer, run: Run) -> None:
if not self.tracing_metadata:
return
metadata = run.metadata
patched = None
for k, v in self.tracing_metadata.items():
if k not in metadata:
if patched is None:
# Copy on first miss to avoid mutating the shared dict.
patched = {**metadata}
run.extra["metadata"] = patched
patched[k] = v