cleveragents-core/robot/helper_llm_trace.py

"""Helper script for LLM trace Robot Framework tests.

Usage:
    python helper_llm_trace.py record-and-retrieve
    python helper_llm_trace.py list-by-plan
    python helper_llm_trace.py compute-metrics
    python helper_llm_trace.py metric-key-count
    python helper_llm_trace.py validation
    python helper_llm_trace.py lifecycle-hooks
"""

from __future__ import annotations

import sys
from pathlib import Path
from typing import Any
from unittest.mock import MagicMock

# Ensure source tree is importable — force the local checkout to win
# over any inherited PYTHONPATH (e.g. /app/src in CI containers).
_SRC = str(Path(__file__).resolve().parents[1] / "src")
sys.path.insert(0, _SRC)

from cleveragents.application.services.trace_service import TraceService  # noqa: E402
from cleveragents.domain.models.observability.llm_trace import LLMTrace  # noqa: E402
from cleveragents.domain.models.observability.metrics import (  # noqa: E402
    OperationalMetricKey,
)

VALID_ULID_1 = "01HXAAAAAAAAAAAAAAAAAAAAAA"
VALID_ULID_2 = "01HXBBBBBBBBBBBBBBBBBBBBBB"
VALID_ULID_3 = "01HXCCCCCCCCCCCCCCCCCCCCCC"
VALID_PLAN_ID = "01HX0000000000PPPPPPPPPPPP"


class InMemoryRepo:
    """Minimal in-memory repo for testing."""

    def __init__(self) -> None:
        self._store: dict[str, LLMTrace] = {}

    def save(self, trace: LLMTrace) -> None:
        self._store[trace.trace_id] = trace

    def get(self, trace_id: str) -> LLMTrace | None:
        return self._store.get(trace_id)

    def list_by_plan(self, plan_id: str) -> list[LLMTrace]:
        return [t for t in self._store.values() if t.plan_id == plan_id]

    def list_by_decision(self, decision_id: str) -> list[LLMTrace]:
        return [t for t in self._store.values() if t.decision_id == decision_id]


def _make_trace(trace_id: str = VALID_ULID_1, **kwargs: Any) -> LLMTrace:
    defaults: dict[str, Any] = {
        "trace_id": trace_id,
        "plan_id": VALID_PLAN_ID,
        "actor": "planner",
        "provider": "openai",
        "model": "gpt-4o",
        "prompt_tokens": 500,
        "completion_tokens": 200,
        "cost_usd": 0.01,
        "latency_ms": 350.0,
    }
    defaults.update(kwargs)
    return LLMTrace(**defaults)


def _make_service() -> TraceService:
    repo = InMemoryRepo()
    return TraceService(settings=MagicMock(), repository=repo)  # type: ignore[arg-type]


def cmd_record_and_retrieve() -> None:
    """Record a trace and retrieve by ID."""
    svc = _make_service()
    trace = _make_trace()
    svc.record_trace(trace)
    result = svc.get_trace(trace.trace_id)
    if result is None:
        print("FAIL: trace not found after recording")
        sys.exit(1)
    if result.trace_id != trace.trace_id:
        print(f"FAIL: trace_id mismatch {result.trace_id} != {trace.trace_id}")
        sys.exit(1)
    print("llm-trace-record-ok")


def cmd_list_by_plan() -> None:
    """Record multiple traces and list by plan."""
    svc = _make_service()
    for uid in (VALID_ULID_1, VALID_ULID_2, VALID_ULID_3):
        svc.record_trace(_make_trace(trace_id=uid))
    results = svc.get_traces(VALID_PLAN_ID)
    if len(results) != 3:
        print(f"FAIL: expected 3 traces, got {len(results)}")
        sys.exit(1)
    print("llm-trace-list-plan-ok")


def cmd_compute_metrics() -> None:
    """Compute metrics from traces."""
    svc = _make_service()
    for uid in (VALID_ULID_1, VALID_ULID_2, VALID_ULID_3):
        svc.record_trace(_make_trace(trace_id=uid))
    metrics = svc.compute_metrics(VALID_PLAN_ID)
    if len(metrics) != 4:
        print(f"FAIL: expected 4 metrics, got {len(metrics)}")
        sys.exit(1)
    keys = {m.key for m in metrics}
    expected = {
        OperationalMetricKey.LLM_CALL_COUNT,
        OperationalMetricKey.LLM_TOTAL_TOKENS,
        OperationalMetricKey.LLM_TOTAL_COST_USD,
        OperationalMetricKey.LLM_AVG_LATENCY_MS,
    }
    if keys != expected:
        print(f"FAIL: unexpected metric keys {keys}")
        sys.exit(1)
    print("llm-trace-metrics-ok")


def cmd_metric_key_count() -> None:
    """Verify all 14 metric keys exist."""
    if len(OperationalMetricKey) != 14:
        print(f"FAIL: expected 14 keys, got {len(OperationalMetricKey)}")
        sys.exit(1)
    print("llm-trace-metric-keys-ok")


def cmd_validation() -> None:
    """Verify model validation rejects bad inputs."""
    from pydantic import ValidationError

    # Invalid trace_id
    try:
        _make_trace(trace_id="INVALID")
        print("FAIL: invalid trace_id accepted")
        sys.exit(1)
    except ValidationError:
        pass

    # Negative prompt tokens
    try:
        _make_trace(prompt_tokens=-1)
        print("FAIL: negative prompt_tokens accepted")
        sys.exit(1)
    except ValidationError:
        pass

    # Negative cost
    try:
        _make_trace(cost_usd=-0.5)
        print("FAIL: negative cost_usd accepted")
        sys.exit(1)
    except ValidationError:
        pass

    print("llm-trace-validation-ok")


def cmd_lifecycle_hooks() -> None:
    """Verify lifecycle hooks produce correct metrics."""
    svc = _make_service()

    # on_plan_start
    entry = svc.on_plan_start(VALID_PLAN_ID)
    if entry is None or entry.key != OperationalMetricKey.PLAN_DECISION_COUNT:
        print("FAIL: on_plan_start incorrect")
        sys.exit(1)

    # on_actor_invocation
    entries = svc.on_actor_invocation(VALID_PLAN_ID, "planner", 250.0)
    keys = {m.key for m in entries}
    if OperationalMetricKey.ACTOR_INVOCATION_COUNT not in keys:
        print("FAIL: on_actor_invocation missing ACTOR_INVOCATION_COUNT")
        sys.exit(1)

    # on_tool_execution
    entries = svc.on_tool_execution(VALID_PLAN_ID, "file_write", errored=True)
    keys = {m.key for m in entries}
    if OperationalMetricKey.TOOL_ERROR_RATE not in keys:
        print("FAIL: on_tool_execution missing TOOL_ERROR_RATE")
        sys.exit(1)

    print("llm-trace-hooks-ok")


def main() -> None:
    """Dispatch subcommand."""
    if len(sys.argv) < 2:
        print("Usage: helper_llm_trace.py <command>")
        sys.exit(1)

    cmd = sys.argv[1]
    commands = {
        "record-and-retrieve": cmd_record_and_retrieve,
        "list-by-plan": cmd_list_by_plan,
        "compute-metrics": cmd_compute_metrics,
        "metric-key-count": cmd_metric_key_count,
        "validation": cmd_validation,
        "lifecycle-hooks": cmd_lifecycle_hooks,
    }
    fn = commands.get(cmd)
    if fn is None:
        print(f"Unknown command: {cmd}")
        sys.exit(1)
    fn()


if __name__ == "__main__":
    main()