Initial commit

skills/telemetry.capture/telemetry_utils.py

@@ -0,0 +1,200 @@
+#!/usr/bin/env python3
+"""
+Betty Framework - Telemetry Utilities
+Decorator and utilities for automatic execution logging
+"""
+
+import functools
+import time
+import inspect
+from typing import Any, Callable, Dict, Optional
+from pathlib import Path
+from .telemetry_capture import TelemetryCapture
+
+
+def capture_telemetry(
+    skill_name: Optional[str] = None,
+    caller: str = "runtime",
+    capture_inputs: bool = False,
+    sanitize_keys: Optional[list] = None
+):
+    """
+    Decorator to automatically capture telemetry for function execution.
+
+    Usage:
+        @capture_telemetry(skill_name="agent.run", caller="CLI", capture_inputs=True)
+        def run_agent(agent_path, task_context=None):
+            # ... implementation
+            return result
+
+    Args:
+        skill_name: Name of the skill/component. If None, uses function name.
+        caller: Source of the call (e.g., 'CLI', 'API', 'workflow')
+        capture_inputs: Whether to capture function arguments (default: False)
+        sanitize_keys: List of parameter names to exclude (e.g., ['password', 'api_key'])
+
+    The decorated function can:
+    - Return a dict with 'status' key to set telemetry status
+    - Raise exceptions (captured as 'error' status)
+    - Return any value (captured as 'success' status)
+    """
+    def decorator(func: Callable) -> Callable:
+        @functools.wraps(func)
+        def wrapper(*args, **kwargs):
+            # Determine skill name
+            skill = skill_name or f"{func.__module__}.{func.__name__}"
+
+            # Capture inputs if requested
+            inputs = {}
+            if capture_inputs:
+                # Get function signature
+                sig = inspect.signature(func)
+                bound_args = sig.bind_partial(*args, **kwargs)
+                bound_args.apply_defaults()
+
+                # Sanitize inputs
+                inputs = dict(bound_args.arguments)
+                if sanitize_keys:
+                    for key in sanitize_keys:
+                        if key in inputs:
+                            inputs[key] = "***REDACTED***"
+
+            # Start timing
+            start_time = time.time()
+            status = "success"
+            error_message = None
+            result = None
+
+            try:
+                # Execute function
+                result = func(*args, **kwargs)
+
+                # Check if result contains status
+                if isinstance(result, dict) and 'status' in result:
+                    status = result['status']
+
+                return result
+
+            except Exception as e:
+                status = "error"
+                error_message = str(e)
+                raise  # Re-raise the exception
+
+            finally:
+                # Calculate duration
+                duration_ms = (time.time() - start_time) * 1000
+
+                # Capture telemetry
+                try:
+                    telemetry = TelemetryCapture()
+                    telemetry.capture(
+                        skill=skill,
+                        status=status,
+                        duration_ms=duration_ms,
+                        caller=caller,
+                        inputs=inputs,
+                        error_message=error_message
+                    )
+                except Exception as telemetry_error:
+                    # Don't let telemetry errors break the actual function
+                    print(f"Warning: Telemetry capture failed: {telemetry_error}", file=sys.stderr)
+
+        return wrapper
+    return decorator
+
+
+class TelemetryContext:
+    """
+    Context manager for capturing telemetry around code blocks.
+
+    Usage:
+        with TelemetryContext(skill="plugin.build", caller="CLI") as ctx:
+            # ... do work
+            ctx.set_inputs({"plugin_path": "./plugin.yaml"})
+            result = build_plugin()
+            ctx.set_status("success")
+    """
+
+    def __init__(
+        self,
+        skill: str,
+        caller: str = "runtime",
+        inputs: Optional[Dict[str, Any]] = None
+    ):
+        self.skill = skill
+        self.caller = caller
+        self.inputs = inputs or {}
+        self.status = "success"
+        self.error_message = None
+        self.start_time = None
+        self.telemetry = TelemetryCapture()
+
+    def __enter__(self):
+        self.start_time = time.time()
+        return self
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        # Calculate duration
+        duration_ms = (time.time() - self.start_time) * 1000
+
+        # Set status based on exception
+        if exc_type is not None:
+            self.status = "error"
+            self.error_message = str(exc_val)
+
+        # Capture telemetry
+        try:
+            self.telemetry.capture(
+                skill=self.skill,
+                status=self.status,
+                duration_ms=duration_ms,
+                caller=self.caller,
+                inputs=self.inputs,
+                error_message=self.error_message
+            )
+        except Exception as telemetry_error:
+            print(f"Warning: Telemetry capture failed: {telemetry_error}", file=sys.stderr)
+
+        # Don't suppress exceptions
+        return False
+
+    def set_status(self, status: str):
+        """Update the execution status."""
+        self.status = status
+
+    def set_inputs(self, inputs: Dict[str, Any]):
+        """Update the captured inputs."""
+        self.inputs.update(inputs)
+
+    def set_error(self, error_message: str):
+        """Set error message."""
+        self.error_message = error_message
+        self.status = "error"
+
+
+# Example usage for testing
+if __name__ == "__main__":
+    # Test decorator
+    @capture_telemetry(skill_name="test.function", caller="CLI", capture_inputs=True)
+    def test_function(name: str, value: int = 10):
+        """Test function for telemetry capture."""
+        print(f"Executing test_function with name={name}, value={value}")
+        time.sleep(0.1)  # Simulate work
+        return {"status": "success", "result": value * 2}
+
+    # Test context manager
+    def test_context():
+        """Test context manager for telemetry capture."""
+        with TelemetryContext(skill="test.context", caller="CLI") as ctx:
+            ctx.set_inputs({"operation": "test"})
+            print("Executing within context")
+            time.sleep(0.05)  # Simulate work
+            ctx.set_status("success")
+
+    print("Testing decorator...")
+    result = test_function("example", value=42)
+    print(f"Result: {result}\n")
+
+    print("Testing context manager...")
+    test_context()
+    print("\nTelemetry tests complete!")