diff --git a/libs/langgraph/langgraph/pregel/__init__.py b/libs/langgraph/langgraph/pregel/__init__.py index 1b28260d1..cb0a89a79 100644 --- a/libs/langgraph/langgraph/pregel/__init__.py +++ b/libs/langgraph/langgraph/pregel/__init__.py @@ -39,11 +39,6 @@ from langchain_core.runnables.utils import ( ConfigurableFieldSpec, get_unique_config_specs, ) - -try: - from langchain_core.tracers._streaming import _StreamingCallbackHandler -except ImportError: - _StreamingCallbackHandler = None from pydantic import BaseModel from typing_extensions import Self @@ -129,6 +124,11 @@ from langgraph.utils.fields import get_enhanced_type_hints from langgraph.utils.pydantic import create_model, is_supported_by_pydantic from langgraph.utils.queue import AsyncQueue, SyncQueue # type: ignore[attr-defined] +try: + from langchain_core.tracers._streaming import _StreamingCallbackHandler +except ImportError: + _StreamingCallbackHandler = None + WriteValue = Union[Callable[[Input], Output], Any] diff --git a/libs/langgraph/langgraph/pregel/messages.py b/libs/langgraph/langgraph/pregel/messages.py index 867012fa6..d53598220 100644 --- a/libs/langgraph/langgraph/pregel/messages.py +++ b/libs/langgraph/langgraph/pregel/messages.py @@ -7,6 +7,7 @@ from typing import ( List, Optional, Sequence, + TypeVar, Union, cast, ) @@ -15,11 +16,16 @@ from uuid import UUID, uuid4 from langchain_core.callbacks import BaseCallbackHandler from langchain_core.messages import BaseMessage from langchain_core.outputs import ChatGenerationChunk, LLMResult -from langchain_core.tracers._streaming import T, _StreamingCallbackHandler from langgraph.constants import NS_SEP, TAG_HIDDEN, TAG_NOSTREAM from langgraph.types import StreamChunk +try: + from langchain_core.tracers._streaming import _StreamingCallbackHandler +except ImportError: + _StreamingCallbackHandler = object + +T = TypeVar("T") Meta = tuple[tuple[str, ...], dict[str, Any]] diff --git a/libs/langgraph/langgraph/utils/runnable.py b/libs/langgraph/langgraph/utils/runnable.py index f21420ad0..f51da12d9 100644 --- a/libs/langgraph/langgraph/utils/runnable.py +++ b/libs/langgraph/langgraph/utils/runnable.py @@ -36,7 +36,6 @@ from langchain_core.runnables.config import ( var_child_runnable_config, ) from langchain_core.runnables.utils import Input, Output -from langchain_core.tracers._streaming import _StreamingCallbackHandler from typing_extensions import TypeGuard from langgraph.constants import ( @@ -54,6 +53,11 @@ from langgraph.utils.config import ( patch_config, ) +try: + from langchain_core.tracers._streaming import _StreamingCallbackHandler +except ImportError: + _StreamingCallbackHandler = None + def _set_config_context( config: RunnableConfig, @@ -683,13 +687,15 @@ class RunnableSeq(Runnable): iterator = step.stream(input, config, **kwargs) else: iterator = step.transform(iterator, config) - if stream_handler := next( - ( - cast(_StreamingCallbackHandler, h) - for h in run_manager.handlers - if isinstance(h, _StreamingCallbackHandler) - ), - None, + if _StreamingCallbackHandler is not None and ( + stream_handler := next( + ( + cast(_StreamingCallbackHandler, h) # type: ignore + for h in run_manager.handlers + if isinstance(h, _StreamingCallbackHandler) + ), + None, + ) ): # populates streamed_output in astream_log() output if needed iterator = stream_handler.tap_output_iter(run_manager.run_id, iterator) @@ -749,13 +755,15 @@ class RunnableSeq(Runnable): aiterator = step.atransform(aiterator, config) if hasattr(aiterator, "aclose"): stack.push_async_callback(aiterator.aclose) - if stream_handler := next( - ( - cast(_StreamingCallbackHandler, h) - for h in run_manager.handlers - if isinstance(h, _StreamingCallbackHandler) - ), - None, + if _StreamingCallbackHandler is not None and ( + stream_handler := next( + ( + cast(_StreamingCallbackHandler, h) # type: ignore + for h in run_manager.handlers + if isinstance(h, _StreamingCallbackHandler) + ), + None, + ) ): # populates streamed_output in astream_log() output if needed aiterator = stream_handler.tap_output_aiter(