mirror of
https://github.com/langchain-ai/langgraph.git
synced 2026-09-13 13:17:52 +02:00
Compare commits
9
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
53172459bb | ||
|
|
df254a7f7f | ||
|
|
46bd949f53 | ||
|
|
2e7a6f088e | ||
|
|
f941f467d5 | ||
|
|
a7aaa81ed8 | ||
|
|
eb9e6e2a0b | ||
|
|
2d304b8a1b | ||
|
|
6b17b12ae4 |
@@ -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:
|
||||||
|
|||||||
@@ -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
|
||||||
|
|||||||
@@ -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
Reference in New Issue
Block a user