codekingpro/portable-devtools

1from __future__ import annotations2 3from collections import ChainMap4from collections.abc import Sequence5from os import getenv6from typing import Any, cast7 8from langchain_core.callbacks import (9    AsyncCallbackManager,10    BaseCallbackManager,11    CallbackManager,12    Callbacks,13)14from langchain_core.runnables import RunnableConfig15from langchain_core.runnables.config import (16    CONFIG_KEYS,17    COPIABLE_KEYS,18    var_child_runnable_config,19)20from langgraph.checkpoint.base import CheckpointMetadata21 22from langgraph._internal._constants import (23    CONF,24    CONFIG_KEY_CHECKPOINT_ID,25    CONFIG_KEY_CHECKPOINT_MAP,26    CONFIG_KEY_CHECKPOINT_NS,27    NS_END,28    NS_SEP,29)30 31DEFAULT_RECURSION_LIMIT = int(getenv("LANGGRAPH_DEFAULT_RECURSION_LIMIT", "10007"))32DELTA_MAX_SUPERSTEPS_SINCE_SNAPSHOT = int(33    getenv("LANGGRAPH_DELTA_MAX_SUPERSTEPS_SINCE_SNAPSHOT", "5000")34)35 36 37def recast_checkpoint_ns(ns: str) -> str:38    """Remove task IDs from checkpoint namespace.39 40    Args:41        ns: The checkpoint namespace with task IDs.42 43    Returns:44        str: The checkpoint namespace without task IDs.45    """46    return NS_SEP.join(47        part.split(NS_END)[0] for part in ns.split(NS_SEP) if not part.isdigit()48    )49 50 51def patch_configurable(52    config: RunnableConfig | None, patch: dict[str, Any]53) -> RunnableConfig:54    if config is None:55        return {CONF: patch}56    elif CONF not in config:57        return {**config, CONF: patch}58    else:59        return {**config, CONF: {**config[CONF], **patch}}60 61 62def patch_checkpoint_map(63    config: RunnableConfig | None, metadata: CheckpointMetadata | None64) -> RunnableConfig:65    if config is None:66        return config67    elif parents := (metadata.get("parents") if metadata else None):68        conf = config[CONF]69        return patch_configurable(70            config,71            {72                CONFIG_KEY_CHECKPOINT_MAP: {73                    **parents,74                    conf[CONFIG_KEY_CHECKPOINT_NS]: conf[CONFIG_KEY_CHECKPOINT_ID],75                },76            },77        )78    else:79        return config80 81 82def merge_configs(*configs: RunnableConfig | None) -> RunnableConfig:83    """Merge multiple configs into one.84 85    Args:86        *configs: The configs to merge.87 88    Returns:89        RunnableConfig: The merged config.90    """91    base: RunnableConfig = {}92    # Even though the keys aren't literals, this is correct93    # because both dicts are the same type94    for config in configs:95        if config is None:96            continue97        for key, value in config.items():98            if not value:99                continue100            if key == "metadata":101                if base_value := base.get(key):102                    base[key] = {**base_value, **value}  # type: ignore103                else:104                    base[key] = value  # type: ignore[literal-required]105            elif key == "tags":106                if base_value := base.get(key):107                    base[key] = [*base_value, *value]  # type: ignore108                else:109                    base[key] = value  # type: ignore[literal-required]110            elif key == CONF:111                if base_value := base.get(key):112                    base[key] = {**base_value, **value}  # type: ignore[dict-item]113                else:114                    base[key] = value115            elif key == "callbacks":116                base_callbacks = base.get("callbacks")117                # callbacks can be either None, list[handler] or manager118                # so merging two callbacks values has 6 cases119                if isinstance(value, list):120                    if base_callbacks is None:121                        base["callbacks"] = value.copy()122                    elif isinstance(base_callbacks, list):123                        base["callbacks"] = base_callbacks + value124                    else:125                        # base_callbacks is a manager126                        mngr = base_callbacks.copy()127                        for callback in value:128                            mngr.add_handler(callback, inherit=True)129                        base["callbacks"] = mngr130                elif isinstance(value, BaseCallbackManager):131                    # value is a manager132                    if base_callbacks is None:133                        base["callbacks"] = value.copy()134                    elif isinstance(base_callbacks, list):135                        mngr = value.copy()136                        for callback in base_callbacks:137                            mngr.add_handler(callback, inherit=True)138                        base["callbacks"] = mngr139                    else:140                        # base_callbacks is also a manager141                        base["callbacks"] = base_callbacks.merge(value)142                else:143                    raise NotImplementedError144            elif key == "recursion_limit":145                if config["recursion_limit"] != DEFAULT_RECURSION_LIMIT:146                    base["recursion_limit"] = config["recursion_limit"]147            else:148                base[key] = config[key]  # type: ignore[literal-required]149    if CONF not in base:150        base[CONF] = {}151    return base152 153 154def patch_config(155    config: RunnableConfig | None,156    *,157    callbacks: Callbacks = None,158    recursion_limit: int | None = None,159    max_concurrency: int | None = None,160    run_name: str | None = None,161    configurable: dict[str, Any] | None = None,162) -> RunnableConfig:163    """Patch a config with new values.164 165    Args:166        config: The config to patch.167        callbacks: The callbacks to set.168        recursion_limit: The recursion limit to set.169        max_concurrency: The max number of concurrent steps to run, which also applies to parallelized steps.170        run_name: The run name to set.171        configurable: The configurable to set.172 173    Returns:174        RunnableConfig: The patched config.175    """176    config = config.copy() if config is not None else {}177    if callbacks is not None:178        # If we're replacing callbacks, we need to unset run_name179        # As that should apply only to the same run as the original callbacks180        config["callbacks"] = callbacks181        if "run_name" in config:182            del config["run_name"]183        if "run_id" in config:184            del config["run_id"]185    if recursion_limit is not None:186        config["recursion_limit"] = recursion_limit187    if max_concurrency is not None:188        config["max_concurrency"] = max_concurrency189    if run_name is not None:190        config["run_name"] = run_name191    if configurable is not None:192        config[CONF] = {**config.get(CONF, {}), **configurable}193    return config194 195 196def get_callback_manager_for_config(197    config: RunnableConfig, tags: Sequence[str] | None = None198) -> CallbackManager:199    """Get a callback manager for a config.200 201    Args:202        config: The config.203 204    Returns:205        CallbackManager: The callback manager.206    """207    from langchain_core.callbacks.manager import CallbackManager208 209    # merge tags210    all_tags = config.get("tags")211    if all_tags is not None and tags is not None:212        all_tags = [*all_tags, *tags]213    elif tags is not None:214        all_tags = list(tags)215    # use existing callbacks if they exist216    if (callbacks := config.get("callbacks")) and isinstance(217        callbacks, CallbackManager218    ):219        if all_tags:220            callbacks.add_tags(all_tags)221        if metadata := config.get("metadata"):222            callbacks.add_metadata(metadata)223        manager = callbacks224    else:225        # otherwise create a new manager226        manager = CallbackManager.configure(227            inheritable_callbacks=config.get("callbacks"),228            inheritable_tags=all_tags,229            inheritable_metadata=config.get("metadata"),230            langsmith_inheritable_metadata=_get_tracing_metadata_defaults(config),231        )232    return manager233 234 235def get_async_callback_manager_for_config(236    config: RunnableConfig,237    tags: Sequence[str] | None = None,238) -> AsyncCallbackManager:239    """Get an async callback manager for a config.240 241    Args:242        config: The config.243 244    Returns:245        AsyncCallbackManager: The async callback manager.246    """247    from langchain_core.callbacks.manager import AsyncCallbackManager248 249    # merge tags250    all_tags = config.get("tags")251    if all_tags is not None and tags is not None:252        all_tags = [*all_tags, *tags]253    elif tags is not None:254        all_tags = list(tags)255    # use existing callbacks if they exist256    if (callbacks := config.get("callbacks")) and isinstance(257        callbacks, AsyncCallbackManager258    ):259        if all_tags:260            callbacks.add_tags(all_tags)261        if metadata := config.get("metadata"):262            callbacks.add_metadata(metadata)263        manager = callbacks264    else:265        # otherwise create a new manager266        manager = AsyncCallbackManager.configure(267            inheritable_callbacks=config.get("callbacks"),268            inheritable_tags=all_tags,269            inheritable_metadata=config.get("metadata"),270            langsmith_inheritable_metadata=_get_tracing_metadata_defaults(config),271        )272    return manager273 274 275def _is_not_empty(value: Any) -> bool:276    if isinstance(value, (list, tuple, dict)):277        return len(value) > 0278    else:279        return value is not None280 281 282def ensure_config(*configs: RunnableConfig | None) -> RunnableConfig:283    """Return a config with all keys, merging any provided configs.284 285    Args:286        *configs: Configs to merge before ensuring defaults.287 288    Returns:289        RunnableConfig: The merged and ensured config.290    """291    empty = RunnableConfig(292        tags=[],293        metadata=ChainMap(),294        callbacks=None,295        recursion_limit=DEFAULT_RECURSION_LIMIT,296        configurable={},297    )298    if var_config := var_child_runnable_config.get():299        empty.update(300            {301                k: v.copy() if k in COPIABLE_KEYS else v  # type: ignore[attr-defined]302                for k, v in var_config.items()303                if _is_not_empty(v)304            },305        )306    for config in configs:307        if config is None:308            continue309        for k, v in config.items():310            if _is_not_empty(v) and k in CONFIG_KEYS:311                if k == CONF:312                    empty[k] = cast(dict, v).copy()313                else:314                    empty[k] = v  # type: ignore[literal-required]315        for k, v in config.items():316            if _is_not_empty(v) and k not in CONFIG_KEYS:317                empty[CONF][k] = v318 319    configurable = empty.get("configurable")320    metadata = empty.get("metadata")321    if configurable and metadata is not None:322        for key in _PROPAGATE_TO_METADATA:323            if key in metadata:324                continue325            value = configurable.get(key)326            if value:327                metadata[key] = value328    return empty329 330 331_OMIT = ("key", "token", "secret", "password", "auth")332 333 334def _exclude_as_metadata(key: str, value: Any) -> bool:335    key_lower = key.casefold()336    return (337        key.startswith("__")338        or not isinstance(value, (str, int, float, bool))339        or any(substr in key_lower for substr in _OMIT)340    )341 342 343def _get_tracing_metadata_defaults(344    config: RunnableConfig,345) -> dict[str, Any] | None:346    """Get tracer-only metadata defaults from configurable values."""347    configurable = config.get("configurable")348    if not configurable:349        return None350    metadata: dict[str, Any] = {}351    for key, value in configurable.items():352        if _exclude_as_metadata(key, value):353            continue354        metadata[key] = value355    return metadata or None356 357 358_PROPAGATE_TO_METADATA = frozenset(359    (360        "thread_id",361        "checkpoint_id",362        "checkpoint_ns",363        "task_id",364        "run_id",365        "assistant_id",366        "graph_id",367    )368)369