Tony20100/langchain_full huggingface dataset:
Source code for langchain_core.runnables.base
from __future__ import annotations
import asyncio
import collections
import inspect
import threading
from abc import ABC, abstractmethod
from concurrent.futures import FIRST_COMPLETED, wait
from contextvars import copy_context
from functools import wraps
from itertools import groupby, tee
from operator import itemgetter
from typing import (
    TYPE_CHECKING,
    Any,
    AsyncIterator,
    Awaitable,
    Callable,
    Coroutine,
    Dict,
    Generic,
    Iterator,
    List,
    Mapping,
    Optional,
    Sequence,
    Set,
    Tuple,
    Type,
    TypeVar,
    Union,
    cast,
    overload,
)
from typing_extensions import Literal, get_args
from langchain_core._api import beta_decorator
from langchain_core.load.dump import dumpd
from langchain_core.load.serializable import (
    Serializable,
    SerializedConstructor,
    SerializedNotImplemented,
)
from langchain_core.pydantic_v1 import BaseModel, Field
from langchain_core.runnables.config import (
    RunnableConfig,
    _set_config_context,
    acall_func_with_variable_args,
    call_func_with_variable_args,
    ensure_config,
    get_async_callback_manager_for_config,
    get_callback_manager_for_config,
    get_config_list,
    get_executor_for_config,
    merge_configs,
    patch_config,
    run_in_executor,
)
from langchain_core.runnables.graph import Graph
from langchain_core.runnables.schema import StreamEvent
from langchain_core.runnables.utils import (
    AddableDict,
    AnyConfigurableField,
    ConfigurableField,
    ConfigurableFieldSpec,
    Input,
    Output,
ConfigurableFieldSpec,
    Input,
    Output,
    accepts_config,
    accepts_context,
    accepts_run_manager,
    create_model,
    gather_with_concurrency,
    get_function_first_arg_dict_keys,
    get_function_nonlocals,
    get_lambda_source,
    get_unique_config_specs,
    indent_lines_after_first,
    is_async_callable,
    is_async_generator,
)
from langchain_core.utils.aiter import atee, py_anext
from langchain_core.utils.iter import safetee
if TYPE_CHECKING:
    from langchain_core.callbacks.manager import (
        AsyncCallbackManagerForChainRun,
        CallbackManagerForChainRun,
    )
    from langchain_core.prompts.base import BasePromptTemplate
    from langchain_core.runnables.fallbacks import (
        RunnableWithFallbacks as RunnableWithFallbacksT,
    )
    from langchain_core.tracers.log_stream import (
        RunLog,
        RunLogPatch,
    )
    from langchain_core.tracers.root_listeners import AsyncListener
    from langchain_core.tracers.schemas import Run
Other = TypeVar("Other")
[docs]class Runnable(Generic[Input, Output], ABC):
    """A unit of work that can be invoked, batched, streamed, transformed and composed.
     Key Methods
     ===========
    - **invoke/ainvoke**: Transforms a single input into an output.
    - **batch/abatch**: Efficiently transforms multiple inputs into outputs.
    - **stream/astream**: Streams output from a single input as it's produced.
    - **astream_log**: Streams output and selected intermediate results from an input.
    Built-in optimizations:
Built-in optimizations:
    - **Batch**: By default, batch runs invoke() in parallel using a thread pool executor.
             Override to optimize batching.
    - **Async**: Methods with "a" suffix are asynchronous. By default, they execute
             the sync counterpart using asyncio's thread pool.
             Override for native async.
    All methods accept an optional config argument, which can be used to configure
    execution, add tags and metadata for tracing and debugging etc.
    Runnables expose schematic information about their input, output and config via
    the input_schema property, the output_schema property and config_schema method.
    LCEL and Composition
    ====================
    The LangChain Expression Language (LCEL) is a declarative way to compose Runnables
    into chains. Any chain constructed this way will automatically have sync, async,
    batch, and streaming support.
    The main composition primitives are RunnableSequence and RunnableParallel.
    RunnableSequence invokes a series of runnables sequentially, with one runnable's
    output serving as the next's input. Construct using the `|` operator or by
    passing a list of runnables to RunnableSequence.
    RunnableParallel invokes runnables concurrently, providing the same input
    to each. Construct it using a dict literal within a sequence or by passing a
    dict to RunnableParallel.
    For example,
    .. code-block:: python
        from langchain_core.runnables import RunnableLambda
        # A RunnableSequence constructed using the `|` operator
        sequence = RunnableLambda(lambda x: x + 1) | RunnableLambda(lambda x: x * 2)
        sequence.invoke(1) # 4
        sequence.batch([1, 2, 3]) # [4, 6, 8]
# A sequence that contains a RunnableParallel constructed using a dict literal
        sequence = RunnableLambda(lambda x: x + 1) | {
            'mul_2': RunnableLambda(lambda x: x * 2),
            'mul_5': RunnableLambda(lambda x: x * 5)
        }
        sequence.invoke(1) # {'mul_2': 4, 'mul_5': 10}
    Standard Methods
    ================
    All Runnables expose additional methods that can be used to modify their behavior
    (e.g., add a retry policy, add lifecycle listeners, make them configurable, etc.).
    These methods will work on any Runnable, including Runnable chains constructed
    by composing other Runnables. See the individual methods for details.
    For example,
    .. code-block:: python
        from langchain_core.runnables import RunnableLambda
        import random
        def add_one(x: int) -> int:
            return x + 1
        def buggy_double(y: int) -> int:
            '''Buggy code that will fail 70% of the time'''
            if random.random() > 0.3:
                print('This code failed, and will probably be retried!')  # noqa: T201
                raise ValueError('Triggered buggy code')
            return y * 2
        sequence = (
            RunnableLambda(add_one) |
            RunnableLambda(buggy_double).with_retry( # Retry on failure
                stop_after_attempt=10,
                wait_exponential_jitter=False
            )
        )
        print(sequence.input_schema.schema()) # Show inferred input schema
        print(sequence.output_schema.schema()) # Show inferred output schema
        print(sequence.invoke(2)) # invoke the sequence (note the retry above!!)
    Debugging and tracing
    =====================
Debugging and tracing
    =====================
    As the chains get longer, it can be useful to be able to see intermediate results
    to debug and trace the chain.
    You can set the global debug flag to True to enable debug output for all chains:
        .. code-block:: python
            from langchain_core.globals import set_debug
            set_debug(True)
    Alternatively, you can pass existing or custom callbacks to any given chain:
        .. code-block:: python
            from langchain_core.tracers import ConsoleCallbackHandler
            chain.invoke(
                ...,
                config={'callbacks': [ConsoleCallbackHandler()]}
            )
    For a UI (and much more) checkout LangSmith: https://docs.smith.langchain.com/
    """  # noqa: E501
    name: Optional[str] = None
    """The name of the runnable. Used for debugging and tracing."""
[docs]    def get_name(
        self, suffix: Optional[str] = None, *, name: Optional[str] = None
    ) -> str:
        """Get the name of the runnable."""
        name = name or self.name or self.__class__.__name__
        if suffix:
            if name[0].isupper():
                return name + suffix.title()
            else:
                return name + "_" + suffix.lower()
        else:
            return name
    @property
    def InputType(self) -> Type[Input]:
        """The type of input this runnable accepts specified as a type annotation."""
        for cls in self.__class__.__orig_bases__:  # type: ignore[attr-defined]
            type_args = get_args(cls)
            if type_args and len(type_args) == 2:
                return type_args[0]
        raise TypeError(
return type_args[0]
        raise TypeError(
            f"Runnable {self.get_name()} doesn't have an inferable InputType. "
            "Override the InputType property to specify the input type."
        )
    @property
    def OutputType(self) -> Type[Output]:
        """The type of output this runnable produces specified as a type annotation."""
        for cls in self.__class__.__orig_bases__:  # type: ignore[attr-defined]
            type_args = get_args(cls)
            if type_args and len(type_args) == 2:
                return type_args[1]
        raise TypeError(
            f"Runnable {self.get_name()} doesn't have an inferable OutputType. "
            "Override the OutputType property to specify the output type."
        )
    @property
    def input_schema(self) -> Type[BaseModel]:
        """The type of input this runnable accepts specified as a pydantic model."""
        return self.get_input_schema()
[docs]    def get_input_schema(
        self, config: Optional[RunnableConfig] = None
    ) -> Type[BaseModel]:
        """Get a pydantic model that can be used to validate input to the runnable.
        Runnables that leverage the configurable_fields and configurable_alternatives
        methods will have a dynamic input schema that depends on which
        configuration the runnable is invoked with.
        This method allows to get an input schema for a specific configuration.
        Args:
            config: A config to use when generating the schema.
        Returns:
            A pydantic model that can be used to validate input.
        """
        root_type = self.InputType
        if inspect.isclass(root_type) and issubclass(root_type, BaseModel):
            return root_type
        return create_model(
return root_type
        return create_model(
            self.get_name("Input"),
            __root__=(root_type, None),
        )
    @property
    def output_schema(self) -> Type[BaseModel]:
        """The type of output this runnable produces specified as a pydantic model."""
        return self.get_output_schema()
[docs]    def get_output_schema(
        self, config: Optional[RunnableConfig] = None
    ) -> Type[BaseModel]:
        """Get a pydantic model that can be used to validate output to the runnable.
        Runnables that leverage the configurable_fields and configurable_alternatives
        methods will have a dynamic output schema that depends on which
        configuration the runnable is invoked with.
        This method allows to get an output schema for a specific configuration.
        Args:
            config: A config to use when generating the schema.
        Returns:
            A pydantic model that can be used to validate output.
        """
        root_type = self.OutputType
        if inspect.isclass(root_type) and issubclass(root_type, BaseModel):
            return root_type
        return create_model(
            self.get_name("Output"),
            __root__=(root_type, None),
        )
    @property
    def config_specs(self) -> List[ConfigurableFieldSpec]:
        """List configurable fields for this runnable."""
        return []
[docs]    def config_schema(
        self, *, include: Optional[Sequence[str]] = None
    ) -> Type[BaseModel]:
        """The type of config this runnable accepts specified as a pydantic model.
        To mark a field as configurable, see the `configurable_fields`
        and `configurable_alternatives` methods.
        Args:
and `configurable_alternatives` methods.
        Args:
            include: A list of fields to include in the config schema.
        Returns:
            A pydantic model that can be used to validate config.
        """
        include = include or []
        config_specs = self.config_specs
        configurable = (
            create_model(  # type: ignore[call-overload]
                "Configurable",
                **{
                    spec.id: (
                        spec.annotation,
                        Field(
                            spec.default, title=spec.name, description=spec.description
                        ),
                    )
                    for spec in config_specs
                },
            )
            if config_specs
            else None
        )
        return create_model(  # type: ignore[call-overload]
            self.get_name("Config"),
            **({"configurable": (configurable, None)} if configurable else {}),
            **{
                field_name: (field_type, None)
                for field_name, field_type in RunnableConfig.__annotations__.items()
                if field_name in [i for i in include if i != "configurable"]
            },
        )
[docs]    def get_graph(self, config: Optional[RunnableConfig] = None) -> Graph:
        """Return a graph representation of this runnable."""
        from langchain_core.runnables.graph import Graph
        graph = Graph()
        try:
            input_node = graph.add_node(self.get_input_schema(config))
        except TypeError:
            input_node = graph.add_node(create_model(self.get_name("Input")))
        runnable_node = graph.add_node(self)
        try:
            output_node = graph.add_node(self.get_output_schema(config))
        except TypeError:
            output_node = graph.add_node(create_model(self.get_name("Output")))
output_node = graph.add_node(create_model(self.get_name("Output")))
        graph.add_edge(input_node, runnable_node)
        graph.add_edge(runnable_node, output_node)
        return graph
[docs]    def get_prompts(
        self, config: Optional[RunnableConfig] = None
    ) -> List[BasePromptTemplate]:
        from langchain_core.prompts.base import BasePromptTemplate
        prompts = []
        for _, node in self.get_graph(config=config).nodes.items():
            if isinstance(node.data, BasePromptTemplate):
                prompts.append(node.data)
        return prompts
    def __or__(
        self,
        other: Union[
            Runnable[Any, Other],
            Callable[[Any], Other],
            Callable[[Iterator[Any]], Iterator[Other]],
            Mapping[str, Union[Runnable[Any, Other], Callable[[Any], Other], Any]],
        ],
    ) -> RunnableSerializable[Input, Other]:
        """Compose this runnable with another object to create a RunnableSequence."""
        return RunnableSequence(self, coerce_to_runnable(other))
    def __ror__(
        self,
        other: Union[
            Runnable[Other, Any],
            Callable[[Other], Any],
            Callable[[Iterator[Other]], Iterator[Any]],
            Mapping[str, Union[Runnable[Other, Any], Callable[[Other], Any], Any]],
        ],
    ) -> RunnableSerializable[Other, Output]:
        """Compose this runnable with another object to create a RunnableSequence."""
        return RunnableSequence(coerce_to_runnable(other), self)
[docs]    def pipe(
        self,
        *others: Union[Runnable[Any, Other], Callable[[Any], Other]],
        name: Optional[str] = None,
    ) -> RunnableSerializable[Input, Other]:
) -> RunnableSerializable[Input, Other]:
        """Compose this Runnable with Runnable-like objects to make a RunnableSequence.
        Equivalent to `RunnableSequence(self, *others)` or `self | others[0] | ...`
        Example:
            .. code-block:: python
                from langchain_core.runnables import RunnableLambda
                def add_one(x: int) -> int:
                    return x + 1
                def mul_two(x: int) -> int:
                    return x * 2
                runnable_1 = RunnableLambda(add_one)
                runnable_2 = RunnableLambda(mul_two)
                sequence = runnable_1.pipe(runnable_2)
                # Or equivalently:
                # sequence = runnable_1 | runnable_2
                # sequence = RunnableSequence(first=runnable_1, last=runnable_2)
                sequence.invoke(1)
                await sequence.ainvoke(1)
                # -> 4
                sequence.batch([1, 2, 3])
                await sequence.abatch([1, 2, 3])
                # -> [4, 6, 8]
        """
        return RunnableSequence(self, *others, name=name)
[docs]    def pick(self, keys: Union[str, List[str]]) -> RunnableSerializable[Any, Any]:
        """Pick keys from the dict output of this runnable.
        Pick single key:
            .. code-block:: python
                import json
                from langchain_core.runnables import RunnableLambda, RunnableMap
                as_str = RunnableLambda(str)
                as_json = RunnableLambda(json.loads)
                chain = RunnableMap(str=as_str, json=as_json)
                chain.invoke("[1, 2, 3]")
chain.invoke("[1, 2, 3]")
                # -> {"str": "[1, 2, 3]", "json": [1, 2, 3]}
                json_only_chain = chain.pick("json")
                json_only_chain.invoke("[1, 2, 3]")
                # -> [1, 2, 3]
        Pick list of keys:
            .. code-block:: python
                from typing import Any
                import json
                from langchain_core.runnables import RunnableLambda, RunnableMap
                as_str = RunnableLambda(str)
                as_json = RunnableLambda(json.loads)
                def as_bytes(x: Any) -> bytes:
                    return bytes(x, "utf-8")
                chain = RunnableMap(
                    str=as_str,
                    json=as_json,
                    bytes=RunnableLambda(as_bytes)
                )
                chain.invoke("[1, 2, 3]")
                # -> {"str": "[1, 2, 3]", "json": [1, 2, 3], "bytes": b"[1, 2, 3]"}
                json_and_bytes_chain = chain.pick(["json", "bytes"])
                json_and_bytes_chain.invoke("[1, 2, 3]")
                # -> {"json": [1, 2, 3], "bytes": b"[1, 2, 3]"}
        """
        from langchain_core.runnables.passthrough import RunnablePick
        return self | RunnablePick(keys)
[docs]    def assign(
        self,
        **kwargs: Union[
            Runnable[Dict[str, Any], Any],
            Callable[[Dict[str, Any]], Any],
            Mapping[
                str,
Callable[[Dict[str, Any]], Any],
            Mapping[
                str,
                Union[Runnable[Dict[str, Any], Any], Callable[[Dict[str, Any]], Any]],
            ],
        ],
    ) -> RunnableSerializable[Any, Any]:
        """Assigns new fields to the dict output of this runnable.
        Returns a new runnable.
        .. code-block:: python
            from langchain_community.llms.fake import FakeStreamingListLLM
            from langchain_core.output_parsers import StrOutputParser
            from langchain_core.prompts import SystemMessagePromptTemplate
            from langchain_core.runnables import Runnable
            from operator import itemgetter
            prompt = (
                SystemMessagePromptTemplate.from_template("You are a nice assistant.")
                + "{question}"
            )
            llm = FakeStreamingListLLM(responses=["foo-lish"])
            chain: Runnable = prompt | llm | {"str": StrOutputParser()}
            chain_with_assign = chain.assign(hello=itemgetter("str") | llm)
            print(chain_with_assign.input_schema.schema())
            # {'title': 'PromptInput', 'type': 'object', 'properties':
            {'question': {'title': 'Question', 'type': 'string'}}}
            print(chain_with_assign.output_schema.schema()) #
            {'title': 'RunnableSequenceOutput', 'type': 'object', 'properties':
            {'str': {'title': 'Str',
            'type': 'string'}, 'hello': {'title': 'Hello', 'type': 'string'}}}
        """
        from langchain_core.runnables.passthrough import RunnableAssign
        return self | RunnableAssign(RunnableParallel(kwargs))
    """ --- Public API --- """
[docs]    @abstractmethod
""" --- Public API --- """
[docs]    @abstractmethod
    def invoke(self, input: Input, config: Optional[RunnableConfig] = None) -> Output:
        """Transform a single input into an output. Override to implement.
        Args:
            input: The input to the runnable.
            config: A config to use when invoking the runnable.
               The config supports standard keys like 'tags', 'metadata' for tracing
               purposes, 'max_concurrency' for controlling how much work to do
               in parallel, and other keys. Please refer to the RunnableConfig
               for more details.
        Returns:
            The output of the runnable.
        """
[docs]    async def ainvoke(
        self, input: Input, config: Optional[RunnableConfig] = None, **kwargs: Any
    ) -> Output:
        """Default implementation of ainvoke, calls invoke from a thread.
        The default implementation allows usage of async code even if
        the runnable did not implement a native async version of invoke.
        Subclasses should override this method if they can run asynchronously.
        """
        return await run_in_executor(config, self.invoke, input, config, **kwargs)
[docs]    def batch(
        self,
        inputs: List[Input],
        config: Optional[Union[RunnableConfig, List[RunnableConfig]]] = None,
        *,
        return_exceptions: bool = False,
        **kwargs: Optional[Any],
    ) -> List[Output]:
        """Default implementation runs invoke in parallel using a thread pool executor.
        The default implementation of batch works well for IO bound runnables.
        Subclasses should override this method if they can batch more efficiently;
        e.g., if the underlying runnable uses an API which supports a batch mode.
        """
"""
        if not inputs:
            return []
        configs = get_config_list(config, len(inputs))
        def invoke(input: Input, config: RunnableConfig) -> Union[Output, Exception]:
            if return_exceptions:
                try:
                    return self.invoke(input, config, **kwargs)
                except Exception as e:
                    return e
            else:
                return self.invoke(input, config, **kwargs)
        # If there's only one input, don't bother with the executor
        if len(inputs) == 1:
            return cast(List[Output], [invoke(inputs[0], configs[0])])
        with get_executor_for_config(configs[0]) as executor:
            return cast(List[Output], list(executor.map(invoke, inputs, configs)))
    @overload
    def batch_as_completed(
        self,
        inputs: Sequence[Input],
        config: Optional[Union[RunnableConfig, Sequence[RunnableConfig]]] = None,
        *,
        return_exceptions: Literal[False] = False,
        **kwargs: Any,
    ) -> Iterator[Tuple[int, Output]]:
        ...
    @overload
    def batch_as_completed(
        self,
        inputs: Sequence[Input],
        config: Optional[Union[RunnableConfig, Sequence[RunnableConfig]]] = None,
        *,
        return_exceptions: Literal[True],
        **kwargs: Any,
    ) -> Iterator[Tuple[int, Union[Output, Exception]]]:
        ...
[docs]    def batch_as_completed(
        self,
        inputs: Sequence[Input],
        config: Optional[Union[RunnableConfig, Sequence[RunnableConfig]]] = None,
        *,
        return_exceptions: bool = False,
        **kwargs: Optional[Any],