Import AITURK IDE 1.0.0-beta.1 from Hermes 63279301; preserve MIT license

agent/moa_trace.py

@@ -0,0 +1,182 @@
+"""Full MoA turn trace persistence (opt-in via config ``moa.save_traces``).
+
+When enabled, every Mixture-of-Agents turn that actually runs the reference
+fan-out (a cache MISS in ``MoAChatCompletions.create``) appends one JSON line
+to ``<hermes_home>/moa-traces/<session_id>.jsonl``. The record is the TRUE
+FULL turn — the exact messages array each reference model received (system
+prompt + advisory view, not the truncated display preview), each reference's
+full output, and the exact messages array the aggregator received (including
+the injected reference-context guidance block) plus its output when available
+— so a run can be audited end-to-end offline: what every model saw, what every
+model said, and what it cost.
+
+This is a side-channel trace. It is NOT the conversation ``messages`` table and
+never enters message history or replay — MoA references are advisory side-calls
+with their own system prompt, not conversation turns, so persisting them as
+message rows would corrupt role alternation / replay. Traces live in their own
+files, keyed by session id, and are safe to delete.
+
+Cost model note: gated OFF by default. When off, the only overhead is the
+``_traces_enabled()`` config read (cheap) — no file I/O, no serialization.
+"""
+
+from __future__ import annotations
+
+import json
+import logging
+import os
+import time
+from pathlib import Path
+from typing import Any, Optional
+
+from hermes_constants import get_hermes_home
+
+logger = logging.getLogger(__name__)
+
+
+def _traces_enabled_and_dir() -> Optional[Path]:
+    """Return the trace directory if ``moa.save_traces`` is on, else None.
+
+    Reads config lazily per call (config is cheap to load and this only runs on
+    a cache-MISS MoA turn, i.e. once per user turn, not per tool iteration).
+    ``moa.trace_dir`` overrides the default ``<hermes_home>/moa-traces/``.
+    """
+    try:
+        from hermes_cli.config import load_config
+
+        moa_cfg = (load_config() or {}).get("moa") or {}
+    except Exception:  # pragma: no cover - defensive: never break a turn over tracing
+        return None
+    if not moa_cfg.get("save_traces"):
+        return None
+    override = moa_cfg.get("trace_dir")
+    if override:
+        base = Path(os.path.expandvars(os.path.expanduser(str(override))))
+    else:
+        base = get_hermes_home() / "moa-traces"
+    return base
+
+
+def _sanitize_session_id(session_id: Optional[str]) -> str:
+    """Make a session id safe as a filename component."""
+    if not session_id:
+        return "unknown-session"
+    return "".join(c if (c.isalnum() or c in "-_.") else "_" for c in str(session_id))
+
+
+def _slot_trace(acct: Any, label: str) -> dict[str, Any]:
+    """Render one reference's _RefAccounting into a full trace dict.
+
+    Includes the FULL input messages the reference received and its FULL
+    output — not the truncated display preview.
+    """
+    usage = getattr(acct, "usage", None)
+    usage_dict: dict[str, Any] = {}
+    if usage is not None:
+        usage_dict = {
+            "input_tokens": getattr(usage, "input_tokens", 0),
+            "output_tokens": getattr(usage, "output_tokens", 0),
+            "cache_read_tokens": getattr(usage, "cache_read_tokens", 0),
+            "cache_write_tokens": getattr(usage, "cache_write_tokens", 0),
+            "reasoning_tokens": getattr(usage, "reasoning_tokens", 0),
+        }
+    return {
+        "label": label,
+        "model": getattr(acct, "model", None),
+        "provider": getattr(acct, "provider", None),
+        "temperature": getattr(acct, "temperature", None),
+        "input_messages": getattr(acct, "messages", None),
+        "output": getattr(acct, "output", None),
+        "usage": usage_dict,
+        "cost_usd": getattr(acct, "cost_usd", None),
+        "cost_status": getattr(acct, "cost_status", None),
+        "cost_source": getattr(acct, "cost_source", None),
+    }
+
+
+def slot_metrics(acct: Any, label: str, output: Any = None) -> dict[str, Any]:
+    """Render one reference's accounting for observability hooks.
+
+    Same fields as ``_slot_trace`` minus ``input_messages``, which is the bulk
+    of a trace record and would cross the plugin-hook boundary for every
+    advisor on every turn. ``output`` comes from the caller because the
+    privacy-redacted advisor text lives alongside the accounting, not on it.
+    """
+    trace = _slot_trace(acct, label)
+    trace.pop("input_messages", None)
+    if output is not None:
+        trace["output"] = output
+    return trace
+
+
+def save_moa_turn(
+    *,
+    session_id: Optional[str],
+    preset_name: str,
+    reference_outputs: list[tuple[str, str, Any]],
+    aggregator_label: str,
+    aggregator_model: Optional[str],
+    aggregator_provider: Optional[str],
+    aggregator_temperature: Any,
+    aggregator_input_messages: Any,
+    aggregator_output: Optional[str],
+    aggregator_streamed: bool,
+) -> None:
+    """Append one full MoA turn record to the session's trace JSONL, if enabled.
+
+    Best-effort: any failure is logged at debug and swallowed — tracing must
+    never break a live turn. Called once per turn on a reference cache MISS.
+
+    ``aggregator_output`` is the aggregator's synthesized text. On the
+    non-streaming path (eval / quiet-mode / subagents) it was captured inline
+    at call time. On the streaming path it is captured after the fact from the
+    caller's resolved assistant text (``aggregator_output_fallback`` in
+    ``consume_and_save_trace``) so the trace is self-contained either way; if
+    that resolved text was unavailable, it falls back to None and the record
+    points at the session store via ``output_location``.
+    """
+    base = _traces_enabled_and_dir()
+    if base is None:
+        return
+    try:
+        base.mkdir(parents=True, exist_ok=True)
+        path = base / f"{_sanitize_session_id(session_id)}.jsonl"
+        # output_location tells an offline reader where the acting text lives:
+        # embedded here when we have it (both non-streaming inline capture and
+        # streaming after-the-fact capture), else the session-db assistant row.
+        _have_output = bool(aggregator_output)
+        if not aggregator_streamed:
+            _output_location = "inline"
+        elif _have_output:
+            _output_location = "inline_from_stream"
+        else:
+            _output_location = "assistant_message_in_session_db"
+        record = {
+            "ts": time.time(),
+            "session_id": session_id,
+            "preset": preset_name,
+            "references": [
+                _slot_trace(acct, label)
+                for label, _text, acct in reference_outputs
+            ],
+            "aggregator": {
+                "label": aggregator_label,
+                "model": aggregator_model,
+                "provider": aggregator_provider,
+                "temperature": aggregator_temperature,
+                "input_messages": aggregator_input_messages,
+                "output": aggregator_output,
+                "streamed": aggregator_streamed,
+                # Where the aggregator's acting output lives for this record.
+                # "inline"             — non-streaming inline capture
+                # "inline_from_stream" — streamed, then captured from the
+                #                        caller's resolved assistant text
+                # "assistant_message_in_session_db" — streamed and the resolved
+                #                        text was unavailable at flush time
+                "output_location": _output_location,
+            },
+        }
+        with path.open("a", encoding="utf-8") as f:
+            f.write(json.dumps(record, ensure_ascii=False, default=str) + "\n")
+    except Exception as exc:  # pragma: no cover - tracing must never break a turn
+        logger.debug("MoA trace write failed (session=%s): %s", session_id, exc)