code: 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)
      )
  
apis: langchain.llms.get_type_to_cls_dict langchain.schema.runnable.passthrough.RunnableAssign langchain.chains.loading.load_chain langchain.schema.runnable.RunnableParallel langchain.schema.runnable.RunnableSequence langchain.llms.loading.load_llm langchain.schema.runnable.RunnableBranch langchain.prompts.loading.load_prompt
extract_api: [((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')]