lang-agent/lang_agent/front_api/build_server_utils.py

from typing import Any, Dict, List, Literal
import os
import os.path as osp
import subprocess
import json

from lang_agent.config.core_config import load_tyro_conf
from lang_agent.config.constants import TY_BUILD_SCRIPT, _PROJECT_ROOT

_DEEP_AGENT_BACKEND_ALIASES = {
    "state_bk": "statebk",
    "statebk": "statebk",
    "local_shell": "localshell",
    "localshell": "localshell",
    "daytona_sandbox": "daytonasandbox",
    "daytonasandbox": "daytonasandbox",
}


def opt_to_config(save_path: str, *nargs):
    os.makedirs(osp.dirname(save_path), exist_ok=True)
    subprocess.run(
        ["python", TY_BUILD_SCRIPT, "--save-path", save_path, *nargs],
        check=True,
        cwd=_PROJECT_ROOT,
    )


def _build_and_load_pipeline_config(
    pipeline_id: str, pipeline_config_dir: str, cmd: List[str]
):
    save_config_f = osp.join(pipeline_config_dir, f"{pipeline_id}.yaml")
    opt_to_config(save_config_f, *cmd)

    # TODO: think if returning the built pipeline is better or just the config obj for front_api
    return load_tyro_conf(save_config_f)


def update_pipeline_registry(
    pipeline_id: str,
    graph_id: str,
    config_file: str,
    llm_name: str,
    enabled: bool = True,
    registry_f: str = "configs/pipeline_registry.json",
):
    if not osp.isabs(registry_f):
        registry_f = osp.join(_PROJECT_ROOT, registry_f)
    os.makedirs(osp.dirname(registry_f), exist_ok=True)
    if not osp.exists(registry_f):
        with open(registry_f, "w", encoding="utf-8") as f:
            json.dump({"pipelines": {}, "api_keys": {}}, f, indent=4)

    with open(registry_f, "r") as f:
        registry = json.load(f)

    pipelines: Dict[str, Dict[str, Any]] = registry.setdefault("pipelines", {})
    pipeline = pipelines.setdefault(pipeline_id, {})
    pipeline["enabled"] = bool(enabled)
    pipeline["config_file"] = config_file
    pipeline["graph_id"] = graph_id
    pipeline["llm_name"] = llm_name

    with open(registry_f, "w", encoding="utf-8") as f:
        json.dump(registry, f, indent=4)


def build_route(
    pipeline_id: str,
    prompt_set: str,
    tool_keys: List[str],
    api_key: str,
    llm_name: str = "qwen-plus",
    pipeline_config_dir: str = "configs/pipelines",
    **_: Any,
):
    cmd_opt = [
        "--pipeline.pipeline-id", pipeline_id,
        "route",  # ------------
        "--llm-name", llm_name,
        "--api-key", api_key,
        "--pipeline-id", pipeline_id,
        "--prompt-set-id", prompt_set,
        "tool_node",  # ------------
        "--llm-name", llm_name,
        "--api-key", api_key,
        "--pipeline-id", pipeline_id,
        "--prompt-set-id", prompt_set,
    ]

    if tool_keys:
        cmd_opt.extend(
            ["--tool-manager-config.client-tool-manager.tool-keys", *tool_keys]
        )

    return _build_and_load_pipeline_config(pipeline_id, pipeline_config_dir, cmd_opt)


def build_react(
    pipeline_id: str,
    prompt_set: str,
    tool_keys: List[str],
    api_key: str,
    llm_name: str = "qwen-plus",
    pipeline_config_dir: str = "configs/pipelines",
    **_: Any,
):
    cmd_opt = [
        "--pipeline.pipeline-id", pipeline_id,
        "react",  # ------------
        "--llm-name", llm_name,
        "--api-key", api_key,
        "--pipeline-id", pipeline_id,
        "--prompt-set-id", prompt_set,
        ]
    if tool_keys:
        cmd_opt.extend(
            ["--tool-manager-config.client-tool-manager.tool-keys", *tool_keys]
        )

    return _build_and_load_pipeline_config(pipeline_id, pipeline_config_dir, cmd_opt)


def build_deep_agent(
    pipeline_id: str,
    prompt_set: str,
    tool_keys: List[str],
    api_key: str,
    llm_name: str = "qwen-plus",
    pipeline_config_dir: str = "configs/pipelines",
    act_bkend: Literal[
        "local_shell",
        "localshell",
        "state_bk",
        "statebk",
        "daytona_sandbox",
        "daytonasandbox",
    ] = "state_bk",
    **_: Any,
):
    backend_subcommand = _DEEP_AGENT_BACKEND_ALIASES.get(act_bkend)
    if backend_subcommand is None:
        raise ValueError(
            "Unsupported deepagent backend "
            f"'{act_bkend}'. Expected one of {sorted(_DEEP_AGENT_BACKEND_ALIASES.keys())}"
        )

    cmd_opt = [
        "--pipeline.pipeline-id", pipeline_id,
        "deepagent",
        "--llm-name", llm_name,
        "--api-key", api_key,
        "--pipeline-id", pipeline_id,
        "--prompt-set-id", prompt_set,
        backend_subcommand,
    ]

    if tool_keys:
        cmd_opt.extend(
            ["--tool-manager-config.client-tool-manager.tool-keys", *tool_keys]
        )

    return _build_and_load_pipeline_config(pipeline_id, pipeline_config_dir, cmd_opt)

# {pipeline_id: build_function}
GRAPH_BUILD_FNCS = {
    "routing": build_route,
    "react": build_react,
    "deepagent": build_deep_agent,
}