observability.py

"""
observability.py — OpenTelemetry Tracing & Monitoring
Provides distributed tracing, metrics, and logging for Manus-class agent
"""
from __future__ import annotations

import asyncio
import functools
import json
import logging
import os
import time
import traceback
import uuid
from contextvars import ContextVar
from dataclasses import dataclass, field
from typing import Any, Callable, Dict, List, Optional
from enum import Enum

logger = logging.getLogger("observability")

# ─── Configuration ─────────────────────────────────────────────────────────────
OTEL_ENABLED = os.environ.get("OTEL_ENABLED", "false").lower() == "true"
OTEL_SERVICE_NAME = os.environ.get("OTEL_SERVICE_NAME", "onehands-agent")
OTEL_EXPORTER_ENDPOINT = os.environ.get("OTEL_EXPORTER_ENDPOINT", "http://localhost:4317")

# ─── Trace Context ──────────────────────────────────────────────────────────────
trace_id_var: ContextVar[str] = ContextVar("trace_id", default="")
span_id_var: ContextVar[str] = ContextVar("span_id", default="")


def generate_trace_id() -> str:
    """Generate a new trace ID."""
    return uuid.uuid4().hex[:32]


def generate_span_id() -> str:
    """Generate a new span ID."""
    return uuid.uuid4().hex[:16]


class SpanKind(Enum):
    """Kind of span."""
    INTERNAL = "internal"
    SERVER = "server"
    CLIENT = "client"
    PRODUCER = "producer"
    CONSUMER = "consumer"


class SpanStatus(Enum):
    """Span status."""
    UNSET = "unset"
    OK = "ok"
    ERROR = "error"


# ─── Span Model ────────────────────────────────────────────────────────────────
@dataclass
class Span:
    """Represents a trace span."""
    trace_id: str
    span_id: str
    name: str
    kind: SpanKind
    start_time: float
    end_time: Optional[float] = None
    status: SpanStatus = SpanStatus.UNSET
    error_message: Optional[str] = None
    attributes: Dict[str, Any] = field(default_factory=dict)
    events: List[Dict] = field(default_factory=list)
    links: List[Dict] = field(default_factory=list)
    parent_span_id: Optional[str] = None
    
    @property
    def duration_ms(self) -> Optional[float]:
        """Calculate span duration in milliseconds."""
        if self.end_time:
            return (self.end_time - self.start_time) * 1000
        return None
    
    def to_dict(self) -> Dict:
        """Convert span to dictionary for export."""
        return {
            "trace_id": self.trace_id,
            "span_id": self.span_id,
            "name": self.name,
            "kind": self.kind.value,
            "start_time": self.start_time,
            "end_time": self.end_time,
            "duration_ms": self.duration_ms,
            "status": self.status.value,
            "error_message": self.error_message,
            "attributes": self.attributes,
            "events": self.events,
            "links": self.links,
            "parent_span_id": self.parent_span_id,
        }


# ─── Tracer ───────────────────────────────────────────────────────────────────
class Tracer:
    """
    OpenTelemetry-compatible tracer for distributed tracing.
    
    Features:
    - Span creation and management
    - Context propagation
    - Event logging
    - Attribute tracking
    - Error recording
    """
    
    def __init__(self, service_name: str = OTEL_SERVICE_NAME):
        self.service_name = service_name
        self._spans: Dict[str, Span] = {}
        self._current_span: Optional[Span] = None
        self._exporter = None
        
        if OTEL_ENABLED:
            self._init_exporter()
    
    def _init_exporter(self):
        """Initialize OTLP exporter."""
        try:
            # Try to use OpenTelemetry SDK
            from opentelemetry import trace
            from opentelemetry.sdk.trace import TracerProvider
            from opentelemetry.sdk.trace.export import BatchSpanProcessor
            from opentelemetry.exporter.otlp.proto.grpc.trace_exporter import OTLPSpanExporter
            
            provider = TracerProvider()
            exporter = OTLPSpanExporter(endpoint=OTEL_EXPORTER_ENDPOINT)
            provider.add_span_processor(BatchSpanProcessor(exporter))
            trace.set_tracer_provider(provider)
            
            self._otel_tracer = trace.get_tracer(self.service_name)
            logger.info(f"OpenTelemetry initialized: {OTEL_EXPORTER_ENDPOINT}")
        except ImportError:
            logger.warning("OpenTelemetry SDK not installed, using simple tracer")
            self._otel_tracer = None
        except Exception as e:
            logger.warning(f"Failed to initialize OTLP exporter: {e}")
            self._otel_tracer = None
    
    def start_span(
        self,
        name: str,
        kind: SpanKind = SpanKind.INTERNAL,
        attributes: Dict[str, Any] = None,
        parent_trace_id: str = None,
        parent_span_id: str = None,
    ) -> Span:
        """Start a new span."""
        trace_id = parent_trace_id or trace_id_var.get() or generate_trace_id()
        parent = parent_span_id or span_id_var.get()
        
        span = Span(
            trace_id=trace_id,
            span_id=generate_span_id(),
            name=name,
            kind=kind,
            start_time=time.time(),
            attributes=attributes or {},
            parent_span_id=parent,
        )
        
        self._spans[span.span_id] = span
        
        # Set context variables
        trace_id_var.set(trace_id)
        span_id_var.set(span.span_id)
        
        # Use OTEL tracer if available
        if hasattr(self, "_otel_tracer") and self._otel_tracer:
            self._otel_span = self._otel_tracer.start_span(name)
        
        return span
    
    def end_span(
        self,
        span: Span,
        status: SpanStatus = SpanStatus.OK,
        error: Exception = None,
    ):
        """End a span."""
        span.end_time = time.time()
        span.status = status
        
        if error:
            span.status = SpanStatus.ERROR
            span.error_message = str(error)
            span.attributes["error.type"] = type(error).__name__
            span.attributes["error.message"] = str(error)
            span.attributes["error.stacktrace"] = traceback.format_exc()
        
        # End OTEL span if available
        if hasattr(self, "_otel_span"):
            try:
                self._otel_span.__exit__(None, None, None)
            except Exception:
                pass
        
        return span
    
    def add_event(self, span: Span, name: str, attributes: Dict[str, Any] = None):
        """Add event to span."""
        span.events.append({
            "name": name,
            "timestamp": time.time(),
            "attributes": attributes or {},
        })
    
    def set_attribute(self, span: Span, key: str, value: Any):
        """Set attribute on span."""
        span.attributes[key] = value
    
    def get_current_span(self) -> Optional[Span]:
        """Get current active span."""
        span_id = span_id_var.get()
        return self._spans.get(span_id)
    
    def get_trace(self, trace_id: str) -> List[Span]:
        """Get all spans for a trace."""
        return [s for s in self._spans.values() if s.trace_id == trace_id]
    
    def export_traces(self) -> List[Dict]:
        """Export all completed spans."""
        return [s.to_dict() for s in self._spans.values() if s.end_time]
    
    def clear(self):
        """Clear all spans (for testing)."""
        self._spans.clear()


# ─── Metrics ──────────────────────────────────────────────────────────────────
@dataclass
class Metric:
    """Represents a metric data point."""
    name: str
    value: float
    timestamp: float
    labels: Dict[str, str] = field(default_factory=dict)
    metric_type: str = "gauge"  # gauge, counter, histogram


class MetricsCollector:
    """
    Simple metrics collector for monitoring.
    
    Features:
    - Counters
    - Gauges
    - Histograms
    - Labels/support
    """
    
    def __init__(self):
        self._counters: Dict[str, float] = {}
        self._gauges: Dict[str, float] = {}
        self._histograms: Dict[str, List[float]] = {}
        self._metrics: List[Metric] = []
    
    def inc_counter(self, name: str, value: float = 1, labels: Dict[str, str] = None):
        """Increment a counter."""
        key = self._make_key(name, labels)
        self._counters[key] = self._counters.get(key, 0) + value
        self._record_metric(name, self._counters[key], "counter", labels)
    
    def set_gauge(self, name: str, value: float, labels: Dict[str, str] = None):
        """Set a gauge value."""
        key = self._make_key(name, labels)
        self._gauges[key] = value
        self._record_metric(name, value, "gauge", labels)
    
    def record_histogram(self, name: str, value: float, labels: Dict[str, str] = None):
        """Record a histogram value."""
        key = self._make_key(name, labels)
        if key not in self._histograms:
            self._histograms[key] = []
        self._histograms[key].append(value)
        self._record_metric(name, value, "histogram", labels)
    
    def _make_key(self, name: str, labels: Dict[str, str] = None) -> str:
        """Create metric key from name and labels."""
        if not labels:
            return name
        label_str = ",".join(f"{k}={v}" for k, v in sorted(labels.items()))
        return f"{name}{{{label_str}}}"
    
    def _record_metric(self, name: str, value: float, metric_type: str, labels: Dict[str, str] = None):
        """Record metric."""
        self._metrics.append(Metric(
            name=name,
            value=value,
            timestamp=time.time(),
            labels=labels or {},
            metric_type=metric_type,
        ))
        # Keep only last 10000 metrics
        if len(self._metrics) > 10000:
            self._metrics = self._metrics[-5000:]
    
    def get_metrics(self, name: str = None) -> List[Dict]:
        """Get current metric values."""
        result = []
        
        if name:
            # Specific metric
            for key, value in self._counters.items():
                if key.startswith(name):
                    result.append({"name": key, "value": value, "type": "counter"})
            for key, value in self._gauges.items():
                if key.startswith(name):
                    result.append({"name": key, "value": value, "type": "gauge"})
        else:
            # All metrics
            for key, value in self._counters.items():
                result.append({"name": key, "value": value, "type": "counter"})
            for key, value in self._gauges.items():
                result.append({"name": key, "value": value, "type": "gauge"})
        
        return result
    
    def get_prometheus_format(self) -> str:
        """Export metrics in Prometheus format."""
        lines = []
        
        for key, value in self._counters.items():
            lines.append(f"# TYPE {key} counter")
            lines.append(f"{key} {value}")
        
        for key, value in self._gauges.items():
            lines.append(f"# TYPE {key} gauge")
            lines.append(f"{key} {value}")
        
        return "\n".join(lines)


# ─── Global instances ──────────────────────────────────────────────────────────
tracer = Tracer()
metrics = MetricsCollector()


# ─── Decorators ────────────────────────────────────────────────────────────────
def traced(
    name: str = None,
    kind: SpanKind = SpanKind.INTERNAL,
    attributes: Dict[str, Any] = None,
):
    """Decorator to trace a function."""
    def decorator(func: Callable):
        @functools.wraps(func)
        async def async_wrapper(*args, **kwargs):
            span_name = name or f"{func.__module__}.{func.__name__}"
            
            span = tracer.start_span(span_name, kind, attributes)
            start_time = time.time()
            
            try:
                result = await func(*args, **kwargs)
                tracer.end_span(span, SpanStatus.OK)
                
                # Record metrics
                duration_ms = (time.time() - start_time) * 1000
                metrics.record_histogram("function.duration_ms", duration_ms, {"function": span_name})
                metrics.inc_counter("function.calls", 1, {"function": span_name, "status": "success"})
                
                return result
            except Exception as e:
                tracer.end_span(span, SpanStatus.ERROR, e)
                
                duration_ms = (time.time() - start_time) * 1000
                metrics.inc_counter("function.calls", 1, {"function": span_name, "status": "error"})
                metrics.inc_counter("function.errors", 1, {"function": span_name, "error_type": type(e).__name__})
                
                raise
        
        @functools.wraps(func)
        def sync_wrapper(*args, **kwargs):
            span_name = name or f"{func.__module__}.{func.__name__}"
            
            span = tracer.start_span(span_name, kind, attributes)
            start_time = time.time()
            
            try:
                result = func(*args, **kwargs)
                tracer.end_span(span, SpanStatus.OK)
                
                duration_ms = (time.time() - start_time) * 1000
                metrics.record_histogram("function.duration_ms", duration_ms, {"function": span_name})
                metrics.inc_counter("function.calls", 1, {"function": span_name, "status": "success"})
                
                return result
            except Exception as e:
                tracer.end_span(span, SpanStatus.ERROR, e)
                
                duration_ms = (time.time() - start_time) * 1000
                metrics.inc_counter("function.calls", 1, {"function": span_name, "status": "error"})
                metrics.inc_counter("function.errors", 1, {"function": span_name, "error_type": type(e).__name__})
                
                raise
        
        if asyncio.iscoroutinefunction(func):
            return async_wrapper
        return sync_wrapper
    
    return decorator


# ─── Agent-specific instrumentation ──────────────────────────────────────────
class AgentTracer:
    """
    Specialized tracer for autonomous agent operations.
    """
    
    def __init__(self, agent_id: str = None):
        self.agent_id = agent_id or str(uuid.uuid4())[:8]
        self.task_spans: Dict[str, Span] = {}
        self.step_traces: List[Dict] = []
    
    async def start_task(self, task_id: str, task_description: str) -> Span:
        """Start tracing a new agent task."""
        span = tracer.start_span(
            name=f"agent.task.{task_id[:8]}",
            kind=SpanKind.INTERNAL,
            attributes={
                "agent.id": self.agent_id,
                "task.id": task_id,
                "task.description": task_description,
            },
        )
        self.task_spans[task_id] = span
        return span
    
    async def end_task(self, task_id: str, status: str, result: Any = None):
        """End tracing a task."""
        span = self.task_spans.get(task_id)
        if span:
            span.attributes["task.status"] = status
            if result:
                span.attributes["task.result_summary"] = str(result)[:200]
            tracer.end_span(span, SpanStatus.OK if status == "completed" else SpanStatus.ERROR)
            del self.task_spans[task_id]
    
    async def trace_step(
        self,
        task_id: str,
        step_number: int,
        thought: str,
        action: str,
        observation: str = None,
    ):
        """Trace a single agent step."""
        span = tracer.start_span(
            name=f"agent.step.{step_number}",
            kind=SpanKind.INTERNAL,
            attributes={
                "agent.id": self.agent_id,
                "task.id": task_id,
                "step.number": step_number,
                "step.thought": thought[:500],
                "step.action": action,
                "step.observation": observation[:500] if observation else None,
            },
        )
        tracer.end_span(span)
        
        self.step_traces.append({
            "task_id": task_id,
            "step": step_number,
            "thought": thought,
            "action": action,
            "observation": observation,
            "timestamp": time.time(),
        })
    
    def get_task_trace(self, task_id: str) -> List[Dict]:
        """Get complete trace for a task."""
        return [s for s in self.step_traces if s["task_id"] == task_id]
    
    def get_all_traces(self) -> Dict[str, List[Dict]]:
        """Get all traces grouped by task."""
        result = {}
        for trace in self.step_traces:
            task_id = trace["task_id"]
            if task_id not in result:
                result[task_id] = []
            result[task_id].append(trace)
        return result


# ─── Observability API Routes ──────────────────────────────────────────────────
def register_observability_routes(app):
    """Register observability routes with FastAPI app."""
    from fastapi import APIRouter, Response
    
    router = APIRouter(prefix="/observability", tags=["observability"])
    
    @router.get("/traces")
    async def get_traces(trace_id: str = None, limit: int = 100):
        """Get trace data."""
        if trace_id:
            return {"traces": tracer.get_trace(trace_id)}
        spans = tracer.export_traces()
        return {"traces": spans[-limit:]}
    
    @router.get("/metrics")
    async def get_metrics(name: str = None):
        """Get metric values."""
        return {"metrics": metrics.get_metrics(name)}
    
    @router.get("/metrics/prometheus")
    async def get_prometheus_metrics():
        """Get metrics in Prometheus format."""
        return Response(
            content=metrics.get_prometheus_format(),
            media_type="text/plain",
        )
    
    @router.get("/health")
    async def health_check():
        """Health check with basic metrics."""
        return {
            "status": "healthy",
            "spans_active": len(tracer._spans),
            "metrics_collected": len(metrics._metrics),
            "service": OTEL_SERVICE_NAME,
            "otel_enabled": OTEL_ENABLED,
        }
    
    app.include_router(router)