rex/automation_registry.py
Rama master · réplica de lectura del repositorio original.
"""
Automation Registry for Rex AI Assistant.
Stores and retrieves automations (trigger + action pairs) persistently.
Automations are saved as JSON and survive restarts.
"""
from __future__ import annotations
import json
import logging
import threading
import uuid
from dataclasses import dataclass, field
from datetime import UTC, datetime
from pathlib import Path
from typing import Any
logger = logging.getLogger(__name__)
_DEFAULT_STORAGE = Path("data/automations/automations.json")
@dataclass
class Automation:
"""A single automation definition."""
automation_id: str
name: str
trigger: dict[str, Any]
action: dict[str, Any]
enabled: bool = True
created_at: str = field(default_factory=lambda: datetime.now(UTC).isoformat())
updated_at: str = field(default_factory=lambda: datetime.now(UTC).isoformat())
metadata: dict[str, Any] = field(default_factory=dict)
def to_dict(self) -> dict[str, Any]:
return {
"automation_id": self.automation_id,
"name": self.name,
"trigger": self.trigger,
"action": self.action,
"enabled": self.enabled,
"created_at": self.created_at,
"updated_at": self.updated_at,
"metadata": self.metadata,
}
@classmethod
def from_dict(cls, data: dict[str, Any]) -> Automation:
return cls(
automation_id=str(data.get("automation_id") or str(uuid.uuid4())),
name=str(data.get("name") or "Unnamed Automation"),
trigger=dict(data.get("trigger") or {}),
action=dict(data.get("action") or {}),
enabled=bool(data.get("enabled", True)),
created_at=str(data.get("created_at") or datetime.now(UTC).isoformat()),
updated_at=str(data.get("updated_at") or datetime.now(UTC).isoformat()),
metadata=dict(data.get("metadata") or {}),
)
class AutomationRegistry:
"""
Persistent store for automation definitions.
Automations are stored in a JSON file and loaded on startup,
so they persist across restarts.
"""
def __init__(self, storage_path: Path | None = None) -> None:
self._path = storage_path or _DEFAULT_STORAGE
self._lock = threading.RLock()
self._automations: dict[str, Automation] = {}
self._load()
# ------------------------------------------------------------------
# Persistence
# ------------------------------------------------------------------
def _load(self) -> None:
if not self._path.exists():
return
try:
raw = json.loads(self._path.read_text(encoding="utf-8"))
except Exception as exc:
logger.error("Failed to read automations file %s: %s", self._path, exc)
return
items: list[dict[str, Any]] = raw if isinstance(raw, list) else []
loaded = 0
with self._lock:
self._automations.clear()
for item in items:
if not isinstance(item, dict):
continue
auto = Automation.from_dict(item)
self._automations[auto.automation_id] = auto
loaded += 1
logger.info("Loaded %d automation(s) from %s", loaded, self._path)
def _save(self) -> None:
with self._lock:
payload = [a.to_dict() for a in self._automations.values()]
try:
self._path.parent.mkdir(parents=True, exist_ok=True)
self._path.write_text(json.dumps(payload, indent=2), encoding="utf-8")
except Exception as exc:
logger.error("Failed to save automations to %s: %s", self._path, exc)
# ------------------------------------------------------------------
# CRUD
# ------------------------------------------------------------------
def store(
self,
name: str,
trigger: dict[str, Any],
action: dict[str, Any],
*,
automation_id: str | None = None,
enabled: bool = True,
metadata: dict[str, Any] | None = None,
) -> Automation:
"""Create and persist a new automation. Returns the stored Automation."""
auto = Automation(
automation_id=str(automation_id or uuid.uuid4()),
name=name,
trigger=dict(trigger),
action=dict(action),
enabled=enabled,
metadata=dict(metadata or {}),
)
with self._lock:
self._automations[auto.automation_id] = auto
self._save()
logger.info("Stored automation '%s' (%s)", name, auto.automation_id)
return auto
def get(self, automation_id: str) -> Automation | None:
"""Retrieve an automation by ID."""
with self._lock:
return self._automations.get(automation_id)
def get_by_name(self, name: str) -> Automation | None:
"""Retrieve the first automation matching the given name."""
with self._lock:
for auto in self._automations.values():
if auto.name == name:
return auto
return None
def list_all(self, *, include_disabled: bool = True) -> list[Automation]:
"""Return all stored automations."""
with self._lock:
autos = list(self._automations.values())
if not include_disabled:
autos = [a for a in autos if a.enabled]
return autos
def update(self, automation_id: str, **updates: Any) -> Automation | None:
"""Update fields on an existing automation. Returns the updated object or None."""
with self._lock:
auto = self._automations.get(automation_id)
if auto is None:
return None
for key, value in updates.items():
if hasattr(auto, key):
setattr(auto, key, value)
auto.updated_at = datetime.now(UTC).isoformat()
self._save()
return auto
def remove(self, automation_id: str) -> bool:
"""Delete an automation. Returns True if it existed."""
with self._lock:
existed = automation_id in self._automations
if existed:
del self._automations[automation_id]
if existed:
self._save()
return existed
def clear(self) -> None:
"""Remove all automations."""
with self._lock:
self._automations.clear()
self._save()
# ------------------------------------------------------------------
# Global registry
# ------------------------------------------------------------------
_REGISTRY: AutomationRegistry | None = None
_REGISTRY_LOCK = threading.Lock()
def get_automation_registry() -> AutomationRegistry:
"""Return the global AutomationRegistry instance."""
global _REGISTRY
if _REGISTRY is None:
with _REGISTRY_LOCK:
if _REGISTRY is None:
_REGISTRY = AutomationRegistry()
return _REGISTRY
def set_automation_registry(registry: AutomationRegistry | None) -> None:
"""Replace the global registry (primarily for testing)."""
global _REGISTRY
with _REGISTRY_LOCK:
_REGISTRY = registry
__all__ = [
"Automation",
"AutomationRegistry",
"get_automation_registry",
"set_automation_registry",
]