201 lines
6.4 KiB
Python
Executable File
201 lines
6.4 KiB
Python
Executable File
#!/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!")
|