aco/trackio_integration.py

"""Trackio integration for live ACO monitoring.

Logs optimization decisions, costs, and quality metrics to Trackio dashboard.
"""

import os
from typing import Dict, Any, Optional


class ACOTrackioLogger:
    """Lightweight Trackio logger for Agent Cost Optimizer runs."""

    def __init__(self, project: str = "aco", space_id: Optional[str] = None):
        self.project = project
        self.space_id = space_id or os.environ.get("TRACKIO_SPACE_ID")
        self._enabled = False
        
        try:
            import trackio
            self.api = trackio.Api()
            self._enabled = True
        except ImportError:
            self.api = None

    def log_decision(self, run_id: str, decision: Any, cost: float, success: bool) -> None:
        """Log an optimization decision to Trackio."""
        if not self._enabled:
            return
        
        try:
            self.api.log_metric(
                project=self.project,
                run=run_id,
                key="optimization_cost",
                value=cost,
            )
            self.api.log_metric(
                project=self.project,
                run=run_id,
                key="task_success",
                value=1.0 if success else 0.0,
            )
            if hasattr(decision, "routing_decision"):
                self.api.log_param(
                    project=self.project,
                    run=run_id,
                    key="model_tier",
                    value=decision.routing_decision.tier,
                )
        except Exception:
            pass

    def alert(self, run_id: str, title: str, text: str, level: str = "INFO") -> None:
        """Send an alert to Trackio."""
        if not self._enabled:
            return
        
        try:
            self.api.alert(
                project=self.project,
                run=run_id,
                title=title,
                text=text,
                level=level,
            )
        except Exception:
            pass

    def log_cost_adjusted_score(self, run_id: str, score: float) -> None:
        if not self._enabled:
            return
        try:
            self.api.log_metric(
                project=self.project,
                run=run_id,
                key="cost_adjusted_score",
                value=score,
            )
        except Exception:
            pass