prela 0.1.0__py3-none-any.whl

prela/instrumentation/multi_agent/models.py

@@ -0,0 +1,203 @@
+"""Shared data models for multi-agent instrumentation.
+
+This module provides common data structures for tracing multi-agent systems
+across different frameworks (CrewAI, AutoGen, LangGraph, Swarm).
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from datetime import datetime
+from enum import Enum
+from typing import Any, Literal, Optional
+import hashlib
+
+
+class AgentRole(str, Enum):
+    """Standard agent roles across frameworks."""
+
+    MANAGER = "manager"
+    WORKER = "worker"
+    SPECIALIST = "specialist"
+    CRITIC = "critic"
+    USER_PROXY = "user_proxy"
+    ASSISTANT = "assistant"
+    CUSTOM = "custom"
+
+
+class MessageType(str, Enum):
+    """Types of inter-agent messages."""
+
+    TASK_ASSIGNMENT = "task_assignment"
+    TASK_RESULT = "task_result"
+    QUESTION = "question"
+    ANSWER = "answer"
+    FEEDBACK = "feedback"
+    DELEGATION = "delegation"
+    HANDOFF = "handoff"
+    SYSTEM = "system"
+
+
+@dataclass
+class AgentDefinition:
+    """Represents an agent in a multi-agent system."""
+
+    agent_id: str
+    name: str
+    role: AgentRole
+    framework: str  # "crewai", "autogen", "langgraph", "swarm"
+
+    # Agent configuration
+    model: Optional[str] = None
+    system_prompt: Optional[str] = None
+    tools: list[str] = field(default_factory=list)
+
+    # Framework-specific metadata
+    metadata: dict[str, Any] = field(default_factory=dict)
+
+    def to_span_attributes(self) -> dict[str, Any]:
+        """Convert to span attributes for tracing."""
+        return {
+            "agent.id": self.agent_id,
+            "agent.name": self.name,
+            "agent.role": self.role.value,
+            "agent.framework": self.framework,
+            "agent.model": self.model,
+            "agent.tools": self.tools,
+        }
+
+
+@dataclass
+class AgentMessage:
+    """A message between agents."""
+
+    message_id: str
+    sender_id: str
+    receiver_id: Optional[str]  # None for broadcast
+    message_type: MessageType
+    content: str
+    timestamp: datetime
+
+    # Optional structured data
+    tool_calls: Optional[list[dict]] = None
+    tool_results: Optional[list[dict]] = None
+    metadata: dict[str, Any] = field(default_factory=dict)
+
+    def to_span_event(self) -> dict[str, Any]:
+        """Convert to span event for tracing."""
+        return {
+            "name": f"agent.message.{self.message_type.value}",
+            "attributes": {
+                "message.id": self.message_id,
+                "message.sender": self.sender_id,
+                "message.receiver": self.receiver_id,
+                "message.type": self.message_type.value,
+                "message.content_length": len(self.content),
+            },
+        }
+
+
+@dataclass
+class TaskAssignment:
+    """A task assigned from one agent to another."""
+
+    task_id: str
+    assigner_id: str
+    assignee_id: str
+    description: str
+    expected_output: Optional[str] = None
+
+    # Execution tracking
+    status: Literal["pending", "in_progress", "completed", "failed"] = "pending"
+    started_at: Optional[datetime] = None
+    completed_at: Optional[datetime] = None
+    result: Optional[str] = None
+    error: Optional[str] = None
+
+    # Token/cost tracking
+    total_tokens: int = 0
+    total_cost_usd: float = 0.0
+
+
+@dataclass
+class CrewExecution:
+    """Top-level execution of a multi-agent crew/team."""
+
+    execution_id: str
+    framework: str
+
+    # Crew configuration
+    agents: list[AgentDefinition] = field(default_factory=list)
+    tasks: list[TaskAssignment] = field(default_factory=list)
+
+    # Execution tracking
+    started_at: Optional[datetime] = None
+    completed_at: Optional[datetime] = None
+    status: Literal["running", "completed", "failed"] = "running"
+
+    # Aggregated metrics
+    total_llm_calls: int = 0
+    total_tokens: int = 0
+    total_cost_usd: float = 0.0
+    total_tool_calls: int = 0
+    total_agent_messages: int = 0
+
+
+@dataclass
+class ConversationTurn:
+    """A single turn in a multi-agent conversation (for AutoGen style)."""
+
+    turn_id: str
+    turn_number: int
+    speaker_id: str
+    content: str
+    timestamp: datetime
+
+    # If this turn triggered an LLM call
+    llm_call_id: Optional[str] = None
+    tokens_used: int = 0
+
+    # If this turn triggered tool use
+    tool_calls: list[dict] = field(default_factory=list)
+
+
+# Helper functions
+
+
+def generate_agent_id(framework: str, name: str) -> str:
+    """Generate a consistent agent ID.
+
+    Args:
+        framework: The multi-agent framework (e.g., "crewai", "autogen")
+        name: The agent name
+
+    Returns:
+        A deterministic 12-character hash based on framework and name
+    """
+    return hashlib.sha256(f"{framework}:{name}".encode()).hexdigest()[:12]
+
+
+def extract_agent_graph(
+    agents: list[AgentDefinition], messages: list[AgentMessage]
+) -> dict:
+    """Extract agent communication graph for visualization.
+
+    Args:
+        agents: List of agent definitions
+        messages: List of messages exchanged between agents
+
+    Returns:
+        A dictionary with "nodes" and "edges" for graph visualization
+    """
+    nodes = {a.agent_id: {"name": a.name, "role": a.role.value} for a in agents}
+    edges = []
+    for msg in messages:
+        if msg.receiver_id:
+            edges.append(
+                {
+                    "from": msg.sender_id,
+                    "to": msg.receiver_id,
+                    "type": msg.message_type.value,
+                }
+            )
+    return {"nodes": nodes, "edges": edges}

prela/instrumentation/multi_agent/swarm.py

@@ -0,0 +1,231 @@
+"""Swarm instrumentation for Prela.
+
+This module provides automatic instrumentation for OpenAI Swarm multi-agent framework,
+capturing agent orchestration, handoffs, and execution flow.
+"""
+
+from __future__ import annotations
+
+import functools
+import uuid
+from typing import Any, Optional
+
+from prela.core.span import SpanType
+from prela.core.tracer import Tracer, get_tracer
+from prela.instrumentation.base import Instrumentor
+from prela.instrumentation.multi_agent.models import generate_agent_id
+from prela.license import require_tier
+
+
+class SwarmInstrumentor(Instrumentor):
+    """Instrumentor for OpenAI Swarm multi-agent framework."""
+
+    FRAMEWORK = "swarm"
+
+    @property
+    def is_instrumented(self) -> bool:
+        """Check if Swarm is currently instrumented."""
+        return self._is_instrumented
+
+    def __init__(self):
+        super().__init__()
+        self._active_swarms: dict[str, dict] = {}
+        self._is_instrumented = False
+        self._tracer: Optional[Tracer] = None
+
+    @require_tier("Swarm instrumentation", "lunch-money")
+    def instrument(self, tracer: Optional[Tracer] = None) -> None:
+        """Patch Swarm classes for tracing.
+
+        Args:
+            tracer: Optional tracer to use. If None, uses global tracer.
+        """
+        if self.is_instrumented:
+            return
+
+        try:
+            from swarm import Agent, Swarm
+        except ImportError:
+            return  # Swarm not installed
+
+        self._tracer = tracer or get_tracer()
+
+        # Patch Swarm.run (main execution entry point)
+        if hasattr(Swarm, "run"):
+            self._patch_swarm_run(Swarm)
+
+        # Patch Agent.__init__ to track agents
+        if hasattr(Agent, "__init__"):
+            self._patch_agent_init(Agent)
+
+        self._is_instrumented = True
+
+    def uninstrument(self) -> None:
+        """Restore original methods."""
+        if not self.is_instrumented:
+            return
+
+        try:
+            from swarm import Agent, Swarm
+        except ImportError:
+            return
+
+        # Restore all patched methods
+        for cls, method_name in [
+            (Swarm, "run"),
+            (Agent, "__init__"),
+        ]:
+            if hasattr(cls, f"_prela_original_{method_name}"):
+                original = getattr(cls, f"_prela_original_{method_name}")
+                setattr(cls, method_name, original)
+                delattr(cls, f"_prela_original_{method_name}")
+
+        self._is_instrumented = False
+        self._tracer = None
+
+    def _patch_swarm_run(self, swarm_cls) -> None:
+        """Patch Swarm.run to create execution spans."""
+        if hasattr(swarm_cls, "_prela_original_run"):
+            return
+
+        original_run = swarm_cls.run
+        swarm_cls._prela_original_run = original_run
+
+        instrumentor = self
+
+        @functools.wraps(original_run)
+        def wrapped_run(swarm_self, agent, messages, *args, **kwargs):
+            tracer = instrumentor._tracer
+            if not tracer:
+                return original_run(swarm_self, agent, messages, *args, **kwargs)
+
+            execution_id = str(uuid.uuid4())
+            instrumentor._active_swarms[execution_id] = {
+                "agents_seen": set(),
+                "handoffs": [],
+                "total_tokens": 0,
+            }
+
+            # Get initial agent name
+            initial_agent_name = agent.name if hasattr(agent, "name") else "unnamed"
+
+            # Extract context variables
+            context_vars = kwargs.get("context_variables", {})
+            context_keys = list(context_vars.keys()) if context_vars else []
+
+            swarm_attributes = {
+                "swarm.execution_id": execution_id,
+                "swarm.framework": instrumentor.FRAMEWORK,
+                "swarm.initial_agent": initial_agent_name,
+                "swarm.num_messages": len(messages) if messages else 0,
+                "swarm.context_variables": context_keys,
+            }
+
+            # NEW: Replay capture if enabled
+            replay_capture = None
+            if tracer.capture_for_replay:
+                from prela.core.replay import ReplayCapture
+
+                replay_capture = ReplayCapture()
+                # Capture agent context
+                replay_capture.set_agent_context(
+                    system_prompt=getattr(agent, "instructions", None),
+                    available_tools=[
+                        {"name": func.__name__, "description": func.__doc__ or ""}
+                        for func in getattr(agent, "functions", [])
+                    ],
+                    config={
+                        "framework": instrumentor.FRAMEWORK,
+                        "execution_id": execution_id,
+                        "agent_name": initial_agent_name,
+                    },
+                    memory={"context_variables": context_keys},
+                )
+
+            with tracer.span(
+                name=f"swarm.run.{initial_agent_name}",
+                span_type=SpanType.AGENT,
+                attributes=swarm_attributes,
+            ) as span:
+                try:
+                    result = original_run(swarm_self, agent, messages, *args, **kwargs)
+
+                    # Add execution statistics
+                    swarm_state = instrumentor._active_swarms.get(execution_id, {})
+                    agents_used = list(swarm_state.get("agents_seen", set()))
+                    if agents_used:
+                        span.set_attribute("swarm.agents_used", agents_used)
+
+                    handoffs = swarm_state.get("handoffs", [])
+                    span.set_attribute("swarm.num_handoffs", len(handoffs))
+
+                    # Capture final agent
+                    if hasattr(result, "agent"):
+                        final_agent = (
+                            result.agent.name
+                            if hasattr(result.agent, "name")
+                            else "unnamed"
+                        )
+                        span.set_attribute("swarm.final_agent", final_agent)
+
+                    # Capture response messages
+                    if hasattr(result, "messages"):
+                        span.set_attribute(
+                            "swarm.response_messages", len(result.messages)
+                        )
+
+                    # Capture context updates
+                    if hasattr(result, "context_variables"):
+                        updated_keys = list(result.context_variables.keys())
+                        span.set_attribute("swarm.updated_context", updated_keys)
+
+                    # NEW: Attach replay snapshot
+                    if replay_capture:
+                        try:
+                            object.__setattr__(span, "replay_snapshot", replay_capture.build())
+                        except Exception as e:
+                            import logging
+                            logger = logging.getLogger(__name__)
+                            logger.debug(f"Failed to capture replay data: {e}")
+
+                    return result
+                except Exception as e:
+                    span.add_event(
+                        "exception",
+                        attributes={
+                            "exception.type": type(e).__name__,
+                            "exception.message": str(e),
+                        },
+                    )
+                    raise
+                finally:
+                    if execution_id in instrumentor._active_swarms:
+                        del instrumentor._active_swarms[execution_id]
+
+        swarm_cls.run = wrapped_run
+
+    def _patch_agent_init(self, agent_cls) -> None:
+        """Patch Agent.__init__ to track agent creation."""
+        if hasattr(agent_cls, "_prela_original___init__"):
+            return
+
+        original_init = agent_cls.__init__
+        agent_cls._prela_original___init__ = original_init
+
+        instrumentor = self
+
+        @functools.wraps(original_init)
+        def wrapped_init(agent_self, *args, **kwargs):
+            result = original_init(agent_self, *args, **kwargs)
+
+            # Generate agent ID for tracking
+            agent_name = (
+                agent_self.name if hasattr(agent_self, "name") else str(id(agent_self))
+            )
+            agent_self._prela_agent_id = generate_agent_id(
+                instrumentor.FRAMEWORK, agent_name
+            )
+
+            return result
+
+        agent_cls.__init__ = wrapped_init

prela/instrumentation/n8n/__init__.py

@@ -0,0 +1,68 @@
+"""
+n8n workflow automation platform instrumentation.
+
+This module provides automatic tracing for n8n workflows, capturing:
+- Workflow executions with status and timing
+- Individual node executions (AI and non-AI)
+- LLM calls with token usage and costs
+- Tool invocations and retrievals
+- Error tracking and retries
+
+Usage:
+    import prela
+    from prela.instrumentation.n8n import N8nInstrumentor
+
+    # Initialize Prela
+    tracer = prela.init(service_name="n8n-workflows")
+
+    # Instrument n8n
+    N8nInstrumentor().instrument(tracer)
+
+    # All n8n workflow executions are now automatically traced
+"""
+
+from __future__ import annotations
+
+# Check tier on module import
+from prela.license import check_tier
+
+if not check_tier("n8n instrumentation", "lunch-money", silent=False):
+    raise ImportError(
+        "n8n instrumentation requires 'lunch-money' subscription or higher. "
+        "Upgrade at https://prela.dev/pricing"
+    )
+
+from prela.instrumentation.n8n.code_node import (
+    PrelaN8nContext,
+    prela_n8n_traced,
+    trace_n8n_code,
+)
+from prela.instrumentation.n8n.models import (
+    N8nAINodeExecution,
+    N8nNodeExecution,
+    N8nSpanType,
+    N8nWorkflowExecution,
+)
+from prela.instrumentation.n8n.webhook import (
+    N8N_AI_NODE_TYPES,
+    N8nWebhookHandler,
+    N8nWebhookPayload,
+    parse_n8n_webhook,
+)
+
+__all__ = [
+    # Models
+    "N8nWorkflowExecution",
+    "N8nNodeExecution",
+    "N8nAINodeExecution",
+    "N8nSpanType",
+    # Webhook
+    "N8nWebhookHandler",
+    "N8nWebhookPayload",
+    "parse_n8n_webhook",
+    "N8N_AI_NODE_TYPES",
+    # Code Node Helpers
+    "PrelaN8nContext",
+    "trace_n8n_code",
+    "prela_n8n_traced",
+]