agentloop-profiler 0.4.0__py3-none-any.whl

agentloop/graph.py

@@ -0,0 +1,312 @@
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from typing import Any
+
+from agentloop.timing import (
+    cumulative_span_time_ms,
+    elapsed_runtime_ms,
+    event_interval_ms,
+)
+from agentloop.timing import (
+    duration_ms as recorded_duration_ms,
+)
+
+
+@dataclass
+class ExecutionNode:
+    node_id: str
+    name: str
+    event_type: str
+    duration_ms: float
+    model: str | None = None
+    input_tokens: int = 0
+    output_tokens: int = 0
+    status: str = "ok"
+    metadata: dict[str, Any] = field(default_factory=dict)
+    started_at: str | None = None
+    ended_at: str | None = None
+    parent_id: str | None = None
+
+    @property
+    def total_tokens(self) -> int:
+        return self.input_tokens + self.output_tokens
+
+    def to_dict(self) -> dict[str, Any]:
+        return {
+            "node_id": self.node_id,
+            "name": self.name,
+            "event_type": self.event_type,
+            "duration_ms": round(self.duration_ms, 3),
+            "started_at": self.started_at,
+            "ended_at": self.ended_at,
+            "parent_id": self.parent_id,
+            "model": self.model,
+            "input_tokens": self.input_tokens,
+            "output_tokens": self.output_tokens,
+            "total_tokens": self.total_tokens,
+            "status": self.status,
+            "metadata": self.metadata,
+        }
+
+
+@dataclass
+class ExecutionEdge:
+    source: str
+    target: str
+    kind: str = "sequence"
+
+    def to_dict(self) -> dict[str, str]:
+        return {"source": self.source, "target": self.target, "kind": self.kind}
+
+
+@dataclass
+class CriticalPath:
+    node_ids: list[str]
+    duration_ms: float
+
+    def to_dict(self) -> dict[str, Any]:
+        return {"node_ids": self.node_ids, "duration_ms": round(self.duration_ms, 3)}
+
+
+@dataclass
+class _PathState:
+    node_ids: list[str]
+    duration_ms: float
+    started_ms: float | None = None
+    ended_ms: float | None = None
+
+
+@dataclass
+class ExecutionGraph:
+    nodes: list[ExecutionNode]
+    edges: list[ExecutionEdge]
+    elapsed_ms: float | None = None
+
+    @classmethod
+    def from_trace(cls, trace: Any) -> "ExecutionGraph":
+        indexed_events = list(enumerate(trace.events))
+        indexed_events.sort(key=_event_sort_key)
+
+        nodes: list[ExecutionNode] = []
+        event_by_node_id: dict[str, Any] = {}
+        for original_index, event in indexed_events:
+            node_id = event.event_id or f"node_{original_index}"
+            nodes.append(
+                ExecutionNode(
+                    node_id=node_id,
+                    name=event.name,
+                    event_type=event.event_type,
+                    duration_ms=event.duration_ms,
+                    model=event.model,
+                    input_tokens=event.input_tokens,
+                    output_tokens=event.output_tokens,
+                    status=event.status,
+                    metadata=event.metadata or {},
+                    started_at=event.started_at,
+                    ended_at=event.ended_at,
+                    parent_id=event.parent_id,
+                )
+            )
+            event_by_node_id[node_id] = event
+
+        node_ids = set(event_by_node_id)
+        edges: list[ExecutionEdge] = []
+        roots: list[ExecutionNode] = []
+        for node in nodes:
+            if node.parent_id and node.parent_id in node_ids and node.parent_id != node.node_id:
+                edges.append(
+                    ExecutionEdge(source=node.parent_id, target=node.node_id, kind="parent")
+                )
+            else:
+                roots.append(node)
+
+        edges.extend(_inferred_sequence_edges(roots))
+        return cls(
+            nodes=nodes, edges=_deduplicate_edges(edges), elapsed_ms=elapsed_runtime_ms(trace)
+        )
+
+    def total_runtime_ms(self) -> float:
+        if self.elapsed_ms is not None:
+            return max(0.0, self.elapsed_ms)
+        intervals = [event_interval_ms(node) for node in self.nodes]
+        if intervals and all(interval is not None for interval in intervals):
+            valid = [interval for interval in intervals if interval is not None]
+            return max(ended for _, ended in valid) - min(started for started, _ in valid)
+        return cumulative_span_time_ms(self.nodes)
+
+    def critical_path(self) -> CriticalPath:
+        if not self.nodes:
+            return CriticalPath([], 0.0)
+
+        node_by_id = {node.node_id: node for node in self.nodes}
+        predecessors: dict[str, list[str]] = {node_id: [] for node_id in node_by_id}
+        successors: dict[str, list[str]] = {node_id: [] for node_id in node_by_id}
+        for edge in self.edges:
+            if edge.source not in node_by_id or edge.target not in node_by_id:
+                continue
+            predecessors[edge.target].append(edge.source)
+            successors[edge.source].append(edge.target)
+
+        ordered_ids = _topological_order(self.nodes, predecessors, successors)
+        states: dict[str, _PathState] = {}
+        for node_id in ordered_ids:
+            node = node_by_id[node_id]
+            interval = event_interval_ms(node)
+            if interval is None:
+                own = _PathState([node_id], recorded_duration_ms(node))
+            else:
+                own = _PathState([node_id], interval[1] - interval[0], interval[0], interval[1])
+
+            candidates = [own]
+            for predecessor_id in predecessors[node_id]:
+                predecessor = states.get(predecessor_id)
+                if predecessor is not None:
+                    candidates.append(_extend_path(predecessor, node, interval))
+            states[node_id] = max(
+                candidates, key=lambda state: (state.duration_ms, len(state.node_ids))
+            )
+
+        best = max(states.values(), key=lambda state: (state.duration_ms, len(state.node_ids)))
+        return CriticalPath(best.node_ids, best.duration_ms)
+
+    def parallelizable_groups(self) -> list[dict[str, Any]]:
+        groups: dict[tuple[str, str], list[ExecutionNode]] = {}
+        for node in self.nodes:
+            key = (node.event_type, node.name)
+            groups.setdefault(key, []).append(node)
+
+        out: list[dict[str, Any]] = []
+        for (event_type, name), items in groups.items():
+            if event_type != "tool_call" or len(items) < 3:
+                continue
+            sequential = sum(item.duration_ms for item in items)
+            parallel = max(item.duration_ms for item in items)
+            out.append(
+                {
+                    "name": name,
+                    "event_type": event_type,
+                    "count": len(items),
+                    "node_ids": [item.node_id for item in items],
+                    "sequential_time_ms": round(sequential, 3),
+                    "estimated_parallel_time_ms": round(parallel, 3),
+                    "estimated_savings_ms": round(max(0.0, sequential - parallel), 3),
+                }
+            )
+        return sorted(out, key=lambda item: item["estimated_savings_ms"], reverse=True)
+
+    def bottlenecks(self, limit: int = 5) -> list[dict[str, Any]]:
+        total = self.total_runtime_ms() or 1.0
+        ranked = sorted(self.nodes, key=lambda node: node.duration_ms, reverse=True)[:limit]
+        return [
+            {
+                **node.to_dict(),
+                "runtime_share": round(node.duration_ms / total, 4),
+            }
+            for node in ranked
+        ]
+
+    def to_dict(self) -> dict[str, Any]:
+        return {
+            "nodes": [node.to_dict() for node in self.nodes],
+            "edges": [edge.to_dict() for edge in self.edges],
+            "total_runtime_ms": round(self.total_runtime_ms(), 3),
+            "cumulative_span_time_ms": round(cumulative_span_time_ms(self.nodes), 3),
+            "critical_path": self.critical_path().to_dict(),
+            "parallelizable_groups": self.parallelizable_groups(),
+            "bottlenecks": self.bottlenecks(),
+        }
+
+
+def _event_sort_key(item: tuple[int, Any]) -> tuple[Any, ...]:
+    original_index, event = item
+    interval = event_interval_ms(event)
+    if interval is None:
+        return (1, original_index)
+    return (0, interval[0], interval[1], str(event.event_id), original_index)
+
+
+def _node_sort_key(node: ExecutionNode) -> tuple[Any, ...]:
+    interval = event_interval_ms(node)
+    if interval is None:
+        return (1, node.node_id)
+    return (0, interval[0], interval[1], node.node_id)
+
+
+def _inferred_sequence_edges(roots: list[ExecutionNode]) -> list[ExecutionEdge]:
+    if len(roots) < 2:
+        return []
+
+    intervals = {node.node_id: event_interval_ms(node) for node in roots}
+    if any(interval is None for interval in intervals.values()):
+        return [
+            ExecutionEdge(source=previous.node_id, target=current.node_id)
+            for previous, current in zip(roots, roots[1:])
+        ]
+
+    ordered = sorted(roots, key=_node_sort_key)
+    edges: list[ExecutionEdge] = []
+    for index, current in enumerate(ordered):
+        current_interval = intervals[current.node_id]
+        if current_interval is None:
+            continue
+        candidates = []
+        for previous in ordered[:index]:
+            previous_interval = intervals[previous.node_id]
+            if previous_interval is not None and previous_interval[1] <= current_interval[0]:
+                candidates.append((previous_interval[1], previous_interval[0], previous))
+        if candidates:
+            predecessor = max(candidates, key=lambda item: (item[0], item[1], item[2].node_id))[2]
+            edges.append(ExecutionEdge(source=predecessor.node_id, target=current.node_id))
+    return edges
+
+
+def _deduplicate_edges(edges: list[ExecutionEdge]) -> list[ExecutionEdge]:
+    seen: set[tuple[str, str, str]] = set()
+    result: list[ExecutionEdge] = []
+    for edge in edges:
+        key = (edge.source, edge.target, edge.kind)
+        if key not in seen:
+            result.append(edge)
+            seen.add(key)
+    return result
+
+
+def _topological_order(
+    nodes: list[ExecutionNode],
+    predecessors: dict[str, list[str]],
+    successors: dict[str, list[str]],
+) -> list[str]:
+    node_by_id = {node.node_id: node for node in nodes}
+    indegree = {node_id: len(set(items)) for node_id, items in predecessors.items()}
+    ready = sorted(
+        (node_by_id[node_id] for node_id, count in indegree.items() if count == 0),
+        key=_node_sort_key,
+    )
+    ordered: list[str] = []
+    while ready:
+        node = ready.pop(0)
+        ordered.append(node.node_id)
+        for target in successors[node.node_id]:
+            indegree[target] -= 1
+            if indegree[target] == 0:
+                ready.append(node_by_id[target])
+                ready.sort(key=_node_sort_key)
+
+    if len(ordered) != len(nodes):
+        remaining = [node for node in nodes if node.node_id not in set(ordered)]
+        ordered.extend(node.node_id for node in sorted(remaining, key=_node_sort_key))
+    return ordered
+
+
+def _extend_path(
+    predecessor: _PathState,
+    node: ExecutionNode,
+    interval: tuple[float, float] | None,
+) -> _PathState:
+    node_ids = [*predecessor.node_ids, node.node_id]
+    if predecessor.started_ms is not None and predecessor.ended_ms is not None and interval:
+        started = min(predecessor.started_ms, interval[0])
+        ended = max(predecessor.ended_ms, interval[1])
+        return _PathState(node_ids, ended - started, started, ended)
+    return _PathState(node_ids, predecessor.duration_ms + recorded_duration_ms(node))

agentloop/integrations/__init__.py

@@ -0,0 +1 @@
+"""Optional integrations for external agent frameworks."""

agentloop/integrations/crewai.py

@@ -0,0 +1,144 @@
+from __future__ import annotations
+
+import inspect
+import time
+from functools import wraps
+from typing import Any
+
+from agentloop.events import utc_now_iso
+from agentloop.tracer import record_tool_call, trace_agent
+
+
+def _patch_method(obj: Any, method_name: str, span_name: str, metadata: dict[str, Any]) -> bool:
+    method = getattr(obj, method_name, None)
+    if method is None or not callable(method):
+        return False
+    if getattr(method, "_agentloop_wrapped", False):
+        return True
+
+    if inspect.iscoroutinefunction(method):
+
+        @wraps(method)
+        async def async_wrapper(*args: Any, **kwargs: Any) -> Any:
+            started_at = utc_now_iso()
+            start = time.perf_counter()
+            status = "ok"
+            error = None
+            try:
+                return await method(*args, **kwargs)
+            except Exception as exc:
+                status = "error"
+                error = str(exc)
+                raise
+            finally:
+                record_tool_call(
+                    span_name,
+                    started_at=started_at,
+                    duration_ms=(time.perf_counter() - start) * 1000,
+                    status=status,
+                    error=error,
+                    metadata=metadata,
+                )
+
+        async_wrapper._agentloop_wrapped = True  # type: ignore[attr-defined]
+        setattr(obj, method_name, async_wrapper)
+        return True
+
+    @wraps(method)
+    def wrapper(*args: Any, **kwargs: Any) -> Any:
+        started_at = utc_now_iso()
+        start = time.perf_counter()
+        status = "ok"
+        error = None
+        try:
+            return method(*args, **kwargs)
+        except Exception as exc:
+            status = "error"
+            error = str(exc)
+            raise
+        finally:
+            record_tool_call(
+                span_name,
+                started_at=started_at,
+                duration_ms=(time.perf_counter() - start) * 1000,
+                status=status,
+                error=error,
+                metadata=metadata,
+            )
+
+    wrapper._agentloop_wrapped = True  # type: ignore[attr-defined]
+    setattr(obj, method_name, wrapper)
+    return True
+
+
+def instrument_crew(crew: Any, *, name: str | None = None) -> Any:
+    """Wrap a CrewAI Crew-like object with AgentLoop tracing.
+
+    This patches `.kickoff(...)` and `.kickoff_async(...)` when present. CrewAI is
+    optional; the integration accepts any object with those methods.
+    """
+
+    run_name = name or getattr(crew, "name", None) or crew.__class__.__name__
+    metadata = {"integration": "crewai", "crew_name": str(run_name)}
+
+    kickoff = getattr(crew, "kickoff", None)
+    if (
+        kickoff is not None
+        and callable(kickoff)
+        and not getattr(kickoff, "_agentloop_wrapped", False)
+    ):
+
+        @wraps(kickoff)
+        def kickoff_wrapper(*args: Any, **kwargs: Any) -> Any:
+            with trace_agent(str(run_name), metadata=metadata):
+                return kickoff(*args, **kwargs)
+
+        kickoff_wrapper._agentloop_wrapped = True  # type: ignore[attr-defined]
+        setattr(crew, "kickoff", kickoff_wrapper)
+
+    kickoff_async = getattr(crew, "kickoff_async", None)
+    if (
+        kickoff_async is not None
+        and callable(kickoff_async)
+        and not getattr(kickoff_async, "_agentloop_wrapped", False)
+    ):
+
+        @wraps(kickoff_async)
+        async def kickoff_async_wrapper(*args: Any, **kwargs: Any) -> Any:
+            with trace_agent(str(run_name), metadata=metadata):
+                return await kickoff_async(*args, **kwargs)
+
+        kickoff_async_wrapper._agentloop_wrapped = True  # type: ignore[attr-defined]
+        setattr(crew, "kickoff_async", kickoff_async_wrapper)
+
+    return crew
+
+
+def instrument_task(task: Any, *, name: str | None = None) -> Any:
+    """Patch common CrewAI task execution methods as tool-call spans."""
+
+    task_name = (
+        name
+        or getattr(task, "description", None)
+        or getattr(task, "name", None)
+        or task.__class__.__name__
+    )
+    metadata = {"integration": "crewai", "task_name": str(task_name)}
+    for method_name in ("execute", "execute_sync", "execute_async", "run"):
+        _patch_method(task, method_name, f"crewai.task.{method_name}", metadata)
+    return task
+
+
+def instrument_agent(agent: Any, *, name: str | None = None) -> Any:
+    """Patch common CrewAI agent execution methods as tool-call spans."""
+
+    agent_name = (
+        name
+        or getattr(agent, "role", None)
+        or getattr(agent, "name", None)
+        or agent.__class__.__name__
+    )
+    metadata = {"integration": "crewai", "agent_name": str(agent_name)}
+    for method_name in ("execute_task", "run", "invoke"):
+        _patch_method(agent, method_name, f"crewai.agent.{method_name}", metadata)
+    return agent

agentloop/integrations/langgraph.py

@@ -0,0 +1,140 @@
+from __future__ import annotations
+
+import inspect
+from collections.abc import Callable
+from functools import wraps
+from typing import Any, TypeVar
+
+from agentloop import trace_agent, trace_tool_call
+
+F = TypeVar("F", bound=Callable[..., Any])
+
+
+def _merge_metadata(
+    metadata: dict[str, Any] | None, extra: dict[str, Any] | None = None
+) -> dict[str, Any]:
+    merged = {"integration": "langgraph"}
+    if metadata:
+        merged.update(metadata)
+    if extra:
+        merged.update(extra)
+    return merged
+
+
+def trace_node(name: str | None = None, metadata: dict[str, Any] | None = None) -> Callable[[F], F]:
+    """Wrap a LangGraph node function as an AgentLoop node span.
+
+    This is dependency-free: it works for normal callables, async callables, or
+    LangGraph node functions without importing LangGraph inside AgentLoop.
+    """
+
+    def decorator(fn: F) -> F:
+        step_name = name or getattr(fn, "__name__", "langgraph_node")
+        step_metadata = _merge_metadata(metadata, {"node_name": step_name})
+
+        if inspect.iscoroutinefunction(fn):
+
+            @wraps(fn)
+            async def async_wrapper(*args: Any, **kwargs: Any) -> Any:
+                with trace_tool_call(step_name, metadata=step_metadata):
+                    return await fn(*args, **kwargs)
+
+            return async_wrapper  # type: ignore[return-value]
+
+        @wraps(fn)
+        def wrapper(*args: Any, **kwargs: Any) -> Any:
+            with trace_tool_call(step_name, metadata=step_metadata):
+                return fn(*args, **kwargs)
+
+        return wrapper  # type: ignore[return-value]
+
+    return decorator
+
+
+def instrument_state_graph(graph: Any, metadata: dict[str, Any] | None = None) -> Any:
+    """Monkey-patch a LangGraph StateGraph-like builder so add_node auto-wraps nodes.
+
+    Usage:
+        builder = StateGraph(State)
+        instrument_state_graph(builder)
+        builder.add_node("retrieve", retrieve_fn)  # automatically traced
+
+    The function intentionally accepts `Any` so AgentLoop does not need LangGraph
+    as a dependency. It expects the object to expose an `add_node(name, action, ...)`
+    method, which matches LangGraph's builder style.
+    """
+
+    if getattr(graph, "_agentloop_instrumented", False):
+        return graph
+    original_add_node = getattr(graph, "add_node", None)
+    if original_add_node is None or not callable(original_add_node):
+        raise TypeError("Expected a LangGraph-like object with callable add_node(...).")
+
+    @wraps(original_add_node)
+    def add_node(name: str, action: Any = None, *args: Any, **kwargs: Any) -> Any:
+        wrapped_action = action
+        if callable(action):
+            wrapped_action = trace_node(
+                str(name), _merge_metadata(metadata, {"auto_instrumented": True})
+            )(action)
+        return original_add_node(name, wrapped_action, *args, **kwargs)
+
+    setattr(graph, "add_node", add_node)
+    setattr(graph, "_agentloop_instrumented", True)
+    return graph
+
+
+class TracedRunnable:
+    """Small wrapper for compiled LangGraph apps/runnables.
+
+    It creates an AgentLoop trace around `.invoke(...)`, `.ainvoke(...)`, `.stream(...)`,
+    and `.astream(...)` while delegating every other attribute to the wrapped app.
+    """
+
+    def __init__(
+        self, app: Any, name: str = "langgraph_app", metadata: dict[str, Any] | None = None
+    ):
+        self.app = app
+        self.name = name
+        self.metadata = _merge_metadata(metadata, {"runnable": True})
+        self.last_trace = None
+
+    def __getattr__(self, item: str) -> Any:
+        return getattr(self.app, item)
+
+    def invoke(self, *args: Any, **kwargs: Any) -> Any:
+        with trace_agent(self.name, metadata=self.metadata) as trace:
+            result = self.app.invoke(*args, **kwargs)
+        self.last_trace = trace
+        return result
+
+    async def ainvoke(self, *args: Any, **kwargs: Any) -> Any:
+        with trace_agent(self.name, metadata=self.metadata) as trace:
+            result = await self.app.ainvoke(*args, **kwargs)
+        self.last_trace = trace
+        return result
+
+    def stream(self, *args: Any, **kwargs: Any) -> Any:
+        with trace_agent(self.name, metadata=self.metadata) as trace:
+            for item in self.app.stream(*args, **kwargs):
+                yield item
+        self.last_trace = trace
+
+    async def astream(self, *args: Any, **kwargs: Any) -> Any:
+        with trace_agent(self.name, metadata=self.metadata) as trace:
+            async for item in self.app.astream(*args, **kwargs):
+                yield item
+        self.last_trace = trace
+
+    def export_last_trace(self, path: str) -> None:
+        if self.last_trace is None:
+            raise RuntimeError("No trace recorded yet. Call invoke/ainvoke/stream/astream first.")
+        self.last_trace.export_json(path)
+
+
+def trace_runnable(
+    app: Any, name: str = "langgraph_app", metadata: dict[str, Any] | None = None
+) -> TracedRunnable:
+    """Wrap a compiled LangGraph app/runnable with AgentLoop tracing."""
+
+    return TracedRunnable(app=app, name=name, metadata=metadata)