HadTavern/agentui/pipeline/storage.py

from __future__ import annotations

from pathlib import Path
from typing import Any, Dict, List
import json
from agentui.pipeline.defaults import default_pipeline


PIPELINE_FILE = Path("pipeline.json")
PRESETS_DIR = Path("presets")
VARS_DIR = Path(".agentui") / "vars"


# DRY нормализация meta/пайплайна: единый источник дефолтов и типов
def normalize_pipeline(pipeline: Dict[str, Any]) -> Dict[str, Any]:
   """
   Приводит верхнеуровневые ключи пайплайна к согласованному виду, заполняет дефолты.
   Безопасно к отсутствующим ключам и неверным типам.
   """
   if not isinstance(pipeline, dict):
       pipeline = {}
   out: Dict[str, Any] = dict(pipeline)

   def _to_int(v, d):
       try:
           n = int(v)
           return n if n > 0 else d
       except Exception:
           return d

   def _to_float(v, d):
       try:
           n = float(v)
           return n if n > 0 else d
       except Exception:
           return d

   # Базовые поля
   out["id"] = str(out.get("id") or "pipeline_editor")
   out["name"] = str(out.get("name") or "Edited Pipeline")
   out["parallel_limit"] = _to_int(out.get("parallel_limit"), 8)
   out["loop_mode"] = str(out.get("loop_mode") or "dag")
   out["loop_max_iters"] = _to_int(out.get("loop_max_iters"), 1000)
   out["loop_time_budget_ms"] = _to_int(out.get("loop_time_budget_ms"), 10000)
   out["clear_var_store"] = bool(out.get("clear_var_store", True))
   out["http_timeout_sec"] = _to_float(out.get("http_timeout_sec"), 60)

   # Глобальные опции извлечения текста для [[OUTx]]
   out["text_extract_strategy"] = str(out.get("text_extract_strategy") or "auto")
   out["text_extract_json_path"] = str(out.get("text_extract_json_path") or "")
   # Поддержка разных написаний text_join_sep
   join_sep = out.get("text_join_sep")
   if join_sep is None:
       for k in list(out.keys()):
           if isinstance(k, str) and k.lower() == "text_join_sep":
               join_sep = out.get(k)
               break
   out["text_join_sep"] = str(join_sep or "\n")

   # Пресеты парсинга
   presets = out.get("text_extract_presets")
   norm_presets: List[Dict[str, Any]] = []
   if isinstance(presets, list):
       for i, it in enumerate(presets):
           if not isinstance(it, dict):
               continue
           norm_presets.append({
               "id": str(it.get("id") or f"p{i}"),
               "name": str(it.get("name") or it.get("json_path") or "Preset"),
               "strategy": str(it.get("strategy") or "auto"),
               "json_path": str(it.get("json_path") or ""),
               "join_sep": str(it.get("join_sep") or "\n"),
           })
   out["text_extract_presets"] = norm_presets

   # Узлы — список
   try:
       nodes = out.get("nodes") or []
       if not isinstance(nodes, list):
           nodes = []
       out["nodes"] = nodes
   except Exception:
       out["nodes"] = []

   return out


def load_pipeline() -> Dict[str, Any]:
   if PIPELINE_FILE.exists():
       try:
           data = json.loads(PIPELINE_FILE.read_text(encoding="utf-8"))
           return normalize_pipeline(data)
       except Exception:
           pass
   return normalize_pipeline(default_pipeline())


def save_pipeline(pipeline: Dict[str, Any]) -> None:
    norm = normalize_pipeline(pipeline or {})
    PIPELINE_FILE.write_text(json.dumps(norm, ensure_ascii=False, indent=2), encoding="utf-8")


def list_presets() -> List[str]:
    PRESETS_DIR.mkdir(parents=True, exist_ok=True)
    return sorted([p.stem for p in PRESETS_DIR.glob("*.json")])


def load_preset(name: str) -> Dict[str, Any]:
    PRESETS_DIR.mkdir(parents=True, exist_ok=True)
    path = PRESETS_DIR / f"{name}.json"
    if not path.exists():
        raise FileNotFoundError(name)
    return json.loads(path.read_text(encoding="utf-8"))


def save_preset(name: str, pipeline: Dict[str, Any]) -> None:
    PRESETS_DIR.mkdir(parents=True, exist_ok=True)
    path = PRESETS_DIR / f"{name}.json"
    path.write_text(json.dumps(pipeline, ensure_ascii=False, indent=2), encoding="utf-8")


# ---------------- Variable Store (per-pipeline) ----------------

def _var_store_path(pipeline_id: str) -> Path:
    pid = pipeline_id or "pipeline_editor"
    VARS_DIR.mkdir(parents=True, exist_ok=True)
    # normalize to safe filename
    safe = "".join(ch if ch.isalnum() or ch in ("-", "_", ".") else "_" for ch in str(pid))
    return VARS_DIR / f"{safe}.json"


def load_var_store(pipeline_id: str) -> Dict[str, Any]:
    """
    Load variable store dictionary for given pipeline id.
    Returns {} if not exists or invalid.
    """
    path = _var_store_path(pipeline_id)
    if not path.exists():
        return {}
    try:
        data = json.loads(path.read_text(encoding="utf-8"))
        return data if isinstance(data, dict) else {}
    except Exception:
        return {}


def save_var_store(pipeline_id: str, data: Dict[str, Any]) -> None:
    """
    Save variable store dictionary for given pipeline id.
    """
    path = _var_store_path(pipeline_id)
    try:
        VARS_DIR.mkdir(parents=True, exist_ok=True)
    except Exception:
        pass
    path.write_text(json.dumps(data or {}, ensure_ascii=False, indent=2), encoding="utf-8")


def clear_var_store(pipeline_id: str) -> None:
    """
    Delete/reset variable store for given pipeline id.
    """
    path = _var_store_path(pipeline_id)
    try:
        if path.exists():
            path.unlink()
    except Exception:
        # ignore failures
        pass