fit-boilerplate-wox/wxo/tools/python/_observable_tool.py

"""Observable tool decorator.

Wraps the ADK `@tool` so every call emits a structured trace
(inputs, outputs, latency, side effects). See `docs/observability-pattern.md`.

Usage:
    from wxo.tools.python._observable_tool import observable_tool

    @observable_tool(name="reset_password", domain="ad")
    def reset_password(username: str) -> dict:
        # your logic
        return {"temp_password": "xxx"}

The decorated function is still a valid ADK `@tool` — the decorator delegates.
The LLM sees only `result`. The full trace goes to the configured sink.

Sinks (controlled by env var TRACE_SINK):
- `sqlite` (default): writes to `traces.db` next to the process
- `http`: POST to $TRACE_SINK_URL (default http://web:8000/api/traces)
- `otlp`: OpenTelemetry OTLP gRPC to $OTEL_EXPORTER_OTLP_ENDPOINT

Set `_bypass_tracing=True` in the decorator to skip tracing for a specific
tool (rare — used for hot-path tools or debug).
"""
from __future__ import annotations

import functools
import json
import os
import sqlite3
import time
import uuid
from datetime import datetime, timezone
from typing import Any, Callable

# === inline compat shim (NO eliminar — TRM lo necesita) ===
try:
    from ibm_watsonx_orchestrate.agent_builder.tools import tool as _adk_tool
except ImportError:
    def _adk_tool(*args, **kwargs):
        def deco(f): return f
        return deco if not args else deco(args[0])
# ============================================================


TRACE_SINK = os.environ.get("TRACE_SINK", "sqlite")
TRACE_SINK_URL = os.environ.get("TRACE_SINK_URL", "http://web:8000/api/traces")
TRACES_DB_PATH = os.environ.get("TRACES_DB_PATH", "/tmp/traces.db")


def _now_iso() -> str:
    return datetime.now(timezone.utc).isoformat()


def _emit_trace(trace: dict[str, Any]) -> None:
    """Send the trace to the configured sink. Failures must NEVER break the tool."""
    try:
        if TRACE_SINK == "sqlite":
            _emit_sqlite(trace)
        elif TRACE_SINK == "http":
            _emit_http(trace)
        elif TRACE_SINK == "otlp":
            _emit_otlp(trace)
        # else: silently drop (TRACE_SINK=off)
    except Exception:
        # Never let observability break the tool.
        pass


def _emit_sqlite(trace: dict[str, Any]) -> None:
    conn = sqlite3.connect(TRACES_DB_PATH)
    conn.execute("""
        CREATE TABLE IF NOT EXISTS traces (
            trace_id TEXT PRIMARY KEY,
            tool TEXT,
            domain TEXT,
            agent_caller TEXT,
            correlation_id TEXT,
            started_at TEXT,
            duration_ms INTEGER,
            payload TEXT
        )
    """)
    conn.execute("CREATE INDEX IF NOT EXISTS ix_traces_started ON traces (started_at)")
    conn.execute(
        "INSERT OR REPLACE INTO traces VALUES (?, ?, ?, ?, ?, ?, ?, ?)",
        (
            trace["trace_id"], trace["tool"], trace.get("domain"),
            trace.get("agent_caller"), trace.get("correlation_id"),
            trace["started_at"], trace["duration_ms"],
            json.dumps(trace),
        ),
    )
    conn.commit()
    conn.close()


def _emit_http(trace: dict[str, Any]) -> None:
    import requests
    requests.post(TRACE_SINK_URL, json=trace, timeout=2)


def _emit_otlp(trace: dict[str, Any]) -> None:
    # Stub. Requires opentelemetry-sdk + exporter. Implement in v2.
    pass


def observable_tool(
    *,
    name: str,
    domain: str = "",
    description: str | None = None,
    _bypass_tracing: bool = False,
):
    """Decorator that wraps ADK `@tool` with structured tracing."""
    def decorator(func: Callable) -> Callable:
        # First wrap the function to capture traces.
        @functools.wraps(func)
        def wrapper(*args, **kwargs):
            if _bypass_tracing:
                return func(*args, **kwargs)

            trace_id = f"{domain or 'tool'}-{name}-{uuid.uuid4().hex[:6]}"
            started = time.perf_counter()
            started_iso = _now_iso()

            inputs = {}
            try:
                inputs = kwargs.copy()
                # positional → harder; we capture the dict for now
                if args:
                    inputs["_positional"] = [repr(a)[:200] for a in args]
            except Exception:
                pass

            result = None
            error = None
            try:
                result = func(*args, **kwargs)
            except Exception as e:
                error = repr(e)
                raise
            finally:
                duration_ms = int((time.perf_counter() - started) * 1000)
                trace = {
                    "trace_id": trace_id,
                    "tool": name,
                    "domain": domain,
                    "started_at": started_iso,
                    "duration_ms": duration_ms,
                    "inputs": inputs,
                    "result_preview": _preview(result),
                    "error": error,
                    "agent_caller": os.environ.get("WXO_CURRENT_AGENT", "unknown"),
                    "correlation_id": os.environ.get("WXO_CORRELATION_ID", ""),
                }
                _emit_trace(trace)
            return result

        # Then wrap with the ADK tool decorator so it's discoverable.
        kwargs = {"name": name}
        if description:
            kwargs["description"] = description
        return _adk_tool(**kwargs)(wrapper)
    return decorator


def _preview(value: Any, max_len: int = 500) -> Any:
    """Truncate big values for trace storage."""
    if value is None:
        return None
    try:
        s = json.dumps(value, default=str)
        if len(s) <= max_len:
            return json.loads(s)
        return {"_truncated": True, "preview": s[:max_len]}
    except Exception:
        return {"_repr": repr(value)[:max_len]}