coco-framework 0.1.0__py3-none-any.whl

coco_core/runnables/config.py

@@ -0,0 +1,712 @@
+"""Configuration utilities for `Runnable` objects."""
+
+from __future__ import annotations
+
+import asyncio
+
+# Cannot move uuid to TYPE_CHECKING as RunnableConfig is used in Pydantic models
+import uuid  # noqa: TC003
+import warnings
+from collections.abc import (
+    Awaitable,
+    Callable,
+    Generator,
+    Iterable,
+    Iterator,
+    Mapping,
+    Sequence,
+)
+from concurrent.futures import Executor, Future, ThreadPoolExecutor
+from contextlib import contextmanager
+from contextvars import Context, ContextVar, Token, copy_context
+from functools import partial
+from typing import (
+    TYPE_CHECKING,
+    Any,
+    ParamSpec,
+    TypeVar,
+    cast,
+)
+
+from typing_extensions import TypedDict
+
+from coco_core.callbacks.manager import AsyncCallbackManager, CallbackManager
+from coco_core.runnables.utils import (
+    Input,
+    Output,
+    accepts_config,
+    accepts_run_manager,
+)
+
+if TYPE_CHECKING:
+    from coco_core.callbacks.base import BaseCallbackManager, Callbacks
+    from coco_core.callbacks.manager import (
+        AsyncCallbackManagerForChainRun,
+        CallbackManagerForChainRun,
+    )
+else:
+    # Pydantic validates through typed dicts, but
+    # the callbacks need forward refs updated
+    Callbacks = list | Any | None
+
+
+class EmptyDict(TypedDict, total=False):
+    """Empty dict type."""
+
+
+class RunnableConfig(TypedDict, total=False):
+    """Configuration for a `Runnable`.
+
+    !!! note Custom values
+
+        The `TypedDict` has `total=False` set intentionally to:
+
+        - Allow partial configs to be created and merged together via `merge_configs`
+        - Support config propagation from parent to child runnables via
+            `var_child_runnable_config` (a `ContextVar` that automatically passes
+            config down the call stack without explicit parameter passing), where
+            configs are merged rather than replaced
+
+        !!! example
+
+            ```python
+            # Parent sets tags
+            chain.invoke(input, config={"tags": ["parent"]})
+            # Child automatically inherits and can add:
+            # ensure_config({"tags": ["child"]}) -> {"tags": ["parent", "child"]}
+            ```
+    """
+
+    tags: list[str]
+    """Tags for this call and any sub-calls (e.g. a Chain calling an LLM).
+
+    You can use these to filter calls.
+    """
+
+    metadata: dict[str, Any]
+    """Metadata for this call and any sub-calls (e.g. a Chain calling an LLM).
+
+    Keys should be strings, values should be JSON-serializable.
+    """
+
+    callbacks: Callbacks
+    """Callbacks for this call and any sub-calls (e.g. a Chain calling an LLM).
+
+    Tags are passed to all callbacks, metadata is passed to handle*Start callbacks.
+    """
+
+    run_name: str
+    """Name for the tracer run for this call.
+
+    Defaults to the name of the class."""
+
+    max_concurrency: int | None
+    """Maximum number of parallel calls to make.
+
+    If not provided, defaults to `ThreadPoolExecutor`'s default.
+    """
+
+    recursion_limit: int
+    """Maximum number of times a call can recurse.
+
+    If not provided, defaults to `25`.
+    """
+
+    configurable: dict[str, Any]
+    """Runtime values for attributes previously made configurable on this `Runnable`,
+    or sub-`Runnable` objects, through `configurable_fields` or
+    `configurable_alternatives`.
+
+    Check `output_schema` for a description of the attributes that have been made
+    configurable.
+    """
+
+    run_id: uuid.UUID | None
+    """Unique identifier for the tracer run for this call.
+
+    If not provided, a new UUID will be generated.
+    """
+
+
+CONFIG_KEYS = [
+    "tags",
+    "metadata",
+    "callbacks",
+    "run_name",
+    "max_concurrency",
+    "recursion_limit",
+    "configurable",
+    "run_id",
+]
+
+COPIABLE_KEYS = [
+    "tags",
+    "metadata",
+    "callbacks",
+    "configurable",
+]
+
+
+# Users are expected to use the `context` API with a context object
+# (which does not get traced)
+CONFIGURABLE_TO_TRACING_METADATA_EXCLUDED_KEYS = frozenset(("api_key",))
+
+
+def _get_langsmith_inheritable_metadata_from_config(
+    config: RunnableConfig,
+) -> dict[str, Any] | None:
+    """Get LangSmith-only inheritable metadata defaults derived from config."""
+    configurable = config.get("configurable") or {}
+    metadata = {
+        key: value
+        for key, value in configurable.items()
+        if not key.startswith("__")
+        and isinstance(value, (str, int, float, bool))
+        and key not in config.get("metadata", {})
+        and key not in CONFIGURABLE_TO_TRACING_METADATA_EXCLUDED_KEYS
+    }
+    return metadata or None
+
+
+DEFAULT_RECURSION_LIMIT = 25
+
+
+var_child_runnable_config: ContextVar[RunnableConfig | None] = ContextVar(
+    "child_runnable_config", default=None
+)
+
+
+# This is imported and used in coco_graph, so don't break.
+def _set_config_context(
+    config: RunnableConfig,
+) -> tuple[Token[RunnableConfig | None], dict[str, Any] | None]:
+    """Set the child Runnable config + tracing context.
+
+    Args:
+        config: The config to set.
+
+    Returns:
+        The token to reset the config and the previous tracing context.
+    """
+    # Deferred to avoid importing langsmith at module level (~132ms).
+    from coco_core.tracers._context import (  # noqa: PLC0415
+        _set_tracing_context,
+        get_tracing_context,
+    )
+
+    from coco_core.tracers.coco import CocoTracer  # noqa: PLC0415
+
+    config_token = var_child_runnable_config.set(config)
+    current_context = None
+    if (
+        (callbacks := config.get("callbacks"))
+        and (
+            parent_run_id := getattr(callbacks, "parent_run_id", None)
+        )  # Is callback manager
+        and (
+            tracer := next(
+                (
+                    handler
+                    for handler in getattr(callbacks, "handlers", [])
+                    if isinstance(handler, CocoTracer)
+                ),
+                None,
+            )
+        )
+        and (run := tracer.run_map.get(str(parent_run_id)))
+    ):
+        current_context = get_tracing_context()
+        _set_tracing_context({"parent": run})
+    return config_token, current_context
+
+
+@contextmanager
+def set_config_context(config: RunnableConfig) -> Generator[Context, None, None]:
+    """Set the child Runnable config + tracing context.
+
+    Args:
+        config: The config to set.
+
+    Yields:
+        The config context.
+    """
+    # Deferred to avoid importing langsmith at module level (~132ms).
+    from coco_core.tracers._context import _set_tracing_context  # noqa: PLC0415
+
+    ctx = copy_context()
+    config_token, _ = ctx.run(_set_config_context, config)
+    try:
+        yield ctx
+    finally:
+        ctx.run(var_child_runnable_config.reset, config_token)
+        ctx.run(
+            _set_tracing_context,
+            {
+                "parent": None,
+                "project_name": None,
+                "tags": None,
+                "metadata": None,
+                "enabled": None,
+                "client": None,
+            },
+        )
+
+
+def ensure_config(config: RunnableConfig | None = None) -> RunnableConfig:
+    """Ensure that a config is a dict with all keys present.
+
+    Args:
+        config: The config to ensure.
+
+    Returns:
+        The ensured config.
+    """
+    empty = RunnableConfig(
+        tags=[],
+        metadata={},
+        callbacks=None,
+        recursion_limit=DEFAULT_RECURSION_LIMIT,
+        configurable={},
+    )
+    if var_config := var_child_runnable_config.get():
+        empty.update(
+            cast(
+                "RunnableConfig",
+                {
+                    k: v.copy() if k in COPIABLE_KEYS else v  # type: ignore[attr-defined]
+                    for k, v in var_config.items()
+                    if v is not None
+                },
+            )
+        )
+    if config is not None:
+        empty.update(
+            cast(
+                "RunnableConfig",
+                {
+                    k: v.copy() if k in COPIABLE_KEYS else v  # type: ignore[attr-defined]
+                    for k, v in config.items()
+                    if v is not None and k in CONFIG_KEYS
+                },
+            )
+        )
+    if config is not None:
+        for k, v in config.items():
+            if k not in CONFIG_KEYS and v is not None:
+                empty["configurable"][k] = v
+    for configurable_key in ("model", "checkpoint_ns"):
+        if (
+            isinstance(
+                configurable_value := empty.get("configurable", {}).get(
+                    configurable_key
+                ),
+                str,
+            )
+            and configurable_key not in empty["metadata"]
+        ):
+            empty["metadata"][configurable_key] = configurable_value
+    return empty
+
+
+def get_config_list(
+    config: RunnableConfig | Sequence[RunnableConfig] | None, length: int
+) -> list[RunnableConfig]:
+    """Get a list of configs from a single config or a list of configs.
+
+     It is useful for subclasses overriding batch() or abatch().
+
+    Args:
+        config: The config or list of configs.
+        length: The length of the list.
+
+    Returns:
+        The list of configs.
+
+    Raises:
+        ValueError: If the length of the list is not equal to the length of the inputs.
+
+    """
+    if length < 0:
+        msg = f"length must be >= 0, but got {length}"
+        raise ValueError(msg)
+    if isinstance(config, Sequence) and len(config) != length:
+        msg = (
+            f"config must be a list of the same length as inputs, "
+            f"but got {len(config)} configs for {length} inputs"
+        )
+        raise ValueError(msg)
+
+    if isinstance(config, Sequence):
+        return list(map(ensure_config, config))
+    if length > 1 and isinstance(config, dict) and config.get("run_id") is not None:
+        warnings.warn(
+            "Provided run_id be used only for the first element of the batch.",
+            category=RuntimeWarning,
+            stacklevel=3,
+        )
+        subsequent = cast(
+            "RunnableConfig", {k: v for k, v in config.items() if k != "run_id"}
+        )
+        return [
+            ensure_config(subsequent) if i else ensure_config(config)
+            for i in range(length)
+        ]
+    return [ensure_config(config) for i in range(length)]
+
+
+def patch_config(
+    config: RunnableConfig | None,
+    *,
+    callbacks: BaseCallbackManager | None = None,
+    recursion_limit: int | None = None,
+    max_concurrency: int | None = None,
+    run_name: str | None = None,
+    configurable: dict[str, Any] | None = None,
+) -> RunnableConfig:
+    """Patch a config with new values.
+
+    Args:
+        config: The config to patch.
+        callbacks: The callbacks to set.
+        recursion_limit: The recursion limit to set.
+        max_concurrency: The max concurrency to set.
+        run_name: The run name to set.
+        configurable: The configurable to set.
+
+    Returns:
+        The patched config.
+    """
+    config = ensure_config(config)
+    if callbacks is not None:
+        # If we're replacing callbacks, we need to unset run_name
+        # As that should apply only to the same run as the original callbacks
+        config["callbacks"] = callbacks
+        if "run_name" in config:
+            del config["run_name"]
+        if "run_id" in config:
+            del config["run_id"]
+    if recursion_limit is not None:
+        config["recursion_limit"] = recursion_limit
+    if max_concurrency is not None:
+        config["max_concurrency"] = max_concurrency
+    if run_name is not None:
+        config["run_name"] = run_name
+    if configurable is not None:
+        config["configurable"] = {**config.get("configurable", {}), **configurable}
+    return config
+
+
+def _merge_metadata_dicts(
+    base: Mapping[str, Any], incoming: Mapping[str, Any]
+) -> dict[str, Any]:
+    """Merge metadata dicts, accumulating only `lc_versions` nested mappings.
+
+    Metadata generally uses last-writer-wins semantics. The Coco-owned
+    `lc_versions` key is the only nested mapping that accumulates across merges
+    so package versions from core, `coco`, and partner packages coexist.
+
+    Args:
+        base: The base metadata dict.
+
+            Values here are kept unless overridden by `incoming`.
+        incoming: The metadata dict to merge on top.
+
+            Its values take precedence on conflict.
+
+    Returns:
+        A new merged dict.
+    """
+    merged = {**base, **incoming}
+    base_versions = base.get("lc_versions")
+    incoming_versions = incoming.get("lc_versions")
+    if isinstance(base_versions, Mapping) and isinstance(incoming_versions, Mapping):
+        merged["lc_versions"] = {**base_versions, **incoming_versions}
+    elif isinstance(incoming_versions, Mapping):
+        merged["lc_versions"] = {**incoming_versions}
+    elif "lc_versions" not in incoming and isinstance(base_versions, Mapping):
+        merged["lc_versions"] = {**base_versions}
+    return merged
+
+
+def merge_configs(*configs: RunnableConfig | None) -> RunnableConfig:
+    """Merge multiple configs into one.
+
+    Args:
+        *configs: The configs to merge.
+
+    Returns:
+        The merged config.
+    """
+    base: RunnableConfig = {}
+    # Even though the keys aren't literals, this is correct
+    # because both dicts are the same type
+    for config in (ensure_config(c) for c in configs if c is not None):
+        for key in config:
+            if key == "metadata":
+                base["metadata"] = _merge_metadata_dicts(
+                    base.get("metadata", {}),
+                    config.get("metadata") or {},
+                )
+            elif key == "tags":
+                base["tags"] = sorted(
+                    set(base.get("tags", []) + (config.get("tags") or [])),
+                )
+            elif key == "configurable":
+                base["configurable"] = {
+                    **base.get("configurable", {}),
+                    **(config.get("configurable") or {}),
+                }
+            elif key == "callbacks":
+                base_callbacks = base.get("callbacks")
+                these_callbacks = config["callbacks"]
+                # callbacks can be either None, list[handler] or manager
+                # so merging two callbacks values has 6 cases
+                if isinstance(these_callbacks, list):
+                    if base_callbacks is None:
+                        base["callbacks"] = these_callbacks.copy()
+                    elif isinstance(base_callbacks, list):
+                        base["callbacks"] = base_callbacks + these_callbacks
+                    else:
+                        # base_callbacks is a manager
+                        mngr = base_callbacks.copy()
+                        for callback in these_callbacks:
+                            mngr.add_handler(callback, inherit=True)
+                        base["callbacks"] = mngr
+                elif these_callbacks is not None:
+                    # these_callbacks is a manager
+                    if base_callbacks is None:
+                        base["callbacks"] = these_callbacks.copy()
+                    elif isinstance(base_callbacks, list):
+                        mngr = these_callbacks.copy()
+                        for callback in base_callbacks:
+                            mngr.add_handler(callback, inherit=True)
+                        base["callbacks"] = mngr
+                    else:
+                        # base_callbacks is also a manager
+                        base["callbacks"] = base_callbacks.merge(these_callbacks)
+            elif key == "recursion_limit":
+                if config["recursion_limit"] != DEFAULT_RECURSION_LIMIT:
+                    base["recursion_limit"] = config["recursion_limit"]
+            elif key in COPIABLE_KEYS and config[key] is not None:  # type: ignore[literal-required]
+                base[key] = config[key].copy()  # type: ignore[literal-required]
+            else:
+                base[key] = config[key] or base.get(key)  # type: ignore[literal-required]
+    return base
+
+
+def call_func_with_variable_args(
+    func: Callable[[Input], Output]
+    | Callable[[Input, RunnableConfig], Output]
+    | Callable[[Input, CallbackManagerForChainRun], Output]
+    | Callable[[Input, CallbackManagerForChainRun, RunnableConfig], Output],
+    input: Input,
+    config: RunnableConfig,
+    run_manager: CallbackManagerForChainRun | None = None,
+    **kwargs: Any,
+) -> Output:
+    """Call function that may optionally accept a run_manager and/or config.
+
+    Args:
+        func: The function to call.
+        input: The input to the function.
+        config: The config to pass to the function.
+        run_manager: The run manager to pass to the function.
+        **kwargs: The keyword arguments to pass to the function.
+
+    Returns:
+        The output of the function.
+    """
+    if accepts_config(func):
+        if run_manager is not None:
+            kwargs["config"] = patch_config(config, callbacks=run_manager.get_child())
+        else:
+            kwargs["config"] = config
+    if run_manager is not None and accepts_run_manager(func):
+        kwargs["run_manager"] = run_manager
+    return func(input, **kwargs)  # type: ignore[call-arg]
+
+
+def acall_func_with_variable_args(
+    func: Callable[[Input], Awaitable[Output]]
+    | Callable[[Input, RunnableConfig], Awaitable[Output]]
+    | Callable[[Input, AsyncCallbackManagerForChainRun], Awaitable[Output]]
+    | Callable[
+        [Input, AsyncCallbackManagerForChainRun, RunnableConfig], Awaitable[Output]
+    ],
+    input: Input,
+    config: RunnableConfig,
+    run_manager: AsyncCallbackManagerForChainRun | None = None,
+    **kwargs: Any,
+) -> Awaitable[Output]:
+    """Async call function that may optionally accept a run_manager and/or config.
+
+    Args:
+        func: The function to call.
+        input: The input to the function.
+        config: The config to pass to the function.
+        run_manager: The run manager to pass to the function.
+        **kwargs: The keyword arguments to pass to the function.
+
+    Returns:
+        The output of the function.
+    """
+    if accepts_config(func):
+        if run_manager is not None:
+            kwargs["config"] = patch_config(config, callbacks=run_manager.get_child())
+        else:
+            kwargs["config"] = config
+    if run_manager is not None and accepts_run_manager(func):
+        kwargs["run_manager"] = run_manager
+    return func(input, **kwargs)  # type: ignore[call-arg]
+
+
+def get_callback_manager_for_config(config: RunnableConfig) -> CallbackManager:
+    """Get a callback manager for a config.
+
+    Args:
+        config: The config.
+
+    Returns:
+        The callback manager.
+    """
+    return CallbackManager.configure(
+        inheritable_callbacks=config.get("callbacks"),
+        inheritable_tags=config.get("tags"),
+        inheritable_metadata=config.get("metadata"),
+        langsmith_inheritable_metadata=_get_langsmith_inheritable_metadata_from_config(
+            config
+        ),
+    )
+
+
+def get_async_callback_manager_for_config(
+    config: RunnableConfig,
+) -> AsyncCallbackManager:
+    """Get an async callback manager for a config.
+
+    Args:
+        config: The config.
+
+    Returns:
+        The async callback manager.
+    """
+    return AsyncCallbackManager.configure(
+        inheritable_callbacks=config.get("callbacks"),
+        inheritable_tags=config.get("tags"),
+        inheritable_metadata=config.get("metadata"),
+        langsmith_inheritable_metadata=_get_langsmith_inheritable_metadata_from_config(
+            config
+        ),
+    )
+
+
+P = ParamSpec("P")
+T = TypeVar("T")
+
+
+class ContextThreadPoolExecutor(ThreadPoolExecutor):
+    """ThreadPoolExecutor that copies the context to the child thread."""
+
+    def submit(  # type: ignore[override]
+        self,
+        func: Callable[P, T],
+        *args: P.args,
+        **kwargs: P.kwargs,
+    ) -> Future[T]:
+        """Submit a function to the executor.
+
+        Args:
+            func: The function to submit.
+            *args: The positional arguments to the function.
+            **kwargs: The keyword arguments to the function.
+
+        Returns:
+            The future for the function.
+        """
+        return super().submit(
+            cast("Callable[..., T]", partial(copy_context().run, func, *args, **kwargs))
+        )
+
+    def map(
+        self,
+        fn: Callable[..., T],
+        *iterables: Iterable[Any],
+        **kwargs: Any,
+    ) -> Iterator[T]:
+        """Map a function to multiple iterables.
+
+        Args:
+            fn: The function to map.
+            *iterables: The iterables to map over.
+            timeout: The timeout for the map.
+            chunksize: The chunksize for the map.
+
+        Returns:
+            The iterator for the mapped function.
+        """
+        contexts = [copy_context() for _ in range(len(iterables[0]))]  # type: ignore[arg-type]
+
+        def _wrapped_fn(*args: Any) -> T:
+            return contexts.pop().run(fn, *args)
+
+        return super().map(
+            _wrapped_fn,
+            *iterables,
+            **kwargs,
+        )
+
+
+@contextmanager
+def get_executor_for_config(
+    config: RunnableConfig | None,
+) -> Generator[Executor, None, None]:
+    """Get an executor for a config.
+
+    Args:
+        config: The config.
+
+    Yields:
+        The executor.
+    """
+    config = config or {}
+    with ContextThreadPoolExecutor(
+        max_workers=config.get("max_concurrency")
+    ) as executor:
+        yield executor
+
+
+async def run_in_executor(
+    executor_or_config: Executor | RunnableConfig | None,
+    func: Callable[P, T],
+    *args: P.args,
+    **kwargs: P.kwargs,
+) -> T:
+    """Run a function in an executor.
+
+    Args:
+        executor_or_config: The executor or config to run in.
+        func: The function.
+        *args: The positional arguments to the function.
+        **kwargs: The keyword arguments to the function.
+
+    Returns:
+        The output of the function.
+    """
+
+    def wrapper() -> T:
+        try:
+            return func(*args, **kwargs)
+        except StopIteration as exc:
+            # StopIteration can't be set on an asyncio.Future
+            # it raises a TypeError and leaves the Future pending forever
+            # so we need to convert it to a RuntimeError
+            raise RuntimeError from exc
+
+    if executor_or_config is None or isinstance(executor_or_config, dict):
+        # Use default executor with context copied from current context
+        return await asyncio.get_running_loop().run_in_executor(
+            None,
+            cast("Callable[..., T]", partial(copy_context().run, wrapper)),
+        )
+
+    return await asyncio.get_running_loop().run_in_executor(executor_or_config, wrapper)