1"""Utility code for `Runnable` objects."""23from __future__ import annotations45import ast6import asyncio7import inspect8import sys9import textwrap1011# Cannot move to TYPE_CHECKING as Mapping and Sequence are needed at runtime by12# RunnableConfigurableFields.13from collections.abc import Mapping, Sequence  # noqa: TC00314from functools import lru_cache15from inspect import signature16from itertools import groupby17from typing import (18    TYPE_CHECKING,19    Any,20    NamedTuple,21    Protocol,22    TypeGuard,23    TypeVar,24)2526from typing_extensions import override2728# Re-export create-model for backwards compatibility29from langchain_core.utils.pydantic import create_model  # noqa: F4013031if TYPE_CHECKING:32    from collections.abc import (33        AsyncIterable,34        AsyncIterator,35        Awaitable,36        Callable,37        Coroutine,38        Iterable,39    )40    from contextvars import Context4142    from langchain_core.runnables.schema import StreamEvent4344Input = TypeVar("Input", contravariant=True)  # noqa: PLC010545# Output type should implement __concat__, as eg str, list, dict do46Output = TypeVar("Output", covariant=True)  # noqa: PLC0105474849async def gated_coro(50    semaphore: asyncio.Semaphore, coro: Coroutine[Any, Any, Any]51) -> Any:52    """Run a coroutine with a semaphore.5354    Args:55        semaphore: The semaphore to use.56        coro: The coroutine to run.5758    Returns:59        The result of the coroutine.60    """61    async with semaphore:62        return await coro636465async def gather_with_concurrency(66    n: int | None, *coros: Coroutine[Any, Any, Any]67) -> list[Any]:68    """Gather coroutines with a limit on the number of concurrent coroutines.6970    Args:71        n: The number of coroutines to run concurrently.72        *coros: The coroutines to run.7374    Returns:75        The results of the coroutines.76    """77    if n is None:78        return await asyncio.gather(*coros)7980    semaphore = asyncio.Semaphore(n)8182    return await asyncio.gather(*(gated_coro(semaphore, c) for c in coros))838485def accepts_run_manager(callable: Callable[..., Any]) -> bool:  # noqa: A00286    """Check if a callable accepts a run_manager argument.8788    Args:89        callable: The callable to check.9091    Returns:92        `True` if the callable accepts a run_manager argument, `False` otherwise.93    """94    try:95        return signature(callable).parameters.get("run_manager") is not None96    except ValueError:97        return False9899100def accepts_config(callable: Callable[..., Any]) -> bool:  # noqa: A002101    """Check if a callable accepts a config argument.102103    Args:104        callable: The callable to check.105106    Returns:107        `True` if the callable accepts a config argument, `False` otherwise.108    """109    try:110        return signature(callable).parameters.get("config") is not None111    except ValueError:112        return False113114115def accepts_context(callable: Callable[..., Any]) -> bool:  # noqa: A002116    """Check if a callable accepts a context argument.117118    Args:119        callable: The callable to check.120121    Returns:122        `True` if the callable accepts a context argument, `False` otherwise.123    """124    try:125        return signature(callable).parameters.get("context") is not None126    except ValueError:127        return False128129130def asyncio_accepts_context() -> bool:131    """Check if asyncio.create_task accepts a `context` arg.132133    Returns:134        True if `asyncio.create_task` accepts a context argument, `False` otherwise.135    """136    return sys.version_info >= (3, 11)137138139_T = TypeVar("_T")140141142def coro_with_context(143    coro: Awaitable[_T], context: Context, *, create_task: bool = False144) -> Awaitable[_T]:145    """Await a coroutine with a context.146147    Args:148        coro: The coroutine to await.149        context: The context to use.150        create_task: Kept for compatibility; this helper always creates a task.151152    Returns:153        The coroutine with the context.154    """155    if asyncio_accepts_context():156        return asyncio.create_task(coro, context=context)  # type: ignore[arg-type,call-arg,unused-ignore]157    del create_task158    return context.run(asyncio.create_task, coro)  # type: ignore[arg-type]159160161class IsLocalDict(ast.NodeVisitor):162    """Check if a name is a local dict."""163164    def __init__(self, name: str, keys: set[str]) -> None:165        """Initialize the visitor.166167        Args:168            name: The name to check.169            keys: The keys to populate.170        """171        self.name = name172        self.keys = keys173174    @override175    def visit_Subscript(self, node: ast.Subscript) -> None:176        """Visit a subscript node.177178        Args:179            node: The node to visit.180        """181        if (182            isinstance(node.ctx, ast.Load)183            and isinstance(node.value, ast.Name)184            and node.value.id == self.name185            and isinstance(node.slice, ast.Constant)186            and isinstance(node.slice.value, str)187        ):188            # we've found a subscript access on the name we're looking for189            self.keys.add(node.slice.value)190191    @override192    def visit_Call(self, node: ast.Call) -> None:193        """Visit a call node.194195        Args:196            node: The node to visit.197        """198        if (199            isinstance(node.func, ast.Attribute)200            and isinstance(node.func.value, ast.Name)201            and node.func.value.id == self.name202            and node.func.attr == "get"203            and len(node.args) in {1, 2}204            and isinstance(node.args[0], ast.Constant)205            and isinstance(node.args[0].value, str)206        ):207            # we've found a .get() call on the name we're looking for208            self.keys.add(node.args[0].value)209210211class IsFunctionArgDict(ast.NodeVisitor):212    """Check if the first argument of a function is a dict."""213214    def __init__(self) -> None:215        """Create a IsFunctionArgDict visitor."""216        self.keys: set[str] = set()217218    @override219    def visit_Lambda(self, node: ast.Lambda) -> None:220        """Visit a lambda function.221222        Args:223            node: The node to visit.224        """225        if not node.args.args:226            return227        input_arg_name = node.args.args[0].arg228        IsLocalDict(input_arg_name, self.keys).visit(node.body)229230    @override231    def visit_FunctionDef(self, node: ast.FunctionDef) -> None:232        """Visit a function definition.233234        Args:235            node: The node to visit.236        """237        if not node.args.args:238            return239        input_arg_name = node.args.args[0].arg240        IsLocalDict(input_arg_name, self.keys).visit(node)241242    @override243    def visit_AsyncFunctionDef(self, node: ast.AsyncFunctionDef) -> None:244        """Visit an async function definition.245246        Args:247            node: The node to visit.248        """249        if not node.args.args:250            return251        input_arg_name = node.args.args[0].arg252        IsLocalDict(input_arg_name, self.keys).visit(node)253254255class NonLocals(ast.NodeVisitor):256    """Get nonlocal variables accessed."""257258    def __init__(self) -> None:259        """Create a NonLocals visitor."""260        self.loads: set[str] = set()261        self.stores: set[str] = set()262263    @override264    def visit_Name(self, node: ast.Name) -> None:265        """Visit a name node.266267        Args:268            node: The node to visit.269        """270        if isinstance(node.ctx, ast.Load):271            self.loads.add(node.id)272        elif isinstance(node.ctx, ast.Store):273            self.stores.add(node.id)274275    @override276    def visit_Attribute(self, node: ast.Attribute) -> None:277        """Visit an attribute node.278279        Args:280            node: The node to visit.281        """282        if isinstance(node.ctx, ast.Load):283            parent = node.value284            attr_expr = node.attr285            while isinstance(parent, ast.Attribute):286                attr_expr = parent.attr + "." + attr_expr287                parent = parent.value288            if isinstance(parent, ast.Name):289                self.loads.add(parent.id + "." + attr_expr)290                self.loads.discard(parent.id)291            elif isinstance(parent, ast.Call):292                if isinstance(parent.func, ast.Name):293                    self.loads.add(parent.func.id)294                else:295                    parent = parent.func296                    attr_expr = ""297                    while isinstance(parent, ast.Attribute):298                        if attr_expr:299                            attr_expr = parent.attr + "." + attr_expr300                        else:301                            attr_expr = parent.attr302                        parent = parent.value303                    if isinstance(parent, ast.Name):304                        self.loads.add(parent.id + "." + attr_expr)305306307class FunctionNonLocals(ast.NodeVisitor):308    """Get the nonlocal variables accessed of a function."""309310    def __init__(self) -> None:311        """Create a FunctionNonLocals visitor."""312        self.nonlocals: set[str] = set()313314    @override315    def visit_FunctionDef(self, node: ast.FunctionDef) -> None:316        """Visit a function definition.317318        Args:319            node: The node to visit.320        """321        visitor = NonLocals()322        visitor.visit(node)323        self.nonlocals.update(visitor.loads - visitor.stores)324325    @override326    def visit_AsyncFunctionDef(self, node: ast.AsyncFunctionDef) -> None:327        """Visit an async function definition.328329        Args:330            node: The node to visit.331        """332        visitor = NonLocals()333        visitor.visit(node)334        self.nonlocals.update(visitor.loads - visitor.stores)335336    @override337    def visit_Lambda(self, node: ast.Lambda) -> None:338        """Visit a lambda function.339340        Args:341            node: The node to visit.342        """343        visitor = NonLocals()344        visitor.visit(node)345        self.nonlocals.update(visitor.loads - visitor.stores)346347348class GetLambdaSource(ast.NodeVisitor):349    """Get the source code of a lambda function."""350351    def __init__(self) -> None:352        """Initialize the visitor."""353        self.source: str | None = None354        self.count = 0355356    @override357    def visit_Lambda(self, node: ast.Lambda) -> None:358        """Visit a lambda function.359360        Args:361            node: The node to visit.362        """363        self.count += 1364        if hasattr(ast, "unparse"):365            self.source = ast.unparse(node)366367368def get_function_first_arg_dict_keys(func: Callable[..., Any]) -> list[str] | None:369    """Get the keys of the first argument of a function if it is a dict.370371    Args:372        func: The function to check.373374    Returns:375        The keys of the first argument if it is a dict, None otherwise.376    """377    try:378        code = inspect.getsource(func)379        tree = ast.parse(textwrap.dedent(code))380        visitor = IsFunctionArgDict()381        visitor.visit(tree)382        return sorted(visitor.keys) if visitor.keys else None383    except (SyntaxError, TypeError, OSError, SystemError):384        return None385386387def get_lambda_source(func: Callable[..., Any]) -> str | None:388    """Get the source code of a lambda function.389390    Args:391        func: a Callable that can be a lambda function.392393    Returns:394        the source code of the lambda function.395    """396    try:397        name = func.__name__ if func.__name__ != "<lambda>" else None398    except AttributeError:399        name = None400    try:401        code = inspect.getsource(func)402        tree = ast.parse(textwrap.dedent(code))403        visitor = GetLambdaSource()404        visitor.visit(tree)405    except (SyntaxError, TypeError, OSError, SystemError):406        return name407    return visitor.source if visitor.count == 1 else name408409410@lru_cache(maxsize=256)411def get_function_nonlocals(func: Callable[..., Any]) -> list[Any]:412    """Get the nonlocal variables accessed by a function.413414    Args:415        func: The function to check.416417    Returns:418        The nonlocal variables accessed by the function.419    """420    try:421        code = inspect.getsource(func)422        tree = ast.parse(textwrap.dedent(code))423        visitor = FunctionNonLocals()424        visitor.visit(tree)425        values: list[Any] = []426        closure = (427            inspect.getclosurevars(func.__wrapped__)428            if hasattr(func, "__wrapped__") and callable(func.__wrapped__)429            else inspect.getclosurevars(func)430        )431        candidates = {**closure.globals, **closure.nonlocals}432        for k, v in candidates.items():433            if k in visitor.nonlocals:434                values.append(v)435            for kk in visitor.nonlocals:436                if "." in kk and kk.startswith(k):437                    vv = v438                    for part in kk.split(".")[1:]:439                        if vv is None:440                            break441                        try:442                            vv = getattr(vv, part)443                        except AttributeError:444                            break445                    else:446                        values.append(vv)447    except (SyntaxError, TypeError, OSError, SystemError):448        return []449450    return values451452453def indent_lines_after_first(text: str, prefix: str) -> str:454    """Indent all lines of text after the first line.455456    Args:457        text: The text to indent.458        prefix: Used to determine the number of spaces to indent.459460    Returns:461        The indented text.462    """463    n_spaces = len(prefix)464    spaces = " " * n_spaces465    lines = text.splitlines()466    return "\n".join([lines[0]] + [spaces + line for line in lines[1:]])467468469class AddableDict(dict[str, Any]):470    """Dictionary that can be added to another dictionary."""471472    def __add__(self, other: AddableDict) -> AddableDict:473        """Add a dictionary to this dictionary.474475        Args:476            other: The other dictionary to add.477478        Returns:479            A dictionary that is the result of adding the two dictionaries.480        """481        chunk = AddableDict(self)482        for key in other:483            if key not in chunk or chunk[key] is None:484                chunk[key] = other[key]485            elif other[key] is not None:486                try:487                    added = chunk[key] + other[key]488                except TypeError:489                    added = other[key]490                chunk[key] = added491        return chunk492493    def __radd__(self, other: AddableDict) -> AddableDict:494        """Add this dictionary to another dictionary.495496        Args:497            other: The other dictionary to be added to.498499        Returns:500            A dictionary that is the result of adding the two dictionaries.501        """502        chunk = AddableDict(other)503        for key in self:504            if key not in chunk or chunk[key] is None:505                chunk[key] = self[key]506            elif self[key] is not None:507                try:508                    added = chunk[key] + self[key]509                except TypeError:510                    added = self[key]511                chunk[key] = added512        return chunk513514515_T_co = TypeVar("_T_co", covariant=True)516_T_contra = TypeVar("_T_contra", contravariant=True)517518519class SupportsAdd(Protocol[_T_contra, _T_co]):520    """Protocol for objects that support addition."""521522    def __add__(self, x: _T_contra, /) -> _T_co:523        """Add the object to another object."""524525526Addable = TypeVar("Addable", bound=SupportsAdd[Any, Any])527528529def add(addables: Iterable[Addable]) -> Addable | None:530    """Add a sequence of addable objects together.531532    Args:533        addables: The addable objects to add.534535    Returns:536        The result of adding the addable objects.537    """538    final: Addable | None = None539    for chunk in addables:540        final = chunk if final is None else final + chunk541    return final542543544async def aadd(addables: AsyncIterable[Addable]) -> Addable | None:545    """Asynchronously add a sequence of addable objects together.546547    Args:548        addables: The addable objects to add.549550    Returns:551        The result of adding the addable objects.552    """553    final: Addable | None = None554    async for chunk in addables:555        final = chunk if final is None else final + chunk556    return final557558559class ConfigurableField(NamedTuple):560    """Field that can be configured by the user."""561562    id: str563    """The unique identifier of the field."""564565    name: str | None = None566    """The name of the field. """567568    description: str | None = None569    """The description of the field. """570571    annotation: Any | None = None572    """The annotation of the field. """573574    is_shared: bool = False575    """Whether the field is shared."""576577    @override578    def __hash__(self) -> int:579        return hash((self.id, self.annotation))580581582class ConfigurableFieldSingleOption(NamedTuple):583    """Field that can be configured by the user with a default value."""584585    id: str586    """The unique identifier of the field."""587588    options: Mapping[str, Any]589    """The options for the field."""590591    default: str592    """The default value for the field."""593594    name: str | None = None595    """The name of the field. """596597    description: str | None = None598    """The description of the field. """599600    is_shared: bool = False601    """Whether the field is shared."""602603    @override604    def __hash__(self) -> int:605        return hash((self.id, tuple(self.options.keys()), self.default))606607608class ConfigurableFieldMultiOption(NamedTuple):609    """Field that can be configured by the user with multiple default values."""610611    id: str612    """The unique identifier of the field."""613614    options: Mapping[str, Any]615    """The options for the field."""616617    default: Sequence[str]618    """The default values for the field."""619620    name: str | None = None621    """The name of the field. """622623    description: str | None = None624    """The description of the field. """625626    is_shared: bool = False627    """Whether the field is shared."""628629    @override630    def __hash__(self) -> int:631        return hash((self.id, tuple(self.options.keys()), tuple(self.default)))632633634AnyConfigurableField = (635    ConfigurableField | ConfigurableFieldSingleOption | ConfigurableFieldMultiOption636)637638639class ConfigurableFieldSpec(NamedTuple):640    """Field that can be configured by the user. It is a specification of a field."""641642    id: str643    """The unique identifier of the field."""644645    annotation: Any646    """The annotation of the field."""647648    name: str | None = None649    """The name of the field. """650651    description: str | None = None652    """The description of the field. """653654    default: Any = None655    """The default value for the field. """656657    is_shared: bool = False658    """Whether the field is shared."""659660    dependencies: list[str] | None = None661    """The dependencies of the field. """662663664def get_unique_config_specs(665    specs: Iterable[ConfigurableFieldSpec],666) -> list[ConfigurableFieldSpec]:667    """Get the unique config specs from a sequence of config specs.668669    Args:670        specs: The config specs.671672    Returns:673        The unique config specs.674675    Raises:676        ValueError: If the runnable sequence contains conflicting config specs.677    """678    grouped = groupby(679        sorted(specs, key=lambda s: (s.id, *(s.dependencies or []))), lambda s: s.id680    )681    unique: list[ConfigurableFieldSpec] = []682    for spec_id, dupes in grouped:683        first = next(dupes)684        others = list(dupes)685        if len(others) == 0 or all(o == first for o in others):686            unique.append(first)687        else:688            msg = (689                "RunnableSequence contains conflicting config specs"690                f"for {spec_id}: {[first, *others]}"691            )692            raise ValueError(msg)693    return unique694695696class _RootEventFilter:697    def __init__(698        self,699        *,700        include_names: Sequence[str] | None = None,701        include_types: Sequence[str] | None = None,702        include_tags: Sequence[str] | None = None,703        exclude_names: Sequence[str] | None = None,704        exclude_types: Sequence[str] | None = None,705        exclude_tags: Sequence[str] | None = None,706    ) -> None:707        """Utility to filter the root event in the astream_events implementation.708709        This is simply binding the arguments to the namespace to make save on710        a bit of typing in the astream_events implementation.711        """712        self.include_names = include_names713        self.include_types = include_types714        self.include_tags = include_tags715        self.exclude_names = exclude_names716        self.exclude_types = exclude_types717        self.exclude_tags = exclude_tags718719    def include_event(self, event: StreamEvent, root_type: str) -> bool:720        """Determine whether to include an event."""721        if (722            self.include_names is None723            and self.include_types is None724            and self.include_tags is None725        ):726            include = True727        else:728            include = False729730        event_tags = event.get("tags") or []731732        if self.include_names is not None:733            include = include or event["name"] in self.include_names734        if self.include_types is not None:735            include = include or root_type in self.include_types736        if self.include_tags is not None:737            include = include or any(tag in self.include_tags for tag in event_tags)738739        if self.exclude_names is not None:740            include = include and event["name"] not in self.exclude_names741        if self.exclude_types is not None:742            include = include and root_type not in self.exclude_types743        if self.exclude_tags is not None:744            include = include and all(745                tag not in self.exclude_tags for tag in event_tags746            )747748        return include749750751def is_async_generator(752    func: Any,753) -> TypeGuard[Callable[..., AsyncIterator[Any]]]:754    """Check if a function is an async generator.755756    Args:757        func: The function to check.758759    Returns:760        `True` if the function is an async generator, `False` otherwise.761    """762    return inspect.isasyncgenfunction(func) or (763        hasattr(func, "__call__")  # noqa: B004764        and inspect.isasyncgenfunction(func.__call__)765    )766767768def is_async_callable(769    func: Any,770) -> TypeGuard[Callable[..., Awaitable[Any]]]:771    """Check if a function is async.772773    Args:774        func: The function to check.775776    Returns:777        `True` if the function is async, `False` otherwise.778    """779    return inspect.iscoroutinefunction(func) or (780        hasattr(func, "__call__")  # noqa: B004781        and inspect.iscoroutinefunction(func.__call__)782    )