import os
from pathlib import Path
from typing import Union

import cloudpickle
import yaml

from mlflow.exceptions import MlflowException
from mlflow.langchain.utils 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,
    _RUNNABLE_LOAD_KEY,
    _UNSUPPORTED_MODEL_ERROR_MESSAGE,
    _load_base_lcs,
    _load_from_json,
    _load_from_pickle,
    _load_from_yaml,
    _save_base_lcs,
    _validate_and_wrap_lc_model,
    base_lc_types,
    custom_type_to_loader_dict,
    lc_runnable_assign_types,
    lc_runnable_branch_types,
    lc_runnable_with_steps_types,
    lc_runnables_types,
    picklable_runnable_types,
)

_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"


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 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.llms.loading import load_llm

        return 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)
    raise MlflowException(f"Unsupported model load key {model_load_fn}")


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

    # Convert file to Path object.
    load_path = Path(file_path)
    if not load_path.exists() or not load_path.is_dir():
        raise MlflowException(
            f"File {load_path} must exist and must be a directory "
            "in order to load runnable with steps."
        )

    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
    if not steps_path.exists() or not steps_path.is_dir():
        raise MlflowException(
            f"Folder {steps_path} must exist and must be a directory "
            "in order to load runnable with steps."
        )

    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

    # Convert file to Path object.
    load_path = Path(file_path)
    if not load_path.exists() or not load_path.is_dir():
        raise MlflowException(
            f"File {load_path} must exist and must be a directory "
            "in order to load runnable with steps."
        )

    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
    if not branches_path.exists() or not branches_path.is_dir():
        raise MlflowException(
            f"Folder {branches_path} must exist and must be a directory "
            "in order to load runnable with steps."
        )

    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

    # Convert file to Path object.
    load_path = Path(file_path)
    if not load_path.exists() or not load_path.is_dir():
        raise MlflowException(
            f"File {load_path} must exist and must be a directory in order to load runnable."
        )

    mapper_file = load_path / _MAPPER_FOLDER_NAME
    if not mapper_file.exists() or not mapper_file.is_dir():
        raise MlflowException(
            f"Folder {mapper_file} must exist and must be a directory "
            "in order to load runnable assign with mapper."
        )
    mapper = _load_runnable_with_steps(mapper_file, "RunnableParallel")
    return RunnableAssign(mapper)


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_wrap_lc_model(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,
        }
        path = path / _MODEL_DATA_YAML_FILE_NAME
        # Save some simple runnables that langchain natively supports.
        if hasattr(runnable, "save"):
            runnable.save(path)
        # TODO: check if `dict` is enough to load it back
        elif hasattr(runnable, "dict"):
            runnable_dict = runnable.dict()
            with open(path, "w") as f:
                yaml.dump(runnable_dict, f, default_flow_style=False)
        else:
            return
    return conf


def _save_runnable_with_steps(model, file_path: Union[Path, str], loader_fn=None, persist_dir=None):
    """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 = model.steps
    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()
        if result := _save_internal_runnables(runnable, save_runnable_path, loader_fn, persist_dir):
            steps_conf[step] = result
        else:
            unsaved_runnables[step] = str(runnable)

    if unsaved_runnables:
        raise MlflowException(
            f"Failed to save runnable sequence: {unsaved_runnables}. "
            "Runnable must have either `save` or `dict` method."
        )

    # 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}"] = {}

            if result := _save_internal_runnables(
                runnable, save_runnable_path, loader_fn, persist_dir
            ):
                branches_conf[f"{index}-{i}"] = result
            else:
                unsaved_runnables[f"{index}-{i}"] = str(runnable)

    # save default branch
    default_branch_path = branches_path / _DEFAULT_BRANCH_NAME
    default_branch_path.mkdir()
    if result := _save_internal_runnables(
        model.default, default_branch_path, loader_fn, persist_dir
    ):
        branches_conf[_DEFAULT_BRANCH_NAME] = result
    else:
        unsaved_runnables[_DEFAULT_BRANCH_NAME] = str(model.default)

    if unsaved_runnables:
        raise MlflowException(
            f"Failed to save runnable branch: {unsaved_runnables}. "
            "Runnable must have either `save` or `dict` method."
        )

    # 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 _save_picklable_runnable(model, path):
    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}
    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
        _save_picklable_runnable(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)
    else:
        raise MlflowException.invalid_parameter_value(
            _UNSUPPORTED_MODEL_ERROR_MESSAGE.format(instance_type=type(model).__name__)
        )
    model_data_kwargs.update({_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))
    raise MlflowException.invalid_parameter_value(
        _UNSUPPORTED_MODEL_ERROR_MESSAGE.format(instance_type=model_type)
    )

langchain.chains.loading.load_chain langchain.prompts.loading.load_prompt langchain.schema.runnable.RunnableSequence langchain.schema.runnable.RunnableParallel langchain.schema.runnable.passthrough.RunnableAssign langchain.schema.runnable.RunnableBranch langchain.llms.get_type_to_cls_dict langchain.llms.loading.load_llm
[((2386, 2443), 'mlflow.exceptions.MlflowException', 'MlflowException', (['f"""Unsupported type {_type} for loading."""'], {}), "(f'Unsupported type {_type} for loading.')\n", (2401, 2443), False, 'from mlflow.exceptions import MlflowException\n'), ((2853, 2915), 'mlflow.exceptions.MlflowException', 'MlflowException', (['f"""Unsupported model load key {model_load_fn}"""'], {}), "(f'Unsupported model load key {model_load_fn}')\n", (2868, 2915), False, 'from mlflow.exceptions import MlflowException\n'), ((3268, 3283), 'pathlib.Path', 'Path', (['file_path'], {}), '(file_path)\n', (3272, 3283), False, 'from pathlib import Path\n'), ((3745, 3777), 'mlflow.langchain.utils._load_from_yaml', '_load_from_yaml', (['steps_conf_file'], {}), '(steps_conf_file)\n', (3760, 3777), False, 'from mlflow.langchain.utils 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, _RUNNABLE_LOAD_KEY, _UNSUPPORTED_MODEL_ERROR_MESSAGE, _load_base_lcs, _load_from_json, _load_from_pickle, _load_from_yaml, _save_base_lcs, _validate_and_wrap_lc_model, base_lc_types, custom_type_to_loader_dict, lc_runnable_assign_types, lc_runnable_branch_types, lc_runnable_with_steps_types, lc_runnables_types, picklable_runnable_types\n'), ((5047, 5102), 'langchain.schema.runnable.RunnableSequence', 'RunnableSequence', ([], {'first': 'first', 'middle': 'middle', 'last': 'last'}), '(first=first, middle=middle, last=last)\n', (5063, 5102), False, 'from langchain.schema.runnable import RunnableSequence\n'), ((5367, 5382), 'pathlib.Path', 'Path', (['file_path'], {}), '(file_path)\n', (5371, 5382), False, 'from pathlib import Path\n'), ((5859, 5894), 'mlflow.langchain.utils._load_from_yaml', '_load_from_yaml', (['branches_conf_file'], {}), '(branches_conf_file)\n', (5874, 5894), False, 'from mlflow.langchain.utils 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, _RUNNABLE_LOAD_KEY, _UNSUPPORTED_MODEL_ERROR_MESSAGE, _load_base_lcs, _load_from_json, _load_from_pickle, _load_from_yaml, _save_base_lcs, _validate_and_wrap_lc_model, base_lc_types, custom_type_to_loader_dict, lc_runnable_assign_types, lc_runnable_branch_types, lc_runnable_with_steps_types, lc_runnables_types, picklable_runnable_types\n'), ((6218, 6243), 'os.listdir', 'os.listdir', (['branches_path'], {}), '(branches_path)\n', (6228, 6243), False, 'import os\n'), ((7003, 7028), 'langchain.schema.runnable.RunnableBranch', 'RunnableBranch', (['*branches'], {}), '(*branches)\n', (7017, 7028), False, 'from langchain.schema.runnable import RunnableBranch\n'), ((7305, 7320), 'pathlib.Path', 'Path', (['file_path'], {}), '(file_path)\n', (7309, 7320), False, 'from pathlib import Path\n'), ((7881, 7903), 'langchain.schema.runnable.passthrough.RunnableAssign', 'RunnableAssign', (['mapper'], {}), '(mapper)\n', (7895, 7903), False, 'from langchain.schema.runnable.passthrough import RunnableAssign\n'), ((10183, 10198), 'pathlib.Path', 'Path', (['file_path'], {}), '(file_path)\n', (10187, 10198), False, 'from pathlib import Path\n'), ((11645, 11660), 'pathlib.Path', 'Path', (['file_path'], {}), '(file_path)\n', (11649, 11660), False, 'from pathlib import Path\n'), ((13438, 13453), 'pathlib.Path', 'Path', (['file_path'], {}), '(file_path)\n', (13442, 13453), False, 'from pathlib import Path\n'), ((1579, 1607), 'mlflow.langchain.utils._load_from_yaml', '_load_from_yaml', (['config_path'], {}), '(config_path)\n', (1594, 1607), False, 'from mlflow.langchain.utils 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, _RUNNABLE_LOAD_KEY, _UNSUPPORTED_MODEL_ERROR_MESSAGE, _load_base_lcs, _load_from_json, _load_from_pickle, _load_from_yaml, _save_base_lcs, _validate_and_wrap_lc_model, base_lc_types, custom_type_to_loader_dict, lc_runnable_assign_types, lc_runnable_branch_types, lc_runnable_with_steps_types, lc_runnables_types, picklable_runnable_types\n'), ((1976, 1999), 'langchain.chains.loading.load_chain', 'load_chain', (['config_path'], {}), '(config_path)\n', (1986, 1999), False, 'from langchain.chains.loading import load_chain\n'), ((2707, 2741), 'mlflow.langchain.utils._load_base_lcs', '_load_base_lcs', (['path', 'model_config'], {}), '(path, model_config)\n', (2721, 2741), False, 'from mlflow.langchain.utils 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, _RUNNABLE_LOAD_KEY, _UNSUPPORTED_MODEL_ERROR_MESSAGE, _load_base_lcs, _load_from_json, _load_from_pickle, _load_from_yaml, _save_base_lcs, _validate_and_wrap_lc_model, base_lc_types, custom_type_to_loader_dict, lc_runnable_assign_types, lc_runnable_branch_types, lc_runnable_with_steps_types, lc_runnables_types, picklable_runnable_types\n'), ((3355, 3474), 'mlflow.exceptions.MlflowException', 'MlflowException', (['f"""File {load_path} must exist and must be a directory in order to load runnable with steps."""'], {}), "(\n    f'File {load_path} must exist and must be a directory in order to load runnable with steps.'\n    )\n", (3370, 3474), False, 'from mlflow.exceptions import MlflowException\n'), ((3614, 3710), 'mlflow.exceptions.MlflowException', 'MlflowException', (['f"""File {steps_conf_file} must exist in order to load runnable with steps."""'], {}), "(\n    f'File {steps_conf_file} must exist in order to load runnable with steps.')\n", (3629, 3710), False, 'from mlflow.exceptions import MlflowException\n'), ((3899, 4021), 'mlflow.exceptions.MlflowException', 'MlflowException', (['f"""Folder {steps_path} must exist and must be a directory in order to load runnable with steps."""'], {}), "(\n    f'Folder {steps_path} must exist and must be a directory in order to load runnable with steps.'\n    )\n", (3914, 4021), False, 'from mlflow.exceptions import MlflowException\n'), ((4624, 4647), 'langchain.schema.runnable.RunnableParallel', 'RunnableParallel', (['steps'], {}), '(steps)\n', (4640, 4647), False, 'from langchain.schema.runnable import RunnableParallel\n'), ((5454, 5573), 'mlflow.exceptions.MlflowException', 'MlflowException', (['f"""File {load_path} must exist and must be a directory in order to load runnable with steps."""'], {}), "(\n    f'File {load_path} must exist and must be a directory in order to load runnable with steps.'\n    )\n", (5469, 5573), False, 'from mlflow.exceptions import MlflowException\n'), ((5722, 5826), 'mlflow.exceptions.MlflowException', 'MlflowException', (['f"""File {branches_conf_file} must exist in order to load runnable with steps."""'], {}), "(\n    f'File {branches_conf_file} must exist in order to load runnable with steps.'\n    )\n", (5737, 5826), False, 'from mlflow.exceptions import MlflowException\n'), ((6028, 6153), 'mlflow.exceptions.MlflowException', 'MlflowException', (['f"""Folder {branches_path} must exist and must be a directory in order to load runnable with steps."""'], {}), "(\n    f'Folder {branches_path} must exist and must be a directory in order to load runnable with steps.'\n    )\n", (6043, 6153), False, 'from mlflow.exceptions import MlflowException\n'), ((7392, 7500), 'mlflow.exceptions.MlflowException', 'MlflowException', (['f"""File {load_path} must exist and must be a directory in order to load runnable."""'], {}), "(\n    f'File {load_path} must exist and must be a directory in order to load runnable.'\n    )\n", (7407, 7500), False, 'from mlflow.exceptions import MlflowException\n'), ((7639, 7770), 'mlflow.exceptions.MlflowException', 'MlflowException', (['f"""Folder {mapper_file} must exist and must be a directory in order to load runnable assign with mapper."""'], {}), "(\n    f'Folder {mapper_file} must exist and must be a directory in order to load runnable assign with mapper.'\n    )\n", (7654, 7770), False, 'from mlflow.exceptions import MlflowException\n'), ((8018, 8038), 'mlflow.langchain.utils.lc_runnables_types', 'lc_runnables_types', ([], {}), '()\n', (8036, 8038), False, 'from mlflow.langchain.utils 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, _RUNNABLE_LOAD_KEY, _UNSUPPORTED_MODEL_ERROR_MESSAGE, _load_base_lcs, _load_from_json, _load_from_pickle, _load_from_yaml, _save_base_lcs, _validate_and_wrap_lc_model, base_lc_types, custom_type_to_loader_dict, lc_runnable_assign_types, lc_runnable_branch_types, lc_runnable_with_steps_types, lc_runnables_types, picklable_runnable_types\n'), ((11186, 11321), 'mlflow.exceptions.MlflowException', 'MlflowException', (['f"""Failed to save runnable sequence: {unsaved_runnables}. Runnable must have either `save` or `dict` method."""'], {}), "(\n    f'Failed to save runnable sequence: {unsaved_runnables}. Runnable must have either `save` or `dict` method.'\n    )\n", (11201, 11321), False, 'from mlflow.exceptions import MlflowException\n'), ((11454, 11504), 'yaml.dump', 'yaml.dump', (['steps_conf', 'f'], {'default_flow_style': '(False)'}), '(steps_conf, f, default_flow_style=False)\n', (11463, 11504), False, 'import yaml\n'), ((12955, 13088), 'mlflow.exceptions.MlflowException', 'MlflowException', (['f"""Failed to save runnable branch: {unsaved_runnables}. Runnable must have either `save` or `dict` method."""'], {}), "(\n    f'Failed to save runnable branch: {unsaved_runnables}. Runnable must have either `save` or `dict` method.'\n    )\n", (12970, 13088), False, 'from mlflow.exceptions import MlflowException\n'), ((13227, 13280), 'yaml.dump', 'yaml.dump', (['branches_conf', 'f'], {'default_flow_style': '(False)'}), '(branches_conf, f, default_flow_style=False)\n', (13236, 13280), False, 'import yaml\n'), ((13679, 13827), 'mlflow.exceptions.MlflowException', 'MlflowException', (['f"""Failed to save model {model} with type {model.__class__.__name__}. RunnableAssign\'s mapper must be a RunnableParallel."""'], {}), '(\n    f"Failed to save model {model} with type {model.__class__.__name__}. RunnableAssign\'s mapper must be a RunnableParallel."\n    )\n', (13694, 13827), False, 'from mlflow.exceptions import MlflowException\n'), ((14126, 14152), 'cloudpickle.dump', 'cloudpickle.dump', (['model', 'f'], {}), '(model, f)\n', (14142, 14152), False, 'import cloudpickle\n'), ((14310, 14340), 'mlflow.langchain.utils.lc_runnable_with_steps_types', 'lc_runnable_with_steps_types', ([], {}), '()\n', (14338, 14340), False, 'from mlflow.langchain.utils 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, _RUNNABLE_LOAD_KEY, _UNSUPPORTED_MODEL_ERROR_MESSAGE, _load_base_lcs, _load_from_json, _load_from_pickle, _load_from_yaml, _save_base_lcs, _validate_and_wrap_lc_model, base_lc_types, custom_type_to_loader_dict, lc_runnable_assign_types, lc_runnable_branch_types, lc_runnable_with_steps_types, lc_runnables_types, picklable_runnable_types\n'), ((16218, 16283), 'mlflow.langchain.utils._UNSUPPORTED_MODEL_ERROR_MESSAGE.format', '_UNSUPPORTED_MODEL_ERROR_MESSAGE.format', ([], {'instance_type': 'model_type'}), '(instance_type=model_type)\n', (16257, 16283), False, 'from mlflow.langchain.utils 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, _RUNNABLE_LOAD_KEY, _UNSUPPORTED_MODEL_ERROR_MESSAGE, _load_base_lcs, _load_from_json, _load_from_pickle, _load_from_yaml, _save_base_lcs, _validate_and_wrap_lc_model, base_lc_types, custom_type_to_loader_dict, lc_runnable_assign_types, lc_runnable_branch_types, lc_runnable_with_steps_types, lc_runnables_types, picklable_runnable_types\n'), ((1665, 1693), 'mlflow.langchain.utils._load_from_json', '_load_from_json', (['config_path'], {}), '(config_path)\n', (1680, 1693), False, 'from mlflow.langchain.utils 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, _RUNNABLE_LOAD_KEY, _UNSUPPORTED_MODEL_ERROR_MESSAGE, _load_base_lcs, _load_from_json, _load_from_pickle, _load_from_yaml, _save_base_lcs, _validate_and_wrap_lc_model, base_lc_types, custom_type_to_loader_dict, lc_runnable_assign_types, lc_runnable_branch_types, lc_runnable_with_steps_types, lc_runnables_types, picklable_runnable_types\n'), ((1718, 1810), 'mlflow.exceptions.MlflowException', 'MlflowException', (['f"""Cannot load runnable without a config file. Got path {config_path}."""'], {}), "(\n    f'Cannot load runnable without a config file. Got path {config_path}.')\n", (1733, 1810), False, 'from mlflow.exceptions import MlflowException\n'), ((2107, 2131), 'langchain.prompts.loading.load_prompt', 'load_prompt', (['config_path'], {}), '(config_path)\n', (2118, 2131), False, 'from langchain.prompts.loading import load_prompt\n'), ((4119, 4141), 'os.listdir', 'os.listdir', (['steps_path'], {}), '(steps_path)\n', (4129, 4141), False, 'import os\n'), ((4297, 4327), 'os.path.join', 'os.path.join', (['steps_path', 'step'], {}), '(steps_path, step)\n', (4309, 4327), False, 'import os\n'), ((8208, 8223), 'mlflow.langchain.utils.base_lc_types', 'base_lc_types', ([], {}), '()\n', (8221, 8223), False, 'from mlflow.langchain.utils 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, _RUNNABLE_LOAD_KEY, _UNSUPPORTED_MODEL_ERROR_MESSAGE, _load_base_lcs, _load_from_json, _load_from_pickle, _load_from_yaml, _save_base_lcs, _validate_and_wrap_lc_model, base_lc_types, custom_type_to_loader_dict, lc_runnable_assign_types, lc_runnable_branch_types, lc_runnable_with_steps_types, lc_runnables_types, picklable_runnable_types\n'), ((8245, 8293), 'mlflow.langchain.utils._validate_and_wrap_lc_model', '_validate_and_wrap_lc_model', (['runnable', 'loader_fn'], {}), '(runnable, loader_fn)\n', (8272, 8293), False, 'from mlflow.langchain.utils 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, _RUNNABLE_LOAD_KEY, _UNSUPPORTED_MODEL_ERROR_MESSAGE, _load_base_lcs, _load_from_json, _load_from_pickle, _load_from_yaml, _save_base_lcs, _validate_and_wrap_lc_model, base_lc_types, custom_type_to_loader_dict, lc_runnable_assign_types, lc_runnable_branch_types, lc_runnable_with_steps_types, lc_runnables_types, picklable_runnable_types\n'), ((14447, 14482), 'os.path.join', 'os.path.join', (['path', 'model_data_path'], {}), '(path, model_data_path)\n', (14459, 14482), False, 'import os\n'), ((14544, 14570), 'mlflow.langchain.utils.picklable_runnable_types', 'picklable_runnable_types', ([], {}), '()\n', (14568, 14570), False, 'from mlflow.langchain.utils 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, _RUNNABLE_LOAD_KEY, _UNSUPPORTED_MODEL_ERROR_MESSAGE, _load_base_lcs, _load_from_json, _load_from_pickle, _load_from_yaml, _save_base_lcs, _validate_and_wrap_lc_model, base_lc_types, custom_type_to_loader_dict, lc_runnable_assign_types, lc_runnable_branch_types, lc_runnable_with_steps_types, lc_runnables_types, picklable_runnable_types\n'), ((15632, 15662), 'os.path.join', 'os.path.join', (['path', 'model_data'], {}), '(path, model_data)\n', (15644, 15662), False, 'import os\n'), ((15847, 15877), 'os.path.join', 'os.path.join', (['path', 'model_data'], {}), '(path, model_data)\n', (15859, 15877), False, 'import os\n'), ((15987, 16017), 'os.path.join', 'os.path.join', (['path', 'model_data'], {}), '(path, model_data)\n', (15999, 16017), False, 'import os\n'), ((16127, 16157), 'os.path.join', 'os.path.join', (['path', 'model_data'], {}), '(path, model_data)\n', (16139, 16157), False, 'import os\n'), ((2150, 2177), 'langchain.llms.get_type_to_cls_dict', 'llms_get_type_to_cls_dict', ([], {}), '()\n', (2175, 2177), True, 'from langchain.llms import get_type_to_cls_dict as llms_get_type_to_cls_dict\n'), ((2247, 2268), 'langchain.llms.loading.load_llm', 'load_llm', (['config_path'], {}), '(config_path)\n', (2255, 2268), False, 'from langchain.llms.loading import load_llm\n'), ((8374, 8428), 'mlflow.langchain.utils._save_base_lcs', '_save_base_lcs', (['lc_model', 'path', 'loader_fn', 'persist_dir'], {}), '(lc_model, path, loader_fn, persist_dir)\n', (8388, 8428), False, 'from mlflow.langchain.utils 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, _RUNNABLE_LOAD_KEY, _UNSUPPORTED_MODEL_ERROR_MESSAGE, _load_base_lcs, _load_from_json, _load_from_pickle, _load_from_yaml, _save_base_lcs, _validate_and_wrap_lc_model, base_lc_types, custom_type_to_loader_dict, lc_runnable_assign_types, lc_runnable_branch_types, lc_runnable_with_steps_types, lc_runnables_types, picklable_runnable_types\n'), ((14665, 14700), 'os.path.join', 'os.path.join', (['path', 'model_data_path'], {}), '(path, model_data_path)\n', (14677, 14700), False, 'import os\n'), ((14729, 14755), 'mlflow.langchain.utils.lc_runnable_branch_types', 'lc_runnable_branch_types', ([], {}), '()\n', (14753, 14755), False, 'from mlflow.langchain.utils 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, _RUNNABLE_LOAD_KEY, _UNSUPPORTED_MODEL_ERROR_MESSAGE, _load_base_lcs, _load_from_json, _load_from_pickle, _load_from_yaml, _save_base_lcs, _validate_and_wrap_lc_model, base_lc_types, custom_type_to_loader_dict, lc_runnable_assign_types, lc_runnable_branch_types, lc_runnable_with_steps_types, lc_runnables_types, picklable_runnable_types\n'), ((15558, 15588), 'mlflow.langchain.utils.lc_runnable_with_steps_types', 'lc_runnable_with_steps_types', ([], {}), '()\n', (15586, 15588), False, 'from mlflow.langchain.utils 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, _RUNNABLE_LOAD_KEY, _UNSUPPORTED_MODEL_ERROR_MESSAGE, _load_base_lcs, _load_from_json, _load_from_pickle, _load_from_yaml, _save_base_lcs, _validate_and_wrap_lc_model, base_lc_types, custom_type_to_loader_dict, lc_runnable_assign_types, lc_runnable_branch_types, lc_runnable_with_steps_types, lc_runnables_types, picklable_runnable_types\n'), ((15921, 15947), 'mlflow.langchain.utils.lc_runnable_branch_types', 'lc_runnable_branch_types', ([], {}), '()\n', (15945, 15947), False, 'from mlflow.langchain.utils 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, _RUNNABLE_LOAD_KEY, _UNSUPPORTED_MODEL_ERROR_MESSAGE, _load_base_lcs, _load_from_json, _load_from_pickle, _load_from_yaml, _save_base_lcs, _validate_and_wrap_lc_model, base_lc_types, custom_type_to_loader_dict, lc_runnable_assign_types, lc_runnable_branch_types, lc_runnable_with_steps_types, lc_runnables_types, picklable_runnable_types\n'), ((16061, 16087), 'mlflow.langchain.utils.lc_runnable_assign_types', 'lc_runnable_assign_types', ([], {}), '()\n', (16085, 16087), False, 'from mlflow.langchain.utils 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, _RUNNABLE_LOAD_KEY, _UNSUPPORTED_MODEL_ERROR_MESSAGE, _load_base_lcs, _load_from_json, _load_from_pickle, _load_from_yaml, _save_base_lcs, _validate_and_wrap_lc_model, base_lc_types, custom_type_to_loader_dict, lc_runnable_assign_types, lc_runnable_branch_types, lc_runnable_with_steps_types, lc_runnables_types, picklable_runnable_types\n'), ((2287, 2315), 'mlflow.langchain.utils.custom_type_to_loader_dict', 'custom_type_to_loader_dict', ([], {}), '()\n', (2313, 2315), False, 'from mlflow.langchain.utils 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, _RUNNABLE_LOAD_KEY, _UNSUPPORTED_MODEL_ERROR_MESSAGE, _load_base_lcs, _load_from_json, _load_from_pickle, _load_from_yaml, _save_base_lcs, _validate_and_wrap_lc_model, base_lc_types, custom_type_to_loader_dict, lc_runnable_assign_types, lc_runnable_branch_types, lc_runnable_with_steps_types, lc_runnables_types, picklable_runnable_types\n'), ((14845, 14880), 'os.path.join', 'os.path.join', (['path', 'model_data_path'], {}), '(path, model_data_path)\n', (14857, 14880), False, 'import os\n'), ((14933, 14959), 'mlflow.langchain.utils.lc_runnable_assign_types', 'lc_runnable_assign_types', ([], {}), '()\n', (14957, 14959), False, 'from mlflow.langchain.utils 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, _RUNNABLE_LOAD_KEY, _UNSUPPORTED_MODEL_ERROR_MESSAGE, _load_base_lcs, _load_from_json, _load_from_pickle, _load_from_yaml, _save_base_lcs, _validate_and_wrap_lc_model, base_lc_types, custom_type_to_loader_dict, lc_runnable_assign_types, lc_runnable_branch_types, lc_runnable_with_steps_types, lc_runnables_types, picklable_runnable_types\n'), ((15728, 15754), 'mlflow.langchain.utils.picklable_runnable_types', 'picklable_runnable_types', ([], {}), '()\n', (15752, 15754), False, 'from mlflow.langchain.utils 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, _RUNNABLE_LOAD_KEY, _UNSUPPORTED_MODEL_ERROR_MESSAGE, _load_base_lcs, _load_from_json, _load_from_pickle, _load_from_yaml, _save_base_lcs, _validate_and_wrap_lc_model, base_lc_types, custom_type_to_loader_dict, lc_runnable_assign_types, lc_runnable_branch_types, lc_runnable_with_steps_types, lc_runnables_types, picklable_runnable_types\n'), ((9016, 9069), 'yaml.dump', 'yaml.dump', (['runnable_dict', 'f'], {'default_flow_style': '(False)'}), '(runnable_dict, f, default_flow_style=False)\n', (9025, 9069), False, 'import yaml\n'), ((15049, 15084), 'os.path.join', 'os.path.join', (['path', 'model_data_path'], {}), '(path, model_data_path)\n', (15061, 15084), False, 'import os\n'), ((2332, 2360), 'mlflow.langchain.utils.custom_type_to_loader_dict', 'custom_type_to_loader_dict', ([], {}), '()\n', (2358, 2360), False, 'from mlflow.langchain.utils 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, _RUNNABLE_LOAD_KEY, _UNSUPPORTED_MODEL_ERROR_MESSAGE, _load_base_lcs, _load_from_json, _load_from_pickle, _load_from_yaml, _save_base_lcs, _validate_and_wrap_lc_model, base_lc_types, custom_type_to_loader_dict, lc_runnable_assign_types, lc_runnable_branch_types, lc_runnable_with_steps_types, lc_runnables_types, picklable_runnable_types\n')]