genesis-flow 1.0.0__py3-none-any.whl

mlflow/langchain/runnables.py

@@ -0,0 +1,527 @@
+from __future__ import annotations
+
+import os
+import re
+import warnings
+from pathlib import Path
+from typing import TYPE_CHECKING, Union
+
+import cloudpickle
+import yaml
+
+from mlflow.exceptions import MlflowException
+from mlflow.langchain.utils.logging import (
+    _BASE_LOAD_KEY,
+    _CONFIG_LOAD_KEY,
+    _MODEL_DATA_FOLDER_NAME,
+    _MODEL_DATA_KEY,
+    _MODEL_DATA_PKL_FILE_NAME,
+    _MODEL_DATA_YAML_FILE_NAME,
+    _MODEL_LOAD_KEY,
+    _MODEL_TYPE_KEY,
+    _PICKLE_LOAD_KEY,
+    _RUNNABLE_LOAD_KEY,
+    _load_base_lcs,
+    _load_from_json,
+    _load_from_pickle,
+    _load_from_yaml,
+    _patch_loader,
+    _save_base_lcs,
+    _validate_and_prepare_lc_model_or_path,
+    base_lc_types,
+    custom_type_to_loader_dict,
+    get_unsupported_model_message,
+    lc_runnable_assign_types,
+    lc_runnable_binding_types,
+    lc_runnable_branch_types,
+    lc_runnable_with_steps_types,
+    lc_runnables_types,
+    patch_langchain_type_to_cls_dict,
+    picklable_runnable_types,
+)
+
+if TYPE_CHECKING:
+    from langchain.schema.runnable import Runnable
+
+_STEPS_FOLDER_NAME = "steps"
+_RUNNABLE_STEPS_FILE_NAME = "steps.yaml"
+_BRANCHES_FOLDER_NAME = "branches"
+_MAPPER_FOLDER_NAME = "mapper"
+_RUNNABLE_BRANCHES_FILE_NAME = "branches.yaml"
+_DEFAULT_BRANCH_NAME = "default"
+_RUNNABLE_BINDING_CONF_FILE_NAME = "binding_conf.yaml"
+
+
+@patch_langchain_type_to_cls_dict
+def _load_model_from_config(path, model_config):
+    from langchain.chains.loading import type_to_loader_dict as chains_type_to_loader_dict
+    from langchain.llms import get_type_to_cls_dict as llms_get_type_to_cls_dict
+
+    try:
+        from langchain.prompts.loading import type_to_loader_dict as prompts_types
+    except ImportError:
+        prompts_types = {"prompt", "few_shot_prompt"}
+
+    config_path = os.path.join(path, model_config.get(_MODEL_DATA_KEY, _MODEL_DATA_YAML_FILE_NAME))
+    # Load runnables from config file
+    if config_path.endswith(".yaml"):
+        config = _load_from_yaml(config_path)
+    elif config_path.endswith(".json"):
+        config = _load_from_json(config_path)
+    else:
+        raise MlflowException(
+            f"Cannot load runnable without a config file. Got path {config_path}."
+        )
+    _type = config.get("_type")
+    if _type in chains_type_to_loader_dict:
+        from langchain.chains.loading import load_chain
+
+        return _patch_loader(load_chain)(config_path)
+    elif _type in prompts_types:
+        from langchain.prompts.loading import load_prompt
+
+        return load_prompt(config_path)
+    elif _type in llms_get_type_to_cls_dict():
+        from langchain_community.llms.loading import load_llm
+
+        return _patch_loader(load_llm)(config_path)
+    elif _type in custom_type_to_loader_dict():
+        return custom_type_to_loader_dict()[_type](config)
+    raise MlflowException(f"Unsupported type {_type} for loading.")
+
+
+def _load_model_from_path(path: str, model_config=None):
+    model_load_fn = model_config.get(_MODEL_LOAD_KEY)
+    if model_load_fn == _RUNNABLE_LOAD_KEY:
+        return _load_runnables(path, model_config)
+    if model_load_fn == _BASE_LOAD_KEY:
+        return _load_base_lcs(path, model_config)
+    if model_load_fn == _CONFIG_LOAD_KEY:
+        return _load_model_from_config(path, model_config)
+    if model_load_fn == _PICKLE_LOAD_KEY:
+        return _load_from_pickle(os.path.join(path, model_config.get(_MODEL_DATA_KEY)))
+    raise MlflowException(f"Unsupported model load key {model_load_fn}")
+
+
+def _validate_path(file_path: Union[str, Path]):
+    load_path = Path(file_path)
+    if not load_path.exists() or not load_path.is_dir():
+        raise MlflowException(
+            f"Path {load_path} must be an existing directory in order to load model."
+        )
+    return load_path
+
+
+def _load_runnable_with_steps(file_path: Union[Path, str], model_type: str):
+    """Load the model
+
+    Args:
+        file_path: Path to file to load the model from.
+        model_type: Type of the model to load.
+    """
+    from langchain.schema.runnable import RunnableParallel, RunnableSequence
+
+    load_path = _validate_path(file_path)
+
+    steps_conf_file = load_path / _RUNNABLE_STEPS_FILE_NAME
+    if not steps_conf_file.exists():
+        raise MlflowException(
+            f"File {steps_conf_file} must exist in order to load runnable with steps."
+        )
+    steps_conf = _load_from_yaml(steps_conf_file)
+    steps_path = load_path / _STEPS_FOLDER_NAME
+    _validate_path(steps_path)
+
+    steps = {}
+    # ignore hidden files
+    for step in (f for f in os.listdir(steps_path) if not f.startswith(".")):
+        config = steps_conf.get(step)
+        # load model from the folder of the step
+        runnable = _load_model_from_path(os.path.join(steps_path, step), config)
+        steps[step] = runnable
+
+    if model_type == RunnableSequence.__name__:
+        steps = [value for _, value in sorted(steps.items(), key=lambda item: int(item[0]))]
+        return runnable_sequence_from_steps(steps)
+    if model_type == RunnableParallel.__name__:
+        return RunnableParallel(steps)
+
+
+def runnable_sequence_from_steps(steps):
+    """Construct a RunnableSequence from steps.
+
+    Args:
+        steps: List of steps to construct the RunnableSequence from.
+    """
+    from langchain.schema.runnable import RunnableSequence
+
+    if len(steps) < 2:
+        raise ValueError(f"RunnableSequence must have at least 2 steps, got {len(steps)}.")
+
+    first, *middle, last = steps
+    return RunnableSequence(first=first, middle=middle, last=last)
+
+
+def _load_runnable_branch(file_path: Union[Path, str]):
+    """Load the model
+
+    Args:
+        file_path: Path to file to load the model from.
+    """
+    from langchain.schema.runnable import RunnableBranch
+
+    load_path = _validate_path(file_path)
+
+    branches_conf_file = load_path / _RUNNABLE_BRANCHES_FILE_NAME
+    if not branches_conf_file.exists():
+        raise MlflowException(
+            f"File {branches_conf_file} must exist in order to load runnable with steps."
+        )
+    branches_conf = _load_from_yaml(branches_conf_file)
+    branches_path = load_path / _BRANCHES_FOLDER_NAME
+    _validate_path(branches_path)
+
+    branches = []
+    for branch in os.listdir(branches_path):
+        # load model from the folder of the branch
+        if branch == _DEFAULT_BRANCH_NAME:
+            default_branch_path = branches_path / _DEFAULT_BRANCH_NAME
+            default = _load_model_from_path(
+                default_branch_path, branches_conf.get(_DEFAULT_BRANCH_NAME)
+            )
+        else:
+            branch_tuple = []
+            for i in range(2):
+                config = branches_conf.get(f"{branch}-{i}")
+                runnable = _load_model_from_path(
+                    os.path.join(branches_path, branch, str(i)), config
+                )
+                branch_tuple.append(runnable)
+            branches.append(tuple(branch_tuple))
+
+    # default branch must be the last branch
+    branches.append(default)
+
+    return RunnableBranch(*branches)
+
+
+def _load_runnable_assign(file_path: Union[Path, str]):
+    """Load the model
+
+    Args:
+        file_path: Path to file to load the model from.
+    """
+    from langchain.schema.runnable.passthrough import RunnableAssign
+
+    load_path = _validate_path(file_path)
+
+    mapper_file = load_path / _MAPPER_FOLDER_NAME
+    _validate_path(mapper_file)
+    mapper = _load_runnable_with_steps(mapper_file, "RunnableParallel")
+    return RunnableAssign(mapper)
+
+
+def _load_runnable_binding(file_path: Union[Path, str]):
+    """
+    Load runnable binding model from the path
+    """
+    from langchain.schema.runnable import RunnableBinding
+
+    load_path = _validate_path(file_path)
+
+    model_conf = _load_from_yaml(load_path / _RUNNABLE_BINDING_CONF_FILE_NAME)
+    for field, value in model_conf.items():
+        if _is_json_primitive(value):
+            model_conf[field] = value
+        # value is dictionary
+        else:
+            model_conf[field] = _load_model_from_path(load_path, value)
+    return RunnableBinding(**model_conf)
+
+
+def _save_internal_runnables(runnable, path, loader_fn, persist_dir):
+    conf = {}
+    if isinstance(runnable, lc_runnables_types()):
+        conf[_MODEL_TYPE_KEY] = runnable.__class__.__name__
+        conf.update(_save_runnables(runnable, path, loader_fn, persist_dir))
+    elif isinstance(runnable, base_lc_types()):
+        lc_model = _validate_and_prepare_lc_model_or_path(runnable, loader_fn)
+        conf[_MODEL_TYPE_KEY] = lc_model.__class__.__name__
+        conf.update(_save_base_lcs(lc_model, path, loader_fn, persist_dir))
+    else:
+        conf = {
+            _MODEL_TYPE_KEY: runnable.__class__.__name__,
+            _MODEL_DATA_KEY: _MODEL_DATA_YAML_FILE_NAME,
+            _MODEL_LOAD_KEY: _CONFIG_LOAD_KEY,
+        }
+        model_path = path / _MODEL_DATA_YAML_FILE_NAME
+
+        _warning_if_imported_from_lc_partner_pkg(runnable)
+
+        # Save some simple runnables that langchain natively supports.
+        if hasattr(runnable, "save"):
+            runnable.save(model_path)
+        elif hasattr(runnable, "dict"):
+            runnable_dict = runnable.dict()
+            with open(model_path, "w") as f:
+                yaml.dump(runnable_dict, f, default_flow_style=False)
+            # if the model cannot be loaded back, then `dict` is not enough for saving.
+            _load_model_from_config(path, conf)
+        else:
+            raise Exception("Cannot save runnable without `save` or `dict` methods.")
+    return conf
+
+
+_LC_PARTNER_MODULE_PATTERN = re.compile(
+    r"langchain_(?!core|community|experimental|cli|text-splitters)([a-z0-9-]+)$"
+)
+
+
+def _warning_if_imported_from_lc_partner_pkg(runnable):
+    """
+    Issues a warning if the model contains LangChain partner packages in its requirements.
+
+    Popular integrations like OpenAI have been migrated from the central langchain-community
+    package to their own partner packages (e.g. langchain-openai). However, the class loading
+    mechanism in MLflow does not handle partner packages and always loads the community version.
+    This can lead to unexpected behavior because the community version is no longer maintained.
+    """
+    module = runnable.__module__
+    root_module = module.split(".")[0]
+    if m := _LC_PARTNER_MODULE_PATTERN.match(root_module):
+        warnings.warn(
+            "Your model contains a class imported from the LangChain partner package "
+            f"`langchain-{m.group(1)}`. When loading the model back, MLflow will use the "
+            "community version of the classes instead of the partner packages, which may "
+            "lead to unexpected behavior. To ensure that the model is loaded correctly, "
+            "it is recommended to save the model with the 'model-from-code' method "
+            "instead: https://mlflow.org/docs/latest/models.html#models-from-code"
+        )
+
+
+def _save_runnable_with_steps(model, file_path: Union[Path, str], loader_fn=None, persist_dir=None):  # noqa: D417
+    """Save the model with steps. Currently it supports saving RunnableSequence and
+    RunnableParallel.
+
+    If saving a RunnableSequence, steps is a list of Runnable objects. We save each step to the
+    subfolder named by the step index.
+    e.g.  - model
+            - steps
+              - 0
+                - model.yaml
+              - 1
+                - model.pkl
+            - steps.yaml
+    If saving a RunnableParallel, steps is a dictionary of key-Runnable pairs. We save each step to
+    the subfolder named by the key.
+    e.g.  - model
+            - steps
+              - context
+                - model.yaml
+              - question
+                - model.pkl
+            - steps.yaml
+
+    We save steps.yaml file to the model folder. It contains each step's model's configuration.
+
+    Args:
+        model: Runnable to be saved.
+        file_path: Path to file to save the model to.
+    """
+    # Convert file to Path object.
+    save_path = Path(file_path)
+    save_path.mkdir(parents=True, exist_ok=True)
+
+    # Save steps into a folder
+    steps_path = save_path / _STEPS_FOLDER_NAME
+    steps_path.mkdir()
+
+    steps = get_runnable_steps(model)
+    if isinstance(steps, list):
+        generator = enumerate(steps)
+    elif isinstance(steps, dict):
+        generator = steps.items()
+    else:
+        raise MlflowException(
+            f"Runnable {model} steps attribute must be either a list or a dictionary. "
+            f"Got {type(steps).__name__}."
+        )
+    unsaved_runnables = {}
+    steps_conf = {}
+    for key, runnable in generator:
+        step = str(key)
+        # Save each step into a subfolder named by step
+        save_runnable_path = steps_path / step
+        save_runnable_path.mkdir()
+        try:
+            steps_conf[step] = _save_internal_runnables(
+                runnable, save_runnable_path, loader_fn, persist_dir
+            )
+        except Exception as e:
+            unsaved_runnables[step] = f"{runnable.get_name()} -- {e}"
+
+    if unsaved_runnables:
+        raise MlflowException(f"Failed to save runnable sequence: {unsaved_runnables}.")
+
+    # save steps configs
+    with save_path.joinpath(_RUNNABLE_STEPS_FILE_NAME).open("w") as f:
+        yaml.dump(steps_conf, f, default_flow_style=False)
+
+
+def _save_runnable_branch(model, file_path, loader_fn, persist_dir):
+    """
+    Save runnable branch in to path.
+    """
+    save_path = Path(file_path)
+    save_path.mkdir(parents=True, exist_ok=True)
+    # save branches into a folder
+    branches_path = save_path / _BRANCHES_FOLDER_NAME
+    branches_path.mkdir()
+
+    unsaved_runnables = {}
+    branches_conf = {}
+    for index, branch_tuple in enumerate(model.branches):
+        # Save each branch into a subfolder named by index
+        # and save condition and runnable into subfolder
+        for i, runnable in enumerate(branch_tuple):
+            save_runnable_path = branches_path / str(index) / str(i)
+            save_runnable_path.mkdir(parents=True)
+            branches_conf[f"{index}-{i}"] = {}
+
+            try:
+                branches_conf[f"{index}-{i}"] = _save_internal_runnables(
+                    runnable, save_runnable_path, loader_fn, persist_dir
+                )
+            except Exception as e:
+                unsaved_runnables[f"{index}-{i}"] = f"{runnable.get_name()} -- {e}"
+
+    # save default branch
+    default_branch_path = branches_path / _DEFAULT_BRANCH_NAME
+    default_branch_path.mkdir()
+    try:
+        branches_conf[_DEFAULT_BRANCH_NAME] = _save_internal_runnables(
+            model.default, default_branch_path, loader_fn, persist_dir
+        )
+    except Exception as e:
+        unsaved_runnables[_DEFAULT_BRANCH_NAME] = f"{model.default.get_name()} -- {e}"
+    if unsaved_runnables:
+        raise MlflowException(f"Failed to save runnable branch: {unsaved_runnables}.")
+
+    # save branches configs
+    with save_path.joinpath(_RUNNABLE_BRANCHES_FILE_NAME).open("w") as f:
+        yaml.dump(branches_conf, f, default_flow_style=False)
+
+
+def _save_runnable_assign(model, file_path, loader_fn=None, persist_dir=None):
+    from langchain.schema.runnable import RunnableParallel
+
+    save_path = Path(file_path)
+    save_path.mkdir(parents=True, exist_ok=True)
+    # save mapper into a folder
+    mapper_path = save_path / _MAPPER_FOLDER_NAME
+    mapper_path.mkdir()
+
+    if not isinstance(model.mapper, RunnableParallel):
+        raise MlflowException(
+            f"Failed to save model {model} with type {model.__class__.__name__}. "
+            "RunnableAssign's mapper must be a RunnableParallel."
+        )
+    _save_runnable_with_steps(model.mapper, mapper_path, loader_fn, persist_dir)
+
+
+def _is_json_primitive(value):
+    return (
+        value is None
+        or isinstance(value, (str, int, float, bool))
+        or (isinstance(value, list) and all(_is_json_primitive(v) for v in value))
+    )
+
+
+def _save_runnable_binding(model, file_path, loader_fn=None, persist_dir=None):
+    save_path = Path(file_path)
+    save_path.mkdir(parents=True, exist_ok=True)
+    model_config = {}
+
+    # runnableBinding bound is the real runnable to be invoked
+    model_config["bound"] = _save_internal_runnables(model.bound, save_path, loader_fn, persist_dir)
+
+    # save other fields
+    for field, value in model.dict().items():
+        if _is_json_primitive(value):
+            model_config[field] = value
+        elif field != "bound":
+            model_config[field] = {
+                _MODEL_LOAD_KEY: _PICKLE_LOAD_KEY,
+                _MODEL_DATA_KEY: f"{field}.pkl",
+            }
+            _pickle_object(value, os.path.join(save_path, f"{field}.pkl"))
+
+    # save fields configs
+    with save_path.joinpath(_RUNNABLE_BINDING_CONF_FILE_NAME).open("w") as f:
+        yaml.dump(model_config, f, default_flow_style=False)
+
+
+def _pickle_object(model, path: str):
+    if not path.endswith(".pkl"):
+        raise ValueError(f"File path must end with .pkl, got {path}.")
+    with open(path, "wb") as f:
+        cloudpickle.dump(model, f)
+
+
+def _save_runnables(model, path, loader_fn=None, persist_dir=None):
+    model_data_kwargs = {
+        _MODEL_LOAD_KEY: _RUNNABLE_LOAD_KEY,
+        _MODEL_TYPE_KEY: model.__class__.__name__,
+    }
+    if isinstance(model, lc_runnable_with_steps_types()):
+        model_data_path = _MODEL_DATA_FOLDER_NAME
+        _save_runnable_with_steps(
+            model, os.path.join(path, model_data_path), loader_fn, persist_dir
+        )
+    elif isinstance(model, picklable_runnable_types()):
+        model_data_path = _MODEL_DATA_PKL_FILE_NAME
+        _pickle_object(model, os.path.join(path, model_data_path))
+    elif isinstance(model, lc_runnable_branch_types()):
+        model_data_path = _MODEL_DATA_FOLDER_NAME
+        _save_runnable_branch(model, os.path.join(path, model_data_path), loader_fn, persist_dir)
+    elif isinstance(model, lc_runnable_assign_types()):
+        model_data_path = _MODEL_DATA_FOLDER_NAME
+        _save_runnable_assign(model, os.path.join(path, model_data_path), loader_fn, persist_dir)
+    elif isinstance(model, lc_runnable_binding_types()):
+        model_data_path = _MODEL_DATA_FOLDER_NAME
+        _save_runnable_binding(model, os.path.join(path, model_data_path), loader_fn, persist_dir)
+    else:
+        raise MlflowException.invalid_parameter_value(
+            get_unsupported_model_message(type(model).__name__)
+        )
+    model_data_kwargs[_MODEL_DATA_KEY] = model_data_path
+    return model_data_kwargs
+
+
+def _load_runnables(path, conf):
+    model_type = conf.get(_MODEL_TYPE_KEY)
+    model_data = conf.get(_MODEL_DATA_KEY, _MODEL_DATA_YAML_FILE_NAME)
+    if model_type in (x.__name__ for x in lc_runnable_with_steps_types()):
+        return _load_runnable_with_steps(os.path.join(path, model_data), model_type)
+    if (
+        model_type in (x.__name__ for x in picklable_runnable_types())
+        or model_data == _MODEL_DATA_PKL_FILE_NAME
+    ):
+        return _load_from_pickle(os.path.join(path, model_data))
+    if model_type in (x.__name__ for x in lc_runnable_branch_types()):
+        return _load_runnable_branch(os.path.join(path, model_data))
+    if model_type in (x.__name__ for x in lc_runnable_assign_types()):
+        return _load_runnable_assign(os.path.join(path, model_data))
+    if model_type in (x.__name__ for x in lc_runnable_binding_types()):
+        return _load_runnable_binding(os.path.join(path, model_data))
+    raise MlflowException.invalid_parameter_value(get_unsupported_model_message(model_type))
+
+
+def get_runnable_steps(model: Runnable):
+    try:
+        return model.steps
+    except AttributeError:
+        # RunnableParallel stores steps as `steps__` attribute since version 0.16.0, while it was
+        # stored as `steps` attribute before that and other runnables like RunnableSequence still
+        # has `steps` property.
+        return model.steps__