Compare commits

...
Author SHA1 Message Date
William Fu-Hinthorn 53172459bb update 2026-01-08 16:08:46 -08:00
William Fu-Hinthorn df254a7f7f more checks 2026-01-08 08:18:38 -08:00
William Fu-Hinthorn 46bd949f53 add cb handler checks 2026-01-08 06:44:24 -08:00
William Fu-Hinthorn 2e7a6f088e name 2026-01-08 06:26:37 -08:00
William Fu-Hinthorn f941f467d5 refac 2026-01-08 04:50:52 -08:00
William Fu-Hinthorn a7aaa81ed8 Support tasks 2026-01-08 04:29:43 -08:00
William Fu-Hinthorn eb9e6e2a0b config 2026-01-07 18:28:14 -08:00
William Fu-Hinthorn 2d304b8a1b cleanup 2026-01-07 16:04:57 -08:00
William Fu-Hinthorn 6b17b12ae4 feat: Honor traceable config 2026-01-07 14:55:53 -08:00
4 changed files with 1645 additions and 201 deletions
+369 -186
View File
@@ -3,6 +3,7 @@ from __future__ import annotations
import asyncio import asyncio
import enum import enum
import inspect import inspect
import logging
import sys import sys
import warnings import warnings
from collections.abc import ( from collections.abc import (
@@ -25,6 +26,7 @@ from typing import (
cast, cast,
) )
from langchain_core.callbacks import AsyncCallbackManager, CallbackManager
from langchain_core.runnables.base import ( from langchain_core.runnables.base import (
Runnable, Runnable,
RunnableConfig, RunnableConfig,
@@ -61,6 +63,86 @@ try:
except ImportError: except ImportError:
_StreamingCallbackHandler = None # type: ignore _StreamingCallbackHandler = None # type: ignore
logger = logging.getLogger(__name__)
def _process_inputs(processor: Callable[[Any], Any] | None, inputs: Any) -> Any:
"""Safely process trace inputs, returning error placeholder on failure."""
if processor is None:
return inputs
try:
return processor(inputs)
except Exception:
logger.exception("trace_inputs filter failed")
return {"error": "<trace_inputs processing failed>"}
def _process_outputs(processor: Callable[[Any], Any] | None, outputs: Any) -> Any:
"""Safely process trace outputs, returning error placeholder on failure."""
if processor is None:
return outputs
try:
return processor(outputs)
except Exception:
logger.exception("trace_outputs filter failed")
return {"error": "<trace_outputs processing failed>"}
def _filter_langsmith_handlers(
callback_manager: CallbackManager,
) -> CallbackManager:
"""Filter out LangChainTracer handlers from a callback manager.
Used when traceable config has enabled=False - we want to fire callbacks
to custom handlers but skip LangSmith tracing.
"""
filtered_handlers = [
h for h in callback_manager.handlers if not isinstance(h, LangChainTracer)
]
filtered_inheritable = [
h
for h in callback_manager.inheritable_handlers
if not isinstance(h, LangChainTracer)
]
# Create a new callback manager with filtered handlers
return CallbackManager(
handlers=filtered_handlers,
inheritable_handlers=filtered_inheritable,
parent_run_id=callback_manager.parent_run_id,
tags=callback_manager.tags,
inheritable_tags=callback_manager.inheritable_tags,
metadata=callback_manager.metadata,
inheritable_metadata=callback_manager.inheritable_metadata,
)
def _filter_langsmith_handlers_async(
callback_manager: AsyncCallbackManager,
) -> AsyncCallbackManager:
"""Filter out LangChainTracer handlers from an async callback manager.
Used when traceable config has enabled=False - we want to fire callbacks
to custom handlers but skip LangSmith tracing.
"""
filtered_handlers = [
h for h in callback_manager.handlers if not isinstance(h, LangChainTracer)
]
filtered_inheritable = [
h
for h in callback_manager.inheritable_handlers
if not isinstance(h, LangChainTracer)
]
# Create a new callback manager with filtered handlers
return AsyncCallbackManager(
handlers=filtered_handlers,
inheritable_handlers=filtered_inheritable,
parent_run_id=callback_manager.parent_run_id,
tags=callback_manager.tags,
inheritable_tags=callback_manager.inheritable_tags,
metadata=callback_manager.metadata,
inheritable_metadata=callback_manager.inheritable_metadata,
)
def _set_config_context( def _set_config_context(
config: RunnableConfig, run: Any = None config: RunnableConfig, run: Any = None
@@ -397,7 +479,9 @@ class RunnableCallable(Runnable):
else: else:
run_manager.on_chain_end(ret) run_manager.on_chain_end(ret)
else: else:
ret = self.func(*args, **kwargs) # Still need to set config context for get_config() to work
with set_config_context(config, None) as context:
ret = context.run(self.func, *args, **kwargs)
if self.recurse and isinstance(ret, Runnable): if self.recurse and isinstance(ret, Runnable):
return ret.invoke(input, config) return ret.invoke(input, config)
return ret return ret
@@ -470,7 +554,13 @@ class RunnableCallable(Runnable):
else: else:
await run_manager.on_chain_end(ret) await run_manager.on_chain_end(ret)
else: else:
ret = await self.afunc(*args, **kwargs) # Still need to set config context for get_config() to work
coro = cast(Coroutine[None, None, Any], self.afunc(*args, **kwargs))
if ASYNCIO_ACCEPTS_CONTEXT:
with set_config_context(config, None) as context:
ret = await asyncio.create_task(coro, context=context)
else:
ret = await coro
if self.recurse and isinstance(ret, Runnable): if self.recurse and isinstance(ret, Runnable):
return await ret.ainvoke(input, config) return await ret.ainvoke(input, config)
return ret return ret
@@ -544,12 +634,20 @@ class RunnableSeq(Runnable):
*steps: RunnableLike, *steps: RunnableLike,
name: str | None = None, name: str | None = None,
trace_inputs: Callable[[Any], Any] | None = None, trace_inputs: Callable[[Any], Any] | None = None,
trace_outputs: Callable[[Any], Any] | None = None,
trace: bool = True,
skip_langsmith: bool = False,
) -> None: ) -> None:
"""Create a new RunnableSeq. """Create a new RunnableSeq.
Args: Args:
steps: The steps to include in the sequence. steps: The steps to include in the sequence.
name: The name of the `Runnable`. name: The name of the `Runnable`.
trace_inputs: Optional function to transform inputs before tracing.
trace_outputs: Optional function to transform outputs before tracing.
trace: Whether to trace this sequence. Defaults to True.
skip_langsmith: If True, filter out LangChainTracer handlers but keep
other callbacks. Used when traceable config has enabled=False.
Raises: Raises:
ValueError: If the sequence has less than 2 steps. ValueError: If the sequence has less than 2 steps.
@@ -569,6 +667,9 @@ class RunnableSeq(Runnable):
self.steps = steps_flat self.steps = steps_flat
self.name = name self.name = name
self.trace_inputs = trace_inputs self.trace_inputs = trace_inputs
self.trace_outputs = trace_outputs
self.trace = trace
self.skip_langsmith = skip_langsmith
def __or__( def __or__(
self, self,
@@ -625,43 +726,58 @@ class RunnableSeq(Runnable):
) -> Any: ) -> Any:
if config is None: if config is None:
config = ensure_config() config = ensure_config()
# setup callbacks and context
callback_manager = get_callback_manager_for_config(config) if self.trace:
# start the root run # setup callbacks and context
run_manager = callback_manager.on_chain_start( callback_manager = get_callback_manager_for_config(config)
None, # Filter out LangChainTracer if skip_langsmith is set
self.trace_inputs(input) if self.trace_inputs is not None else input, if self.skip_langsmith:
name=config.get("run_name") or self.get_name(), callback_manager = _filter_langsmith_handlers(callback_manager)
run_id=config.pop("run_id", None), # start the root run
) run_manager = callback_manager.on_chain_start(
# invoke all steps in sequence None,
try: _process_inputs(self.trace_inputs, input),
for i, step in enumerate(self.steps): name=config.get("run_name") or self.get_name(),
# mark each step as a child run run_id=config.pop("run_id", None),
config = patch_config( )
config, callbacks=run_manager.get_child(f"seq:step:{i + 1}") # invoke all steps in sequence
) try:
# 1st step is the actual node, for i, step in enumerate(self.steps):
# others are writers which don't need to be run in context # mark each step as a child run
if i == 0: config = patch_config(
# get the run object config, callbacks=run_manager.get_child(f"seq:step:{i + 1}")
for h in run_manager.handlers: )
if isinstance(h, LangChainTracer): # 1st step is the actual node,
run = h.run_map.get(str(run_manager.run_id)) # others are writers which don't need to be run in context
break if i == 0:
# get the run object
for h in run_manager.handlers:
if isinstance(h, LangChainTracer):
run = h.run_map.get(str(run_manager.run_id))
break
else:
run = None
# run in context
with set_config_context(config, run) as context:
input = context.run(step.invoke, input, config, **kwargs)
else: else:
run = None input = step.invoke(input, config)
# run in context # finish the root run
with set_config_context(config, run) as context: except BaseException as e:
input = context.run(step.invoke, input, config, **kwargs) run_manager.on_chain_error(e)
else: raise
input = step.invoke(input, config) else:
# finish the root run run_manager.on_chain_end(_process_outputs(self.trace_outputs, input))
except BaseException as e: return input
run_manager.on_chain_error(e)
raise
else: else:
run_manager.on_chain_end(input) # Still need to set config context for get_config() to work
with set_config_context(config, None) as context:
for i, step in enumerate(self.steps):
input = (
context.run(step.invoke, input, config, **kwargs)
if i == 0
else step.invoke(input, config)
)
return input return input
async def ainvoke( async def ainvoke(
@@ -672,49 +788,75 @@ class RunnableSeq(Runnable):
) -> Any: ) -> Any:
if config is None: if config is None:
config = ensure_config() config = ensure_config()
# setup callbacks
callback_manager = get_async_callback_manager_for_config(config)
# start the root run
run_manager = await callback_manager.on_chain_start(
None,
self.trace_inputs(input) if self.trace_inputs is not None else input,
name=config.get("run_name") or self.get_name(),
run_id=config.pop("run_id", None),
)
# invoke all steps in sequence if self.trace:
try: # setup callbacks
for i, step in enumerate(self.steps): callback_manager = get_async_callback_manager_for_config(config)
# mark each step as a child run # Filter out LangChainTracer if skip_langsmith is set
config = patch_config( if self.skip_langsmith:
config, callbacks=run_manager.get_child(f"seq:step:{i + 1}") callback_manager = _filter_langsmith_handlers_async(callback_manager)
) # start the root run
# 1st step is the actual node, run_manager = await callback_manager.on_chain_start(
# others are writers which don't need to be run in context None,
if i == 0: _process_inputs(self.trace_inputs, input),
if ASYNCIO_ACCEPTS_CONTEXT: name=config.get("run_name") or self.get_name(),
# get the run object run_id=config.pop("run_id", None),
for h in run_manager.handlers: )
if isinstance(h, LangChainTracer):
run = h.run_map.get(str(run_manager.run_id)) # invoke all steps in sequence
break try:
for i, step in enumerate(self.steps):
# mark each step as a child run
config = patch_config(
config, callbacks=run_manager.get_child(f"seq:step:{i + 1}")
)
# 1st step is the actual node,
# others are writers which don't need to be run in context
if i == 0:
if ASYNCIO_ACCEPTS_CONTEXT:
# get the run object
for h in run_manager.handlers:
if isinstance(h, LangChainTracer):
run = h.run_map.get(str(run_manager.run_id))
break
else:
run = None
# run in context
with set_config_context(config, run) as context:
input = await asyncio.create_task(
step.ainvoke(input, config, **kwargs),
context=context,
)
else: else:
run = None input = await step.ainvoke(input, config, **kwargs)
# run in context else:
with set_config_context(config, run) as context: input = await step.ainvoke(input, config)
# finish the root run
except BaseException as e:
await run_manager.on_chain_error(e)
raise
else:
await run_manager.on_chain_end(
_process_outputs(self.trace_outputs, input)
)
return input
else:
# Still need to set config context for get_config() to work
if ASYNCIO_ACCEPTS_CONTEXT:
with set_config_context(config, None) as context:
for i, step in enumerate(self.steps):
if i == 0:
input = await asyncio.create_task( input = await asyncio.create_task(
step.ainvoke(input, config, **kwargs), context=context step.ainvoke(input, config, **kwargs), context=context
) )
else: else:
input = await step.ainvoke(input, config)
else:
for i, step in enumerate(self.steps):
if i == 0:
input = await step.ainvoke(input, config, **kwargs) input = await step.ainvoke(input, config, **kwargs)
else: else:
input = await step.ainvoke(input, config) input = await step.ainvoke(input, config)
# finish the root run
except BaseException as e:
await run_manager.on_chain_error(e)
raise
else:
await run_manager.on_chain_end(input)
return input return input
def stream( def stream(
@@ -725,80 +867,20 @@ class RunnableSeq(Runnable):
) -> Iterator[Any]: ) -> Iterator[Any]:
if config is None: if config is None:
config = ensure_config() config = ensure_config()
# setup callbacks
callback_manager = get_callback_manager_for_config(config)
# start the root run
run_manager = callback_manager.on_chain_start(
None,
self.trace_inputs(input) if self.trace_inputs is not None else input,
name=config.get("run_name") or self.get_name(),
run_id=config.pop("run_id", None),
)
# get the run object
for h in run_manager.handlers:
if isinstance(h, LangChainTracer):
run = h.run_map.get(str(run_manager.run_id))
break
else:
run = None
# create first step config
config = patch_config(
config,
callbacks=run_manager.get_child(f"seq:step:{1}"),
)
# run all in context
with set_config_context(config, run) as context:
try:
# stream the last steps
# transform the input stream of each step with the next
# steps that don't natively support transforming an input stream will
# buffer input in memory until all available, and then start emitting output
for idx, step in enumerate(self.steps):
if idx == 0:
iterator = step.stream(input, config, **kwargs)
else:
config = patch_config(
config,
callbacks=run_manager.get_child(f"seq:step:{idx + 1}"),
)
iterator = step.transform(iterator, config)
# populates streamed_output in astream_log() output if needed
if _StreamingCallbackHandler is not None:
for h in run_manager.handlers:
if isinstance(h, _StreamingCallbackHandler):
iterator = h.tap_output_iter(run_manager.run_id, iterator)
# consume into final output
output = context.run(_consume_iter, iterator)
# sequence doesn't emit output, yield to mark as generator
yield
except BaseException as e:
run_manager.on_chain_error(e)
raise
else:
run_manager.on_chain_end(output)
async def astream( if self.trace:
self, # setup callbacks
input: Input, callback_manager = get_callback_manager_for_config(config)
config: RunnableConfig | None = None, # Filter out LangChainTracer if skip_langsmith is set
**kwargs: Any | None, if self.skip_langsmith:
) -> AsyncIterator[Any]: callback_manager = _filter_langsmith_handlers(callback_manager)
if config is None: # start the root run
config = ensure_config() run_manager = callback_manager.on_chain_start(
# setup callbacks None,
callback_manager = get_async_callback_manager_for_config(config) _process_inputs(self.trace_inputs, input),
# start the root run name=config.get("run_name") or self.get_name(),
run_manager = await callback_manager.on_chain_start( run_id=config.pop("run_id", None),
None, )
self.trace_inputs(input) if self.trace_inputs is not None else input,
name=config.get("run_name") or self.get_name(),
run_id=config.pop("run_id", None),
)
# stream the last steps
# transform the input stream of each step with the next
# steps that don't natively support transforming an input stream will
# buffer input in memory until all available, and then start emitting output
if ASYNCIO_ACCEPTS_CONTEXT:
# get the run object # get the run object
for h in run_manager.handlers: for h in run_manager.handlers:
if isinstance(h, LangChainTracer): if isinstance(h, LangChainTracer):
@@ -813,18 +895,136 @@ class RunnableSeq(Runnable):
) )
# run all in context # run all in context
with set_config_context(config, run) as context: with set_config_context(config, run) as context:
try:
# stream the last steps
# transform the input stream of each step with the next
# steps that don't natively support transforming an input stream will
# buffer input in memory until all available, and then start emitting output
for idx, step in enumerate(self.steps):
if idx == 0:
iterator = step.stream(input, config, **kwargs)
else:
config = patch_config(
config,
callbacks=run_manager.get_child(f"seq:step:{idx + 1}"),
)
iterator = step.transform(iterator, config)
# populates streamed_output in astream_log() output if needed
if _StreamingCallbackHandler is not None:
for h in run_manager.handlers:
if isinstance(h, _StreamingCallbackHandler):
iterator = h.tap_output_iter(
run_manager.run_id, iterator
)
# consume into final output
output = context.run(_consume_iter, iterator)
# sequence doesn't emit output, yield to mark as generator
yield
except BaseException as e:
run_manager.on_chain_error(e)
raise
else:
run_manager.on_chain_end(
_process_outputs(self.trace_outputs, output)
)
else:
# No tracing - still need to set config context for get_config() to work
with set_config_context(config, None) as context:
for idx, step in enumerate(self.steps):
if idx == 0:
iterator = step.stream(input, config, **kwargs)
else:
iterator = step.transform(iterator, config)
context.run(_consume_iter, iterator)
yield
async def astream(
self,
input: Input,
config: RunnableConfig | None = None,
**kwargs: Any | None,
) -> AsyncIterator[Any]:
if config is None:
config = ensure_config()
if self.trace:
# setup callbacks
callback_manager = get_async_callback_manager_for_config(config)
# Filter out LangChainTracer if skip_langsmith is set
if self.skip_langsmith:
callback_manager = _filter_langsmith_handlers_async(callback_manager)
# start the root run
run_manager = await callback_manager.on_chain_start(
None,
_process_inputs(self.trace_inputs, input),
name=config.get("run_name") or self.get_name(),
run_id=config.pop("run_id", None),
)
# stream the last steps
# transform the input stream of each step with the next
# steps that don't natively support transforming an input stream will
# buffer input in memory until all available, and then start emitting output
if ASYNCIO_ACCEPTS_CONTEXT:
# get the run object
for h in run_manager.handlers:
if isinstance(h, LangChainTracer):
run = h.run_map.get(str(run_manager.run_id))
break
else:
run = None
# create first step config
config = patch_config(
config,
callbacks=run_manager.get_child(f"seq:step:{1}"),
)
# run all in context
with set_config_context(config, run) as context:
try:
async with AsyncExitStack() as stack:
for idx, step in enumerate(self.steps):
if idx == 0:
aiterator = step.astream(input, config, **kwargs)
else:
config = patch_config(
config,
callbacks=run_manager.get_child(
f"seq:step:{idx + 1}"
),
)
aiterator = step.atransform(aiterator, config)
if hasattr(aiterator, "aclose"):
stack.push_async_callback(aiterator.aclose)
# populates streamed_output in astream_log() output if needed
if _StreamingCallbackHandler is not None:
for h in run_manager.handlers:
if isinstance(h, _StreamingCallbackHandler):
aiterator = h.tap_output_aiter(
run_manager.run_id, aiterator
)
# consume into final output
output = await asyncio.create_task(
_consume_aiter(aiterator), context=context
)
# sequence doesn't emit output, yield to mark as generator
yield
except BaseException as e:
await run_manager.on_chain_error(e)
raise
else:
await run_manager.on_chain_end(
_process_outputs(self.trace_outputs, output)
)
else:
try: try:
async with AsyncExitStack() as stack: async with AsyncExitStack() as stack:
for idx, step in enumerate(self.steps): for idx, step in enumerate(self.steps):
config = patch_config(
config,
callbacks=run_manager.get_child(f"seq:step:{idx + 1}"),
)
if idx == 0: if idx == 0:
aiterator = step.astream(input, config, **kwargs) aiterator = step.astream(input, config, **kwargs)
else: else:
config = patch_config(
config,
callbacks=run_manager.get_child(
f"seq:step:{idx + 1}"
),
)
aiterator = step.atransform(aiterator, config) aiterator = step.atransform(aiterator, config)
if hasattr(aiterator, "aclose"): if hasattr(aiterator, "aclose"):
stack.push_async_callback(aiterator.aclose) stack.push_async_callback(aiterator.aclose)
@@ -836,46 +1036,29 @@ class RunnableSeq(Runnable):
run_manager.run_id, aiterator run_manager.run_id, aiterator
) )
# consume into final output # consume into final output
output = await asyncio.create_task( output = await _consume_aiter(aiterator)
_consume_aiter(aiterator), context=context
)
# sequence doesn't emit output, yield to mark as generator # sequence doesn't emit output, yield to mark as generator
yield yield
except BaseException as e: except BaseException as e:
await run_manager.on_chain_error(e) await run_manager.on_chain_error(e)
raise raise
else: else:
await run_manager.on_chain_end(output) await run_manager.on_chain_end(
_process_outputs(self.trace_outputs, output)
)
else: else:
try: # No tracing - just execute the steps directly
async with AsyncExitStack() as stack: async with AsyncExitStack() as stack:
for idx, step in enumerate(self.steps): for idx, step in enumerate(self.steps):
config = patch_config( if idx == 0:
config, aiterator = step.astream(input, config, **kwargs)
callbacks=run_manager.get_child(f"seq:step:{idx + 1}"), else:
) aiterator = step.atransform(aiterator, config)
if idx == 0: if hasattr(aiterator, "aclose"):
aiterator = step.astream(input, config, **kwargs) stack.push_async_callback(aiterator.aclose)
else: # consume into final output
aiterator = step.atransform(aiterator, config) await _consume_aiter(aiterator)
if hasattr(aiterator, "aclose"): yield
stack.push_async_callback(aiterator.aclose)
# populates streamed_output in astream_log() output if needed
if _StreamingCallbackHandler is not None:
for h in run_manager.handlers:
if isinstance(h, _StreamingCallbackHandler):
aiterator = h.tap_output_aiter(
run_manager.run_id, aiterator
)
# consume into final output
output = await _consume_aiter(aiterator)
# sequence doesn't emit output, yield to mark as generator
yield
except BaseException as e:
await run_manager.on_chain_error(e)
raise
else:
await run_manager.on_chain_end(output)
def _consume_iter(it: Iterator[Any]) -> Any: def _consume_iter(it: Iterator[Any]) -> Any:
+59 -8
View File
@@ -21,6 +21,7 @@ from langgraph._internal._runnable import (
run_in_executor, run_in_executor,
) )
from langgraph.config import get_config from langgraph.config import get_config
from langgraph.pregel._read import _validate_traceable_config
from langgraph.pregel._write import ChannelWrite, ChannelWriteEntry from langgraph.pregel._write import ChannelWrite, ChannelWriteEntry
from langgraph.types import CachePolicy, RetryPolicy from langgraph.types import CachePolicy, RetryPolicy
@@ -192,6 +193,19 @@ def get_runnable_for_entrypoint(func: Callable[..., Any]) -> Runnable:
return CACHE.setdefault(key, run) return CACHE.setdefault(key, run)
def _compose_trace_inputs(
base_fn: Callable[[Any], Any],
process_fn: Callable[[Any], Any],
) -> Callable[[Any], Any]:
"""Compose trace_inputs with process_inputs from traceable config."""
def composed(inputs: Any) -> Any:
exploded = base_fn(inputs)
return process_fn(exploded)
return composed
def get_runnable_for_task(func: Callable[..., Any]) -> Runnable: def get_runnable_for_task(func: Callable[..., Any]) -> Runnable:
key = (func, True) key = (func, True)
if key in CACHE: if key in CACHE:
@@ -206,10 +220,22 @@ def get_runnable_for_task(func: Callable[..., Any]) -> Runnable:
else: else:
name = str(func) name = str(func)
if is_async_callable(func): # Check for traceable config early to handle __unwrapped__
raw_config = getattr(func, "__traceable_config__", None)
traceable_config = _validate_traceable_config(raw_config)
# Use unwrapped function if available to avoid double-tracing
# when @traceable and @task are used together
func_to_run: Callable[..., Any] = func
if traceable_config:
unwrapped = traceable_config.get("__unwrapped__")
if unwrapped is not None:
func_to_run = unwrapped
if is_async_callable(func_to_run):
run = RunnableCallable( run = RunnableCallable(
None, None,
func, func_to_run,
explode_args=True, explode_args=True,
name=name, name=name,
trace=False, trace=False,
@@ -217,20 +243,45 @@ def get_runnable_for_task(func: Callable[..., Any]) -> Runnable:
) )
else: else:
run = RunnableCallable( run = RunnableCallable(
func, func_to_run,
functools.wraps(func)(functools.partial(run_in_executor, None, func)), functools.wraps(func_to_run)(
functools.partial(run_in_executor, None, func_to_run)
),
explode_args=True, explode_args=True,
name=name, name=name,
trace=False, trace=False,
recurse=False, recurse=False,
) )
# Build trace_inputs - compose with process_inputs if provided
# Use original func for signature to match expected args
base_trace_inputs = functools.partial(
_explode_args_trace_inputs, inspect.signature(func)
)
process_inputs = (
traceable_config.get("process_inputs") if traceable_config else None
)
if process_inputs is not None:
trace_inputs = _compose_trace_inputs(base_trace_inputs, process_inputs)
else:
trace_inputs = base_trace_inputs
# Build seq_kwargs
seq_kwargs: dict[str, Any] = {
"name": name,
"trace_inputs": trace_inputs,
}
if traceable_config:
# Filter out LangChainTracer (skip LangSmith) but keep other callbacks
if traceable_config.get("enabled") is False:
seq_kwargs["skip_langsmith"] = True
if traceable_config.get("process_outputs"):
seq_kwargs["trace_outputs"] = traceable_config["process_outputs"]
seq = RunnableSeq( seq = RunnableSeq(
run, run,
ChannelWrite([ChannelWriteEntry(RETURN)]), ChannelWrite([ChannelWriteEntry(RETURN)]),
name=name, **seq_kwargs,
trace_inputs=functools.partial(
_explode_args_trace_inputs, inspect.signature(func)
),
) )
if not _lookup_module_and_qualname(func): if not _lookup_module_and_qualname(func):
return seq return seq
+80 -7
View File
@@ -7,10 +7,15 @@ from typing import (
) )
from langchain_core.runnables import Runnable, RunnableConfig from langchain_core.runnables import Runnable, RunnableConfig
from typing_extensions import TypedDict
from langgraph._internal._config import merge_configs from langgraph._internal._config import merge_configs
from langgraph._internal._constants import CONF, CONFIG_KEY_READ from langgraph._internal._constants import CONF, CONFIG_KEY_READ
from langgraph._internal._runnable import RunnableCallable, RunnableSeq from langgraph._internal._runnable import (
RunnableCallable,
RunnableSeq,
coerce_to_runnable,
)
from langgraph.pregel._utils import find_subgraph_pregel from langgraph.pregel._utils import find_subgraph_pregel
from langgraph.pregel._write import ChannelWrite from langgraph.pregel._write import ChannelWrite
from langgraph.pregel.protocol import PregelProtocol from langgraph.pregel.protocol import PregelProtocol
@@ -20,6 +25,32 @@ READ_TYPE = Callable[[str | Sequence[str], bool], Any | dict[str, Any]]
INPUT_CACHE_KEY_TYPE = tuple[Callable[..., Any], tuple[str, ...]] INPUT_CACHE_KEY_TYPE = tuple[Callable[..., Any], tuple[str, ...]]
class TraceableConfig(TypedDict, total=False):
"""Configuration extracted from @traceable decorated functions."""
process_inputs: Callable[[Any], Any] | None
process_outputs: Callable[[Any], Any] | None
enabled: bool | None
__unwrapped__: Callable[[Any], Any] | None
def _validate_traceable_config(raw: Any) -> TraceableConfig | None:
"""Validate __traceable_config__ has expected structure.
Returns validated config dict or None if invalid.
"""
if not isinstance(raw, dict):
return None
# Support both "__unwrapped__" (langsmith) and "wrapped" (legacy) keys
unwrapped = raw.get("__unwrapped__") or raw.get("wrapped")
return {
"process_inputs": raw.get("process_inputs"),
"process_outputs": raw.get("process_outputs"),
"enabled": raw.get("enabled"), # None means use external context
"__unwrapped__": unwrapped,
}
class ChannelRead(RunnableCallable): class ChannelRead(RunnableCallable):
"""Implements the logic for reading state from CONFIG_KEY_READ. """Implements the logic for reading state from CONFIG_KEY_READ.
Usable both as a runnable as well as a static method to call imperatively.""" Usable both as a runnable as well as a static method to call imperatively."""
@@ -176,6 +207,7 @@ class PregelNode:
attrs = {**self.__dict__, **update} attrs = {**self.__dict__, **update}
# Drop the cached properties # Drop the cached properties
attrs.pop("flat_writers", None) attrs.pop("flat_writers", None)
attrs.pop("_traceable_config", None)
attrs.pop("node", None) attrs.pop("node", None)
attrs.pop("input_cache_key", None) attrs.pop("input_cache_key", None)
return PregelNode(**attrs) return PregelNode(**attrs)
@@ -197,6 +229,31 @@ class PregelNode:
writers.pop() writers.pop()
return writers return writers
@cached_property
def _traceable_config(self) -> TraceableConfig | None:
"""Extract @traceable config if bound wraps a traceable function.
Returns validated TraceableConfig with:
- process_inputs: Optional input processing function
- process_outputs: Optional output processing function
- enabled: Whether tracing is enabled (None = honor external context)
- __unwrapped__: The original unwrapped function (if available)
"""
if not isinstance(self.bound, RunnableCallable):
return None
func = self.bound.func or self.bound.afunc
if func is None:
return None
raw_config = getattr(func, "__traceable_config__", None)
if raw_config is None:
return None
config = _validate_traceable_config(raw_config)
if config is None:
return None
return config
@cached_property @cached_property
def node(self) -> Runnable[Any, Any] | None: def node(self) -> Runnable[Any, Any] | None:
"""Get a runnable that combines `bound` and `writers`.""" """Get a runnable that combines `bound` and `writers`."""
@@ -207,10 +264,24 @@ class PregelNode:
return writers[0] return writers[0]
elif self.bound is DEFAULT_BOUND: elif self.bound is DEFAULT_BOUND:
return RunnableSeq(*writers) return RunnableSeq(*writers)
elif writers: elif not writers:
return RunnableSeq(self.bound, *writers)
else:
return self.bound return self.bound
tc = self._traceable_config
seq_kwargs: dict[str, Any] = {}
if tc:
# Filter out LangChainTracer (skip LangSmith) but keep other callbacks
if tc.get("enabled") is False:
seq_kwargs["skip_langsmith"] = True
seq_kwargs["trace_inputs"] = tc.get("process_inputs")
seq_kwargs["trace_outputs"] = tc.get("process_outputs")
bound = self.bound
if tc and (unwrapped := tc.get("__unwrapped__")):
# We want to avoid double-tracing.
bound = coerce_to_runnable(unwrapped, name=None, trace=False)
return RunnableSeq(bound, *writers, **seq_kwargs)
@cached_property @cached_property
def input_cache_key(self) -> INPUT_CACHE_KEY_TYPE: def input_cache_key(self) -> INPUT_CACHE_KEY_TYPE:
@@ -218,9 +289,11 @@ class PregelNode:
This is used to avoid calculating the same input multiple times.""" This is used to avoid calculating the same input multiple times."""
return ( return (
self.mapper, self.mapper,
tuple(self.channels) (
if isinstance(self.channels, list) tuple(self.channels)
else (self.channels,), if isinstance(self.channels, list)
else (self.channels,)
),
) )
def invoke( def invoke(
File diff suppressed because it is too large Load Diff