zeno-cli 0.3.4__py3-none-any.whl

zeno_sdk/client.py

@@ -0,0 +1,87 @@
+from __future__ import annotations
+
+import uuid
+from dataclasses import dataclass
+from datetime import UTC, datetime
+from pathlib import Path
+from typing import Any
+
+from ._runtime import run_coro
+from .config import ZenoConfig, load_config
+from .daemon import DaemonManager
+from .privacy import redact_metadata
+from .session import SessionContext
+from .storage import ZenoStorage
+from .types import SupervisionEvent
+
+
+@dataclass(slots=True)
+class Zeno:
+    project: str
+    project_root: Path | None = None
+    config: ZenoConfig | None = None
+    _storage: ZenoStorage | None = None
+    _daemon: DaemonManager | None = None
+
+    def __post_init__(self) -> None:
+        self.config = self.config or load_config(self.project_root)
+        self._storage = ZenoStorage(self.config.db_path)
+        self._daemon = DaemonManager(
+            db_path=self.config.db_path,
+            flush_interval_seconds=self.config.flush_interval_seconds,
+        )
+        run_coro(self._storage.migrate())
+
+    @property
+    def storage(self) -> ZenoStorage:
+        if self._storage is None:
+            raise RuntimeError("Storage is not initialized.")
+        return self._storage
+
+    @property
+    def daemon(self) -> DaemonManager:
+        if self._daemon is None:
+            raise RuntimeError("Daemon is not initialized.")
+        return self._daemon
+
+    def _project_id(self) -> uuid.UUID:
+        return uuid.uuid5(uuid.NAMESPACE_URL, self.project)
+
+    def session(self, *, harness: str) -> SessionContext:
+        self.daemon.start()
+        return SessionContext(
+            storage=self.storage,
+            project_id=self._project_id(),
+            harness=harness,
+            probes_enabled=self.config.probes_enabled,
+            record_content=self.config.record_content,
+        )
+
+    def record_event(
+        self,
+        *,
+        session_id: uuid.UUID,
+        type: str,
+        agent_run_id: uuid.UUID | None = None,
+        latency_ms_to_decide: int | None = None,
+        metadata: dict[str, Any] | None = None,
+    ) -> None:
+        self.daemon.start()
+        event = SupervisionEvent(
+            id=uuid.uuid4(),
+            session_id=session_id,
+            agent_run_id=agent_run_id,
+            type=type,
+            timestamp=datetime.now(tz=UTC),
+            latency_ms_to_decide=latency_ms_to_decide,
+            metadata=redact_metadata(metadata, record_content=self.config.record_content),
+        )
+        run_coro(self.storage.insert_event(event))
+
+    def prompt_load_probe(
+        self, *, session: SessionContext, subscales: dict[str, int] | None, skipped: bool = False
+    ) -> None:
+        session.prompt_load_probe(subscales=subscales, skipped=skipped)
+
+    def shutdown(self) -> None:
+        self.daemon.stop()

zeno_sdk/config.py

@@ -0,0 +1,61 @@
+from __future__ import annotations
+
+import os
+import tomllib
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Any
+
+
+@dataclass(slots=True)
+class ZenoConfig:
+    base_dir: Path
+    db_path: Path
+    probes_enabled: bool
+    flush_interval_seconds: int
+    redaction_enabled: bool
+    record_content: bool
+
+
+def _read_toml(path: Path) -> dict[str, Any]:
+    if not path.exists():
+        return {}
+    return tomllib.loads(path.read_text(encoding="utf-8"))
+
+
+def load_config(project_root: Path | None = None) -> ZenoConfig:
+    home = Path.home()
+    base_dir = Path(os.getenv("ZENO_HOME", home / ".zeno")).expanduser()
+    base_dir.mkdir(parents=True, exist_ok=True)
+
+    user_config = _read_toml(base_dir / "config.toml")
+    project_config: dict[str, Any] = {}
+    if project_root is not None:
+        project_config = _read_toml(project_root / ".zeno.toml")
+
+    cfg = {}
+    cfg.update(user_config.get("zeno", {}))
+    cfg.update(project_config.get("zeno", {}))
+
+    probes_enabled = str(
+        os.getenv("ZENO_PROBES_ENABLED", cfg.get("probes_enabled", "true"))
+    ).lower()
+    redaction_enabled = str(
+        os.getenv("ZENO_REDACTION_ENABLED", cfg.get("redaction_enabled", "true"))
+    ).lower()
+    record_content = str(
+        os.getenv("ZENO_RECORD_CONTENT", cfg.get("record_content", "false"))
+    ).lower()
+    flush_interval = int(
+        os.getenv("ZENO_FLUSH_INTERVAL_SECONDS", cfg.get("flush_interval_seconds", 5))
+    )
+    db_path = Path(os.getenv("ZENO_DB_PATH", str(base_dir / "zeno.db"))).expanduser()
+
+    return ZenoConfig(
+        base_dir=base_dir,
+        db_path=db_path,
+        probes_enabled=probes_enabled in {"1", "true", "yes", "on"},
+        flush_interval_seconds=max(flush_interval, 1),
+        redaction_enabled=redaction_enabled in {"1", "true", "yes", "on"},
+        record_content=record_content in {"1", "true", "yes", "on"},
+    )

zeno_sdk/daemon.py

@@ -0,0 +1,72 @@
+from __future__ import annotations
+
+import multiprocessing as mp
+import signal
+import time
+import uuid
+from dataclasses import dataclass
+from pathlib import Path
+
+from ._runtime import run_coro
+from .storage import ZenoStorage, make_tax_point
+
+
+def _daemon_main(db_path: str, flush_interval: int, stop_event: mp.Event) -> None:
+    running = True
+
+    def _handle_sigterm(_signum: int, _frame: object) -> None:
+        nonlocal running
+        running = False
+
+    signal.signal(signal.SIGTERM, _handle_sigterm)
+    storage = ZenoStorage(Path(db_path))
+    while running and not stop_event.is_set():
+        # Keep process alive and ready for future sync/aggregation jobs.
+        # Current v1 loop only runs periodic liveness cycles.
+        time.sleep(max(flush_interval, 1))
+        _ = storage
+
+
+@dataclass(slots=True)
+class DaemonManager:
+    db_path: Path
+    flush_interval_seconds: int
+    _process: mp.Process | None = None
+    _stop_event: mp.Event | None = None
+
+    def start(self) -> None:
+        if self._process and self._process.is_alive():
+            return
+        self._stop_event = mp.Event()
+        self._process = mp.Process(
+            target=_daemon_main,
+            args=(str(self.db_path), self.flush_interval_seconds, self._stop_event),
+            daemon=True,
+        )
+        self._process.start()
+
+    def stop(self, timeout_seconds: float = 5.0) -> None:
+        if not self._process:
+            return
+        if self._stop_event:
+            self._stop_event.set()
+        self._process.join(timeout_seconds)
+        if self._process.is_alive():
+            self._process.terminate()
+            self._process.join(timeout_seconds)
+        self._process = None
+        self._stop_event = None
+
+    @property
+    def pid(self) -> int | None:
+        if self._process and self._process.is_alive():
+            return self._process.pid
+        return None
+
+
+def compute_incremental_tax_point(storage: ZenoStorage, session_id: str) -> None:
+    stats = run_coro(storage.compute_session_stats(session_id))
+    point = make_tax_point(session_id, stats)
+    # Use stable deterministic id to keep updates idempotent.
+    point.id = uuid.uuid5(uuid.UUID(session_id), "tax-point")
+    run_coro(storage.upsert_tax_point(point))

zeno_sdk/privacy.py

@@ -0,0 +1,46 @@
+from __future__ import annotations
+
+import hashlib
+import json
+from collections.abc import Mapping
+from typing import Any
+
+BLOCKED_KEYS = {"prompt", "prompt_text", "output", "output_text", "completion", "response"}
+
+
+def hash_text(value: str) -> str:
+    return hashlib.sha256(value.encode("utf-8")).hexdigest()
+
+
+def redact_metadata(
+    metadata: Mapping[str, Any] | None,
+    *,
+    record_content: bool = False,
+    record_tool_args: bool = False,
+) -> dict[str, Any]:
+    if metadata is None:
+        return {}
+
+    redacted: dict[str, Any] = {}
+    for key, value in metadata.items():
+        lowered = key.lower()
+        if lowered in BLOCKED_KEYS and not record_content:
+            text = str(value)
+            redacted[f"{key}_hash"] = hash_text(text)
+            redacted[f"{key}_len"] = len(text)
+            continue
+
+        if lowered in {"tool_args", "arguments"} and not record_tool_args:
+            if isinstance(value, Mapping):
+                redacted["tool_arg_keys"] = sorted(value.keys())
+            else:
+                redacted["tool_arg_keys"] = []
+            continue
+
+        if isinstance(value, (str, int, float, bool)) or value is None:
+            redacted[key] = value
+            continue
+
+        # Keep complex values serializable and compact.
+        redacted[key] = json.loads(json.dumps(value, default=str))
+    return redacted

zeno_sdk/session.py

@@ -0,0 +1,179 @@
+from __future__ import annotations
+
+import uuid
+from dataclasses import dataclass
+from datetime import UTC, datetime
+from typing import Any
+
+from zeno_core.tlx_s import run_tlx_survey_tui, should_prompt_probe
+
+from ._runtime import run_coro
+from .daemon import compute_incremental_tax_point
+from .privacy import redact_metadata
+from .storage import ZenoStorage
+from .types import AgentRun, LoadProbe, LoadProbeSubscales, Session, SupervisionEvent
+
+
+def _now() -> datetime:
+    return datetime.now(tz=UTC)
+
+
+_probe_state: dict[str, Any] = {
+    "sessions_since_last_probe": 0,
+    "last_probe_at": None,
+}
+
+
+@dataclass(slots=True)
+class AgentRunContext:
+    storage: ZenoStorage
+    session_id: uuid.UUID
+    harness: str
+    model: str
+    run_id: uuid.UUID
+    _outcome: str = "unknown"
+    _started_at: datetime | None = None
+    _ended_at: datetime | None = None
+
+    def outcome(self, value: str) -> None:
+        self._outcome = value
+
+    def __enter__(self) -> AgentRunContext:
+        self._started_at = _now()
+        run = AgentRun(
+            id=self.run_id,
+            session_id=self.session_id,
+            harness=self.harness,
+            model=self.model,
+            started_at=self._started_at,
+            ended_at=None,
+            outcome="unknown",
+        )
+        run_coro(self.storage.insert_agent_run(run))
+        return self
+
+    def __exit__(self, _exc_type: object, _exc: object, _tb: object) -> None:
+        self._ended_at = _now()
+        if self._started_at is None:
+            self._started_at = self._ended_at
+        final = AgentRun(
+            id=self.run_id,
+            session_id=self.session_id,
+            harness=self.harness,
+            model=self.model,
+            started_at=self._started_at,
+            ended_at=self._ended_at,
+            outcome=self._outcome,
+        )
+        run_coro(self.storage.insert_agent_run(final))
+
+
+@dataclass(slots=True)
+class SessionContext:
+    storage: ZenoStorage
+    project_id: uuid.UUID
+    harness: str
+    probes_enabled: bool
+    record_content: bool
+    session_id: uuid.UUID | None = None
+    _started_at: datetime | None = None
+
+    def __enter__(self) -> SessionContext:
+        self.session_id = uuid.uuid4()
+        self._started_at = _now()
+        session = Session(
+            id=self.session_id,
+            start_at=self._started_at,
+            end_at=None,
+            agent_count_max=1,
+            harness=self.harness,
+            project_id=self.project_id,
+        )
+        run_coro(self.storage.insert_session(session))
+        return self
+
+    def __exit__(self, _exc_type: object, _exc: object, _tb: object) -> None:
+        if self.session_id is None or self._started_at is None:
+            return
+
+        stats = run_coro(self.storage.compute_session_stats(str(self.session_id)))
+        finished = Session(
+            id=self.session_id,
+            start_at=self._started_at,
+            end_at=_now(),
+            agent_count_max=max(1, int(round(stats.avg_n_agents)) or 1),
+            harness=self.harness,
+            project_id=self.project_id,
+        )
+        run_coro(self.storage.insert_session(finished))
+        compute_incremental_tax_point(self.storage, str(self.session_id))
+        _probe_state["sessions_since_last_probe"] += 1
+        last_probe_at = _probe_state.get("last_probe_at")
+        seconds_since = None
+        if isinstance(last_probe_at, datetime):
+            seconds_since = (_now() - last_probe_at).total_seconds()
+        should_prompt = should_prompt_probe(
+            session_duration_minutes=(finished.end_at - finished.start_at).total_seconds() / 60,
+            sessions_since_last_probe=int(_probe_state["sessions_since_last_probe"]),
+            probes_enabled=self.probes_enabled,
+            seconds_since_last_probe=seconds_since,
+            active_agent_runs=0,
+        )
+        if should_prompt:
+            result = run_tlx_survey_tui()
+            self.prompt_load_probe(result.subscales, skipped=result.skipped)
+            if result.comment_hash:
+                self.record(type="intervene", metadata={"tlx_comment_hash": result.comment_hash})
+            _probe_state["sessions_since_last_probe"] = 0
+            _probe_state["last_probe_at"] = _now()
+
+    def agent_run(self, *, model: str) -> AgentRunContext:
+        if self.session_id is None:
+            raise RuntimeError("Session is not active.")
+        return AgentRunContext(
+            storage=self.storage,
+            session_id=self.session_id,
+            harness=self.harness,
+            model=model,
+            run_id=uuid.uuid4(),
+        )
+
+    def record(
+        self,
+        *,
+        type: str,
+        agent_run_id: uuid.UUID | None = None,
+        latency_ms_to_decide: int | None = None,
+        metadata: dict[str, Any] | None = None,
+        **extra_metadata: Any,
+    ) -> None:
+        if self.session_id is None:
+            raise RuntimeError("Session is not active.")
+        merged_metadata = dict(metadata or {})
+        merged_metadata.update(extra_metadata)
+        event = SupervisionEvent(
+            id=uuid.uuid4(),
+            session_id=self.session_id,
+            agent_run_id=agent_run_id,
+            type=type,
+            timestamp=_now(),
+            latency_ms_to_decide=latency_ms_to_decide,
+            metadata=redact_metadata(merged_metadata, record_content=self.record_content),
+        )
+        run_coro(self.storage.insert_event(event))
+
+    def prompt_load_probe(self, subscales: dict[str, int] | None, skipped: bool = False) -> None:
+        if self.session_id is None:
+            raise RuntimeError("Session is not active.")
+        if not self.probes_enabled and not skipped:
+            return
+        prompted = _now()
+        probe = LoadProbe(
+            id=uuid.uuid4(),
+            session_id=self.session_id,
+            prompted_at=prompted,
+            responded_at=None if skipped else prompted,
+            skipped=skipped,
+            subscales=LoadProbeSubscales(**subscales) if subscales is not None else None,
+        )
+        run_coro(self.storage.insert_load_probe(probe))