118 lines
3.3 KiB
Python
118 lines
3.3 KiB
Python
from typing import Any, Dict, List
|
|
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
|
|
|
|
|
|
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["overrides"] = {"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="configs/pipelines",
|
|
):
|
|
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="configs/pipelines",
|
|
):
|
|
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)
|
|
|
|
|
|
# {pipeline_id: build_function}
|
|
GRAPH_BUILD_FNCS = {
|
|
"routing": build_route,
|
|
"react": build_react,
|
|
}
|