aeonic-agentguard-sdk-python 0.1.0__py3-none-any.whl

aeonic/core/assessment/executor.py

@@ -0,0 +1,207 @@
+"""Assessment execution against local agent endpoints."""
+
+import json
+import logging
+import os
+import time
+import urllib.request
+from typing import Any, Dict, List
+
+logger = logging.getLogger("aeonic")
+
+
+def _get_endpoint() -> str:
+    """Get the Aeonic endpoint from environment variable."""
+    return os.environ.get("AEONIC_ENDPOINT", "http://45.79.111.106:3291")
+
+
+def execute_assessment(job: Dict[str, Any]) -> List[Dict[str, Any]]:
+    """
+    Execute an assessment job and return results.
+
+    Args:
+        job: Assessment job containing scenario_steps, agent_endpoint, and playbook.
+
+    Returns:
+        List of assessment results for each step.
+    """
+    results: List[Dict[str, Any]] = []
+    start_time = time.time()
+
+    scenario_steps = job.get("scenario_steps", [])
+    agent_endpoint = job.get("agent_endpoint", {})
+    playbook = job.get("playbook", {})
+    timeout = playbook.get("timeout", 30)
+
+    logger.info(
+        f"[Aeonic] Executing assessment job {job.get('job_id')} with {len(scenario_steps)} steps"
+    )
+
+    for step in scenario_steps:
+        step_start_time = time.time()
+
+        try:
+            response = _call_agent_endpoint(
+                agent_endpoint,
+                step.get("input"),
+                timeout,
+            )
+
+            results.append(
+                {
+                    "step": step.get("step"),
+                    "input": step.get("input"),
+                    "output": {
+                        "body": response.get("data"),
+                        "status": response.get("status"),
+                        "headers": response.get("headers", {}),
+                    },
+                    "timestamp": int(time.time() * 1000),
+                    "status": "success",
+                    "latency_ms": int((time.time() - step_start_time) * 1000),
+                }
+            )
+        except Exception as error:
+            logger.warning(f"[Aeonic] Assessment step {step.get('step')} failed: {error}")
+            results.append(
+                {
+                    "step": step.get("step"),
+                    "input": step.get("input"),
+                    "error": str(error),
+                    "timestamp": int(time.time() * 1000),
+                    "status": "error",
+                    "latency_ms": int((time.time() - step_start_time) * 1000),
+                }
+            )
+
+    # Submit results to T-LAIOR
+    execution_time = int((time.time() - start_time) * 1000)
+    _submit_assessment_results(job, results, execution_time)
+
+    return results
+
+
+def _call_agent_endpoint(
+    endpoint: Dict[str, Any],
+    input_data: Any,
+    timeout: int,
+) -> Dict[str, Any]:
+    """
+    Call customer's agent endpoint locally.
+
+    Args:
+        endpoint: Agent endpoint configuration (url, method, headers, auth).
+        input_data: Input data for the agent.
+        timeout: Request timeout in seconds.
+
+    Returns:
+        Response dictionary with data, status, and headers.
+    """
+    api_url = endpoint.get("url")
+    if not api_url:
+        raise ValueError("Agent endpoint URL is required")
+
+    method = endpoint.get("method", "POST")
+    headers: Dict[str, str] = {
+        "Content-Type": "application/json",
+        **endpoint.get("headers", {}),
+    }
+
+    # Add authentication
+    auth_type = endpoint.get("auth_type")
+    if auth_type == "bearer_token" and endpoint.get("token"):
+        headers["Authorization"] = f"Bearer {endpoint['token']}"
+    elif auth_type == "api_key" and endpoint.get("api_key"):
+        headers["X-API-Key"] = endpoint["api_key"]
+
+    data = json.dumps(input_data).encode("utf-8")
+    port = os.environ.get("PORT") or os.environ.get("APP_PORT") or "8000"
+
+    url = f"http://localhost:{port}{api_url}"
+    print("URL=================>", url)
+
+    req = urllib.request.Request(url, data=data, headers=headers, method=method)
+
+    with urllib.request.urlopen(req, timeout=timeout) as response:
+        response_body = response.read()
+        try:
+            response_data = json.loads(response_body)
+        except json.JSONDecodeError:
+            # If response is not JSON, return as string
+            response_data = response_body.decode("utf-8")
+
+        return {
+            "data": response_data,
+            "status": response.getcode(),
+            "headers": dict(response.headers),
+        }
+
+
+def _submit_assessment_results(
+    job: Dict[str, Any],
+    results: List[Dict[str, Any]],
+    execution_time: int,
+) -> None:
+    """
+    Submit assessment results to T-LAIOR.
+
+    Args:
+        job: Original assessment job.
+        results: List of assessment results.
+        execution_time: Total execution time in milliseconds.
+    """
+    from aeonic.core.agentguard.registration import get_agent_guard_credentials
+
+    credentials = get_agent_guard_credentials()
+    if not credentials:
+        logger.error("[Aeonic] Cannot submit results: AgentGuard not registered")
+        return
+
+    endpoint = _get_endpoint()
+    url = f"{endpoint}/assessments/results"
+
+    payload = {
+        "job_id": job["job_id"],
+        "agentguard_id": credentials["agentguard_id"],
+        "agent_id": job.get("agent_id"),
+        "scenario_id": job.get("scenario_id"),
+        "results": results,
+        "metadata": {
+            "execution_time_ms": execution_time,
+            "total_steps": len(results),
+            "successful_steps": sum(1 for r in results if r.get("status") == "success"),
+            "failed_steps": sum(1 for r in results if r.get("status") == "error"),
+        },
+    }
+
+    print("PAYLOAD=================>", payload)
+    print("URL=================>", url)
+    print("CREDENTIALS=================>", credentials)
+
+    try:
+        data = json.dumps(payload).encode("utf-8")
+        req = urllib.request.Request(
+            url,
+            data=data,
+            headers={
+                "Authorization": f"Bearer {credentials['auth_token']}",
+                "Content-Type": "application/json",
+            },
+            method="POST",
+        )
+
+        with urllib.request.urlopen(req, timeout=30) as response:
+            response_data = json.loads(response.read())
+            if response_data.get("success"):
+                logger.info(
+                    f"[Aeonic] Assessment results submitted successfully "
+                    f"for job {job.get('job_id')}"
+                )
+            else:
+                logger.warning(
+                    f"[Aeonic] Assessment results submission failed: {response_data.get('message')}"
+                )
+    except Exception as error:
+        logger.error(
+            f"[Aeonic] Failed to submit assessment results for job {job.get('job_id')}: {error}"
+        )

aeonic/core/buffer.py

@@ -0,0 +1,27 @@
+from typing import List
+
+from aeonic.core.types import AgentContext, RouteContext, Sample
+
+
+class SampleBuffer:
+    def __init__(
+        self,
+        agent_key: str,
+        agent_context: AgentContext,
+        route: RouteContext,
+        max_samples: int,
+    ):
+        self.agent_key = agent_key
+        self.agent_context = agent_context
+        self.route = route
+        self.max_samples = max_samples
+        self._samples: List[Sample] = []
+
+    def add(self, sample: Sample) -> bool:
+        self._samples.append(sample)
+        return len(self._samples) >= self.max_samples
+
+    def flush(self) -> List[Sample]:
+        samples = self._samples
+        self._samples = []
+        return samples

aeonic/core/init.py

@@ -0,0 +1,218 @@
+# aeonic/core/init.py
+
+import logging
+import threading
+import time
+from typing import Any, Dict, Optional
+
+logger = logging.getLogger("aeonic")
+
+_CONFIG: Optional[Dict[str, Any]] = None
+
+
+def init(config: Dict[str, Any]) -> None:
+    """
+    Initialize Aeonic SDK.
+
+    Required:
+      - api_key
+
+    Optional:
+      - app: FastAPI or Django app instance (for route introspection)
+      - service_name: Service identifier
+      - capture_payloads.max_samples: Number of samples to buffer
+      - introspection_delay_ms: Delay before route introspection (default: 2000ms)
+    """
+    global _CONFIG
+
+    if "api_key" not in config:
+        raise ValueError("Aeonic init requires `api_key`")
+
+    _CONFIG = {
+        "api_key": config["api_key"],
+        "service_name": config.get("service_name"),
+        "max_samples": (config.get("capture_payloads", {}) or {}).get("max_samples", 10),
+        "app": config.get("app"),
+        "introspection_delay_ms": config.get("introspection_delay_ms", 2000),
+        "agentguard_enabled": config.get("agentguard_enabled", True),
+    }
+
+    logger.info("[Aeonic] SDK initialized.")
+
+    # Register as AgentGuard (non-blocking, optional)
+    if _CONFIG.get("agentguard_enabled", True):
+        _register_agentguard_async()
+
+    # Schedule deferred route introspection
+    # For Django, we can introspect even without app (uses get_resolver())
+    # For FastAPI, app is required
+    delay_ms = _CONFIG.get("introspection_delay_ms", 2000)
+
+    # Check if we should run introspection
+    app = _CONFIG.get("app")
+    should_introspect = app is not None
+
+    # For Django, try to detect and run introspection even without app
+    if app is None:
+        try:
+            import django  # noqa: F401
+
+            # Django is installed, try to introspect
+            should_introspect = True
+        except ImportError:
+            pass
+
+    if should_introspect:
+        _schedule_introspection(delay_ms)
+    else:
+        # Fallback: emit inventory if agents already registered (legacy behavior)
+        _emit_inventory_if_ready()
+
+
+def get_config() -> Dict[str, Any]:
+    if _CONFIG is None:
+        raise RuntimeError("Aeonic SDK not initialized. Call init() before using the SDK.")
+    return _CONFIG
+
+
+def _schedule_introspection(delay_ms: int) -> None:
+    """
+    Schedule deferred route introspection in a background thread.
+
+    Args:
+        delay_ms: Delay in milliseconds before introspection runs
+    """
+
+    def deferred_introspection() -> None:
+        try:
+            # Wait for routes to be defined
+            time.sleep(delay_ms / 1000.0)
+
+            app = _CONFIG.get("app") if _CONFIG else None
+
+            # For Django, we can introspect even without app (uses get_resolver())
+            # For FastAPI, app is required
+            if app is None:
+                # Try Django introspection
+                try:
+                    import django  # noqa: F401
+
+                    from aeonic.core.introspection import introspect_routes
+
+                    # Pass None for Django - introspection will use get_resolver()
+                    introspect_routes(None)
+                except ImportError:
+                    logger.error(
+                        "[Aeonic] No app provided and Django not detected - skipping route introspection"
+                    )
+                    return
+                except Exception as e:
+                    logger.error(f"[Aeonic] Django introspection failed: {e}")
+                    _emit_inventory_if_ready()
+                    return
+            else:
+                # Perform introspection with app instance
+                from aeonic.core.introspection import introspect_routes
+
+                introspect_routes(app)
+
+            # Emit inventory after introspection
+            _emit_inventory_if_ready()
+
+        except Exception as e:
+            logger.warning(f"[Aeonic] Route introspection failed: {e}")
+            import traceback
+
+            logger.error(f"[Aeonic] Traceback: {traceback.format_exc()}")
+            # Fallback: emit inventory with whatever agents were registered manually
+            _emit_inventory_if_ready()
+
+    # Run in background daemon thread (non-blocking)
+    thread = threading.Thread(target=deferred_introspection, daemon=True)
+    thread.start()
+
+
+def _emit_inventory_if_ready() -> None:
+    """
+    Emit agent inventory if:
+    - SDK is initialized
+    - Agents are registered
+    - Inventory has not been emitted yet
+
+    This function never throws.
+    """
+    try:
+        # Avoid circular import
+        from aeonic.core.agent_registry import (
+            get_all_agents,
+            mark_inventory_emitted,
+            should_emit_inventory,
+        )
+        from aeonic.core.sender import emit_agent_inventory
+
+        if _CONFIG is None:
+            return
+
+        if not should_emit_inventory():
+            return
+
+        agents = get_all_agents()
+        if not agents:
+            logger.warning("[Aeonic] No agents found during introspection")
+            return
+
+        # Detect framework from registered agents
+        framework = agents[0]["framework"] if agents else "unknown"
+
+        # Convert to inventory format
+        inventory_agents = []
+        for agent in agents:
+            inventory_agents.append(
+                {
+                    "name": agent.get("name"),
+                    "type": agent.get("type"),
+                    "model": agent.get("model", []),
+                    "route_path": agent.get("route_path"),
+                    "http_method": agent.get("http_method"),
+                    "framework": agent.get("framework"),
+                    "declared_at": int(agent.get("declared_at", 0) * 1000),  # Convert to ms
+                }
+            )
+
+        payload: Dict[str, Any] = {
+            "event": "agent_inventory",
+            "api_key": _CONFIG["api_key"],
+            "service_name": _CONFIG.get("service_name"),
+            "framework": framework,
+            "agents": inventory_agents,
+        }
+
+        if emit_agent_inventory(payload):
+            mark_inventory_emitted()
+
+    except Exception as e:
+        # Never throw - silent failure
+        logger.warning(f"[Aeonic] Failed to emit inventory: {e}")
+
+
+def _register_agentguard_async() -> None:
+    """
+    Register SDK as AgentGuard instance with T-LAIOR.
+    Runs in background thread and never blocks or throws.
+    """
+
+    def register_async() -> None:
+        try:
+            from aeonic.core.agentguard.registration import register_agent_guard
+
+            register_agent_guard()
+        except Exception as error:
+            logger.warning(
+                f"[Aeonic] AgentGuard registration failed: {error}. "
+                "SDK will continue to function for sample collection."
+            )
+            # Don't raise - SDK can still function for sample collection
+
+    # Run in background daemon thread (non-blocking)
+    thread = threading.Thread(target=register_async, daemon=True)
+    thread.start()