kotaemon/knowledgehub/contribs/promptui/config.py

"""Get config from Pipeline"""
import inspect
from pathlib import Path
from typing import Any, Dict, Optional, Type, Union

import yaml

from ...base import BaseComponent
from .base import DEFAULT_COMPONENT_BY_TYPES


def config_from_value(value: Any) -> dict:
    """Get the config from default value

    Args:
        value (Any): default value

    Returns:
        dict: config
    """
    component = DEFAULT_COMPONENT_BY_TYPES.get(type(value).__name__, "text")
    return {
        "component": component,
        "params": {
            "value": value,
        },
    }


def handle_param(param: dict) -> dict:
    """Convert param definition into promptui-compliant config

    Supported gradio's UI components are (https://www.gradio.app/docs/components)
        - CheckBoxGroup: list (multi select)
        - DropDown: list (single select)
        - File
        - Image
        - Number: int / float
        - Radio: list (single select)
        - Slider: int / float
        - TextBox: str
    """
    params = {}
    default = param.get("default", None)
    if isinstance(default, str) and default.startswith("{{") and default.endswith("}}"):
        default = None
    if default is not None:
        params["value"] = default

    type_: str = type(default).__name__ if default is not None else ""
    ui_component = DEFAULT_COMPONENT_BY_TYPES.get(type_, "text")

    return {
        "component": ui_component,
        "params": params,
    }


def handle_node(node: dict) -> dict:
    """Convert node definition into promptui-compliant config"""
    config = {}
    for name, param_def in node.get("params", {}).items():
        if isinstance(param_def["default_callback"], str):
            continue
        config[name] = handle_param(param_def)
    for name, node_def in node.get("nodes", {}).items():
        if isinstance(node_def["default_callback"], str):
            continue
        for key, value in handle_node(node_def["default"]).items():
            config[f"{name}.{key}"] = value
        for key, value in node_def["default_kwargs"].items():
            config[f"{name}.{key}"] = config_from_value(value)

    return config


def handle_input(pipeline: Union[BaseComponent, Type[BaseComponent]]) -> dict:
    """Get the input from the pipeline"""
    if not hasattr(pipeline, "run_raw"):
        return {}
    signature = inspect.signature(pipeline.run_raw)
    inputs: Dict[str, Dict] = {}
    for name, param in signature.parameters.items():
        if name in ["self", "args", "kwargs"]:
            continue
        input_def: Dict[str, Optional[Any]] = {"component": "text"}
        default = param.default
        if default is param.empty:
            inputs[name] = input_def
            continue

        params = {}
        params["value"] = default
        type_ = type(default).__name__ if default is not None else None
        ui_component = None
        if type_ is not None:
            ui_component = "text"

        input_def["component"] = ui_component
        input_def["params"] = params

        inputs[name] = input_def

    return inputs


def export_pipeline_to_config(
    pipeline: Union[BaseComponent, Type[BaseComponent]],
    path: Optional[str] = None,
) -> dict:
    """Export a pipeline to a promptui-compliant config dict"""
    if inspect.isclass(pipeline):
        pipeline = pipeline()

    pipeline_def = pipeline.describe()
    config = {
        f"{pipeline.__module__}.{pipeline.__class__.__name__}": {
            "params": handle_node(pipeline_def),
            "inputs": handle_input(pipeline),
            "outputs": [{"step": ".", "component": "text"}],
        }
    }
    if path is not None:
        old_config = config
        if Path(path).is_file():
            with open(path) as f:
                old_config = yaml.safe_load(f)
                old_config.update(config)
        with open(path, "w") as f:
            yaml.safe_dump(old_config, f)

    return config