From 12a04358b3a6b7fd671e9f2aa7f8473526a0f080 Mon Sep 17 00:00:00 2001 From: MervinPraison <454862+MervinPraison@users.noreply.github.com> Date: Mon, 26 Oct 2026 11:22:07 +0000 Subject: [PATCH 01/12] feat(trace): replay schema 1.1 with run_id, parent_seq and node_id Context events gain three optional fields (omitted from to_dict when unset) so a replay can tell which run, which enclosing span and which graph node an event belongs to. 1.0 traces still load, and from_dict keeps an unknown event_type as its raw string (with .value) instead of relabelling it. Branch id and the open-span stack move into contextvars keyed by emitter, so concurrent Parallel/Loop branches no longer clobber each other; sequence numbers come from itertools.count under a lock so append order matches sequence order. Adds workflow/task/step/approval start/end events, a kind on agent_handoff, and private _run_scope/_workflow_scope helpers. Every new method is a no-op when the emitter is disabled. Co-Authored-By: Claude Opus 5.5 (1M context) --- .../praisonaiagents/trace/context_events.py | 618 +++++++++++++++++- 1 file changed, 589 insertions(+), 29 deletions(-) diff --git a/src/praisonai-agents/praisonaiagents/trace/context_events.py b/src/praisonai-agents/praisonaiagents/trace/context_events.py index 4215a2b2e0..967fe5b3bb 100644 --- a/src/praisonai-agents/praisonaiagents/trace/context_events.py +++ b/src/praisonai-agents/praisonaiagents/trace/context_events.py @@ -9,10 +9,12 @@ - Disabled emitter has near-zero overhead - Uses __slots__ for memory efficiency -Schema Version: 1.0 +Schema Version: 1.1 (adds optional run_id / parent_seq / node_id; 1.0 files still load) """ import contextvars +import itertools +import threading import time from contextlib import contextmanager from dataclasses import dataclass, field @@ -21,7 +23,7 @@ import json -CONTEXT_SCHEMA_VERSION = "1.0" +CONTEXT_SCHEMA_VERSION = "1.1" # ============================================================================= @@ -37,6 +39,28 @@ # Singleton disabled emitter for default case (zero allocation per call) _default_disabled_emitter: Optional["ContextTraceEmitter"] = None +# Per-context run state. These live in contextvars, not on the emitter, so +# concurrent Parallel/Loop branches (each running in its own copied context, see +# copy_context_to_callable) and gathered asyncio tasks cannot clobber each +# other. Values are immutable tuples and are only ever replaced via .set(), so +# a copied context never shares a mutable object with its parent. +# +# Entries are keyed by the owning emitter's ``_uid`` so state set through one +# emitter never leaks into events of another emitter in the same context. +_branch_var: contextvars.ContextVar = contextvars.ContextVar( + 'praisonai_trace_branch', default=None # (emitter_uid, branch_id) +) +_span_var: contextvars.ContextVar = contextvars.ContextVar( + 'praisonai_trace_spans', default=() # ((emitter_uid, seq, kind, key), ...) +) +_run_id_var: contextvars.ContextVar = contextvars.ContextVar( + 'praisonai_trace_run_id', default=None +) +_llm_open_var: contextvars.ContextVar = contextvars.ContextVar( + 'praisonai_trace_llm_open', default=None # (emitter_uid, parent_seq, previous) +) +_emitter_uids = itertools.count(1) + def get_context_emitter() -> "ContextTraceEmitter": """ @@ -114,12 +138,126 @@ def trace_context(emitter: "ContextTraceEmitter") -> Generator["ContextTraceEmit reset_context_emitter(token) +@contextmanager +def _run_scope(run_id: Optional[str] = None) -> Generator[str, None, None]: + """Stamp ``run_id`` on every event emitted inside the block. + + Nested teams and flows inherit the OUTER run id: when a run id is already + in scope it is kept and ``run_id`` is ignored. A fresh uuid is used when no + id is given. Yields the effective run id. + """ + outer = _run_id_var.get() + if outer is not None: + yield outer + return + if not run_id: + import uuid + run_id = uuid.uuid4().hex + token = _run_id_var.set(run_id) + try: + yield run_id + finally: + try: + _run_id_var.reset(token) + except ValueError: + # Exited from a different context than it was entered in. + _run_id_var.set(None) + + +@contextmanager +def _workflow_scope( + kind: str, + name: Optional[str], + process: str, + graph_fn=None, + run_id: Optional[str] = None, +) -> Generator[Optional[Dict[str, Any]], None, None]: + """Emit ``workflow_start``/``workflow_end`` around a team or flow run. + + Yields ``None`` (and does nothing else, not even building the graph) when + the current emitter is disabled. Otherwise yields a mutable dict whose + ``"status"`` the caller may set to ``"failed"``; an exception escaping the + block always records ``"failed"``. + """ + emitter = get_context_emitter() + if not emitter._enabled: + yield None + return + with _run_scope(run_id): + graph = None + if graph_fn is not None: + try: + graph = graph_fn() + except Exception: + graph = None # A diagram must never break a run. + emitter.workflow_start(kind, name, process, graph=graph) + state: Dict[str, Any] = {"status": "completed"} + started = time.perf_counter() + try: + yield state + except BaseException: + state["status"] = "failed" + raise + finally: + emitter.workflow_end( + status=state.get("status") or "completed", + duration_ms=(time.perf_counter() - started) * 1000, + ) + + +def _approval_begin(request: Any, kind: str = "tool_call", node_id: Optional[str] = None): + """Emit ``approval_requested`` for an ``ApprovalRequest``; token or None.""" + emitter = get_context_emitter() + if not emitter._enabled: + return None + agent_name = getattr(request, "agent_name", None) + tool = getattr(request, "tool_name", None) or "" + if node_id is None and kind == "tool_call": + node_id = f"tool:{agent_name}/{tool}" if agent_name else None + emitter.approval_requested( + getattr(request, "approval_id", None) or "", + tool, + args=getattr(request, "arguments", None), + risk=getattr(request, "risk_level", None), + kind=kind, + agent_name=agent_name, + node_id=node_id, + ) + return emitter, request, node_id + + +def _approval_end(token, decision: Any = None, error: Optional[BaseException] = None) -> None: + """Emit ``approval_decided`` for a token from :func:`_approval_begin`.""" + if token is None: + return + emitter, request, node_id = token + if decision is None: + approved, reason, approver, modified = False, f"approval failed: {error}", None, False + else: + approved = bool(getattr(decision, "approved", False)) + reason = getattr(decision, "reason", None) + approver = getattr(decision, "approver", None) + modified = bool(getattr(decision, "modified_args", None)) + emitter.approval_decided( + getattr(request, "approval_id", None) or "", + approved, + reason=reason, + approver=approver, + modified=modified, + agent_name=getattr(request, "agent_name", None), + node_id=node_id, + ) + + def copy_context_to_callable(func): """ Wrap a callable to copy the current context to the new thread/executor. This is needed because contextvars are NOT automatically propagated to thread pool executors. Use this when calling run_in_executor(). + + The whole context is copied, so the replay state (emitter, run id, branch + id and the open-span stack that yields ``parent_seq``) reaches the thread. Usage: # Instead of: @@ -159,6 +297,53 @@ class ContextEventType(str, Enum): # Knowledge events for knowledge utilization tracking KNOWLEDGE_SEARCH = "knowledge_search" KNOWLEDGE_ADD = "knowledge_add" + # Run structure (schema 1.1): teams/flows, tasks, flow steps, approvals + WORKFLOW_START = "workflow_start" + WORKFLOW_END = "workflow_end" + TASK_START = "task_start" + TASK_END = "task_end" + STEP_START = "step_start" + STEP_END = "step_end" + APPROVAL_REQUESTED = "approval_requested" + APPROVAL_DECIDED = "approval_decided" + + +class _UnknownEventType(str): + """An ``event_type`` string this version does not know. + + ``from_dict`` keeps the raw value instead of relabelling it, so a trace + written by a newer version stays readable. ``.value`` mirrors the Enum API + so existing ``event.event_type.value`` callers keep working. + """ + + __slots__ = () + + @property + def value(self) -> str: + return str(self) + + +# Span bookkeeping: which event types open/close a span, and of which kind. +_SPAN_START = { + "session_start": "session", + "agent_start": "agent", + "tool_call_start": "tool_call", + "workflow_start": "workflow", + "task_start": "task", + "step_start": "step", + "approval_requested": "approval", +} +_SPAN_END = { + "session_end": "session", + "agent_end": "agent", + "tool_call_end": "tool_call", + "workflow_end": "workflow", + "task_end": "task", + "step_end": "step", + "approval_decided": "approval", +} +_PREVIEW_LIMIT = 500 +_UNSET = object() @dataclass @@ -184,6 +369,10 @@ class ContextEvent: completion_tokens: Number of tokens in the completion/response cost_usd: Estimated cost in USD for this event content_hash: SHA256 hash of content for duplicate detection + run_id: The team/flow/recorded run this event belongs to (1.1) + parent_seq: sequence_num of the enclosing (or, on *_end, the closed) + *_start event (1.1) + node_id: Graph node the event concerns, e.g. ``agent:Writer`` (1.1) """ event_type: ContextEventType timestamp: float @@ -199,6 +388,9 @@ class ContextEvent: completion_tokens: int = 0 cost_usd: float = 0.0 content_hash: Optional[str] = None + run_id: Optional[str] = None + parent_seq: Optional[int] = None + node_id: Optional[str] = None def to_dict(self) -> Dict[str, Any]: """Convert to dictionary for serialization.""" @@ -221,6 +413,12 @@ def to_dict(self) -> Dict[str, Any]: result["branch_id"] = self.branch_id if self.content_hash is not None: result["content_hash"] = self.content_hash + if self.run_id is not None: + result["run_id"] = self.run_id + if self.parent_seq is not None: + result["parent_seq"] = self.parent_seq + if self.node_id is not None: + result["node_id"] = self.node_id return result def to_json(self) -> str: @@ -232,12 +430,12 @@ def from_dict(cls, data: Dict[str, Any]) -> "ContextEvent": """Create ContextEvent from dictionary.""" event_type_str = data.get("event_type", "") - # Convert string to enum + # Convert string to enum. An unknown type (written by a newer version) + # keeps its raw string rather than being relabelled as something else. try: event_type = ContextEventType(event_type_str) except ValueError: - # Fallback for unknown event types - event_type = ContextEventType.CONTEXT_SNAPSHOT + event_type = _UnknownEventType(event_type_str) return cls( event_type=event_type, @@ -254,6 +452,9 @@ def from_dict(cls, data: Dict[str, Any]) -> "ContextEvent": completion_tokens=data.get("completion_tokens", 0), cost_usd=data.get("cost_usd", 0.0), content_hash=data.get("content_hash"), + run_id=data.get("run_id"), + parent_seq=data.get("parent_seq"), + node_id=data.get("node_id"), ) @@ -349,6 +550,28 @@ def __iter__(self): return iter(self._events) +def _innermost(stack: tuple, uid: int) -> Optional[int]: + """Sequence number of the innermost open span owned by emitter ``uid``.""" + for entry in reversed(stack): + if entry[0] == uid: + return entry[1] + return None + + +def _find_span(stack: tuple, uid: int, kind: str, key: Any) -> Optional[int]: + """Index of the innermost open span matching kind+key (else kind only).""" + fallback = None + for i in range(len(stack) - 1, -1, -1): + entry = stack[i] + if entry[0] != uid or entry[2] != kind: + continue + if entry[3] == key: + return i + if fallback is None: + fallback = i + return fallback + + class ContextTraceEmitter: """ Emitter for context trace events. @@ -367,7 +590,7 @@ class ContextTraceEmitter: emitter.session_end() """ - __slots__ = ("_sink", "_session_id", "_enabled", "_redact", "_sequence", "_branch_id", "_full_content") + __slots__ = ("_sink", "_session_id", "_enabled", "_redact", "_full_content", "_uid", "_counter", "_lock") # Default limits for truncation (can be overridden with full_content=True) # Increased significantly to capture full search results and avoid false truncation @@ -396,40 +619,124 @@ def __init__( self._session_id = session_id self._enabled = enabled self._redact = redact - self._sequence = 0 - self._branch_id: Optional[str] = None self._full_content = full_content - - def _emit(self, event: ContextEvent) -> None: - """Internal emit with enabled check and sequence assignment. + self._uid = next(_emitter_uids) + # itertools.count is atomic under the GIL; the lock additionally keeps + # the sink's append order equal to sequence order when Parallel/Loop + # branches emit from several threads at once (live tailers resume by + # sequence number, so an out-of-order line would be skipped). + self._counter = itertools.count() + self._lock = threading.RLock() + + def _emit(self, event: ContextEvent, span_key: Any = _UNSET) -> None: + """Internal emit with enabled check, sequence and span assignment. + + Stamps ``sequence_num``, ``branch_id``, ``run_id`` and ``parent_seq``. + A ``*_start`` event (and ``approval_requested``) opens a span on the + per-context stack; the matching ``*_end`` closes it and points its + ``parent_seq`` at the start it closes. Every other event points at the + innermost open span. Exception-safe: sink errors are silently caught to prevent tracing from crashing agent execution. """ if not self._enabled: return - event.sequence_num = self._sequence - event.branch_id = self._branch_id # Include branch context - self._sequence += 1 - try: - self._sink.emit(event) - except Exception: - # Tracing should never crash agent execution - pass + uid = self._uid + branch = _branch_var.get() + event.branch_id = branch[1] if branch is not None and branch[0] == uid else None + if event.run_id is None: + event.run_id = _run_id_var.get() + etype = event.event_type + value = etype.value if isinstance(etype, ContextEventType) else str(etype) + if span_key is _UNSET: + span_key = event.node_id or event.agent_name + with self._lock: + seq = next(self._counter) + event.sequence_num = seq + stack = _span_var.get() + kind = _SPAN_START.get(value) + if kind is not None: + if event.parent_seq is None: + event.parent_seq = _innermost(stack, uid) + _span_var.set(stack + ((uid, seq, kind, span_key),)) + else: + kind = _SPAN_END.get(value) + if kind is not None: + idx = _find_span(stack, uid, kind, span_key) + if idx is not None: + event.parent_seq = stack[idx][1] + # Close it and anything of ours left open above it + # (a start whose end was skipped by an exception). + _span_var.set(stack[:idx] + tuple(e for e in stack[idx + 1:] if e[0] != uid)) + elif event.parent_seq is None: + event.parent_seq = _innermost(stack, uid) + try: + self._sink.emit(event) + except Exception: + # Tracing should never crash agent execution + pass def set_branch(self, branch_id: str) -> None: """Set current branch context for parallel execution tracking. - All subsequent events will include this branch_id until cleared. + All subsequent events emitted from the current context (thread or + asyncio task) include this branch_id until cleared. Concurrent + branches each run in their own copied context, so they cannot + overwrite each other's branch. Args: branch_id: Identifier for the parallel branch (e.g., 'parallel_0', 'loop_5') """ - self._branch_id = branch_id + _branch_var.set((self._uid, branch_id)) def clear_branch(self) -> None: """Clear the current branch context.""" - self._branch_id = None + _branch_var.set(None) + + def _preview(self, text: Any) -> Optional[str]: + """Truncate to 500 chars and pass through redaction (for previews).""" + if text is None: + return None + text = text if isinstance(text, str) else str(text) + if self._redact: + from .redact import redact_string + text = redact_string(text[: _PREVIEW_LIMIT * 4]) + return text[:_PREVIEW_LIMIT] + + def _redact_args(self, args: Optional[Dict[str, Any]]) -> Dict[str, Any]: + """Approval args: sensitive keys redacted; text values previewed + redacted.""" + if not args: + return {} + if not isinstance(args, dict): + return {"value": self._preview(args)} + args = self._redact_dict(args) + return {k: self._preview(v) if isinstance(v, str) else v for k, v in args.items()} + + def _current_agent(self) -> Optional[str]: + """Name of the innermost open agent span in this context, if any.""" + uid = self._uid + for entry in reversed(_span_var.get()): + if entry[0] == uid and entry[2] == "agent": + key = entry[3] + if isinstance(key, str) and key.startswith("agent:"): + return key[len("agent:"):] + return key + return None + + def _llm_call_open(self) -> bool: + """True while a caller-emitted ``llm_request`` awaits its response. + + The agent's OpenAI-client path emits llm_request/llm_response around + its whole dispatch, which may itself reach the LiteLLM path; the LLM + class checks this so the same call is not counted twice. The marker is + only honoured inside the span it was opened in, so one left behind by + an exception goes stale as soon as that span closes. + """ + mark = _llm_open_var.get() + if mark is None or mark[0] != self._uid: + return False + return mark[1] == _innermost(_span_var.get(), self._uid) def _redact_dict(self, data: Optional[Dict[str, Any]]) -> Optional[Dict[str, Any]]: """Redact sensitive data from dictionary if enabled.""" @@ -475,6 +782,7 @@ def agent_start( session_id=self._session_id, agent_name=agent_name, data=metadata or {}, + node_id=f"agent:{agent_name}" if agent_name else None, )) def agent_end( @@ -491,6 +799,7 @@ def agent_end( session_id=self._session_id, agent_name=agent_name, data=metadata or {}, + node_id=f"agent:{agent_name}" if agent_name else None, )) def message_added( @@ -540,6 +849,7 @@ def tool_call_start( "tool_name": tool_name, "tool_args": self._redact_dict(tool_args) if tool_args else {}, }, + node_id=f"tool:{agent_name}/{tool_name}", )) def tool_call_end( @@ -587,6 +897,7 @@ def tool_call_end( "duration_ms": duration_ms, "error": error, }, + node_id=f"tool:{agent_name}/{tool_name}", )) def _smart_truncate_result(self, result: str, tool_name: str) -> str: @@ -686,11 +997,27 @@ def llm_request( """ if not self._enabled: return + event = self._llm_request_event( + agent_name, messages_count, tokens_used, tokens_budget, model, messages + ) + self._emit(event) + # Mark the call open so the LLM class does not record it a second time. + # The previous marker is kept so a nested agent's call restores it. + _llm_open_var.set((self._uid, event.parent_seq, _llm_open_var.get())) + + def _llm_request_event( + self, + agent_name: Optional[str], + messages_count: int, + tokens_used: int, + tokens_budget: int = 0, + model: Optional[str] = None, + messages: Optional[List[Dict[str, Any]]] = None, + ) -> ContextEvent: data = {"model": model} if messages is not None: data["messages"] = messages - - self._emit(ContextEvent( + return ContextEvent( event_type=ContextEventType.LLM_REQUEST, timestamp=time.time(), session_id=self._session_id, @@ -699,7 +1026,8 @@ def llm_request( tokens_used=tokens_used, tokens_budget=tokens_budget, data=data, - )) + node_id=f"agent:{agent_name}" if agent_name else None, + ) def llm_response( self, @@ -711,6 +1039,7 @@ def llm_response( prompt_tokens: int = 0, completion_tokens: int = 0, cost_usd: float = 0.0, + model: Optional[str] = None, ) -> None: """Emit LLM response event. @@ -723,9 +1052,30 @@ def llm_response( prompt_tokens: Number of tokens in the prompt/request completion_tokens: Number of tokens in the completion/response cost_usd: Estimated cost in USD for this LLM call + model: Model that served the call (stored as ``data.model``) """ if not self._enabled: return + self._emit(self._llm_response_event( + agent_name, response_tokens, duration_ms, finish_reason, + response_content, prompt_tokens, completion_tokens, cost_usd, model, + )) + mark = _llm_open_var.get() + if mark is not None and mark[0] == self._uid: + _llm_open_var.set(mark[2]) + + def _llm_response_event( + self, + agent_name: Optional[str], + response_tokens: int = 0, + duration_ms: float = 0.0, + finish_reason: Optional[str] = None, + response_content: Optional[str] = None, + prompt_tokens: int = 0, + completion_tokens: int = 0, + cost_usd: float = 0.0, + model: Optional[str] = None, + ) -> ContextEvent: # Use completion_tokens if provided, fallback to response_tokens for backward compat actual_completion_tokens = completion_tokens if completion_tokens > 0 else response_tokens @@ -748,8 +1098,10 @@ def llm_response( data["response_content"] = f"{head}\n...[{len(response_content):,} chars, showing first/last portions]...\n{tail}" else: data["response_content"] = response_content + if model is not None: + data["model"] = model - self._emit(ContextEvent( + return ContextEvent( event_type=ContextEventType.LLM_RESPONSE, timestamp=time.time(), session_id=self._session_id, @@ -758,7 +1110,8 @@ def llm_response( completion_tokens=actual_completion_tokens, cost_usd=cost_usd, data=data, - )) + node_id=f"agent:{agent_name}" if agent_name else None, + ) def context_snapshot( self, @@ -790,8 +1143,16 @@ def agent_handoff( to_agent: str, reason: Optional[str] = None, context_passed: Optional[Dict[str, Any]] = None, + kind: str = "handoff", + node_id: Optional[str] = None, ) -> None: - """Emit agent handoff event for tracking agent-to-agent flow.""" + """Emit agent handoff event for tracking agent-to-agent flow. + + Args: + kind: ``handoff`` (agent hands over), ``delegate`` (the hierarchical + manager assigns a task), ``route`` or ``spawn``. + node_id: Graph node concerned; defaults to the receiving agent. + """ if not self._enabled: return self._emit(ContextEvent( @@ -803,10 +1164,209 @@ def agent_handoff( "from_agent": from_agent, "to_agent": to_agent, "reason": reason, - "context_passed": context_passed or {}, + "context_passed": self._redact_dict(context_passed) or {}, + "kind": kind, + }, + node_id=node_id or (f"agent:{to_agent}" if to_agent else None), + )) + + # ------------------------------------------------------------------ + # Run structure (schema 1.1). Each method is a no-op when disabled. + # ------------------------------------------------------------------ + + def workflow_start( + self, + kind: str, + name: Optional[str], + process: str, + graph: Optional[Dict[str, Any]] = None, + metadata: Optional[Dict[str, Any]] = None, + ) -> None: + """A team, flow or agent run begins. ``graph`` is its static graph.""" + if not self._enabled: + return + data: Dict[str, Any] = {"kind": kind, "name": name, "process": process} + if graph is not None: + data["graph"] = graph + if metadata: + data.update(metadata) + self._emit(ContextEvent( + event_type=ContextEventType.WORKFLOW_START, + timestamp=time.time(), + session_id=self._session_id, + data=data, + )) + + def workflow_end(self, status: str = "completed", duration_ms: float = 0.0) -> None: + """The innermost open team/flow run ends.""" + if not self._enabled: + return + self._emit(ContextEvent( + event_type=ContextEventType.WORKFLOW_END, + timestamp=time.time(), + session_id=self._session_id, + data={"status": status, "duration_ms": duration_ms}, + )) + + def task_start( + self, + name: Optional[str], + description: Optional[str] = None, + agent: Optional[str] = None, + node_id: Optional[str] = None, + ) -> None: + """A team task starts (node_id ``task:``).""" + if not self._enabled: + return + self._emit(ContextEvent( + event_type=ContextEventType.TASK_START, + timestamp=time.time(), + session_id=self._session_id, + agent_name=agent, + data={ + "name": name, + "description_preview": self._preview(description), + "agent": agent, }, + node_id=node_id, )) + def task_end( + self, + name: Optional[str], + status: str, + retries: int = 0, + duration_ms: float = 0.0, + output: Any = None, + agent: Optional[str] = None, + node_id: Optional[str] = None, + ) -> None: + """A team task ends: ``completed``, ``failed`` or ``skipped``.""" + if not self._enabled: + return + self._emit(ContextEvent( + event_type=ContextEventType.TASK_END, + timestamp=time.time(), + session_id=self._session_id, + agent_name=agent, + data={ + "name": name, + "status": status, + "retries": retries, + "duration_ms": duration_ms, + "output_preview": self._preview(output), + }, + node_id=node_id, + )) + + def step_start( + self, + name: Optional[str], + pattern: Optional[str] = None, + node_id: Optional[str] = None, + agent: Optional[str] = None, + ) -> None: + """An AgentFlow step (or a Route/Parallel/Loop/... pattern) starts.""" + if not self._enabled: + return + data: Dict[str, Any] = {"name": name} + if pattern: + data["pattern"] = pattern + self._emit(ContextEvent( + event_type=ContextEventType.STEP_START, + timestamp=time.time(), + session_id=self._session_id, + agent_name=agent, + data=data, + node_id=node_id, + )) + + def step_end( + self, + name: Optional[str], + status: str, + duration_ms: float = 0.0, + output: Any = None, + route: Optional[str] = None, + node_id: Optional[str] = None, + agent: Optional[str] = None, + ) -> None: + """A flow step ends: ``completed``, ``failed``, ``skipped`` or ``cached``.""" + if not self._enabled: + return + data: Dict[str, Any] = { + "name": name, + "status": status, + "duration_ms": duration_ms, + "output_preview": self._preview(output), + } + if route is not None: + data["route"] = route + self._emit(ContextEvent( + event_type=ContextEventType.STEP_END, + timestamp=time.time(), + session_id=self._session_id, + agent_name=agent, + data=data, + node_id=node_id, + )) + + def approval_requested( + self, + approval_id: str, + tool: str, + args: Optional[Dict[str, Any]] = None, + risk: Optional[str] = None, + kind: str = "tool_call", + agent_name: Optional[str] = None, + node_id: Optional[str] = None, + ) -> None: + """A person is asked to approve a tool call or a task's output.""" + if not self._enabled: + return + self._emit(ContextEvent( + event_type=ContextEventType.APPROVAL_REQUESTED, + timestamp=time.time(), + session_id=self._session_id, + agent_name=agent_name, + data={ + "approval_id": approval_id, + "tool": tool, + "args": self._redact_args(args), + "risk": risk, + "kind": kind, + }, + node_id=node_id, + ), span_key=approval_id) + + def approval_decided( + self, + approval_id: str, + approved: bool, + reason: Optional[str] = None, + approver: Optional[str] = None, + modified: bool = False, + agent_name: Optional[str] = None, + node_id: Optional[str] = None, + ) -> None: + """The decision for ``approval_id`` (parent_seq = the request).""" + if not self._enabled: + return + self._emit(ContextEvent( + event_type=ContextEventType.APPROVAL_DECIDED, + timestamp=time.time(), + session_id=self._session_id, + agent_name=agent_name, + data={ + "approval_id": approval_id, + "approved": bool(approved), + "reason": reason, + "approver": approver, + "modified": bool(modified), + }, + node_id=node_id, + ), span_key=approval_id) + def memory_store( self, agent_name: str, From 662edacc9ef9169cd64b3251f1a5270a0d54be4e Mon Sep 17 00:00:00 2001 From: MervinPraison <454862+MervinPraison@users.noreply.github.com> Date: Mon, 26 Oct 2026 11:22:07 +0000 Subject: [PATCH 02/12] feat(workflows): graph data for agents, teams and flows One builder walks a definition and produces praisonai.graph/1 data with ids derived from names and positions only (agent:, task:, step:, tool:, human:you, manager:Manager, start:begin, end:done), so the ids are the same in every process and are what replay events carry as node_id. Covers task context/next_tasks/conditions, the hierarchical manager, hand-offs, tools, human review and Route/Parallel/Loop/Repeat/If/Include/Discussion. AgentFlow.to_mermaid() output is byte-identical (the builder records the legacy draw order). Adds to_graph(), graph_to_mermaid(), AgentFlow.to_graph() and AgentTeam.to_graph()/to_mermaid(). Stdlib only and never imported by import praisonaiagents; golden JSON and a two-subprocess stability check pin the shape. Co-Authored-By: Claude Opus 5.5 (1M context) --- .../praisonaiagents/agents/agents.py | 15 + .../praisonaiagents/workflows/diagram.py | 567 ++++++++++++++++-- .../praisonaiagents/workflows/workflows.py | 9 + .../tests/unit/test_lazy_imports.py | 58 ++ .../tests/unit/workflows/_graph_fixtures.py | 107 ++++ .../unit/workflows/golden/graph_flow.json | 431 +++++++++++++ .../golden/graph_hierarchical_team.json | 126 ++++ .../golden/graph_sequential_team.json | 126 ++++ .../unit/workflows/golden/graph_team.json | 186 ++++++ .../tests/unit/workflows/test_graph.py | 110 ++++ 10 files changed, 1678 insertions(+), 57 deletions(-) create mode 100644 src/praisonai-agents/tests/unit/workflows/_graph_fixtures.py create mode 100644 src/praisonai-agents/tests/unit/workflows/golden/graph_flow.json create mode 100644 src/praisonai-agents/tests/unit/workflows/golden/graph_hierarchical_team.json create mode 100644 src/praisonai-agents/tests/unit/workflows/golden/graph_sequential_team.json create mode 100644 src/praisonai-agents/tests/unit/workflows/golden/graph_team.json create mode 100644 src/praisonai-agents/tests/unit/workflows/test_graph.py diff --git a/src/praisonai-agents/praisonaiagents/agents/agents.py b/src/praisonai-agents/praisonaiagents/agents/agents.py index f785f08288..27cfbc4f96 100644 --- a/src/praisonai-agents/praisonaiagents/agents/agents.py +++ b/src/praisonai-agents/praisonaiagents/agents/agents.py @@ -2441,6 +2441,21 @@ def run_all_tasks(self): except StopIteration: pass + def to_graph(self) -> Dict[str, Any]: + """This team's definition as graph data (``praisonai.graph/1``). + + Covers agents, tasks (context/depends_on, next_tasks, conditions), + tools, hand-offs, human review and the hierarchical manager. Node ids + are stable across processes and match replay events' ``node_id``. + """ + from ..workflows.diagram import to_graph + return to_graph(self) + + def to_mermaid(self) -> str: + """Render this team's definition as a mermaid diagram.""" + from ..workflows.diagram import graph_to_mermaid, to_graph + return graph_to_mermaid(to_graph(self)) + def get_task_status(self, task_id): if task_id in self.tasks: return self.tasks[task_id].status diff --git a/src/praisonai-agents/praisonaiagents/workflows/diagram.py b/src/praisonai-agents/praisonaiagents/workflows/diagram.py index af8e211763..7fc4da4cec 100644 --- a/src/praisonai-agents/praisonaiagents/workflows/diagram.py +++ b/src/praisonai-agents/praisonaiagents/workflows/diagram.py @@ -1,13 +1,25 @@ -"""Render a workflow's DEFINITION as a mermaid diagram. +"""Render a run's DEFINITION as graph data and as a mermaid diagram. PraisonAI could already diagram a captured execution trace (``telemetry/performance_utils.py``), but that only exists once a flow has run -- so a flow could only be pictured after it had already cost money. This -renders the structure from the step list, before anything executes, which is +renders the structure from the definition, before anything executes, which is what CrewAI's ``flow.plot()`` gives its users. +One builder walks the definition and produces graph data +(``schema: "praisonai.graph/1"``) for an ``Agent``, an ``AgentTeam`` or an +``AgentFlow``. Node ids are derived from names and positions only (never a +uuid), so they are identical across processes and runs and the same ids appear +as ``node_id`` on replay trace events: + +- ``agent:``, ``task:``, ``step:`` + (duplicates become ``#2``, ``#3`` in definition order), + ``tool:/``, ``human:you``, ``manager:Manager`` +- ``start:begin`` ("Brief") and ``end:done`` ("Delivered") + Mermaid rather than graphviz on purpose: it needs no new dependency, and both -GitHub and the docs render it inline. +GitHub and the docs render it inline. Stdlib only; nothing here is imported +when ``praisonaiagents`` is. >>> print(flow.to_mermaid()) graph TD @@ -16,9 +28,25 @@ ... """ -from typing import Any, List, Optional, Tuple +from typing import Any, Dict, List, Optional, Tuple + +__all__ = [ + "to_mermaid", + "flow_to_mermaid", + "to_graph", + "graph_to_mermaid", + "GRAPH_SCHEMA", +] + +GRAPH_SCHEMA = "praisonai.graph/1" + +START_ID = "start:begin" +END_ID = "end:done" +HUMAN_ID = "human:you" +MANAGER_ID = "manager:Manager" -__all__ = ["to_mermaid", "flow_to_mermaid"] +_FLOW_CLASSES = {"AgentFlow"} +_TEAM_CLASSES = {"AgentTeam"} def _label(step: Any) -> str: @@ -41,23 +69,170 @@ def _escape(text: str) -> str: return str(text).replace('"', "'").replace("\n", " ") -class _Builder: - def __init__(self) -> None: - self.lines: List[str] = ["graph TD"] - self._n = 0 +def _class_names(obj: Any) -> set: + return {klass.__name__ for klass in type(obj).__mro__} + + +def _tool_name(tool: Any) -> str: + if isinstance(tool, str): + return tool + if isinstance(tool, dict): + function = tool.get("function") if isinstance(tool.get("function"), dict) else {} + return str(function.get("name") or tool.get("name") or "tool") + for attr in ("name", "__name__"): + value = getattr(tool, attr, None) + if isinstance(value, str) and value: + return value + return type(tool).__name__ + + +def _is_agent(obj: Any) -> bool: + return hasattr(obj, "chat") and isinstance(getattr(obj, "name", None), str) + + +# --------------------------------------------------------------------------- +# Graph builder +# --------------------------------------------------------------------------- + + +class _Graph: + """Accumulates nodes/edges, and the mermaid draw order for flows. + + ``ops`` records exactly the nodes and edges the legacy flow diagram drew, + in the order it drew them, so ``flow_to_mermaid`` stays byte-identical + while the graph data can carry more (agents, tools, pattern groups). + """ + + def __init__(self, kind: str, name: Optional[str], process: str) -> None: + self.kind = kind + self.name = name + self.process = process + self.nodes: Dict[str, Dict[str, Any]] = {} + self.edges: Dict[str, Dict[str, Any]] = {} + self.ops: List[tuple] = [] + # id(definition object) -> node id, first occurrence wins. Used at run + # time to stamp node_id on step events. + self.ids: Dict[int, str] = {} + self._seen: Dict[str, int] = {} + + def unique(self, base: str) -> str: + """``base`` the first time, then ``base#2``, ``base#3`` ...""" + count = self._seen.get(base, 0) + 1 + self._seen[base] = count + return base if count == 1 else f"{base}#{count}" - def node(self, shape: str) -> str: - node_id = f"n{self._n}" - self._n += 1 - self.lines.append(f" {node_id}{shape}") + def node( + self, + node_id: str, + kind: str, + label: str, + parent: Optional[str] = None, + meta: Optional[Dict[str, Any]] = None, + shape: Optional[str] = None, + ) -> str: + if node_id not in self.nodes: + self.nodes[node_id] = { + "id": node_id, + "kind": kind, + "label": label, + "parent": parent, + "meta": dict(meta or {}), + } + if shape is not None: + self.ops.append(("node", node_id, shape)) return node_id - def edge(self, src: str, dst: str, label: Optional[str] = None) -> None: - arrow = f' -->|"{_escape(label)}"| ' if label else " --> " - self.lines.append(f" {src}{arrow}{dst}") + def edge( + self, + source: str, + target: str, + kind: str, + label: Optional[str] = None, + draw: bool = False, + ) -> None: + edge_id = f"{source}->{target}:{kind}" + if edge_id not in self.edges: + self.edges[edge_id] = { + "id": edge_id, + "source": source, + "target": target, + "kind": kind, + "label": label, + } + if draw: + self.ops.append(("edge", source, target, label)) + def has_edge(self, source: str, target: str, kind: str) -> bool: + return f"{source}->{target}:{kind}" in self.edges -def _branch(steps: Any, b: "_Builder") -> Optional[Tuple[str, str]]: + def register(self, obj: Any, node_id: str) -> None: + self.ids.setdefault(id(obj), node_id) + + def to_dict(self) -> Dict[str, Any]: + return { + "schema": GRAPH_SCHEMA, + "kind": self.kind, + "name": self.name, + "process": self.process, + "nodes": list(self.nodes.values()), + "edges": list(self.edges.values()), + } + + def ops_to_mermaid(self) -> str: + lines = ["graph TD"] + short: Dict[str, str] = {} + for op in self.ops: + if op[0] == "node": + ref = f"n{len(short)}" + short[op[1]] = ref + lines.append(f" {ref}{op[2]}") + else: + _, source, target, label = op + arrow = f' -->|"{_escape(label)}"| ' if label else " --> " + lines.append(f" {short[source]}{arrow}{short[target]}") + return "\n".join(lines) + + +def _add_agent(g: _Graph, agent: Any, parent: Optional[str] = None, _seen: Optional[set] = None) -> str: + """Add an agent node with its tools and hand-off targets; return its id.""" + name = getattr(agent, "name", None) or type(agent).__name__ + agent_id = f"agent:{name}" + seen = _seen if _seen is not None else set() + if agent_id in seen: + return agent_id + seen.add(agent_id) + g.node(agent_id, "agent", name, parent=parent, meta={"role": getattr(agent, "role", None)}) + targets = [] + handoff_tools = set() + for handoff in list(getattr(agent, "handoffs", None) or []): + target = handoff if _is_agent(handoff) else getattr(handoff, "agent", None) + if target is None or not getattr(target, "name", None): + continue + targets.append(target) + tool_name = getattr(handoff, "tool_name", None) if not _is_agent(handoff) else None + if tool_name is None: + from ..agent.handoff import Handoff + tool_name = Handoff(agent=target).tool_name + handoff_tools.add(tool_name) + for tool in list(getattr(agent, "tools", None) or []): + tool_name = _tool_name(tool) + if tool_name in handoff_tools: + continue # drawn as the hand-off edge, not as a tool + tool_id = f"tool:{name}/{tool_name}" + g.node(tool_id, "tool", tool_name, parent=parent) + g.edge(agent_id, tool_id, "uses") + for target in targets: + target_id = _add_agent(g, target, parent=parent, _seen=seen) + g.edge(agent_id, target_id, "handoff") + return agent_id + + +# --------------------------------------------------------------------------- +# AgentFlow +# --------------------------------------------------------------------------- + + +def _flow_branch(steps: Any, g: _Graph, parent: Optional[str]) -> Optional[Tuple[str, str]]: """Render a sub-sequence of steps in isolation. Returns the ``(entry, exit)`` pair for the whole sub-sequence, or ``None`` @@ -68,7 +243,7 @@ def _branch(steps: Any, b: "_Builder") -> Optional[Tuple[str, str]]: entry: Optional[str] = None tail: Optional[str] = None for inner in steps or []: - rendered = _render(inner, b, tail) + rendered = _flow_step(inner, g, tail, parent) if rendered is None: continue inner_entry, inner_exit = rendered @@ -80,7 +255,15 @@ def _branch(steps: Any, b: "_Builder") -> Optional[Tuple[str, str]]: return entry, tail -def _render(step: Any, b: _Builder, prev: Optional[str]) -> Optional[Tuple[str, str]]: +def _pattern_group(g: _Graph, step: Any, pattern: str, parent: Optional[str], label: str, meta: Dict[str, Any]) -> Tuple[str, str]: + """The ``step:`` node that groups a pattern's parts; returns (id, suffix).""" + group = g.unique(f"step:{label}") + g.node(group, "step", label, parent=parent, meta={"pattern": pattern, **meta}) + g.register(step, group) + return group, group[len("step:"):] + + +def _flow_step(step: Any, g: _Graph, prev: Optional[str], parent: Optional[str]) -> Optional[Tuple[str, str]]: """Render one step, returning its ``(entry, exit)`` nodes. ``entry`` is what an incoming edge should point at; ``exit`` is what the @@ -89,82 +272,93 @@ def _render(step: Any, b: _Builder, prev: Optional[str]) -> Optional[Tuple[str, its join). """ cls = type(step).__name__ - name = _escape(_label(step)) if cls == "If": - decision = b.node(f'{{"{_escape(getattr(step, "condition", "") or "if")}"}}') + condition = getattr(step, "condition", "") or "if" + group, suffix = _pattern_group(g, step, "if", parent, "if", {"condition": str(condition)}) + decision = g.node(f"decision:{suffix}", "decision", str(condition), parent=group, + shape=f'{{"{_escape(condition)}"}}') if prev: - b.edge(prev, decision) + g.edge(prev, decision, "next", draw=True) ends = [] for branch, edge_label in (("then_steps", "true"), ("else_steps", "false")): - rendered = _branch(getattr(step, branch, None), b) + rendered = _flow_branch(getattr(step, branch, None), g, group) if rendered: branch_entry, branch_exit = rendered - b.edge(decision, branch_entry, edge_label) + g.edge(decision, branch_entry, "routes", edge_label, draw=True) ends.append(branch_exit) - join = b.node('(("join"))') + join = g.node(f"join:{suffix}", "join", "join", parent=group, shape='(("join"))') for end in ends or [decision]: - b.edge(end, join) + g.edge(end, join, "next", draw=True) return decision, join if cls == "Route": - router = b.node(f'{{"route"}}') + group, suffix = _pattern_group(g, step, "route", parent, "route", {}) + router = g.node(f"decision:{suffix}", "decision", "route", parent=group, shape='{"route"}') if prev: - b.edge(prev, router) + g.edge(prev, router, "next", draw=True) ends = [] routes = getattr(step, "routes", None) or {} for key, steps in routes.items(): - rendered = _branch(steps, b) + rendered = _flow_branch(steps, g, group) if rendered: branch_entry, branch_exit = rendered - b.edge(router, branch_entry, key) + g.edge(router, branch_entry, "routes", key, draw=True) ends.append(branch_exit) # ``Route.__init__`` mirrors ``routes["default"]`` into ``.default``; # only render ``.default`` when it was not already drawn as a key above. if "default" not in routes: - rendered = _branch(getattr(step, "default", None), b) + rendered = _flow_branch(getattr(step, "default", None), g, group) if rendered: branch_entry, branch_exit = rendered - b.edge(router, branch_entry, "default") + g.edge(router, branch_entry, "routes", "default", draw=True) ends.append(branch_exit) - join = b.node('(("join"))') + join = g.node(f"join:{suffix}", "join", "join", parent=group, shape='(("join"))') for end in ends or [router]: - b.edge(end, join) + g.edge(end, join, "next", draw=True) return router, join if cls == "Parallel": - fork = b.node('(("fork"))') + group, suffix = _pattern_group(g, step, "parallel", parent, "parallel", {}) + fork = g.node(f"fork:{suffix}", "fork", "fork", parent=group, shape='(("fork"))') if prev: - b.edge(prev, fork) + g.edge(prev, fork, "next", draw=True) ends = [] for inner in getattr(step, "steps", None) or []: - rendered = _render(inner, b, None) + rendered = _flow_step(inner, g, None, group) if rendered: branch_entry, branch_exit = rendered - b.edge(fork, branch_entry) + g.edge(fork, branch_entry, "parallel", draw=True) ends.append(branch_exit) - join = b.node('(("join"))') + join = g.node(f"join:{suffix}", "join", "join", parent=group, shape='(("join"))') for end in ends or [fork]: - b.edge(end, join) + g.edge(end, join, "next", draw=True) return fork, join if cls in ("Loop", "Repeat"): + pattern = cls.lower() inner_steps = getattr(step, "steps", None) or ( [getattr(step, "step")] if getattr(step, "step", None) is not None else [] ) + over = getattr(step, "over", None) + iterations = getattr(step, "max_iterations", None) + meta = {} + if over: + meta["over"] = str(over) + if iterations: + meta["max_iterations"] = iterations + group, _ = _pattern_group(g, step, pattern, parent, pattern, meta) tail = prev first = None for inner in inner_steps: - rendered = _render(inner, b, tail) + rendered = _flow_step(inner, g, tail, group) if rendered: inner_entry, inner_exit = rendered first = first or inner_entry tail = inner_exit if first and tail: - over = getattr(step, "over", None) - iterations = getattr(step, "max_iterations", None) back = f"over {over}" if over else (f"up to {iterations}x" if iterations else "repeat") - b.edge(tail, first, back) + g.edge(tail, first, "loop", back, draw=True) if first is None or tail is None: return None return first, tail @@ -174,33 +368,292 @@ def _render(step: Any, b: _Builder, prev: Optional[str]) -> Optional[Tuple[str, # them inline rather than as an opaque node. A recipe include (by name) # is only resolvable at runtime, so it stays a single node. workflow = getattr(step, "workflow", None) + recipe = getattr(step, "recipe", None) + meta = {"recipe": recipe} if recipe else {} + if workflow is not None and getattr(workflow, "name", None): + meta["workflow"] = workflow.name + group, _ = _pattern_group(g, step, "include", parent, "include", meta) if workflow is not None: - rendered = _branch(getattr(workflow, "steps", None), b) + rendered = _flow_branch(getattr(workflow, "steps", None), g, group) if rendered: branch_entry, branch_exit = rendered if prev: - b.edge(prev, branch_entry) + g.edge(prev, branch_entry, "next", draw=True) return branch_entry, branch_exit + # Unexpanded: the group node itself is the drawn step. + g.ops.append(("node", group, f'["{_escape(_label(step))}"]')) + if prev: + g.edge(prev, group, "next", draw=True) + return group, group - node = b.node(f'["{name}"]') + label = _label(step) + node_id = g.unique(f"step:{label}") + meta: Dict[str, Any] = {} + if cls == "Discussion": + meta["pattern"] = "discussion" + g.node(node_id, "step", label, parent=parent, meta=meta, shape=f'["{_escape(label)}"]') + g.register(step, node_id) if prev: - b.edge(prev, node) - return node, node + g.edge(prev, node_id, "next", draw=True) + # Graph-only detail: who does the step, and with which tools. + if cls == "Discussion": + for speaker in getattr(step, "agents", None) or []: + if _is_agent(speaker): + g.edge(node_id, _add_agent(g, speaker), "assigned") + g.register(speaker, node_id) + elif _is_agent(step): + g.edge(node_id, _add_agent(g, step), "assigned") + elif _is_agent(getattr(step, "agent", None)): + g.edge(node_id, _add_agent(g, step.agent), "assigned") + return node_id, node_id -def to_mermaid(steps: Any, name: Optional[str] = None) -> str: - """Render a list of workflow steps as a mermaid ``graph TD``.""" - b = _Builder() - start = b.node('([" start "])' if not name else f'(["{_escape(name)}"])') + +def _build_flow(steps: Any, name: Optional[str], kind: str = "flow", process: str = "flow") -> _Graph: + g = _Graph(kind, name, process) + start = g.node(START_ID, "start", "Brief", + shape='([" start "])' if not name else f'(["{_escape(name)}"])') prev: Optional[str] = start for step in (steps or []): - rendered = _render(step, b, prev) + rendered = _flow_step(step, g, prev, None) if rendered: prev = rendered[1] - end = b.node('([" end "])') + end = g.node(END_ID, "end", "Delivered", shape='([" end "])') if prev: - b.edge(prev, end) - return "\n".join(b.lines) + g.edge(prev, end, "next", draw=True) + return g + + +def _flow_process(flow: Any) -> str: + return "hierarchical" if getattr(flow, "process", None) == "hierarchical" else "flow" + + +def _flow_graph(flow: Any) -> _Graph: + return _build_flow(getattr(flow, "steps", None), getattr(flow, "name", None), + process=_flow_process(flow)) + + +# --------------------------------------------------------------------------- +# AgentTeam +# --------------------------------------------------------------------------- + + +def _is_synthetic_manager(task: Any) -> bool: + """The hierarchical process injects a ``manager_task`` run by "Manager".""" + agent = getattr(task, "agent", None) + return getattr(task, "name", None) == "manager_task" and getattr(agent, "name", None) == "Manager" + + +def _team_tasks(team: Any) -> List[Any]: + tasks = getattr(team, "tasks", None) or [] + if isinstance(tasks, dict): + return list(tasks.values()) + return list(tasks) + + +def task_node_ids(tasks: List[Any]) -> List[str]: + """Stable node ids for a team's tasks, aligned with ``tasks``. + + ``task:`` (or ``task:`` for an unnamed task), duplicates + suffixed ``#2``... The hierarchical synthetic manager task maps to + ``manager:Manager``. + """ + seen: Dict[str, int] = {} + ids = [] + for index, task in enumerate(tasks): + if _is_synthetic_manager(task): + ids.append(MANAGER_ID) + continue + name = getattr(task, "name", None) + base = f"task:{name}" if name else f"task:{index}" + count = seen.get(base, 0) + 1 + seen[base] = count + ids.append(base if count == 1 else f"{base}#{count}") + return ids + + +def task_node_id(tasks: Any, task: Any) -> Optional[str]: + """Node id of ``task`` within a team's ``tasks`` (dict or list).""" + if isinstance(tasks, dict): + tasks = list(tasks.values()) + tasks = list(tasks or []) + for task_obj, node_id in zip(tasks, task_node_ids(tasks)): + if task_obj is task: + return node_id + name = getattr(task, "name", None) + return f"task:{name}" if name else None + + +def _team_graph(team: Any) -> _Graph: + process = getattr(team, "process", None) or "sequential" + hierarchical = process == "hierarchical" + g = _Graph("team", getattr(team, "name", None), process) + all_tasks = _team_tasks(team) + ids = task_node_ids(all_tasks) + pairs = [(t, i) for t, i in zip(all_tasks, ids) if not _is_synthetic_manager(t)] + by_name = {} + for task, node_id in pairs: + if getattr(task, "name", None): + by_name.setdefault(task.name, node_id) + + g.node(START_ID, "start", "Brief") + if hierarchical: + g.node(MANAGER_ID, "manager", "Manager") + for agent in getattr(team, "agents", None) or []: + if getattr(agent, "name", None): + _add_agent(g, agent) + if any(getattr(t, "human_input", False) for t, _ in pairs): + g.node(HUMAN_ID, "human", "You") + + for index, (task, node_id) in enumerate(pairs): + label = getattr(task, "name", None) or ( + (getattr(task, "description", None) or "")[:60] or f"Task {index + 1}" + ) + agent = getattr(task, "agent", None) + meta: Dict[str, Any] = {} + if getattr(task, "task_type", None) not in (None, "task"): + meta["task_type"] = task.task_type + if getattr(task, "human_input", False): + meta["human_input"] = True + g.node(node_id, "task", str(label), meta=meta) + g.register(task, node_id) + if agent is not None and getattr(agent, "name", None): + g.edge(node_id, _add_agent(g, agent), "assigned") + if getattr(task, "human_input", False): + g.edge(HUMAN_ID, node_id, "reviews") + + if hierarchical: + g.edge(START_ID, MANAGER_ID, "next") + for agent in getattr(team, "agents", None) or []: + if getattr(agent, "name", None): + g.edge(MANAGER_ID, f"agent:{agent.name}", "delegates") + g.edge(MANAGER_ID, END_ID, "next") + elif process == "workflow" and pairs: + starts = [n for t, n in pairs if getattr(t, "is_start", False)] or [pairs[0][1]] + for node_id in starts: + g.edge(START_ID, node_id, "next") + for task, node_id in pairs: + condition = getattr(task, "condition", None) or {} + followers = [x for x in (getattr(task, "next_tasks", None) or []) if x in by_name] + routed = False + if isinstance(condition, dict): + for key, targets in condition.items(): + for target in (targets if isinstance(targets, (list, tuple)) else [targets]): + if target in by_name: + g.edge(node_id, by_name[target], "routes", str(key)) + routed = True + elif target in ("exit", "end", "done"): + g.edge(node_id, END_ID, "routes", str(key)) + routed = True + for target in followers: + g.edge(node_id, by_name[target], "next") + if not followers and not routed: + g.edge(node_id, END_ID, "next") + else: + prev = START_ID + for _, node_id in pairs: + g.edge(prev, node_id, "next") + prev = node_id + g.edge(prev, END_ID, "next") + + # Explicit data dependencies (context=/depends_on=), unless the same pair + # is already joined by a control edge (sequential teams add both). + task_ids = {id(t): n for t, n in pairs} + for task, node_id in pairs: + for dep in getattr(task, "context", None) or []: + dep_id = task_ids.get(id(dep)) + if dep_id and dep_id != node_id and not g.has_edge(dep_id, node_id, "next"): + g.edge(dep_id, node_id, "depends_on") + + g.node(END_ID, "end", "Delivered") + return g + + +def _agent_graph(agent: Any) -> _Graph: + g = _Graph("agent", getattr(agent, "name", None), "sequential") + g.node(START_ID, "start", "Brief") + agent_id = _add_agent(g, agent) + g.register(agent, agent_id) + g.edge(START_ID, agent_id, "next") + g.node(END_ID, "end", "Delivered") + g.edge(agent_id, END_ID, "next") + return g + + +def _graph_for(obj: Any) -> _Graph: + names = _class_names(obj) + if names & _FLOW_CLASSES: + return _flow_graph(obj) + if names & _TEAM_CLASSES: + return _team_graph(obj) + if isinstance(obj, (list, tuple)): + return _build_flow(obj, None) + if _is_agent(obj): + return _agent_graph(obj) + if hasattr(obj, "steps"): + return _flow_graph(obj) + if hasattr(obj, "tasks") and hasattr(obj, "agents"): + return _team_graph(obj) + raise TypeError(f"to_graph() expects an Agent, AgentTeam or AgentFlow, got {type(obj).__name__}") + + +def to_graph(obj: Any) -> Dict[str, Any]: + """Graph data (``praisonai.graph/1``) for an Agent, AgentTeam or AgentFlow.""" + return _graph_for(obj).to_dict() + + +# --------------------------------------------------------------------------- +# Mermaid +# --------------------------------------------------------------------------- + +_SHAPES = { + "start": '(["{}"])', + "end": '(["{}"])', + "decision": '{{"{}"}}', + "fork": '(("{}"))', + "join": '(("{}"))', + "manager": '{{{{"{}"}}}}', + "human": '>"{}"]', + "tool": '[["{}"]]', +} + + +def graph_to_mermaid(graph: Dict[str, Any]) -> str: + """Render ``praisonai.graph/1`` data as a mermaid ``graph TD``. + + Group nodes (a node other nodes name as ``parent``) are containers, not + boxes, so they and their edges are left out of the picture. + """ + nodes = list(graph.get("nodes") or []) + parents = {n.get("parent") for n in nodes if n.get("parent")} + short: Dict[str, str] = {} + lines = ["graph TD"] + for node in nodes: + if node["id"] in parents: + continue + ref = f"n{len(short)}" + short[node["id"]] = ref + shape = _SHAPES.get(node.get("kind"), '["{}"]') + lines.append(f" {ref}{shape.format(_escape(node.get('label') or node['id']))}") + for edge in graph.get("edges") or []: + source, target = short.get(edge["source"]), short.get(edge["target"]) + if source is None or target is None: + continue + kind = edge.get("kind") + label = edge.get("label") + if kind in ("next", "parallel"): + arrow = f' -->|"{_escape(label)}"| ' if label else " --> " + elif kind in ("routes", "loop", "handoff", "delegates"): + arrow = f' -->|"{_escape(label or kind)}"| ' + else: + arrow = f' -.->|"{_escape(label or kind)}"| ' + lines.append(f" {source}{arrow}{target}") + return "\n".join(lines) + + +def to_mermaid(steps: Any, name: Optional[str] = None) -> str: + """Render a list of workflow steps as a mermaid ``graph TD``.""" + return _build_flow(steps, name).ops_to_mermaid() def flow_to_mermaid(flow: Any) -> str: diff --git a/src/praisonai-agents/praisonaiagents/workflows/workflows.py b/src/praisonai-agents/praisonaiagents/workflows/workflows.py index f7af765c36..2dcbf16d1e 100644 --- a/src/praisonai-agents/praisonaiagents/workflows/workflows.py +++ b/src/praisonai-agents/praisonaiagents/workflows/workflows.py @@ -1498,6 +1498,15 @@ def to_mermaid(self) -> str: from .diagram import flow_to_mermaid return flow_to_mermaid(self) + def to_graph(self) -> Dict[str, Any]: + """This flow's definition as graph data (``praisonai.graph/1``). + + Node ids are stable across processes (derived from step names and + positions), and are the ``node_id`` values replay trace events carry. + """ + from .diagram import to_graph + return to_graph(self) + @property def state(self): """The flow's variables as the declared ``state_model``. diff --git a/src/praisonai-agents/tests/unit/test_lazy_imports.py b/src/praisonai-agents/tests/unit/test_lazy_imports.py index 66fc17cf04..04d66faf38 100644 --- a/src/praisonai-agents/tests/unit/test_lazy_imports.py +++ b/src/praisonai-agents/tests/unit/test_lazy_imports.py @@ -278,3 +278,61 @@ def test_do_not_track_takes_precedence(self): if __name__ == "__main__": pytest.main([__file__, "-v"]) + + +def _fresh_python(code): + """Run ``code`` in a fresh interpreter against this checkout; return stdout.""" + import os + import subprocess + from pathlib import Path + + root = str(Path(__file__).resolve().parents[2]) + env = dict(os.environ, PYTHONPATH=root, PYTHONDONTWRITEBYTECODE="1") + out = subprocess.run([sys.executable, "-c", code], env=env, capture_output=True, + text=True, timeout=120, check=True) + return out.stdout.strip().splitlines()[-1] + + +class TestReplayStaysLazy: + """PraisonAI Replay's graph builder must cost nothing until a trace is on.""" + + def test_package_import_does_not_load_the_graph_builder(self): + loaded = _fresh_python( + "import sys, praisonaiagents;" + "from praisonaiagents import Agent, AgentTeam, AgentFlow;" + "print('praisonaiagents.workflows.diagram' in sys.modules)" + ) + assert loaded == "False" + + def test_an_untraced_run_does_not_load_the_graph_builder(self): + loaded = _fresh_python( + "import sys\n" + "from praisonaiagents import AgentFlow\n" + "def step(ctx): return 'ok'\n" + "AgentFlow(steps=[step]).run('go');" + "print('praisonaiagents.workflows.diagram' in sys.modules)" + ) + assert loaded == "False" + + def test_a_traced_run_builds_the_graph(self): + has_graph = _fresh_python( + "from praisonaiagents import AgentFlow\n" + "from praisonaiagents.trace.context_events import ContextListSink, ContextTraceEmitter, trace_context\n" + "sink = ContextListSink()\n" + "def step(ctx): return 'ok'\n" + "with trace_context(ContextTraceEmitter(sink=sink, session_id='s')):\n" + " AgentFlow(steps=[step]).run('go')\n" + "print(sink.get_events()[0].data['graph']['schema'])" + ) + assert has_graph == "praisonai.graph/1" + + def test_package_import_stays_under_budget(self): + """``import praisonaiagents`` < 200 ms (best of three fresh interpreters).""" + best = min( + float(_fresh_python( + "import time; t = time.perf_counter(); import praisonaiagents;" + "print((time.perf_counter() - t) * 1000)" + )) + for _ in range(3) + ) + assert best < 200, f"import praisonaiagents took {best:.0f} ms" diff --git a/src/praisonai-agents/tests/unit/workflows/_graph_fixtures.py b/src/praisonai-agents/tests/unit/workflows/_graph_fixtures.py new file mode 100644 index 0000000000..68d747f074 --- /dev/null +++ b/src/praisonai-agents/tests/unit/workflows/_graph_fixtures.py @@ -0,0 +1,107 @@ +"""Definitions whose graphs are pinned as golden JSON (see test_graph.py). + +Kept in their own module so a subprocess can rebuild them and prove the node +ids do not depend on the process (hash seed, uuids, object ids). +""" + + +def search_web(query: str) -> str: + """Search the web.""" + return query + + +def score(ctx): + return "approved" + + +def fetch(ctx): + return "a,b" + + +def tidy(ctx): + return "ok" + + +def team(): + from praisonaiagents import Agent, Task + from praisonaiagents.agents.agents import AgentTeam + + researcher = Agent(name="Researcher", instructions="Find facts", tools=[search_web]) + billing = Agent(name="Billing", instructions="Handle billing") + writer = Agent(name="Writer", instructions="Write", handoffs=[billing]) + research = Task(name="research", description="Find facts", agent=researcher, + is_start=True, next_tasks=["draft"]) + draft = Task(name="draft", description="Write the piece", agent=writer, + context=[research], human_input=True, + condition={"approved": ["publish"], "rejected": ["research"]}) + publish = Task(name="publish", description="Ship it", agent=writer) + unnamed = Task(description="Archive everything afterwards", agent=writer) + return AgentTeam(agents=[researcher, writer], tasks=[research, draft, publish, unnamed], + process="workflow", name="newsroom") + + +def sequential_team(): + from praisonaiagents import Agent, Task + from praisonaiagents.agents.agents import AgentTeam + + researcher = Agent(name="Researcher", instructions="Find facts", tools=[search_web]) + writer = Agent(name="Writer", instructions="Write") + return AgentTeam( + agents=[researcher, writer], + tasks=[ + Task(name="research", description="Find facts", agent=researcher), + Task(name="write", description="Write it up", agent=writer), + Task(name="write", description="Write it again", agent=writer), + ], + name="pair", + ) + + +def hierarchical_team(): + from praisonaiagents import Agent, Task + from praisonaiagents.agents.agents import AgentTeam + + researcher = Agent(name="Researcher", instructions="Find facts", tools=[search_web]) + writer = Agent(name="Writer", instructions="Write") + return AgentTeam( + agents=[researcher, writer], + tasks=[ + Task(name="research", description="Find facts", agent=researcher), + Task(name="write", description="Write it up", agent=writer), + ], + process="hierarchical", + manager_llm="gpt-4o-mini", + name="managed", + ) + + +def flow(): + from praisonaiagents import Agent + from praisonaiagents.workflows.workflows import ( + AgentFlow, Discussion, If, Include, Loop, Parallel, Repeat, Route, + ) + + researcher = Agent(name="Researcher", instructions="Find facts", tools=[search_web]) + critic = Agent(name="Critic", instructions="Critique") + return AgentFlow( + name="pipeline", + steps=[ + researcher, + score, + Route(routes={"approved": [tidy], "default": [fetch]}), + Parallel(steps=[fetch, tidy]), + Loop(step=tidy, over="items"), + If(condition="{{score}} > 80", then_steps=[tidy], else_steps=[Repeat(fetch, max_iterations=2)]), + Discussion(agents=[researcher, critic], rounds=1), + researcher, + Include(recipe="publisher"), + ], + ) + + +BUILDERS = { + "team": team, + "sequential_team": sequential_team, + "hierarchical_team": hierarchical_team, + "flow": flow, +} diff --git a/src/praisonai-agents/tests/unit/workflows/golden/graph_flow.json b/src/praisonai-agents/tests/unit/workflows/golden/graph_flow.json new file mode 100644 index 0000000000..47e58289eb --- /dev/null +++ b/src/praisonai-agents/tests/unit/workflows/golden/graph_flow.json @@ -0,0 +1,431 @@ +{ + "edges": [ + { + "id": "start:begin->step:Researcher:next", + "kind": "next", + "label": null, + "source": "start:begin", + "target": "step:Researcher" + }, + { + "id": "agent:Researcher->tool:Researcher/search_web:uses", + "kind": "uses", + "label": null, + "source": "agent:Researcher", + "target": "tool:Researcher/search_web" + }, + { + "id": "step:Researcher->agent:Researcher:assigned", + "kind": "assigned", + "label": null, + "source": "step:Researcher", + "target": "agent:Researcher" + }, + { + "id": "step:Researcher->step:score:next", + "kind": "next", + "label": null, + "source": "step:Researcher", + "target": "step:score" + }, + { + "id": "step:score->decision:route:next", + "kind": "next", + "label": null, + "source": "step:score", + "target": "decision:route" + }, + { + "id": "decision:route->step:tidy:routes", + "kind": "routes", + "label": "approved", + "source": "decision:route", + "target": "step:tidy" + }, + { + "id": "decision:route->step:fetch:routes", + "kind": "routes", + "label": "default", + "source": "decision:route", + "target": "step:fetch" + }, + { + "id": "step:tidy->join:route:next", + "kind": "next", + "label": null, + "source": "step:tidy", + "target": "join:route" + }, + { + "id": "step:fetch->join:route:next", + "kind": "next", + "label": null, + "source": "step:fetch", + "target": "join:route" + }, + { + "id": "join:route->fork:parallel:next", + "kind": "next", + "label": null, + "source": "join:route", + "target": "fork:parallel" + }, + { + "id": "fork:parallel->step:fetch#2:parallel", + "kind": "parallel", + "label": null, + "source": "fork:parallel", + "target": "step:fetch#2" + }, + { + "id": "fork:parallel->step:tidy#2:parallel", + "kind": "parallel", + "label": null, + "source": "fork:parallel", + "target": "step:tidy#2" + }, + { + "id": "step:fetch#2->join:parallel:next", + "kind": "next", + "label": null, + "source": "step:fetch#2", + "target": "join:parallel" + }, + { + "id": "step:tidy#2->join:parallel:next", + "kind": "next", + "label": null, + "source": "step:tidy#2", + "target": "join:parallel" + }, + { + "id": "join:parallel->step:tidy#3:next", + "kind": "next", + "label": null, + "source": "join:parallel", + "target": "step:tidy#3" + }, + { + "id": "step:tidy#3->step:tidy#3:loop", + "kind": "loop", + "label": "over items", + "source": "step:tidy#3", + "target": "step:tidy#3" + }, + { + "id": "step:tidy#3->decision:if:next", + "kind": "next", + "label": null, + "source": "step:tidy#3", + "target": "decision:if" + }, + { + "id": "decision:if->step:tidy#4:routes", + "kind": "routes", + "label": "true", + "source": "decision:if", + "target": "step:tidy#4" + }, + { + "id": "step:fetch#3->step:fetch#3:loop", + "kind": "loop", + "label": "up to 2x", + "source": "step:fetch#3", + "target": "step:fetch#3" + }, + { + "id": "decision:if->step:fetch#3:routes", + "kind": "routes", + "label": "false", + "source": "decision:if", + "target": "step:fetch#3" + }, + { + "id": "step:tidy#4->join:if:next", + "kind": "next", + "label": null, + "source": "step:tidy#4", + "target": "join:if" + }, + { + "id": "step:fetch#3->join:if:next", + "kind": "next", + "label": null, + "source": "step:fetch#3", + "target": "join:if" + }, + { + "id": "join:if->step:discussion:next", + "kind": "next", + "label": null, + "source": "join:if", + "target": "step:discussion" + }, + { + "id": "step:discussion->agent:Researcher:assigned", + "kind": "assigned", + "label": null, + "source": "step:discussion", + "target": "agent:Researcher" + }, + { + "id": "step:discussion->agent:Critic:assigned", + "kind": "assigned", + "label": null, + "source": "step:discussion", + "target": "agent:Critic" + }, + { + "id": "step:discussion->step:Researcher#2:next", + "kind": "next", + "label": null, + "source": "step:discussion", + "target": "step:Researcher#2" + }, + { + "id": "step:Researcher#2->agent:Researcher:assigned", + "kind": "assigned", + "label": null, + "source": "step:Researcher#2", + "target": "agent:Researcher" + }, + { + "id": "step:Researcher#2->step:include:next", + "kind": "next", + "label": null, + "source": "step:Researcher#2", + "target": "step:include" + }, + { + "id": "step:include->end:done:next", + "kind": "next", + "label": null, + "source": "step:include", + "target": "end:done" + } + ], + "kind": "flow", + "name": "pipeline", + "nodes": [ + { + "id": "start:begin", + "kind": "start", + "label": "Brief", + "meta": {}, + "parent": null + }, + { + "id": "step:Researcher", + "kind": "step", + "label": "Researcher", + "meta": {}, + "parent": null + }, + { + "id": "agent:Researcher", + "kind": "agent", + "label": "Researcher", + "meta": { + "role": "Assistant" + }, + "parent": null + }, + { + "id": "tool:Researcher/search_web", + "kind": "tool", + "label": "search_web", + "meta": {}, + "parent": null + }, + { + "id": "step:score", + "kind": "step", + "label": "score", + "meta": {}, + "parent": null + }, + { + "id": "step:route", + "kind": "step", + "label": "route", + "meta": { + "pattern": "route" + }, + "parent": null + }, + { + "id": "decision:route", + "kind": "decision", + "label": "route", + "meta": {}, + "parent": "step:route" + }, + { + "id": "step:tidy", + "kind": "step", + "label": "tidy", + "meta": {}, + "parent": "step:route" + }, + { + "id": "step:fetch", + "kind": "step", + "label": "fetch", + "meta": {}, + "parent": "step:route" + }, + { + "id": "join:route", + "kind": "join", + "label": "join", + "meta": {}, + "parent": "step:route" + }, + { + "id": "step:parallel", + "kind": "step", + "label": "parallel", + "meta": { + "pattern": "parallel" + }, + "parent": null + }, + { + "id": "fork:parallel", + "kind": "fork", + "label": "fork", + "meta": {}, + "parent": "step:parallel" + }, + { + "id": "step:fetch#2", + "kind": "step", + "label": "fetch", + "meta": {}, + "parent": "step:parallel" + }, + { + "id": "step:tidy#2", + "kind": "step", + "label": "tidy", + "meta": {}, + "parent": "step:parallel" + }, + { + "id": "join:parallel", + "kind": "join", + "label": "join", + "meta": {}, + "parent": "step:parallel" + }, + { + "id": "step:loop", + "kind": "step", + "label": "loop", + "meta": { + "over": "items", + "pattern": "loop" + }, + "parent": null + }, + { + "id": "step:tidy#3", + "kind": "step", + "label": "tidy", + "meta": {}, + "parent": "step:loop" + }, + { + "id": "step:if", + "kind": "step", + "label": "if", + "meta": { + "condition": "{{score}} > 80", + "pattern": "if" + }, + "parent": null + }, + { + "id": "decision:if", + "kind": "decision", + "label": "{{score}} > 80", + "meta": {}, + "parent": "step:if" + }, + { + "id": "step:tidy#4", + "kind": "step", + "label": "tidy", + "meta": {}, + "parent": "step:if" + }, + { + "id": "step:repeat", + "kind": "step", + "label": "repeat", + "meta": { + "max_iterations": 2, + "pattern": "repeat" + }, + "parent": "step:if" + }, + { + "id": "step:fetch#3", + "kind": "step", + "label": "fetch", + "meta": {}, + "parent": "step:repeat" + }, + { + "id": "join:if", + "kind": "join", + "label": "join", + "meta": {}, + "parent": "step:if" + }, + { + "id": "step:discussion", + "kind": "step", + "label": "discussion", + "meta": { + "pattern": "discussion" + }, + "parent": null + }, + { + "id": "agent:Critic", + "kind": "agent", + "label": "Critic", + "meta": { + "role": "Assistant" + }, + "parent": null + }, + { + "id": "step:Researcher#2", + "kind": "step", + "label": "Researcher", + "meta": {}, + "parent": null + }, + { + "id": "step:include", + "kind": "step", + "label": "include", + "meta": { + "pattern": "include", + "recipe": "publisher" + }, + "parent": null + }, + { + "id": "end:done", + "kind": "end", + "label": "Delivered", + "meta": {}, + "parent": null + } + ], + "process": "flow", + "schema": "praisonai.graph/1" +} diff --git a/src/praisonai-agents/tests/unit/workflows/golden/graph_hierarchical_team.json b/src/praisonai-agents/tests/unit/workflows/golden/graph_hierarchical_team.json new file mode 100644 index 0000000000..7b9acbe085 --- /dev/null +++ b/src/praisonai-agents/tests/unit/workflows/golden/graph_hierarchical_team.json @@ -0,0 +1,126 @@ +{ + "edges": [ + { + "id": "agent:Researcher->tool:Researcher/search_web:uses", + "kind": "uses", + "label": null, + "source": "agent:Researcher", + "target": "tool:Researcher/search_web" + }, + { + "id": "task:research->agent:Researcher:assigned", + "kind": "assigned", + "label": null, + "source": "task:research", + "target": "agent:Researcher" + }, + { + "id": "task:write->agent:Writer:assigned", + "kind": "assigned", + "label": null, + "source": "task:write", + "target": "agent:Writer" + }, + { + "id": "start:begin->manager:Manager:next", + "kind": "next", + "label": null, + "source": "start:begin", + "target": "manager:Manager" + }, + { + "id": "manager:Manager->agent:Researcher:delegates", + "kind": "delegates", + "label": null, + "source": "manager:Manager", + "target": "agent:Researcher" + }, + { + "id": "manager:Manager->agent:Writer:delegates", + "kind": "delegates", + "label": null, + "source": "manager:Manager", + "target": "agent:Writer" + }, + { + "id": "manager:Manager->end:done:next", + "kind": "next", + "label": null, + "source": "manager:Manager", + "target": "end:done" + }, + { + "id": "task:research->task:write:depends_on", + "kind": "depends_on", + "label": null, + "source": "task:research", + "target": "task:write" + } + ], + "kind": "team", + "name": "managed", + "nodes": [ + { + "id": "start:begin", + "kind": "start", + "label": "Brief", + "meta": {}, + "parent": null + }, + { + "id": "manager:Manager", + "kind": "manager", + "label": "Manager", + "meta": {}, + "parent": null + }, + { + "id": "agent:Researcher", + "kind": "agent", + "label": "Researcher", + "meta": { + "role": "Assistant" + }, + "parent": null + }, + { + "id": "tool:Researcher/search_web", + "kind": "tool", + "label": "search_web", + "meta": {}, + "parent": null + }, + { + "id": "agent:Writer", + "kind": "agent", + "label": "Writer", + "meta": { + "role": "Assistant" + }, + "parent": null + }, + { + "id": "task:research", + "kind": "task", + "label": "research", + "meta": {}, + "parent": null + }, + { + "id": "task:write", + "kind": "task", + "label": "write", + "meta": {}, + "parent": null + }, + { + "id": "end:done", + "kind": "end", + "label": "Delivered", + "meta": {}, + "parent": null + } + ], + "process": "hierarchical", + "schema": "praisonai.graph/1" +} diff --git a/src/praisonai-agents/tests/unit/workflows/golden/graph_sequential_team.json b/src/praisonai-agents/tests/unit/workflows/golden/graph_sequential_team.json new file mode 100644 index 0000000000..fa48c7dddd --- /dev/null +++ b/src/praisonai-agents/tests/unit/workflows/golden/graph_sequential_team.json @@ -0,0 +1,126 @@ +{ + "edges": [ + { + "id": "agent:Researcher->tool:Researcher/search_web:uses", + "kind": "uses", + "label": null, + "source": "agent:Researcher", + "target": "tool:Researcher/search_web" + }, + { + "id": "task:research->agent:Researcher:assigned", + "kind": "assigned", + "label": null, + "source": "task:research", + "target": "agent:Researcher" + }, + { + "id": "task:write->agent:Writer:assigned", + "kind": "assigned", + "label": null, + "source": "task:write", + "target": "agent:Writer" + }, + { + "id": "task:write#2->agent:Writer:assigned", + "kind": "assigned", + "label": null, + "source": "task:write#2", + "target": "agent:Writer" + }, + { + "id": "start:begin->task:research:next", + "kind": "next", + "label": null, + "source": "start:begin", + "target": "task:research" + }, + { + "id": "task:research->task:write:next", + "kind": "next", + "label": null, + "source": "task:research", + "target": "task:write" + }, + { + "id": "task:write->task:write#2:next", + "kind": "next", + "label": null, + "source": "task:write", + "target": "task:write#2" + }, + { + "id": "task:write#2->end:done:next", + "kind": "next", + "label": null, + "source": "task:write#2", + "target": "end:done" + } + ], + "kind": "team", + "name": "pair", + "nodes": [ + { + "id": "start:begin", + "kind": "start", + "label": "Brief", + "meta": {}, + "parent": null + }, + { + "id": "agent:Researcher", + "kind": "agent", + "label": "Researcher", + "meta": { + "role": "Assistant" + }, + "parent": null + }, + { + "id": "tool:Researcher/search_web", + "kind": "tool", + "label": "search_web", + "meta": {}, + "parent": null + }, + { + "id": "agent:Writer", + "kind": "agent", + "label": "Writer", + "meta": { + "role": "Assistant" + }, + "parent": null + }, + { + "id": "task:research", + "kind": "task", + "label": "research", + "meta": {}, + "parent": null + }, + { + "id": "task:write", + "kind": "task", + "label": "write", + "meta": {}, + "parent": null + }, + { + "id": "task:write#2", + "kind": "task", + "label": "write", + "meta": {}, + "parent": null + }, + { + "id": "end:done", + "kind": "end", + "label": "Delivered", + "meta": {}, + "parent": null + } + ], + "process": "sequential", + "schema": "praisonai.graph/1" +} diff --git a/src/praisonai-agents/tests/unit/workflows/golden/graph_team.json b/src/praisonai-agents/tests/unit/workflows/golden/graph_team.json new file mode 100644 index 0000000000..538185b803 --- /dev/null +++ b/src/praisonai-agents/tests/unit/workflows/golden/graph_team.json @@ -0,0 +1,186 @@ +{ + "edges": [ + { + "id": "agent:Researcher->tool:Researcher/search_web:uses", + "kind": "uses", + "label": null, + "source": "agent:Researcher", + "target": "tool:Researcher/search_web" + }, + { + "id": "agent:Writer->agent:Billing:handoff", + "kind": "handoff", + "label": null, + "source": "agent:Writer", + "target": "agent:Billing" + }, + { + "id": "task:research->agent:Researcher:assigned", + "kind": "assigned", + "label": null, + "source": "task:research", + "target": "agent:Researcher" + }, + { + "id": "task:draft->agent:Writer:assigned", + "kind": "assigned", + "label": null, + "source": "task:draft", + "target": "agent:Writer" + }, + { + "id": "human:you->task:draft:reviews", + "kind": "reviews", + "label": null, + "source": "human:you", + "target": "task:draft" + }, + { + "id": "task:publish->agent:Writer:assigned", + "kind": "assigned", + "label": null, + "source": "task:publish", + "target": "agent:Writer" + }, + { + "id": "task:3->agent:Writer:assigned", + "kind": "assigned", + "label": null, + "source": "task:3", + "target": "agent:Writer" + }, + { + "id": "start:begin->task:research:next", + "kind": "next", + "label": null, + "source": "start:begin", + "target": "task:research" + }, + { + "id": "task:research->task:draft:next", + "kind": "next", + "label": null, + "source": "task:research", + "target": "task:draft" + }, + { + "id": "task:draft->task:publish:routes", + "kind": "routes", + "label": "approved", + "source": "task:draft", + "target": "task:publish" + }, + { + "id": "task:draft->task:research:routes", + "kind": "routes", + "label": "rejected", + "source": "task:draft", + "target": "task:research" + }, + { + "id": "task:publish->end:done:next", + "kind": "next", + "label": null, + "source": "task:publish", + "target": "end:done" + }, + { + "id": "task:3->end:done:next", + "kind": "next", + "label": null, + "source": "task:3", + "target": "end:done" + } + ], + "kind": "team", + "name": "newsroom", + "nodes": [ + { + "id": "start:begin", + "kind": "start", + "label": "Brief", + "meta": {}, + "parent": null + }, + { + "id": "agent:Researcher", + "kind": "agent", + "label": "Researcher", + "meta": { + "role": "Assistant" + }, + "parent": null + }, + { + "id": "tool:Researcher/search_web", + "kind": "tool", + "label": "search_web", + "meta": {}, + "parent": null + }, + { + "id": "agent:Writer", + "kind": "agent", + "label": "Writer", + "meta": { + "role": "Assistant" + }, + "parent": null + }, + { + "id": "agent:Billing", + "kind": "agent", + "label": "Billing", + "meta": { + "role": "Assistant" + }, + "parent": null + }, + { + "id": "human:you", + "kind": "human", + "label": "You", + "meta": {}, + "parent": null + }, + { + "id": "task:research", + "kind": "task", + "label": "research", + "meta": {}, + "parent": null + }, + { + "id": "task:draft", + "kind": "task", + "label": "draft", + "meta": { + "human_input": true + }, + "parent": null + }, + { + "id": "task:publish", + "kind": "task", + "label": "publish", + "meta": {}, + "parent": null + }, + { + "id": "task:3", + "kind": "task", + "label": "Archive everything afterwards", + "meta": {}, + "parent": null + }, + { + "id": "end:done", + "kind": "end", + "label": "Delivered", + "meta": {}, + "parent": null + } + ], + "process": "workflow", + "schema": "praisonai.graph/1" +} diff --git a/src/praisonai-agents/tests/unit/workflows/test_graph.py b/src/praisonai-agents/tests/unit/workflows/test_graph.py new file mode 100644 index 0000000000..ebd28ff02a --- /dev/null +++ b/src/praisonai-agents/tests/unit/workflows/test_graph.py @@ -0,0 +1,110 @@ +"""Graph data (``praisonai.graph/1``) for teams, flows and agents. + +The replay UI draws the Map from this graph and joins it to trace events by +``node_id``, so the ids must be derived from the definition alone: identical +in every process and on every run. The goldens pin the full shape; regenerate +them deliberately with ``PRAISONAI_UPDATE_GOLDEN=1``. +""" +import json +import os +import subprocess +import sys +from pathlib import Path + +import pytest + +HERE = Path(__file__).parent +GOLDEN = HERE / "golden" +sys.path.insert(0, str(HERE)) +import _graph_fixtures as fixtures # noqa: E402 + +from praisonaiagents.workflows.diagram import graph_to_mermaid, to_graph # noqa: E402 + + +@pytest.mark.parametrize("name", ["team", "sequential_team", "hierarchical_team", "flow"]) +def test_graph_matches_golden(name): + graph = to_graph(fixtures.BUILDERS[name]()) + path = GOLDEN / f"graph_{name}.json" + if os.environ.get("PRAISONAI_UPDATE_GOLDEN") == "1": + path.write_text(json.dumps(graph, indent=2, sort_keys=True) + "\n") + assert graph == json.loads(path.read_text()) + + +def test_ids_are_stable_across_processes(): + """Two fresh interpreters with different hash seeds produce the same graphs.""" + code = ( + "import json, sys; sys.path.insert(0, sys.argv[1]);" + "import _graph_fixtures as f;" + "from praisonaiagents.workflows.diagram import to_graph;" + "print(json.dumps({k: to_graph(b()) for k, b in sorted(f.BUILDERS.items())}, sort_keys=True))" + ) + root = str(Path(__import__("praisonaiagents").__file__).parent.parent) + outputs = [] + for seed in ("1", "2"): + env = dict(os.environ, PYTHONHASHSEED=seed, PYTHONPATH=root, PYTHONDONTWRITEBYTECODE="1") + out = subprocess.run([sys.executable, "-c", code, str(HERE)], env=env, + capture_output=True, text=True, timeout=120, check=True) + outputs.append(json.loads(out.stdout.strip().splitlines()[-1])) + assert outputs[0] == outputs[1] + assert outputs[0]["flow"] == to_graph(fixtures.flow()) + + +def test_every_edge_joins_known_nodes_and_ids_are_unique(): + for build in fixtures.BUILDERS.values(): + graph = to_graph(build()) + ids = [n["id"] for n in graph["nodes"]] + assert len(ids) == len(set(ids)) + for edge in graph["edges"]: + assert edge["source"] in ids and edge["target"] in ids, edge + assert edge["id"] == f'{edge["source"]}->{edge["target"]}:{edge["kind"]}' + for node in graph["nodes"]: + assert node["parent"] is None or node["parent"] in ids + + +def test_no_uuid_reaches_the_graph(): + import re + + uuid = re.compile(r"[0-9a-f]{8}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{12}") + for build in fixtures.BUILDERS.values(): + assert not uuid.search(json.dumps(to_graph(build()))) + + +def test_team_and_flow_methods(): + team = fixtures.hierarchical_team() + assert team.to_graph() == to_graph(team) + diagram = team.to_mermaid() + assert diagram.startswith("graph TD") + assert '{{"Manager"}}' in diagram and '-->|"delegates"|' in diagram + flow = fixtures.flow() + assert flow.to_graph() == to_graph(flow) + + +def test_agent_graph_has_tools_and_handoffs(): + from praisonaiagents import Agent + + billing = Agent(name="Billing", instructions="b") + triage = Agent(name="Triage", instructions="t", tools=[fixtures.search_web], handoffs=[billing]) + graph = to_graph(triage) + assert graph["kind"] == "agent" + edges = {e["id"] for e in graph["edges"]} + assert { + "start:begin->agent:Triage:next", + "agent:Triage->tool:Triage/search_web:uses", + "agent:Triage->agent:Billing:handoff", + "agent:Triage->end:done:next", + } <= edges + # The hand-off's generated transfer tool is the hand-off edge, not a tool node. + assert not any(n["id"].startswith("tool:Triage/transfer_to") for n in graph["nodes"]) + + +def test_graph_to_mermaid_draws_leaf_nodes_not_groups(): + graph = to_graph(fixtures.flow()) + diagram = graph_to_mermaid(graph) + assert diagram.startswith("graph TD") + assert '(("fork"))' in diagram and '{"route"}' in diagram + assert '["parallel"]' not in diagram # a pattern group is a container + + +def test_unsupported_object_is_rejected(): + with pytest.raises(TypeError): + to_graph(42) From 08ca3aca1879a0cd020e2b71457e4d57768119b9 Mon Sep 17 00:00:00 2001 From: MervinPraison <454862+MervinPraison@users.noreply.github.com> Date: Mon, 26 Oct 2026 11:22:42 +0000 Subject: [PATCH 03/12] feat(replay): record hand-offs, team tasks, flow steps, manager choices and approvals - agent/handoff.py: every place a hand-off actually happens (programmatic, async, LLM tool, typed sync/async) emits one agent_handoff(kind="handoff") with the context policy and redacted args. - agents/agents.py: task_start/task_end (completed, failed after retries, skipped by should_run) on the sync and async paths in every process mode; workflow_start/workflow_end with the team graph inside _run_scope(run_id); approval_requested/approval_decided around human review. - workflows/workflows.py: step_start/step_end for top-level and nested steps (completed/failed/skipped/cached) via helpers that wrap the existing on_step_start/on_step_complete calls; Route and If record the branch taken; pattern executors are wrapped once so every nesting level is covered; workflow_start/workflow_end with the flow graph and a fresh run id. - process/process.py: the hierarchical manager's choice is an agent_handoff from "Manager" with kind="delegate", the task node id and a task-status summary (sync and async). - approval/registry.py: approval_requested/approval_decided around the real backend call only (not the fast paths), args redacted. All of it is skipped when the context emitter is disabled; the graph builder is imported only for a traced run. Co-Authored-By: Claude Opus 5.5 (1M context) --- .../praisonaiagents/agent/handoff.py | 26 ++ .../praisonaiagents/agents/agents.py | 98 ++++++- .../praisonaiagents/approval/registry.py | 108 +++++--- .../praisonaiagents/process/process.py | 27 ++ .../praisonaiagents/workflows/workflows.py | 254 ++++++++++++++---- 5 files changed, 419 insertions(+), 94 deletions(-) diff --git a/src/praisonai-agents/praisonaiagents/agent/handoff.py b/src/praisonai-agents/praisonaiagents/agent/handoff.py index bd2f850650..2c113c9a74 100644 --- a/src/praisonai-agents/praisonaiagents/agent/handoff.py +++ b/src/praisonai-agents/praisonaiagents/agent/handoff.py @@ -630,6 +630,27 @@ async def _execute_with_runtime_resolution_async( None, lambda: self.agent.chat(prompt, tools=effective_tools) ) + def _trace_handoff(self, source_agent: 'Agent', context: Optional[Dict[str, Any]], reason: Optional[str] = None) -> None: + """Record the hand-off on the replay trace (no-op when tracing is off).""" + from ..trace.context_events import get_context_emitter + emitter = get_context_emitter() + if not emitter._enabled: + return + context = dict(context or {}) + if reason is None: + reason = context.get("reason") or context.get("task") or context.get("prompt") + policy = getattr(self.config, "context_policy", None) + emitter.agent_handoff( + getattr(source_agent, "name", None) or "agent", + getattr(self.agent, "name", None), + reason=emitter._preview(reason) if reason is not None else None, + context_passed={ + "context_policy": getattr(policy, "value", policy), + "args": context, + }, + kind="handoff", + ) + def _check_safety(self, source_agent: 'Agent') -> None: """ Check safety constraints before handoff. @@ -778,6 +799,7 @@ def execute_programmatic( # Track handoff chain _push_handoff(source_agent.name) pushed = True + self._trace_handoff(source_agent, kwargs, reason=prompt) # Execute on_handoff callback self._execute_callback(self.config.on_handoff or self.on_handoff, source_agent, kwargs) @@ -888,6 +910,7 @@ async def _execute(): self._check_safety(source_agent) _push_handoff(source_agent.name) pushed = True + self._trace_handoff(source_agent, kwargs, reason=prompt) # Execute callback self._execute_callback(self.config.on_handoff or self.on_handoff, source_agent, kwargs) @@ -1016,6 +1039,7 @@ def handoff_tool(**kwargs): # Track handoff chain _push_handoff(source_agent.name) pushed = True + self._trace_handoff(source_agent, kwargs, reason=None) # Execute on_handoff callback self._execute_callback(self.config.on_handoff or self.on_handoff, source_agent, kwargs) @@ -1632,6 +1656,7 @@ def execute_programmatic( # Track handoff chain _push_handoff(source_agent.name) pushed = True + self._trace_handoff(source_agent, kwargs, reason=None) # Execute on_handoff callback self._execute_callback(self.config.on_handoff or self.on_handoff, source_agent, kwargs) @@ -1747,6 +1772,7 @@ async def _execute(): self._check_safety(source_agent) _push_handoff(source_agent.name) pushed = True + self._trace_handoff(source_agent, kwargs, reason=None) # Execute callback self._execute_callback(self.config.on_handoff or self.on_handoff, source_agent, kwargs) diff --git a/src/praisonai-agents/praisonaiagents/agents/agents.py b/src/praisonai-agents/praisonaiagents/agents/agents.py index 27cfbc4f96..add920ad7a 100644 --- a/src/praisonai-agents/praisonaiagents/agents/agents.py +++ b/src/praisonai-agents/praisonaiagents/agents/agents.py @@ -356,6 +356,39 @@ def _record_skipped_task(agents_instance, task): task.status = "completed" +def _trace_task_begin(agents_instance, task): + """Emit ``task_start``; returns a span for ``_trace_task_end`` (None when off).""" + from ..trace.context_events import get_context_emitter + emitter = get_context_emitter() + if not emitter._enabled: + return None + from ..workflows.diagram import task_node_id + node_id = task_node_id(agents_instance.tasks, task) + agent = getattr(getattr(task, "agent", None), "name", None) + emitter.task_start( + getattr(task, "name", None), getattr(task, "description", None), + agent=agent, node_id=node_id, + ) + return emitter, node_id, agent, time.perf_counter() + + +def _trace_task_end(span, task, status): + """Emit ``task_end`` (completed / failed / skipped) for a ``_trace_task_begin`` span.""" + if span is None: + return + emitter, node_id, agent, started = span + result = getattr(task, "result", None) + emitter.task_end( + getattr(task, "name", None), + status, + retries=getattr(task, "retry_count", 0) or 0, + duration_ms=(time.perf_counter() - started) * 1000, + output=getattr(result, "raw", None) if result is not None else None, + agent=agent, + node_id=node_id, + ) + + def _task_should_skip(agents_instance, task): """Return True when a task's should_run gate evaluates falsy. @@ -1539,7 +1572,19 @@ def _apply_human_review(self, task, task_id, task_output): agent_name=getattr(getattr(task, "agent", None), "name", None), context={"kind": "task_output_review", "prompt": prompt}, ) - decision = backend.request_approval_sync(request) + from ..trace.context_events import _approval_begin, _approval_end, get_context_emitter + _approval = None + if get_context_emitter()._enabled: + from ..workflows.diagram import task_node_id + _approval = _approval_begin( + request, kind="task_output_review", node_id=task_node_id(self.tasks, task), + ) + try: + decision = backend.request_approval_sync(request) + except BaseException as exc: + _approval_end(_approval, error=exc) + raise + _approval_end(_approval, decision) if getattr(decision, "approved", False): return task_output, False @@ -1740,6 +1785,18 @@ async def arun_task(self, task_id): logger.info(f"Task with ID {task_id} is already completed") return + _span = _trace_task_begin(self, task) + _status = "failed" + try: + if await self._arun_task_body(task_id, task) == "skipped": + _status = "skipped" + elif task.status == "completed": + _status = "completed" + finally: + _trace_task_end(_span, task, _status) + + async def _arun_task_body(self, task_id, task): + """Body of :meth:`arun_task`; returns ``"skipped"`` for a gated-out task.""" # Call on_task_start callback and propagate variables (async-aware, mirrors run_task) await self._arun_task_start_hook(task, task_id) @@ -1747,7 +1804,7 @@ async def arun_task(self, task_id): if await _atask_should_skip(self, task): logger.info(f"Task {task_id} skipped by should_run gate") _record_skipped_task(self, task) - return + return "skipped" # Use per-task max_retries if available task_max = getattr(task, "max_retries", self.max_retries) @@ -2201,7 +2258,12 @@ async def astart(self, content=None, return_dict=False, **kwargs): # Re-entrant on the owning thread so batch runs (astart_for_each) can hold # the guard across the whole batch while still calling astart() per item. with self._guard_execution(): - return await self._astart_impl(content=content, return_dict=return_dict, **kwargs) + from ..trace.context_events import _workflow_scope + with _workflow_scope("team", self.name, self.process, + graph_fn=self.to_graph, run_id=self.run_id) as _run: + result = await self._astart_impl(content=content, return_dict=return_dict, **kwargs) + self._mark_run_status(_run) + return result async def _astart_impl(self, content=None, return_dict=False, **kwargs): # Track execution via telemetry @@ -2318,6 +2380,18 @@ def run_task(self, task_id): logger.info(f"Task with ID {task_id} is already completed") return + _span = _trace_task_begin(self, task) + _status = "failed" + try: + if self._run_task_body(task_id, task) == "skipped": + _status = "skipped" + elif task.status == "completed": + _status = "completed" + finally: + _trace_task_end(_span, task, _status) + + def _run_task_body(self, task_id, task): + """Body of :meth:`run_task`; returns ``"skipped"`` for a gated-out task.""" # Call on_task_start callback and propagate variables (shared with arun_task) self._run_task_start_hook(task, task_id) @@ -2325,7 +2399,7 @@ def run_task(self, task_id): if _task_should_skip(self, task): logger.info(f"Task {task_id} skipped by should_run gate") _record_skipped_task(self, task) - return + return "skipped" # Use per-task max_retries if available task_max = getattr(task, "max_retries", self.max_retries) @@ -2534,8 +2608,20 @@ def start(self, content=None, return_dict=False, output=None, **kwargs): # Re-entrant on the owning thread so batch runs (start_for_each) can hold # the guard across the whole batch while still calling start() per item. with self._guard_execution(): - return self._start_impl(content=content, return_dict=return_dict, - output=output, **kwargs) + from ..trace.context_events import _workflow_scope + with _workflow_scope("team", self.name, self.process, + graph_fn=self.to_graph, run_id=self.run_id) as _run: + result = self._start_impl(content=content, return_dict=return_dict, + output=output, **kwargs) + self._mark_run_status(_run) + return result + + def _mark_run_status(self, run_state): + """Record a failed replay run when any task ended failed (tracing on only).""" + if run_state is not None and any( + getattr(t, "status", None) == "failed" for t in self.tasks.values() + ): + run_state["status"] = "failed" def _start_impl(self, content=None, return_dict=False, output=None, **kwargs): # Track execution via telemetry diff --git a/src/praisonai-agents/praisonaiagents/approval/registry.py b/src/praisonai-agents/praisonaiagents/approval/registry.py index 3214c4efc9..28be129b48 100644 --- a/src/praisonai-agents/praisonaiagents/approval/registry.py +++ b/src/praisonai-agents/praisonaiagents/approval/registry.py @@ -689,29 +689,41 @@ def approve_sync( liveness=liveness, ) - with _human_wait_scope(): - # Prefer sync method if available - if hasattr(backend, "request_approval_sync"): - decision = backend.request_approval_sync(request) - else: - # Use shared utility for consistent async-to-sync bridging - from .utils import run_coroutine_safely - decision = run_coroutine_safely( - backend.request_approval(request), - timeout=self.timeout - ) - - # Live-authority gate: a resolution that lands after the originating - # turn was stopped/superseded is dropped fail-closed — the tool does not - # run and no durable session/always grant is persisted. - stale = self._reject_if_stale(request) - if stale is not None: - return stale - - if decision.approved: - self.mark_approved(tool_name, arguments, agent_name, scope_id) - self._persist_scoped_decision(agent_name, tool_name, arguments, decision, scope_id) - return decision + # Replay: a person is actually being asked (the fast paths above are not). + from ..trace.context_events import _approval_begin, _approval_end + _trace = _approval_begin(request) + try: + with _human_wait_scope(): + # Prefer sync method if available + if hasattr(backend, "request_approval_sync"): + decision = backend.request_approval_sync(request) + else: + # Use shared utility for consistent async-to-sync bridging + from .utils import run_coroutine_safely + decision = run_coroutine_safely( + backend.request_approval(request), + timeout=self.timeout + ) + + # Live-authority gate: a resolution that lands after the originating + # turn was stopped/superseded is dropped fail-closed — the tool does not + # run and no durable session/always grant is persisted. + stale = self._reject_if_stale(request) + if stale is not None: + decision = stale + return stale + + if decision.approved: + self.mark_approved(tool_name, arguments, agent_name, scope_id) + self._persist_scoped_decision(agent_name, tool_name, arguments, decision, scope_id) + return decision + except BaseException as exc: + _approval_end(_trace, error=exc) + _trace = None + raise + finally: + if _trace is not None: + _approval_end(_trace, decision) async def approve_async( self, @@ -764,26 +776,38 @@ async def approve_async( liveness=liveness, ) + # Replay: a person is actually being asked (the fast paths above are not). + from ..trace.context_events import _approval_begin, _approval_end + _trace = _approval_begin(request) try: - async with _human_wait_scope(): - decision = await asyncio.wait_for( - backend.request_approval(request), - timeout=self.timeout, - ) - except asyncio.TimeoutError: - decision = ApprovalDecision(approved=False, reason="Approval timed out") - - # Live-authority gate: a resolution that lands after the originating - # turn was stopped/superseded is dropped fail-closed — the tool does not - # run and no durable session/always grant is persisted. - stale = self._reject_if_stale(request) - if stale is not None: - return stale - - if decision.approved: - self.mark_approved(tool_name, arguments, agent_name, scope_id) - self._persist_scoped_decision(agent_name, tool_name, arguments, decision, scope_id) - return decision + try: + async with _human_wait_scope(): + decision = await asyncio.wait_for( + backend.request_approval(request), + timeout=self.timeout, + ) + except asyncio.TimeoutError: + decision = ApprovalDecision(approved=False, reason="Approval timed out") + + # Live-authority gate: a resolution that lands after the originating + # turn was stopped/superseded is dropped fail-closed — the tool does not + # run and no durable session/always grant is persisted. + stale = self._reject_if_stale(request) + if stale is not None: + decision = stale + return stale + + if decision.approved: + self.mark_approved(tool_name, arguments, agent_name, scope_id) + self._persist_scoped_decision(agent_name, tool_name, arguments, decision, scope_id) + return decision + except BaseException as exc: + _approval_end(_trace, error=exc) + _trace = None + raise + finally: + if _trace is not None: + _approval_end(_trace, decision) @staticmethod def _reject_if_stale(request: ApprovalRequest) -> Optional[ApprovalDecision]: diff --git a/src/praisonai-agents/praisonaiagents/process/process.py b/src/praisonai-agents/praisonaiagents/process/process.py index 09b5b4d27b..d75611183c 100644 --- a/src/praisonai-agents/praisonaiagents/process/process.py +++ b/src/praisonai-agents/praisonaiagents/process/process.py @@ -80,6 +80,31 @@ async def _get_state_lock(self): self._state_lock = asyncio.Lock() return self._state_lock + def _trace_delegate(self, task_id, tasks_summary): + """Record the manager's choice as a ``delegate`` hand-off (tracing on only).""" + from ..trace.context_events import get_context_emitter + emitter = get_context_emitter() + if not emitter._enabled: + return + from ..workflows.diagram import task_node_id + task = self.tasks[task_id] + agent = getattr(getattr(task, "agent", None), "name", None) + emitter.agent_handoff( + "Manager", + agent, + reason=f"Manager delegated task {getattr(task, 'name', None) or task_id}", + context_passed={ + "task_id": task_id, + "task_status": [ + {"task_id": t.get("task_id"), "name": t.get("name"), + "status": t.get("status"), "agent": t.get("agent")} + for t in tasks_summary + ], + }, + kind="delegate", + node_id=task_node_id(self.tasks, task), + ) + def _task_by_name(self, name): """Return the task whose name matches, or None. First-match, mirrors prior inline lookup.""" return next((t for t in self.tasks.values() if t.name == name), None) @@ -1256,6 +1281,7 @@ async def ahierarchical(self) -> AsyncGenerator[str, None]: if self.tasks[selected_task_id].status != "completed": logging.info(f"Starting execution of task {selected_task_id}") + self._trace_delegate(selected_task_id, tasks_summary) yield selected_task_id logging.info(f"Finished execution of task {selected_task_id}, status: {self.tasks[selected_task_id].status}") # Count a task toward completion only the first time it actually @@ -1916,6 +1942,7 @@ def hierarchical(self): if self.tasks[selected_task_id].status != "completed": logging.info(f"Starting execution of task {selected_task_id}") + self._trace_delegate(selected_task_id, tasks_summary) yield selected_task_id logging.info(f"Finished execution of task {selected_task_id}, status: {self.tasks[selected_task_id].status}") # Count a task toward completion only the first time it actually diff --git a/src/praisonai-agents/praisonaiagents/workflows/workflows.py b/src/praisonai-agents/praisonaiagents/workflows/workflows.py index 2dcbf16d1e..5aec581a0d 100644 --- a/src/praisonai-agents/praisonaiagents/workflows/workflows.py +++ b/src/praisonai-agents/praisonaiagents/workflows/workflows.py @@ -725,6 +725,97 @@ def include(recipe: Optional[str] = None, workflow: Optional[Any] = None, input: # Maximum nesting depth for patterns (prevents stack overflow) MAX_NESTING_DEPTH = 5 + +# --------------------------------------------------------------------------- +# Replay step events (step_start / step_end). Every helper is a no-op returning +# None when the context emitter is disabled, so an untraced run pays one +# attribute check per step. +# --------------------------------------------------------------------------- + +# The innermost pattern span (Route/If) so it can record the branch it chose. +_replay_step_span: "contextvars.ContextVar[Optional[dict]]" = contextvars.ContextVar( + "praisonai_replay_step_span", default=None +) + + +def _replay_step_begin(flow, step, name=None, pattern=None): + """Emit ``step_start`` for a definition object; returns a span or None.""" + from ..trace.context_events import get_context_emitter + emitter = get_context_emitter() + if not emitter._enabled: + return None + node_id = (getattr(flow, "_replay_node_ids", None) or {}).get(id(step)) + if name is None: + if node_id is not None: + name = node_id[len("step:"):] if node_id.startswith("step:") else node_id + else: + from .diagram import _label + name = pattern or _label(step) + if node_id is None: + node_id = f"step:{name}" + agent = getattr(step, "agent", None) if not hasattr(step, "chat") else step + agent = getattr(agent, "name", None) if agent is not None else None + emitter.step_start(name, pattern=pattern, node_id=node_id, agent=agent) + return {"emitter": emitter, "name": name, "node_id": node_id, "agent": agent, + "started": time.perf_counter()} + + +def _replay_step_end(span, status, output=None): + """Emit ``step_end`` for a span from :func:`_replay_step_begin`.""" + if span is None: + return + span["emitter"].step_end( + span["name"], + status, + duration_ms=(time.perf_counter() - span["started"]) * 1000, + output=output, + route=span.get("route"), + node_id=span["node_id"], + agent=span["agent"], + ) + + +def _replay_note_route(label): + """Record which branch the enclosing Route/If chose (tracing on only).""" + span = _replay_step_span.get() + if span is not None: + span["route"] = label + + +def _replay_pattern(pattern): + """Wrap a pattern executor (``_execute_route`` ...) in step_start/step_end. + + Top-level and nested patterns both reach the executor, so wrapping it here + covers every nesting level once. + """ + def decorate(execute): + def wrapper(self, step, *args, **kwargs): + name = getattr(step, "name", None) if pattern == "discussion" else None + span = _replay_step_begin(self, step, name, pattern) + if span is None: + return execute(self, step, *args, **kwargs) + token = _replay_step_span.set(span) + try: + result = execute(self, step, *args, **kwargs) + except BaseException: + _replay_step_end(span, "failed") + raise + finally: + _replay_step_span.reset(token) + is_dict = isinstance(result, dict) + _replay_step_end( + span, + "failed" if is_dict and result.get("stop") else "completed", + result.get("output") if is_dict else None, + ) + return result + wrapper.__name__ = execute.__name__ + wrapper.__qualname__ = execute.__qualname__ + wrapper.__doc__ = execute.__doc__ + wrapper.__wrapped__ = execute + return wrapper + return decorate + @dataclass class If: """ @@ -1464,19 +1555,26 @@ def run( "separate Workflow instance per concurrent run." ) try: - if self.tools_run_on is None: - return self._run_impl(input, llm, verbose, stream) - # One sandbox for the whole run; torn down even if a step raises. - with self._shared_compute() as shared: - shared.attach(self._collect_agents()) - # Expose the live sandbox so agents created *during* the run - # (e.g. an action-only step's temporary agent) can be attached - # too, instead of silently executing tools locally. - self._active_shared_compute = shared - try: - return self._run_impl(input, llm, verbose, stream) - finally: - self._active_shared_compute = None + from ..trace.context_events import _workflow_scope + _process = "hierarchical" if self.process == "hierarchical" else "flow" + with _workflow_scope("flow", self.name, _process, graph_fn=self._replay_graph) as _run: + if self.tools_run_on is None: + result = self._run_impl(input, llm, verbose, stream) + else: + # One sandbox for the whole run; torn down even if a step raises. + with self._shared_compute() as shared: + shared.attach(self._collect_agents()) + # Expose the live sandbox so agents created *during* the run + # (e.g. an action-only step's temporary agent) can be attached + # too, instead of silently executing tools locally. + self._active_shared_compute = shared + try: + result = self._run_impl(input, llm, verbose, stream) + finally: + self._active_shared_compute = None + if _run is not None and isinstance(result, dict) and result.get("status") == "failed": + _run["status"] = "failed" + return result finally: # Restore the generation the caller's context carried before this # run stamped its own, so a nested Workflow.run() (Include) does not @@ -1610,6 +1708,34 @@ def visit(step: Any) -> None: visit(self.steps) return found + def _notify_step_start(self, raw_step: Any, step: Any, context: Any) -> Optional[dict]: + """Emit replay ``step_start`` then call ``on_step_start``; returns the span.""" + span = _replay_step_begin(self, raw_step, step.name) + if self.on_step_start: + try: + self.on_step_start(step.name, context) + except Exception as e: + logger.error(f"on_step_start callback failed: {e}") + return span + + def _notify_step_complete( + self, span: Optional[dict], step_name: str, step_result: Any, status: str, output: Any = None, + ) -> None: + """Call ``on_step_complete`` then emit replay ``step_end``.""" + if self.on_step_complete: + try: + self.on_step_complete(step_name, step_result) + except Exception as e: + logger.error(f"on_step_complete callback failed: {e}") + _replay_step_end(span, status, output) + + def _replay_graph(self) -> Dict[str, Any]: + """Graph for ``workflow_start``; also maps step objects to node ids.""" + from .diagram import _flow_graph + graph = _flow_graph(self) + self._replay_node_ids = graph.ids + return graph.to_dict() + def _run_impl( self, input: str = "", @@ -1830,6 +1956,7 @@ def _run_impl( continue # Normalize single step + raw_step = step step = self._normalize_single_step(step, i) # Create context for handlers @@ -1848,11 +1975,7 @@ def _run_impl( # Preserve the preparation hook before the gate: it may populate # context variables used by should_run, including on cached runs. - if self.on_step_start: - try: - self.on_step_start(step.name, context) - except Exception as e: - logger.error(f"on_step_start callback failed: {e}") + _replay = self._notify_step_start(raw_step, step, context) # Check should_run condition if step.should_run: @@ -1863,11 +1986,9 @@ def _run_impl( if hasattr(step, 'status'): step.status = "skipped" self.step_statuses[step.name] = "skipped" - if self.on_step_complete: - try: - self.on_step_complete(step.name, StepResult(output="", skipped=True)) - except Exception as e: - logger.error(f"on_step_complete callback failed: {e}") + self._notify_step_complete( + _replay, step.name, StepResult(output="", skipped=True), "skipped", + ) i += 1 continue except Exception as e: @@ -1913,16 +2034,12 @@ def _run_impl( step.status = cached_record["status"] if _cached.get("variables"): all_variables.update(_cached["variables"]) - if self.on_step_complete: - try: - self.on_step_complete( - step.name, StepResult( - output=previous_output or "", - stop_workflow=bool(_cached.get("stop")), - ), - ) - except Exception as e: - logger.error(f"on_step_complete callback failed: {e}") + self._notify_step_complete( + _replay, step.name, StepResult( + output=previous_output or "", + stop_workflow=bool(_cached.get("stop")), + ), "cached", previous_output, + ) if _cached.get("stop"): if verbose: print(f"🛑 Workflow stopped at: {step.name}") @@ -2205,14 +2322,10 @@ def _run_impl( }) # Still notify the step-complete callback so lifecycle observers # see the (failed) terminal step before the workflow aborts. - if self.on_step_complete: - try: - self.on_step_complete( - step.name, - StepResult(output=output or "", stop_workflow=True), - ) - except Exception as e: - logger.error(f"on_step_complete callback failed: {e}") + self._notify_step_complete( + _replay, step.name, + StepResult(output=output or "", stop_workflow=True), "failed", output, + ) if verbose: print(f"🛑 Workflow stopped: step '{step.name}' failed (on_error='stop')") break @@ -2221,11 +2334,10 @@ def _run_impl( step_result = StepResult(output=output or "", stop_workflow=stop) # Call on_step_complete callback - if self.on_step_complete: - try: - self.on_step_complete(step.name, step_result) - except Exception as e: - logger.error(f"on_step_complete callback failed: {e}") + self._notify_step_complete( + _replay, step.name, step_result, + "failed" if step_failed else "completed", output, + ) # Handle output_file - save output to file if hasattr(step, 'output_file') and step.output_file and output: @@ -2993,10 +3105,47 @@ def _execute_single_step_internal( # them on every run so their conditions and nested Task gates execute; # their individual leaves can still reuse cached results. if isinstance(step, (Loop, Parallel, Route, Repeat, Discussion, If, Include)): + # Pattern executors emit their own replay step events. return self._execute_single_step_uncached( step, previous_output, input, all_variables, model, verbose, index, stream=stream, depth=depth, ) + span = _replay_step_begin(self, step) + if span is None: + return self._execute_leaf_step( + step, previous_output, input, all_variables, model, verbose, index, + stream=stream, depth=depth, + ) + try: + result = self._execute_leaf_step( + step, previous_output, input, all_variables, model, verbose, index, + stream=stream, depth=depth, _replay=span, + ) + except BaseException: + _replay_step_end(span, "failed") + raise + is_dict = isinstance(result, dict) + status = span.get("status") or ( + "skipped" if is_dict and result.get("skipped") + else "failed" if is_dict and result.get("error") else "completed" + ) + _replay_step_end(span, status, result.get("output") if is_dict else None) + return result + + def _execute_leaf_step( + self, + step: Any, + previous_output: Optional[str], + input: str, + all_variables: Dict[str, Any], + model: str, + verbose: bool, + index: int = 0, + stream: bool = True, + depth: int = 0, + _replay: Optional[dict] = None, + ) -> Dict[str, Any]: + """Leaf-step body of :meth:`_execute_single_step_internal` (gate + cache).""" if depth > MAX_NESTING_DEPTH: raise ValueError( f"Maximum nesting depth ({MAX_NESTING_DEPTH}) exceeded. " @@ -3045,6 +3194,8 @@ def _execute_single_step_internal( if hasattr(step, "status"): step.status = "completed" self.step_statuses[step.name] = "completed" + if _replay is not None: + _replay["status"] = "cached" # Copy: a caller mutating a returned result must not edit the cache. return dict(hit) result = self._execute_single_step_uncached( @@ -3432,6 +3583,7 @@ def _apply_step_policies(self, step, run_body, all_variables, verbose): return output, step_error + @_replay_pattern("route") def _execute_route( self, route_step: Route, @@ -3462,6 +3614,9 @@ def _execute_route( if matched_route is None: matched_route = route_step.default or [] + _replay_note_route("default") + else: + _replay_note_route(key) if verbose: route_name = next((k for k in route_step.routes if route_step.routes[k] == matched_route), "default") @@ -3713,6 +3868,7 @@ def _merge_branch_variables( winners[key] = idx all_variables[key] = value + @_replay_pattern("parallel") def _execute_parallel( self, parallel_step: Parallel, @@ -3965,6 +4121,7 @@ def execute_with_branch(step=step, idx=idx, opt_prev=optimized_previous, branch_ return {"steps": results, "output": combined_output, "variables": all_variables, "stop": parallel_stopped} + @_replay_pattern("loop") def _execute_loop( self, loop_step: Loop, @@ -4351,6 +4508,7 @@ def _parse_list_from_string(self, text: str) -> List[Any]: # 4. Fallback: wrap as single item return [text] + @_replay_pattern("discussion") def _execute_discussion( self, discussion, previous_output, input, all_variables, model, verbose, stream=True, depth=0 ): @@ -4434,6 +4592,7 @@ def _execute_discussion( "stop": discussion_stopped, } + @_replay_pattern("repeat") def _execute_repeat( self, repeat_step: Repeat, @@ -4508,6 +4667,7 @@ def _execute_repeat( all_variables["repeat_iterations"] = iteration + 1 if repeat_step.max_iterations > 0 else 0 return {"steps": results, "output": output, "variables": all_variables, "stop": repeat_stopped} + @_replay_pattern("if") def _execute_if( self, if_step: If, @@ -4539,6 +4699,7 @@ def _execute_if( # Evaluate the condition condition_result = self._evaluate_condition(if_step.condition, all_variables, previous_output) + _replay_note_route("true" if condition_result else "false") if verbose: branch = "then" if condition_result else "else" @@ -4592,6 +4753,7 @@ def _evaluate_condition( from praisonaiagents.conditions.evaluator import evaluate_condition return evaluate_condition(condition, variables, previous_output) + @_replay_pattern("include") def _execute_include( self, include_step: Include, From 9ba8c09bd97e90867965019df4346427863a1ede Mon Sep 17 00:00:00 2001 From: MervinPraison <454862+MervinPraison@users.noreply.github.com> Date: Mon, 26 Oct 2026 11:22:42 +0000 Subject: [PATCH 04/12] feat(llm): record llm_request/llm_response on the LiteLLM path Only the agent's OpenAI-client path recorded model calls, so a team or flow running on any LiteLLM model showed no LLM activity in a replay. The two LiteLLM choke points (_completion_with_retry and its async twin) now emit the pair with model, duration_ms, prompt/completion tokens and cost when litellm knows the price. Streams are wrapped in a pass-through that records the response when the stream ends; with tracing off the stream is untouched. A call the agent already recorded (its OpenAI-client dispatch can reach the LiteLLM adapter) is not counted twice: the agent's llm_request leaves a context marker the LLM class honours inside the same span. chat_mixin.py is not modified. Co-Authored-By: Claude Opus 5.5 (1M context) --- .../praisonaiagents/llm/llm.py | 190 ++++++++++++++++-- 1 file changed, 173 insertions(+), 17 deletions(-) diff --git a/src/praisonai-agents/praisonaiagents/llm/llm.py b/src/praisonai-agents/praisonaiagents/llm/llm.py index 1ae66b4420..7befbbfff3 100644 --- a/src/praisonai-agents/praisonaiagents/llm/llm.py +++ b/src/praisonai-agents/praisonaiagents/llm/llm.py @@ -43,6 +43,142 @@ def mark_success(self, profile: "AuthProfile") -> None: ... # Bounded shrink-and-retry steps after a provider context-overflow rejection. _CONTEXT_OVERFLOW_SHRINK_STEPS = 3 + +# --------------------------------------------------------------------------- +# Replay llm_request / llm_response for the LiteLLM path. The agent's +# OpenAI-client path records its own pair around the whole dispatch (which can +# reach this class through the LiteLLM adapter); ``_llm_call_open`` makes sure +# that call is not counted twice. A no-op returning None when tracing is off. +# --------------------------------------------------------------------------- + +def _replay_llm_begin(llm: Any, params: Dict[str, Any]): + from ..trace.context_events import get_context_emitter + emitter = get_context_emitter() + if not emitter._enabled or emitter._llm_call_open(): + return None + agent = getattr(llm, "current_agent_name", None) or emitter._current_agent() + model = params.get("model") or getattr(llm, "model", None) + messages = params.get("messages") + messages = messages if isinstance(messages, list) else None + emitter._emit(emitter._llm_request_event( + agent, len(messages) if messages else 0, 0, model=model, messages=messages, + )) + return {"emitter": emitter, "agent": agent, "model": model, + "started": time.perf_counter(), "content": [], "usage": None, "finish": None} + + +def _replay_usage_tokens(usage: Any) -> tuple: + if not usage: + return 0, 0 + get = usage.get if isinstance(usage, dict) else (lambda k, d=None: getattr(usage, k, d)) + return int(get("prompt_tokens", 0) or 0), int(get("completion_tokens", 0) or 0) + + +def _replay_observe(span: Dict[str, Any], response: Any, chunk: bool) -> None: + """Collect content / finish reason / usage from a response or a stream chunk.""" + try: + get = response.get if isinstance(response, dict) else (lambda k, d=None: getattr(response, k, d)) + usage = get("usage", None) + if usage: + span["usage"] = usage + choices = get("choices", None) or [] + if not choices: + return + first = choices[0] + fget = first.get if isinstance(first, dict) else (lambda k, d=None: getattr(first, k, d)) + if fget("finish_reason", None): + span["finish"] = fget("finish_reason", None) + body = fget("delta" if chunk else "message", None) + if body is not None: + content = body.get("content") if isinstance(body, dict) else getattr(body, "content", None) + if isinstance(content, str) and content: + span["content"].append(content) + except Exception: + pass # Tracing must never break a model call. + + +def _replay_llm_end(span: Optional[Dict[str, Any]], error: Optional[BaseException] = None) -> None: + if span is None or span.get("done"): + return + span["done"] = True + emitter = span["emitter"] + prompt_tokens, completion_tokens = _replay_usage_tokens(span.get("usage")) + cost = 0.0 + if prompt_tokens or completion_tokens: + try: + import litellm + prompt_cost, completion_cost = litellm.cost_per_token( + model=span["model"], prompt_tokens=prompt_tokens, + completion_tokens=completion_tokens, + ) + cost = float(prompt_cost or 0) + float(completion_cost or 0) + except Exception: + cost = 0.0 # Unknown price: leave it unknown rather than guess. + emitter._emit(emitter._llm_response_event( + span["agent"], + duration_ms=(time.perf_counter() - span["started"]) * 1000, + finish_reason="error" if error is not None else span.get("finish"), + response_content=str(error) if error is not None else "".join(span["content"]), + prompt_tokens=prompt_tokens, + completion_tokens=completion_tokens, + cost_usd=cost, + model=span["model"], + )) + + +class _ReplayStream: + """Pass-through for a LiteLLM stream that records ``llm_response`` when it ends.""" + + __slots__ = ("_stream", "_iter", "_span") + + def __init__(self, stream: Any, span: Dict[str, Any]) -> None: + self._stream = stream + self._iter = None + self._span = span + + def __iter__(self): + return self + + def __next__(self): + if self._iter is None: + self._iter = iter(self._stream) + try: + chunk = next(self._iter) + except StopIteration: + _replay_llm_end(self._span) + raise + except BaseException as exc: + _replay_llm_end(self._span, exc) + raise + _replay_observe(self._span, chunk, chunk=True) + return chunk + + def __aiter__(self): + return self + + async def __anext__(self): + if self._iter is None: + self._iter = self._stream.__aiter__() + try: + chunk = await self._iter.__anext__() + except StopAsyncIteration: + _replay_llm_end(self._span) + raise + except BaseException as exc: + _replay_llm_end(self._span, exc) + raise + _replay_observe(self._span, chunk, chunk=True) + return chunk + + def close(self): + _replay_llm_end(self._span) + close = getattr(self._stream, "close", None) + if callable(close): + return close() + + def __getattr__(self, name: str) -> Any: + return getattr(self._stream, name) + # Default per-request timeout (seconds) when ``timeout`` is not passed. Without # it litellm falls back to its 6000s ceiling, so a provider that accepts the # connection but never answers blocks the turn for ~100 minutes per attempt. @@ -1736,15 +1872,25 @@ def _completion_with_retry(self, **completion_params): # entirely, so scripted turns never reach (or trip) the guard. check_model_request(self.model, "litellm.completion") import litellm - for step in range(_CONTEXT_OVERFLOW_SHRINK_STEPS + 1): - try: - response = self._call_with_retry(litellm.completion, **completion_params) - break - except Exception as e: - if not self._shrink_on_context_overflow(e, completion_params, step): - raise + _replay = _replay_llm_begin(self, completion_params) + try: + for step in range(_CONTEXT_OVERFLOW_SHRINK_STEPS + 1): + try: + response = self._call_with_retry(litellm.completion, **completion_params) + break + except Exception as e: + if not self._shrink_on_context_overflow(e, completion_params, step): + raise + except BaseException as exc: + _replay_llm_end(_replay, exc) + raise if not completion_params.get("stream"): self._track_token_usage(response, self._response_model_for_tracking(response)) + if _replay is not None: + _replay_observe(_replay, response, chunk=False) + _replay_llm_end(_replay) + elif _replay is not None: + response = _ReplayStream(response, _replay) return response async def _acompletion_with_retry(self, **completion_params): @@ -1765,18 +1911,28 @@ async def _acompletion_with_retry(self, **completion_params): """ check_model_request(self.model, "litellm.acompletion") import litellm - for step in range(_CONTEXT_OVERFLOW_SHRINK_STEPS + 1): - try: - response = await self._call_with_retry_async( - litellm.acompletion, - **completion_params, - ) - break - except Exception as e: - if not self._shrink_on_context_overflow(e, completion_params, step): - raise + _replay = _replay_llm_begin(self, completion_params) + try: + for step in range(_CONTEXT_OVERFLOW_SHRINK_STEPS + 1): + try: + response = await self._call_with_retry_async( + litellm.acompletion, + **completion_params, + ) + break + except Exception as e: + if not self._shrink_on_context_overflow(e, completion_params, step): + raise + except BaseException as exc: + _replay_llm_end(_replay, exc) + raise if not completion_params.get("stream"): self._track_token_usage(response, self._response_model_for_tracking(response)) + if _replay is not None: + _replay_observe(_replay, response, chunk=False) + _replay_llm_end(_replay) + elif _replay is not None: + response = _ReplayStream(response, _replay) return response def _supports_web_search(self) -> bool: From fea0c826809797ab6a91ba4fca86a25eb07ead10 Mon Sep 17 00:00:00 2001 From: MervinPraison <454862+MervinPraison@users.noreply.github.com> Date: Mon, 26 Oct 2026 11:22:43 +0000 Subject: [PATCH 05/12] test(replay): cover run events, spans, branches and every emission site Schema round-trip and 1.0 back-compat, unknown event types, parent_seq and run_id inheritance, 8 concurrent Parallel branches (branch_id, parent_seq, unique sequence numbers), flow step statuses and chosen routes, team task completed/skipped/failed (sync and async), the manager's delegate hand-off, human-review and tool-call approval pairs, every hand-off path emitting once, the LiteLLM path (mocked: sync, async, streaming, errors, no double count), and a micro-benchmark that the disabled emitter adds next to nothing. Co-Authored-By: Claude Opus 5.5 (1M context) --- .../tests/unit/trace/test_run_events.py | 769 ++++++++++++++++++ 1 file changed, 769 insertions(+) create mode 100644 src/praisonai-agents/tests/unit/trace/test_run_events.py diff --git a/src/praisonai-agents/tests/unit/trace/test_run_events.py b/src/praisonai-agents/tests/unit/trace/test_run_events.py new file mode 100644 index 0000000000..7857b3768b --- /dev/null +++ b/src/praisonai-agents/tests/unit/trace/test_run_events.py @@ -0,0 +1,769 @@ +"""PraisonAI Replay: run-structure events on the context trace (schema 1.1). + +A replay can only explain a run if the trace says which run, which task or +step, which graph node and which enclosing span every event belongs to, and if +concurrent branches cannot scramble that. These cover the event schema, the +span/branch bookkeeping, and every place the SDK now records what happened: +hand-offs, team tasks, flow steps, manager decisions, approvals and the +LiteLLM call path. No test here reaches a model. +""" + +import asyncio +import time +import types + +import pytest + +from praisonaiagents.trace.context_events import ( + ContextEvent, + ContextEventType, + ContextListSink, + ContextTraceEmitter, + _run_scope, + copy_context_to_callable, + get_context_emitter, + trace_context, +) + + +@pytest.fixture +def recorder(): + """A live emitter installed for the test; yields its sink.""" + sink = ContextListSink() + emitter = ContextTraceEmitter(sink=sink, session_id="replay-test", enabled=True) + with trace_context(emitter): + yield sink + + +def _types(sink): + return [e.event_type.value for e in sink.get_events()] + + +def _of(sink, event_type): + return [e for e in sink.get_events() if e.event_type.value == event_type] + + +def _by_seq(sink): + return {e.sequence_num: e for e in sink.get_events()} + + +# --------------------------------------------------------------------------- +# Schema +# --------------------------------------------------------------------------- + + +class TestSchema: + def test_schema_version_is_1_1(self): + from praisonaiagents.trace.context_events import CONTEXT_SCHEMA_VERSION + assert CONTEXT_SCHEMA_VERSION == "1.1" + + def test_new_fields_round_trip(self): + event = ContextEvent( + event_type=ContextEventType.TASK_START, + timestamp=1.5, + session_id="s", + run_id="run-1", + parent_seq=3, + node_id="task:research", + ) + data = event.to_dict() + assert data["run_id"] == "run-1" + assert data["parent_seq"] == 3 + assert data["node_id"] == "task:research" + back = ContextEvent.from_dict(data) + assert (back.run_id, back.parent_seq, back.node_id) == ("run-1", 3, "task:research") + assert back.event_type == ContextEventType.TASK_START + + def test_new_fields_are_omitted_when_none(self): + data = ContextEvent(ContextEventType.SESSION_START, 1.0, "s").to_dict() + for key in ("run_id", "parent_seq", "node_id", "branch_id"): + assert key not in data + + def test_a_1_0_event_still_loads(self): + legacy = { + "schema_version": "1.0", + "event_type": "agent_start", + "timestamp": 10.0, + "session_id": "old", + "agent_name": "Writer", + "sequence_num": 4, + "messages_count": 0, + "tokens_used": 0, + "tokens_budget": 0, + "data": {"role": "w"}, + "prompt_tokens": 0, + "completion_tokens": 0, + "cost_usd": 0.0, + } + event = ContextEvent.from_dict(legacy) + assert event.event_type == ContextEventType.AGENT_START + assert event.sequence_num == 4 + assert event.run_id is None and event.parent_seq is None and event.node_id is None + + def test_an_unknown_event_type_keeps_its_raw_value(self): + event = ContextEvent.from_dict({"event_type": "hologram_beamed", "session_id": "s"}) + assert event.event_type == "hologram_beamed" + assert event.event_type.value == "hologram_beamed" # Enum-style callers still work + assert event.to_dict()["event_type"] == "hologram_beamed" + + def test_control_a_known_event_type_is_still_the_enum(self): + event = ContextEvent.from_dict({"event_type": "tool_call_end", "session_id": "s"}) + assert event.event_type is ContextEventType.TOOL_CALL_END + + def test_new_event_types_have_contract_values(self): + assert {t.value for t in ContextEventType} >= { + "workflow_start", "workflow_end", "task_start", "task_end", + "step_start", "step_end", "approval_requested", "approval_decided", + } + + +# --------------------------------------------------------------------------- +# Spans, sequence numbers, branches, run ids +# --------------------------------------------------------------------------- + + +class TestSpans: + def test_parent_seq_nests_and_end_points_at_its_start(self, recorder): + emitter = get_context_emitter() + emitter.workflow_start("team", "t", "sequential") # 0 + emitter.task_start("research", node_id="task:research") # 1 + emitter.agent_start("Writer") # 2 + emitter.message_added("Writer", "user", "hi", 1, 1) # 3 point + emitter.agent_end("Writer") # 4 + emitter.task_end("research", "completed", node_id="task:research") # 5 + emitter.workflow_end() # 6 + e = _by_seq(recorder) + assert e[0].parent_seq is None + assert e[1].parent_seq == 0 + assert e[2].parent_seq == 1 + assert e[3].parent_seq == 2 + assert e[4].parent_seq == 2 + assert e[5].parent_seq == 1 + assert e[6].parent_seq == 0 + + def test_a_start_left_open_is_closed_by_its_parent_end(self, recorder): + emitter = get_context_emitter() + emitter.workflow_start("team", "t", "sequential") + emitter.tool_call_start("Writer", "search") # never ended (exception path) + emitter.workflow_end() + emitter.message_added("x", "user", "after", 1, 1) + assert recorder.get_events()[-1].parent_seq is None + + def test_spans_of_another_emitter_do_not_leak(self, recorder): + other_sink = ContextListSink() + other = ContextTraceEmitter(sink=other_sink, session_id="other") + get_context_emitter().workflow_start("team", "t", "sequential") + other.message_added("x", "user", "hi", 1, 1) + assert other_sink.get_events()[0].parent_seq is None + + def test_run_scope_stamps_run_id_and_nested_scopes_inherit(self, recorder): + emitter = get_context_emitter() + with _run_scope("outer") as outer: + with _run_scope("inner") as inner: + emitter.message_added("a", "user", "x", 1, 1) + emitter.message_added("a", "user", "y", 1, 1) + assert outer == inner == "outer" + first, second = recorder.get_events() + assert first.run_id == "outer" + assert second.run_id is None + + def test_eight_concurrent_branches_keep_their_own_branch_and_parent(self, recorder): + import concurrent.futures + + emitter = get_context_emitter() + emitter.step_start("parallel", pattern="parallel", node_id="step:parallel") + fork_seq = recorder.get_events()[-1].sequence_num + + def branch(idx): + trace = get_context_emitter() + trace.set_branch(f"parallel_{idx}") + try: + for _ in range(20): + trace.step_start(f"b{idx}", node_id=f"step:b{idx}") + time.sleep(0.0005) + trace.step_end(f"b{idx}", "completed", node_id=f"step:b{idx}") + finally: + trace.clear_branch() + + with concurrent.futures.ThreadPoolExecutor(max_workers=8) as pool: + futures = [pool.submit(copy_context_to_callable(lambda i=i: branch(i))) for i in range(8)] + for future in futures: + future.result() + emitter.step_end("parallel", "completed", node_id="step:parallel") + + events = recorder.get_events() + seqs = [e.sequence_num for e in events] + assert seqs == sorted(seqs) == list(range(len(events))) # unique, append order + by_seq = {e.sequence_num: e for e in events} + for e in events[1:-1]: + idx = int(e.node_id.split(":b")[1]) + assert e.branch_id == f"parallel_{idx}" + if e.event_type == ContextEventType.STEP_START: + assert e.parent_seq == fork_seq + else: + start = by_seq[e.parent_seq] + assert start.node_id == e.node_id and start.branch_id == e.branch_id + assert events[-1].parent_seq == fork_seq + assert events[-1].branch_id is None + + +# --------------------------------------------------------------------------- +# AgentFlow steps +# --------------------------------------------------------------------------- + + +def _handler(name, output=None): + def run(ctx): + return output if output is not None else name + run.__name__ = name + return run + + +class TestFlowSteps: + def test_flow_records_workflow_and_step_events_with_graph(self, recorder): + from praisonaiagents.workflows.workflows import AgentFlow + + flow = AgentFlow(name="demo", steps=[_handler("draft"), _handler("polish")]) + flow.run("go") + types_ = _types(recorder) + assert types_[0] == "workflow_start" and types_[-1] == "workflow_end" + start = recorder.get_events()[0] + assert start.data["kind"] == "flow" and start.data["name"] == "demo" + assert start.data["graph"]["schema"] == "praisonai.graph/1" + steps = _of(recorder, "step_start") + assert [s.node_id for s in steps] == ["step:draft", "step:polish"] + assert all(s.parent_seq == start.sequence_num for s in steps) + run_ids = {e.run_id for e in recorder.get_events()} + assert len(run_ids) == 1 and None not in run_ids + assert [e.data["status"] for e in _of(recorder, "step_end")] == ["completed", "completed"] + assert recorder.get_events()[-1].data["status"] == "completed" + + def test_parallel_with_eight_branches(self, recorder): + from praisonaiagents.workflows.workflows import AgentFlow, Parallel + + def slow(name): + def run(ctx): + time.sleep(0.002) + return name + run.__name__ = name + return run + + flow = AgentFlow(steps=[Parallel(steps=[slow(f"b{i}") for i in range(8)])]) + flow.run("go") + events = recorder.get_events() + seqs = [e.sequence_num for e in events] + assert seqs == list(range(len(events))) + parallel = [e for e in _of(recorder, "step_start") if e.data.get("pattern") == "parallel"] + assert len(parallel) == 1 and parallel[0].node_id == "step:parallel" + branches = [e for e in _of(recorder, "step_start") if e.node_id != "step:parallel"] + assert sorted(e.branch_id for e in branches) == sorted(f"parallel_{i}" for i in range(8)) + assert {e.parent_seq for e in branches} == {parallel[0].sequence_num} + by_seq = _by_seq(recorder) + for end in _of(recorder, "step_end"): + assert by_seq[end.parent_seq].node_id == end.node_id + assert {e.node_id for e in branches} == {f"step:b{i}" for i in range(8)} + + def test_route_and_if_record_the_branch_taken(self, recorder): + from praisonaiagents.workflows.workflows import AgentFlow, If, Route + + flow = AgentFlow(steps=[ + _handler("classify", "this is urgent"), + Route(routes={"urgent": [_handler("page")], "default": [_handler("queue")]}), + If(condition="{{missing}} == yes", then_steps=[_handler("a")], else_steps=[_handler("b")]), + ]) + flow.run("go") + ends = {e.node_id: e for e in _of(recorder, "step_end")} + assert ends["step:route"].data["route"] == "urgent" + assert ends["step:if"].data["route"] == "false" + assert "step:page" in ends and "step:queue" not in ends + assert "step:b" in ends and "step:a" not in ends + + def test_skipped_failed_and_cached_steps(self, recorder): + from praisonaiagents.task.task import Task + from praisonaiagents.workflows.workflows import AgentFlow + + def boom(ctx): + raise RuntimeError("nope") + + skip = Task(name="gated", handler=_handler("gated"), should_run=lambda ctx: False) + fail = Task(name="boom", handler=boom, max_retries=0, on_error="continue") + flow = AgentFlow(steps=[_handler("cached_step"), skip, fail], cache=True) + flow.run("go") + statuses = {e.node_id: e.data["status"] for e in _of(recorder, "step_end")} + assert statuses == {"step:cached_step": "completed", "step:gated": "skipped", "step:boom": "failed"} + recorder.clear() + flow.run("go") + statuses = {e.node_id: e.data["status"] for e in _of(recorder, "step_end")} + assert statuses["step:cached_step"] == "cached" + + def test_nested_flow_inherits_the_outer_run_id(self, recorder): + from praisonaiagents.workflows.workflows import AgentFlow + + inner = AgentFlow(name="inner", steps=[_handler("fetch")]) + + def nested(ctx): + return inner.run("sub")["output"] + + outer = AgentFlow(name="outer", steps=[_handler("first"), nested]) + outer.run("go") + starts = _of(recorder, "workflow_start") + assert [s.data["name"] for s in starts] == ["outer", "inner"] + assert len({e.run_id for e in recorder.get_events()}) == 1 + step = [e for e in _of(recorder, "step_start") if e.node_id == "step:nested"][0] + assert starts[1].parent_seq == step.sequence_num + + def test_team_inside_a_flow_inherits_the_flow_run_id(self, recorder, monkeypatch): + from praisonaiagents import Agent, Task + from praisonaiagents.agents.agents import AgentTeam + from praisonaiagents.workflows.workflows import AgentFlow + + monkeypatch.setattr(AgentTeam, "execute_task", _fake_execute({})) + writer = Agent(name="Writer", instructions="write") + team = _team([Task(name="draft", description="d", agent=writer)], agents=[writer]) + + def run_team(ctx): + return team.start() + + AgentFlow(name="outer", steps=[run_team]).run("go") + run_ids = {e.run_id for e in recorder.get_events()} + assert len(run_ids) == 1 and team.run_id not in run_ids + + +# --------------------------------------------------------------------------- +# AgentTeam tasks, manager decisions, human review +# --------------------------------------------------------------------------- + + +def _team(tasks, process="sequential", agents=None): + from praisonaiagents import Agent + from praisonaiagents.agents.agents import AgentTeam + + agents = agents or [Agent(name="Writer", instructions="write")] + return AgentTeam(agents=agents, tasks=tasks, process=process, output="silent") + + +def _fake_execute(outputs): + """Stand-in for AgentTeam.execute_task: no model, scripted results.""" + from praisonaiagents.main import TaskOutput + + def execute(self, task_id): + task = self.tasks[task_id] + raw = outputs.get(task.name, "done") + if raw is None: + return None + out = TaskOutput(description=task.description, raw=raw, agent=task.agent.name) + task.result = out + return out + return execute + + +class TestTeamTasks: + def test_task_completed_skipped_and_failed(self, recorder, monkeypatch): + from praisonaiagents import Agent, Task + from praisonaiagents.agents.agents import AgentTeam + + writer = Agent(name="Writer", instructions="write") + done = Task(name="draft", description="Draft it", agent=writer) + skipped = Task(name="optional", description="Maybe", agent=writer, should_run=lambda ctx: False) + failed = Task(name="flaky", description="Never finishes", agent=writer, max_retries=1, retry_delay=0) + monkeypatch.setattr(AgentTeam, "execute_task", _fake_execute({"flaky": None})) + team = _team([done, skipped, failed], agents=[writer]) + team.start() + + ends = {e.node_id: e for e in _of(recorder, "task_end")} + assert ends["task:draft"].data["status"] == "completed" + assert ends["task:draft"].data["output_preview"] == "done" + assert ends["task:optional"].data["status"] == "skipped" + assert ends["task:flaky"].data["status"] == "failed" + starts = {e.node_id: e for e in _of(recorder, "task_start")} + assert starts["task:draft"].data["description_preview"] == "Draft it" + assert starts["task:draft"].data["agent"] == "Writer" + for node_id, end in ends.items(): + assert end.parent_seq == starts[node_id].sequence_num + wf_start = _of(recorder, "workflow_start")[0] + assert wf_start.data["kind"] == "team" and wf_start.data["graph"]["kind"] == "team" + assert {s.parent_seq for s in starts.values()} == {wf_start.sequence_num} + assert _of(recorder, "workflow_end")[0].data["status"] == "failed" + assert {e.run_id for e in recorder.get_events()} == {team.run_id} + + def test_async_team_records_tasks(self, recorder, monkeypatch): + from praisonaiagents import Agent, Task + from praisonaiagents.agents.agents import AgentTeam + + async def aexecute(self, task_id): + return _fake_execute({})(self, task_id) + + writer = Agent(name="Writer", instructions="write") + task = Task(name="draft", description="Draft", agent=writer, async_execution=True) + monkeypatch.setattr(AgentTeam, "aexecute_task", aexecute) + team = _team([task], agents=[writer]) + asyncio.run(team.astart()) + assert [e.data["status"] for e in _of(recorder, "task_end")] == ["completed"] + assert _types(recorder)[0] == "workflow_start" + assert _types(recorder)[-1] == "workflow_end" + + def test_manager_choice_is_a_delegate_handoff(self, recorder, monkeypatch): + from praisonaiagents import Agent, Task + from praisonaiagents.agents.agents import AgentTeam + from praisonaiagents.process.manager_schema import ManagerInstructions + from praisonaiagents.process.process import Process + + writer = Agent(name="Writer", instructions="write") + editor = Agent(name="Editor", instructions="edit") + t1 = Task(name="draft", description="Draft", agent=writer) + t2 = Task(name="review", description="Review", agent=writer) + script = iter([ + ManagerInstructions(task_id=0, agent_name="Writer", action="execute"), + ManagerInstructions(task_id=1, agent_name="Editor", action="execute"), + ]) + monkeypatch.setattr( + Process, "_get_manager_instructions_with_fallback", + lambda self, *a: next(script), + ) + monkeypatch.setattr(AgentTeam, "execute_task", _fake_execute({})) + team = _team([t1, t2], process="hierarchical", agents=[writer, editor]) + team.start() + + delegates = [e for e in _of(recorder, "agent_handoff") if e.data["kind"] == "delegate"] + assert [(d.agent_name, d.data["to_agent"], d.node_id) for d in delegates] == [ + ("Manager", "Writer", "task:draft"), + ("Manager", "Editor", "task:review"), + ] + summary = delegates[1].data["context_passed"]["task_status"] + assert [(s["name"], s["status"]) for s in summary] == [("draft", "completed"), ("review", "not started")] + manager = [e for e in _of(recorder, "task_start") if e.node_id == "manager:Manager"] + assert len(manager) == 1 + + def test_human_review_records_an_approval_pair(self, recorder): + from praisonaiagents import Agent, Task + from praisonaiagents.approval import ApprovalDecision, get_approval_registry + from praisonaiagents.main import TaskOutput + + class Reviewer: + def request_approval_sync(self, request): + return ApprovalDecision(approved=True, reason="looks good", approver="ana") + + registry = get_approval_registry() + previous = registry.get_backend() + registry.set_backend(Reviewer()) + try: + writer = Agent(name="Writer", instructions="write") + task = Task(name="draft", description="Draft", agent=writer, human_input=True) + team = _team([task], agents=[writer]) + out = TaskOutput(description="Draft", raw="api_key=sk-secret123456 text", agent="Writer") + team._apply_human_review(task, 0, out) + finally: + registry.set_backend(previous) + requested, decided = _of(recorder, "approval_requested"), _of(recorder, "approval_decided") + assert len(requested) == len(decided) == 1 + assert requested[0].data["kind"] == "task_output_review" + assert requested[0].node_id == "task:draft" + assert "sk-secret123456" not in str(requested[0].data["args"]) + assert decided[0].parent_seq == requested[0].sequence_num + assert decided[0].data["approved"] is True and decided[0].data["approver"] == "ana" + assert decided[0].data["approval_id"] == requested[0].data["approval_id"] + + +# --------------------------------------------------------------------------- +# Tool-call approvals +# --------------------------------------------------------------------------- + + +class TestToolApprovals: + def _registry(self, backend): + from praisonaiagents.approval.registry import ApprovalRegistry + + registry = ApprovalRegistry() + registry.set_backend(backend) + registry.add_requirement("delete_file", risk_level="high") + return registry + + def test_sync_approval_records_request_and_decision(self, recorder, monkeypatch): + from praisonaiagents.approval import ApprovalDecision + + monkeypatch.delenv("PRAISONAI_AUTO_APPROVE", raising=False) + + class Deny: + def request_approval_sync(self, request): + return ApprovalDecision(approved=False, reason="not today", approver="ops") + + registry = self._registry(Deny()) + decision = registry.approve_sync("Writer", "delete_file", {"path": "/x", "password": "hunter2"}) + assert decision.approved is False + requested, decided = _of(recorder, "approval_requested"), _of(recorder, "approval_decided") + assert len(requested) == len(decided) == 1 + assert requested[0].data["tool"] == "delete_file" + assert requested[0].data["risk"] == "high" + assert requested[0].data["kind"] == "tool_call" + assert requested[0].node_id == "tool:Writer/delete_file" + assert "hunter2" not in str(requested[0].data["args"]) + assert decided[0].parent_seq == requested[0].sequence_num + assert decided[0].data == { + "approval_id": requested[0].data["approval_id"], + "approved": False, "reason": "not today", "approver": "ops", "modified": False, + } + + def test_async_approval_records_request_and_decision(self, recorder, monkeypatch): + from praisonaiagents.approval import ApprovalDecision + + monkeypatch.delenv("PRAISONAI_AUTO_APPROVE", raising=False) + + class Allow: + async def request_approval(self, request): + return ApprovalDecision(approved=True, reason="ok", modified_args={"path": "/y"}) + + registry = self._registry(Allow()) + asyncio.run(registry.approve_async("Writer", "delete_file", {"path": "/x"})) + decided = _of(recorder, "approval_decided") + assert len(decided) == 1 and decided[0].data["modified"] is True + assert decided[0].parent_seq == _of(recorder, "approval_requested")[0].sequence_num + + def test_control_the_fast_path_records_nothing(self, recorder): + from praisonaiagents.approval.registry import ApprovalRegistry + + ApprovalRegistry().approve_sync("Writer", "read_file", {}) # not required + assert _of(recorder, "approval_requested") == [] + + +# --------------------------------------------------------------------------- +# Hand-offs +# --------------------------------------------------------------------------- + + +class TestHandoffs: + def _pair(self): + from praisonaiagents import Agent + + source = Agent(name="Router", instructions="route") + target = Agent(name="Billing", instructions="bill") + target.chat = lambda *a, **k: "handled" + + async def achat(*a, **k): + return "handled" + target.achat = achat + return source, target + + def _one_handoff(self, sink): + handoffs = _of(sink, "agent_handoff") + assert len(handoffs) == 1, [h.data for h in handoffs] + h = handoffs[0] + assert (h.agent_name, h.data["from_agent"], h.data["to_agent"]) == ("Router", "Router", "Billing") + assert h.data["kind"] == "handoff" + assert h.node_id == "agent:Billing" + return h + + def test_programmatic(self, recorder): + from praisonaiagents.agent.handoff import Handoff + + source, target = self._pair() + Handoff(agent=target).execute_programmatic(source, "refund order 7") + assert self._one_handoff(recorder).data["reason"] == "refund order 7" + + def test_async(self, recorder): + from praisonaiagents.agent.handoff import Handoff + + source, target = self._pair() + asyncio.run(Handoff(agent=target).execute_async(source, "refund order 7")) + self._one_handoff(recorder) + + def test_tool_function(self, recorder): + from praisonaiagents.agent.handoff import Handoff + + source, target = self._pair() + tool = Handoff(agent=target).to_tool_function(source) + tool(task="refund order 7") + h = self._one_handoff(recorder) + assert h.data["context_passed"]["args"] == {"task": "refund order 7"} + + def test_typed_programmatic_and_async(self, recorder): + from pydantic import BaseModel + from praisonaiagents.agent.handoff import TypedHandoff + + class Refund(BaseModel): + order: int + + source, target = self._pair() + TypedHandoff(agent=target, input_schema=Refund).execute_programmatic(source, Refund(order=7)) + self._one_handoff(recorder) + recorder.clear() + asyncio.run(TypedHandoff(agent=target, input_schema=Refund).execute_async(source, Refund(order=7))) + self._one_handoff(recorder) + + def test_control_a_refused_handoff_records_nothing(self, recorder): + from praisonaiagents.agent.handoff import Handoff, HandoffConfig, _clear_handoff_chain, _push_handoff + + source, target = self._pair() + _push_handoff("Billing") # Billing already in the chain -> cycle + try: + with pytest.raises(Exception): + Handoff(agent=target, config=HandoffConfig(detect_cycles=True)).execute_programmatic(source, "x") + finally: + _clear_handoff_chain() + assert _of(recorder, "agent_handoff") == [] + + +# --------------------------------------------------------------------------- +# LiteLLM path +# --------------------------------------------------------------------------- + + +def _response(text="hello", prompt=11, completion=3): + message = types.SimpleNamespace(content=text, tool_calls=None) + choice = types.SimpleNamespace(message=message, finish_reason="stop", delta=None) + usage = types.SimpleNamespace(prompt_tokens=prompt, completion_tokens=completion, total_tokens=prompt + completion) + return types.SimpleNamespace(choices=[choice], usage=usage, model="gpt-4o-mini") + + +def _chunk(text=None, finish=None, usage=None): + delta = types.SimpleNamespace(content=text, tool_calls=None) + choice = types.SimpleNamespace(delta=delta, finish_reason=finish) + return types.SimpleNamespace(choices=[choice], usage=usage) + + +@pytest.fixture +def fake_litellm(monkeypatch): + import litellm + import praisonaiagents.llm.llm as llm_module + + monkeypatch.setattr(llm_module, "check_model_request", lambda *a, **k: None) + calls = [] + + def completion(**params): + calls.append(params) + if params.get("stream"): + usage = types.SimpleNamespace(prompt_tokens=5, completion_tokens=2) + return iter([_chunk("hel"), _chunk("lo", finish="stop"), _chunk(usage=usage)]) + return _response() + + async def acompletion(**params): + calls.append(params) + if params.get("stream"): + async def gen(): + for c in [_chunk("a"), _chunk("b", finish="stop")]: + yield c + return gen() + return _response("async") + + monkeypatch.setattr(litellm, "completion", completion) + monkeypatch.setattr(litellm, "acompletion", acompletion) + return calls + + +class TestLiteLLMPath: + def _llm(self): + from praisonaiagents.llm.llm import LLM + return LLM(model="gpt-4o-mini") + + def test_non_streaming_call_records_request_and_response(self, recorder, fake_litellm): + llm = self._llm() + emitter = get_context_emitter() + emitter.agent_start("Writer") + llm._completion_with_retry(model="gpt-4o-mini", messages=[{"role": "user", "content": "hi"}]) + emitter.agent_end("Writer") + request, response = _of(recorder, "llm_request"), _of(recorder, "llm_response") + assert len(request) == len(response) == 1 + assert request[0].agent_name == "Writer" and request[0].node_id == "agent:Writer" + assert request[0].messages_count == 1 + r = response[0] + assert (r.prompt_tokens, r.completion_tokens) == (11, 3) + assert r.data["model"] == "gpt-4o-mini" + assert r.data["response_content"] == "hello" + assert r.data["finish_reason"] == "stop" + assert r.data["duration_ms"] >= 0 + assert r.cost_usd > 0 # litellm knows gpt-4o-mini's price + + def test_streaming_call_records_response_when_the_stream_ends(self, recorder, fake_litellm): + llm = self._llm() + stream = llm._completion_with_retry(model="gpt-4o-mini", messages=[], stream=True) + assert _of(recorder, "llm_response") == [] + text = "".join(c.choices[0].delta.content or "" for c in stream) + assert text == "hello" + response = _of(recorder, "llm_response") + assert len(response) == 1 + assert response[0].data["response_content"] == "hello" + assert (response[0].prompt_tokens, response[0].completion_tokens) == (5, 2) + + def test_async_calls_record_events(self, recorder, fake_litellm): + llm = self._llm() + + async def run(): + await llm._acompletion_with_retry(model="gpt-4o-mini", messages=[]) + stream = await llm._acompletion_with_retry(model="gpt-4o-mini", messages=[], stream=True) + return "".join([c.choices[0].delta.content async for c in stream]) + + assert asyncio.run(run()) == "ab" + responses = _of(recorder, "llm_response") + assert [r.data["response_content"] for r in responses] == ["async", "ab"] + assert len(_of(recorder, "llm_request")) == 2 + + def test_a_call_already_recorded_by_the_agent_is_not_counted_twice(self, recorder, fake_litellm): + llm = self._llm() + emitter = get_context_emitter() + emitter.agent_start("Writer") + emitter.llm_request("Writer", messages_count=1, tokens_used=0, model="gpt-4o-mini") + llm._completion_with_retry(model="gpt-4o-mini", messages=[]) + emitter.llm_response("Writer", duration_ms=1.0) + llm._completion_with_retry(model="gpt-4o-mini", messages=[]) # a fresh call counts + emitter.agent_end("Writer") + assert len(_of(recorder, "llm_request")) == 2 + assert len(_of(recorder, "llm_response")) == 2 + + def test_a_failed_call_records_an_error_response(self, recorder, monkeypatch): + import litellm + import praisonaiagents.llm.llm as llm_module + + monkeypatch.setattr(llm_module, "check_model_request", lambda *a, **k: None) + + def explode(**params): + raise ValueError("bad request") + monkeypatch.setattr(litellm, "completion", explode) + llm = self._llm() + llm._max_retries = 0 + with pytest.raises(Exception): + llm._completion_with_retry(model="gpt-4o-mini", messages=[]) + response = _of(recorder, "llm_response") + assert len(response) == 1 and response[0].data["finish_reason"] == "error" + + def test_control_disabled_tracing_leaves_the_stream_untouched(self, fake_litellm): + llm = self._llm() + stream = llm._completion_with_retry(model="gpt-4o-mini", messages=[], stream=True) + assert type(stream).__name__ == "list_iterator" + + +# --------------------------------------------------------------------------- +# The disabled emitter costs (almost) nothing +# --------------------------------------------------------------------------- + + +def test_disabled_emitter_adds_next_to_nothing(): + emitter = ContextTraceEmitter(enabled=False) + n = 20000 + + def baseline(): + pass + + t0 = time.perf_counter() + for _ in range(n): + baseline() + base = time.perf_counter() - t0 + + t0 = time.perf_counter() + for _ in range(n): + emitter.task_start("t", "d", "a", "task:t") + emitter.step_end("s", "completed", 1.0, "out") + emitter.agent_handoff("a", "b", kind="delegate") + emitter.approval_requested("id", "tool", {"x": 1}) + emitter.workflow_start("team", "t", "sequential", graph={}) + cost = time.perf_counter() - t0 + # Five disabled calls per iteration: one attribute check each. + assert cost / (n * 5) < 2e-6, f"{cost / (n * 5) * 1e9:.0f} ns per disabled call" + assert cost < base * 5 * 20 + 0.05 + + +def test_disabled_tracing_never_builds_the_graph(): + from praisonaiagents.trace.context_events import _workflow_scope + + built = [] + with _workflow_scope("team", "t", "sequential", graph_fn=lambda: built.append(1)) as run: + assert run is None + assert built == [] From da00df8e69a1e94df723eb24eac1339da5ca718f Mon Sep 17 00:00:00 2001 From: MervinPraison <454862+MervinPraison@users.noreply.github.com> Date: Mon, 26 Oct 2026 11:25:54 +0000 Subject: [PATCH 06/12] feat(llm): record the Responses API calls on the LiteLLM path too OpenAI models reached through LiteLLM (e.g. llm="openai/gpt-4o-mini") use litellm.responses, not litellm.completion, so the real agentic run showed the Writer agent with no LLM events. The four Responses API entry points (sync, async, streamed sync, streamed async) are now thin traced wrappers around the unchanged bodies, recording input/output tokens, text and cost. Co-Authored-By: Claude Opus 5.5 (1M context) --- .../praisonaiagents/llm/llm.py | 91 ++++++++++++++++++- .../tests/unit/trace/test_run_events.py | 25 +++++ 2 files changed, 112 insertions(+), 4 deletions(-) diff --git a/src/praisonai-agents/praisonaiagents/llm/llm.py b/src/praisonai-agents/praisonaiagents/llm/llm.py index 7befbbfff3..da1c475659 100644 --- a/src/praisonai-agents/praisonaiagents/llm/llm.py +++ b/src/praisonai-agents/praisonaiagents/llm/llm.py @@ -58,7 +58,8 @@ def _replay_llm_begin(llm: Any, params: Dict[str, Any]): return None agent = getattr(llm, "current_agent_name", None) or emitter._current_agent() model = params.get("model") or getattr(llm, "model", None) - messages = params.get("messages") + # Chat Completions sends ``messages``; the Responses API sends ``input``. + messages = params.get("messages", params.get("input")) messages = messages if isinstance(messages, list) else None emitter._emit(emitter._llm_request_event( agent, len(messages) if messages else 0, 0, model=model, messages=messages, @@ -71,7 +72,10 @@ def _replay_usage_tokens(usage: Any) -> tuple: if not usage: return 0, 0 get = usage.get if isinstance(usage, dict) else (lambda k, d=None: getattr(usage, k, d)) - return int(get("prompt_tokens", 0) or 0), int(get("completion_tokens", 0) or 0) + # Chat Completions names, then the Responses API names. + prompt = get("prompt_tokens", None) or get("input_tokens", 0) + completion = get("completion_tokens", None) or get("output_tokens", 0) + return int(prompt or 0), int(completion or 0) def _replay_observe(span: Dict[str, Any], response: Any, chunk: bool) -> None: @@ -126,6 +130,19 @@ def _replay_llm_end(span: Optional[Dict[str, Any]], error: Optional[BaseExceptio )) +def _replay_stream_result(span: Dict[str, Any], result: Any) -> None: + """Finish a span from a streamed Responses API ``(text, tools, final)`` tuple.""" + try: + text, final = result[0], result[2] if len(result) > 2 else None + if isinstance(text, str) and text: + span["content"].append(text) + if final is not None: + span["usage"] = getattr(final, "usage", None) + except Exception: + pass + _replay_llm_end(span) + + class _ReplayStream: """Pass-through for a LiteLLM stream that records ``llm_response`` when it ends.""" @@ -7792,6 +7809,32 @@ def _build_responses_params( return params def _call_responses_api(self, **params) -> Any: + """``litellm.responses()`` with retry, recorded on the replay trace.""" + span = _replay_llm_begin(self, params) + if span is None: + return self._call_responses_api_impl(**params) + try: + response = self._call_responses_api_impl(**params) + except BaseException as exc: + _replay_llm_end(span, exc) + raise + self._replay_observe_responses(span, response) + _replay_llm_end(span) + return response + + def _replay_observe_responses(self, span: Dict[str, Any], response: Any) -> None: + """Usage and text of a Responses API result, for the replay trace.""" + try: + span["usage"] = getattr(response, "usage", None) or ( + response.get("usage") if isinstance(response, dict) else None + ) + text = self._extract_from_responses_output(response)[0] + if isinstance(text, str) and text: + span["content"].append(text) + except Exception: + pass # Tracing must never break a model call. + + def _call_responses_api_impl(self, **params) -> Any: """ Call ``litellm.responses()`` synchronously with retry support. Returns a ``ResponsesAPIResponse`` object. @@ -7805,6 +7848,20 @@ def _call_responses_api(self, **params) -> Any: return self._call_with_retry(litellm.responses, **params) async def _call_responses_api_async(self, **params) -> Any: + """``litellm.aresponses()`` with retry, recorded on the replay trace.""" + span = _replay_llm_begin(self, params) + if span is None: + return await self._call_responses_api_async_impl(**params) + try: + response = await self._call_responses_api_async_impl(**params) + except BaseException as exc: + _replay_llm_end(span, exc) + raise + self._replay_observe_responses(span, response) + _replay_llm_end(span) + return response + + async def _call_responses_api_async_impl(self, **params) -> Any: """ Call ``litellm.aresponses()`` asynchronously with retry support. Returns a ``ResponsesAPIResponse`` object. @@ -7934,7 +7991,20 @@ async def _aclose_responses_stream(stream) -> None: except Exception: logging.debug("Responses stream cleanup failed", exc_info=True) - def _stream_responses_api( + def _stream_responses_api(self, params: Dict[str, Any], *args: Any, **kwargs: Any) -> tuple: + """Streamed ``litellm.responses()``, recorded on the replay trace.""" + span = _replay_llm_begin(self, params) + if span is None: + return self._stream_responses_api_impl(params, *args, **kwargs) + try: + result = self._stream_responses_api_impl(params, *args, **kwargs) + except BaseException as exc: + _replay_llm_end(span, exc) + raise + _replay_stream_result(span, result) + return result + + def _stream_responses_api_impl( self, params: Dict[str, Any], verbose: bool = True, @@ -8087,7 +8157,20 @@ def _stream_responses_api( return response_text, tool_calls if tool_calls else None, _final_response - async def _stream_responses_api_async( + async def _stream_responses_api_async(self, params: Dict[str, Any], *args: Any, **kwargs: Any) -> tuple: + """Async streamed ``litellm.aresponses()``, recorded on the replay trace.""" + span = _replay_llm_begin(self, params) + if span is None: + return await self._stream_responses_api_async_impl(params, *args, **kwargs) + try: + result = await self._stream_responses_api_async_impl(params, *args, **kwargs) + except BaseException as exc: + _replay_llm_end(span, exc) + raise + _replay_stream_result(span, result) + return result + + async def _stream_responses_api_async_impl( self, params: Dict[str, Any], stream_callback: Optional[Callable] = None, diff --git a/src/praisonai-agents/tests/unit/trace/test_run_events.py b/src/praisonai-agents/tests/unit/trace/test_run_events.py index 7857b3768b..4e584f2dcd 100644 --- a/src/praisonai-agents/tests/unit/trace/test_run_events.py +++ b/src/praisonai-agents/tests/unit/trace/test_run_events.py @@ -724,6 +724,31 @@ def explode(**params): response = _of(recorder, "llm_response") assert len(response) == 1 and response[0].data["finish_reason"] == "error" + def test_responses_api_calls_record_events(self, recorder, monkeypatch): + """OpenAI models on the LiteLLM path go through the Responses API.""" + from praisonaiagents.llm.llm import LLM + + usage = types.SimpleNamespace(input_tokens=7, output_tokens=2) + monkeypatch.setattr(LLM, "_call_responses_api_impl", lambda self, **p: types.SimpleNamespace(usage=usage)) + monkeypatch.setattr(LLM, "_extract_from_responses_output", lambda self, r: ("answer", None, None)) + monkeypatch.setattr(LLM, "_stream_responses_api_impl", + lambda self, params, *a, **k: ("streamed", None, types.SimpleNamespace(usage=usage))) + + async def acall(self, **p): + return types.SimpleNamespace(usage=usage) + monkeypatch.setattr(LLM, "_call_responses_api_async_impl", acall) + + llm = self._llm() + params = {"model": "gpt-4o-mini", "input": [{"role": "user", "content": "hi"}]} + llm._call_responses_api(**params) + llm._stream_responses_api(dict(params), verbose=False) + asyncio.run(llm._call_responses_api_async(**params)) + requests, responses = _of(recorder, "llm_request"), _of(recorder, "llm_response") + assert len(requests) == len(responses) == 3 + assert requests[0].messages_count == 1 + assert [r.data["response_content"] for r in responses] == ["answer", "streamed", "answer"] + assert all((r.prompt_tokens, r.completion_tokens) == (7, 2) for r in responses) + def test_control_disabled_tracing_leaves_the_stream_untouched(self, fake_litellm): llm = self._llm() stream = llm._completion_with_retry(model="gpt-4o-mini", messages=[], stream=True) From 82680b097085dd425fe5dd61cdb7537666ef25cf Mon Sep 17 00:00:00 2001 From: MervinPraison <454862+MervinPraison@users.noreply.github.com> Date: Mon, 26 Oct 2026 11:25:54 +0000 Subject: [PATCH 07/12] test(replay): real agentic test for team, LLM and hand-off events A two-agent AgentTeam on gpt-4o-mini (one agent on the OpenAI-client path, one on the LiteLLM path) plus a real Triage -> Billing hand-off, recorded with a ContextListSink. Asserts workflow/task/agent/llm events, the hand-off, the team graph, run ids and resolvable parent links. Gated on RUN_REAL_KEY_TESTS=1, PRAISONAI_ALLOW_NETWORK=1 and a real OPENAI_API_KEY; about five small calls. Co-Authored-By: Claude Opus 5.5 (1M context) --- .../integration/test_replay_events_real.py | 113 ++++++++++++++++++ 1 file changed, 113 insertions(+) create mode 100644 src/praisonai-agents/tests/integration/test_replay_events_real.py diff --git a/src/praisonai-agents/tests/integration/test_replay_events_real.py b/src/praisonai-agents/tests/integration/test_replay_events_real.py new file mode 100644 index 0000000000..ed9881153b --- /dev/null +++ b/src/praisonai-agents/tests/integration/test_replay_events_real.py @@ -0,0 +1,113 @@ +"""Real agentic test for PraisonAI Replay events (AGENTS.md section 9.4). + +Runs a real two-agent AgentTeam and a real hand-off against gpt-4o-mini while +recording with a ContextListSink, then checks the trace a replay needs: +workflow, task, agent and LLM events (on both the OpenAI-client path and the +LiteLLM path), a real hand-off, one run id and parent links that resolve. + +About five small model calls. Opt in with: + + RUN_REAL_KEY_TESTS=1 PRAISONAI_ALLOW_NETWORK=1 OPENAI_API_KEY=... \ + pytest tests/integration/test_replay_events_real.py -s +""" + +import os + +import pytest + +pytestmark = [ + pytest.mark.network, + pytest.mark.skipif( + os.getenv("RUN_REAL_KEY_TESTS") != "1" + or not os.getenv("OPENAI_API_KEY") + or os.getenv("OPENAI_API_KEY") == "test-key", + reason="Real model calls: set RUN_REAL_KEY_TESTS=1, PRAISONAI_ALLOW_NETWORK=1 and a real OPENAI_API_KEY", + ), +] + + +def _types(events): + return [e.event_type.value for e in events] + + +def test_team_and_handoff_record_a_replayable_trace(): + from praisonaiagents import Agent, Task + from praisonaiagents.agents.agents import AgentTeam + from praisonaiagents.trace.context_events import ( + ContextListSink, ContextTraceEmitter, trace_context, + ) + + sink = ContextListSink() + emitter = ContextTraceEmitter(sink=sink, session_id="replay-real", enabled=True) + + # Researcher uses the OpenAI-client path, Writer the LiteLLM path. + researcher = Agent(name="Researcher", instructions="Answer in one short sentence.", + llm="gpt-4o-mini") + writer = Agent(name="Writer", instructions="Rewrite the given text as a headline of at most eight words.", + llm="openai/gpt-4o-mini") + team = AgentTeam( + agents=[researcher, writer], + tasks=[ + Task(name="research", description="State one fact about the Moon.", + expected_output="One sentence", agent=researcher), + Task(name="headline", description="Turn the fact into a headline.", + expected_output="A headline", agent=writer), + ], + output="silent", + ) + + billing = Agent(name="Billing", instructions="Reply exactly: Refund logged.", llm="gpt-4o-mini") + triage = Agent( + name="Triage", + instructions=("You never answer yourself. For any request, immediately call the " + "transfer_to_billing tool, then repeat its answer."), + llm="gpt-4o-mini", + handoffs=[billing], + ) + + with trace_context(emitter): + team_output = team.start() + handoff_output = triage.start("Please refund order 7.") + + events = sink.get_events() + print("\n--- team output ---\n", team_output) + print("--- hand-off output ---\n", handoff_output) + print("--- trace ---") + for e in events: + print(f"{e.sequence_num:>3} parent={e.parent_seq!s:>4} {e.event_type.value:<18} " + f"{(e.agent_name or ''):<11} {e.node_id or ''}") + + assert team_output and handoff_output + types = _types(events) + for needed in ("workflow_start", "workflow_end", "task_start", "task_end", + "agent_start", "agent_end", "llm_request", "llm_response", "agent_handoff"): + assert needed in types, needed + + start = next(e for e in events if e.event_type.value == "workflow_start") + assert start.data["kind"] == "team" and start.data["graph"]["schema"] == "praisonai.graph/1" + graph_ids = {n["id"] for n in start.data["graph"]["nodes"]} + task_ends = [e for e in events if e.event_type.value == "task_end"] + assert [e.data["status"] for e in task_ends] == ["completed", "completed"] + assert {e.node_id for e in task_ends} == {"task:research", "task:headline"} <= graph_ids + + team_events = [e for e in events if e.run_id == team.run_id] + assert team_events and team_events[0] is start + + # LLM activity on both paths, attributed to the right agent. + responses = [e for e in events if e.event_type.value == "llm_response"] + by_agent = {e.agent_name for e in responses} + assert {"Researcher", "Writer", "Triage", "Billing"} <= by_agent + writer_calls = [e for e in responses if e.agent_name == "Writer"] + assert "gpt-4o-mini" in writer_calls[0].data.get("model") # recorded by the LLM class + assert writer_calls[0].prompt_tokens > 0 and writer_calls[0].completion_tokens > 0 + assert len([e for e in events if e.event_type.value == "llm_request" and e.agent_name == "Writer"]) == len(writer_calls) + + handoffs = [e for e in events if e.event_type.value == "agent_handoff"] + assert any(h.data["from_agent"] == "Triage" and h.data["to_agent"] == "Billing" + and h.data["kind"] == "handoff" for h in handoffs) + + # Every parent link resolves to an earlier event; sequence numbers are unique. + seqs = [e.sequence_num for e in events] + assert seqs == list(range(len(events))) + for e in events: + assert e.parent_seq is None or e.parent_seq < e.sequence_num From 12d9ee07444c4250ab362870096dbdadd443cebd Mon Sep 17 00:00:00 2001 From: MervinPraison <454862+MervinPraison@users.noreply.github.com> Date: Mon, 26 Oct 2026 12:03:34 +0000 Subject: [PATCH 08/12] fix(replay): label an unnamed team run by its agents AgentTeam.name defaults to None, so workflow_start carried no name and viewers fell back to the first agent's name ('Researcher') as the run's title. An unnamed team is now labelled by its agents ('Researcher + Writer'); a named team keeps its name. Co-Authored-By: Claude Opus 5.5 (1M context) --- .../praisonaiagents/agents/agents.py | 11 +++++++++-- .../tests/unit/trace/test_run_events.py | 16 ++++++++++++++++ 2 files changed, 25 insertions(+), 2 deletions(-) diff --git a/src/praisonai-agents/praisonaiagents/agents/agents.py b/src/praisonai-agents/praisonaiagents/agents/agents.py index add920ad7a..3d7622b2ab 100644 --- a/src/praisonai-agents/praisonaiagents/agents/agents.py +++ b/src/praisonai-agents/praisonaiagents/agents/agents.py @@ -2259,7 +2259,7 @@ async def astart(self, content=None, return_dict=False, **kwargs): # the guard across the whole batch while still calling astart() per item. with self._guard_execution(): from ..trace.context_events import _workflow_scope - with _workflow_scope("team", self.name, self.process, + with _workflow_scope("team", self._replay_name(), self.process, graph_fn=self.to_graph, run_id=self.run_id) as _run: result = await self._astart_impl(content=content, return_dict=return_dict, **kwargs) self._mark_run_status(_run) @@ -2609,13 +2609,20 @@ def start(self, content=None, return_dict=False, output=None, **kwargs): # the guard across the whole batch while still calling start() per item. with self._guard_execution(): from ..trace.context_events import _workflow_scope - with _workflow_scope("team", self.name, self.process, + with _workflow_scope("team", self._replay_name(), self.process, graph_fn=self.to_graph, run_id=self.run_id) as _run: result = self._start_impl(content=content, return_dict=return_dict, output=output, **kwargs) self._mark_run_status(_run) return result + def _replay_name(self) -> str: + """Run label for Replay: the team's name, else its agents joined ("A + B").""" + if self.name: + return self.name + names = [getattr(a, "name", None) or "Agent" for a in (self.agents or [])] + return " + ".join(names) if names else "Agent team" + def _mark_run_status(self, run_state): """Record a failed replay run when any task ended failed (tracing on only).""" if run_state is not None and any( diff --git a/src/praisonai-agents/tests/unit/trace/test_run_events.py b/src/praisonai-agents/tests/unit/trace/test_run_events.py index 4e584f2dcd..f842b70d03 100644 --- a/src/praisonai-agents/tests/unit/trace/test_run_events.py +++ b/src/praisonai-agents/tests/unit/trace/test_run_events.py @@ -792,3 +792,19 @@ def test_disabled_tracing_never_builds_the_graph(): with _workflow_scope("team", "t", "sequential", graph_fn=lambda: built.append(1)) as run: assert run is None assert built == [] + + +class TestTeamRunLabel: + """An unnamed team is labelled by its agents, not by whichever agent comes first.""" + + def test_unnamed_team_is_labelled_by_its_agents(self): + from praisonaiagents import Agent, AgentTeam + + team = AgentTeam(agents=[Agent(name="Researcher", output="silent"), Agent(name="Writer", output="silent")]) + assert team._replay_name() == "Researcher + Writer" + + def test_named_team_keeps_its_name(self): + from praisonaiagents import Agent, AgentTeam + + team = AgentTeam(agents=[Agent(name="Researcher", output="silent")], name="EV brief") + assert team._replay_name() == "EV brief" From 7f92faff5eb2a88e07ab3003b020c4b1e9fc7a3b Mon Sep 17 00:00:00 2001 From: MervinPraison <454862+MervinPraison@users.noreply.github.com> Date: Mon, 26 Oct 2026 12:28:29 +0000 Subject: [PATCH 09/12] fix(replay): address review findings on graph export and LLM events - llm: snapshot the messages list in llm_request; record the model that served the response (failover) for name and cost - diagram: draw when/then_task/else_task team routes; keep both labels when two conditions share a target; expand Repeat(step=[...]) members - tests: replace wall-clock import budget with a deterministic check; add regression tests for each fix Co-Authored-By: Claude Opus 5.5 --- .../praisonaiagents/llm/llm.py | 8 ++++- .../praisonaiagents/workflows/diagram.py | 19 +++++++++- .../tests/unit/test_lazy_imports.py | 15 ++++---- .../tests/unit/trace/test_run_events.py | 13 +++++++ .../tests/unit/workflows/test_graph.py | 35 +++++++++++++++++++ 5 files changed, 79 insertions(+), 11 deletions(-) diff --git a/src/praisonai-agents/praisonaiagents/llm/llm.py b/src/praisonai-agents/praisonaiagents/llm/llm.py index da1c475659..5661ecd5fb 100644 --- a/src/praisonai-agents/praisonaiagents/llm/llm.py +++ b/src/praisonai-agents/praisonaiagents/llm/llm.py @@ -60,7 +60,8 @@ def _replay_llm_begin(llm: Any, params: Dict[str, Any]): model = params.get("model") or getattr(llm, "model", None) # Chat Completions sends ``messages``; the Responses API sends ``input``. messages = params.get("messages", params.get("input")) - messages = messages if isinstance(messages, list) else None + # Copy: the caller keeps appending turns to (and shrinking) the live list. + messages = list(messages) if isinstance(messages, list) else None emitter._emit(emitter._llm_request_event( agent, len(messages) if messages else 0, 0, model=model, messages=messages, )) @@ -85,6 +86,11 @@ def _replay_observe(span: Dict[str, Any], response: Any, chunk: bool) -> None: usage = get("usage", None) if usage: span["usage"] = usage + # The model that actually served the call (a failover profile may + # differ from the requested one), so name and price match it. + served = get("model", None) + if isinstance(served, str) and served: + span["model"] = served choices = get("choices", None) or [] if not choices: return diff --git a/src/praisonai-agents/praisonaiagents/workflows/diagram.py b/src/praisonai-agents/praisonaiagents/workflows/diagram.py index 7fc4da4cec..a3b00744aa 100644 --- a/src/praisonai-agents/praisonaiagents/workflows/diagram.py +++ b/src/praisonai-agents/praisonaiagents/workflows/diagram.py @@ -159,6 +159,12 @@ def edge( "kind": kind, "label": label, } + elif label: + # Two conditions to the same target share one edge; keep both labels. + existing = self.edges[edge_id] + labels = (existing["label"] or "").split(" / ") if existing["label"] else [] + if label not in labels: + existing["label"] = " / ".join(labels + [label]) if draw: self.ops.append(("edge", source, target, label)) @@ -337,8 +343,11 @@ def _flow_step(step: Any, g: _Graph, prev: Optional[str], parent: Optional[str]) if cls in ("Loop", "Repeat"): pattern = cls.lower() + inner = getattr(step, "step", None) + # ``Repeat(step=[a, b])`` runs each member, as ``_execute_repeat`` does. inner_steps = getattr(step, "steps", None) or ( - [getattr(step, "step")] if getattr(step, "step", None) is not None else [] + list(inner) if isinstance(inner, (list, tuple)) + else [inner] if inner is not None else [] ) over = getattr(step, "over", None) iterations = getattr(step, "max_iterations", None) @@ -536,6 +545,14 @@ def _team_graph(team: Any) -> _Graph: condition = getattr(task, "condition", None) or {} followers = [x for x in (getattr(task, "next_tasks", None) or []) if x in by_name] routed = False + # Unified when/then_task/else_task routing, which execution + # checks before condition/next_tasks. + for attr, branch in (("then_task", str(getattr(task, "when", None) or "then")), + ("else_task", "else")): + target = getattr(task, attr, None) + if isinstance(target, str) and target in by_name: + g.edge(node_id, by_name[target], "routes", branch) + routed = True if isinstance(condition, dict): for key, targets in condition.items(): for target in (targets if isinstance(targets, (list, tuple)) else [targets]): diff --git a/src/praisonai-agents/tests/unit/test_lazy_imports.py b/src/praisonai-agents/tests/unit/test_lazy_imports.py index 04d66faf38..7a96faa6b6 100644 --- a/src/praisonai-agents/tests/unit/test_lazy_imports.py +++ b/src/praisonai-agents/tests/unit/test_lazy_imports.py @@ -326,13 +326,10 @@ def test_a_traced_run_builds_the_graph(self): ) assert has_graph == "praisonai.graph/1" - def test_package_import_stays_under_budget(self): - """``import praisonaiagents`` < 200 ms (best of three fresh interpreters).""" - best = min( - float(_fresh_python( - "import time; t = time.perf_counter(); import praisonaiagents;" - "print((time.perf_counter() - t) * 1000)" - )) - for _ in range(3) + def test_package_import_does_not_load_heavy_optional_modules(self): + """Deterministic stand-in for the <200 ms budget: no heavy optional deps.""" + loaded = _fresh_python( + "import sys, praisonaiagents;" + "print(sorted(m for m in ('litellm', 'chromadb') if m in sys.modules))" ) - assert best < 200, f"import praisonaiagents took {best:.0f} ms" + assert loaded == "[]" diff --git a/src/praisonai-agents/tests/unit/trace/test_run_events.py b/src/praisonai-agents/tests/unit/trace/test_run_events.py index f842b70d03..722e6f1774 100644 --- a/src/praisonai-agents/tests/unit/trace/test_run_events.py +++ b/src/praisonai-agents/tests/unit/trace/test_run_events.py @@ -672,6 +672,19 @@ def test_non_streaming_call_records_request_and_response(self, recorder, fake_li assert r.data["duration_ms"] >= 0 assert r.cost_usd > 0 # litellm knows gpt-4o-mini's price + def test_response_records_the_serving_model_and_a_message_snapshot(self, recorder, fake_litellm, monkeypatch): + import litellm + + served = _response() + served.model = "gpt-4o" # a failover profile answered + monkeypatch.setattr(litellm, "completion", lambda **params: served) + messages = [{"role": "user", "content": "hi"}] + self._llm()._completion_with_retry(model="gpt-4o-mini", messages=messages) + messages.append({"role": "assistant", "content": "later turn"}) + assert _of(recorder, "llm_response")[0].data["model"] == "gpt-4o" + assert _of(recorder, "llm_request")[0].messages_count == 1 + assert len(_of(recorder, "llm_request")[0].data["messages"]) == 1 + def test_streaming_call_records_response_when_the_stream_ends(self, recorder, fake_litellm): llm = self._llm() stream = llm._completion_with_retry(model="gpt-4o-mini", messages=[], stream=True) diff --git a/src/praisonai-agents/tests/unit/workflows/test_graph.py b/src/praisonai-agents/tests/unit/workflows/test_graph.py index ebd28ff02a..2f47eda604 100644 --- a/src/praisonai-agents/tests/unit/workflows/test_graph.py +++ b/src/praisonai-agents/tests/unit/workflows/test_graph.py @@ -105,6 +105,41 @@ def test_graph_to_mermaid_draws_leaf_nodes_not_groups(): assert '["parallel"]' not in diagram # a pattern group is a container +def test_team_graph_keeps_when_routes_and_shared_target_labels(): + from praisonaiagents import Agent, Task + from praisonaiagents.agents.agents import AgentTeam + + a = Agent(name="A", instructions="a") + tasks = [ + Task(name="check", description="c", agent=a, is_start=True, + when="score > 80", then_task="publish", else_task="revise"), + Task(name="revise", description="r", agent=a, + condition={"approved": ["publish"], "forced": ["publish"]}), + Task(name="publish", description="p", agent=a), + ] + graph = to_graph(AgentTeam(agents=[a], tasks=tasks, process="workflow")) + edges = {e["id"]: e for e in graph["edges"]} + assert edges["task:check->task:publish:routes"]["label"] == "score > 80" + assert edges["task:check->task:revise:routes"]["label"] == "else" + assert "task:check->end:done:next" not in edges + assert edges["task:revise->task:publish:routes"]["label"] == "approved / forced" + + +def test_repeat_with_a_step_list_draws_each_member(): + from praisonaiagents.workflows.workflows import AgentFlow, Repeat + + def draft(ctx): + return "d" + + def polish(ctx): + return "p" + + graph = to_graph(AgentFlow(steps=[Repeat(step=[draft, polish], until=lambda ctx: True)])) + ids = {n["id"] for n in graph["nodes"]} + assert {"step:draft", "step:polish"} <= ids + assert "step:list" not in ids + + def test_unsupported_object_is_rejected(): with pytest.raises(TypeError): to_graph(42) From d5b87bb0520cca38915c39cdb06caa64284687e8 Mon Sep 17 00:00:00 2001 From: MervinPraison <454862+MervinPraison@users.noreply.github.com> Date: Mon, 26 Oct 2026 13:10:10 +0000 Subject: [PATCH 10/12] chore(parity): record AgentTeam.to_graph/to_mermaid as Python-only The signature gate flags new Python public methods without a TypeScript counterpart. Replay's graph export is Python-only for now (the TS SDK has no Replay viewer or run-event schema 1.1 yet), so the two methods are recorded in the inventory baseline; the regenerated parity reports only move line numbers. Co-Authored-By: Claude Opus 5.5 (1M context) --- src/praisonai-ts/SIGNATURE_PARITY.md | 6 +++--- src/praisonai-ts/signature-parity.json | 6 +++--- .../_dev/parity/signatures/inventory-baseline.json | 2 ++ 3 files changed, 8 insertions(+), 6 deletions(-) diff --git a/src/praisonai-ts/SIGNATURE_PARITY.md b/src/praisonai-ts/SIGNATURE_PARITY.md index cfa1974cb7..6f65a2ddf1 100644 --- a/src/praisonai-ts/SIGNATURE_PARITY.md +++ b/src/praisonai-ts/SIGNATURE_PARITY.md @@ -99,7 +99,7 @@ TS-only members: `config`, `pretty`?, `fetch`?, `outputSchema`?, `outputSchemaNa ### `AgentTeam.__init__` -- Python: `src/praisonai-agents/praisonaiagents/agents/agents.py:804` +- Python: `src/praisonai-agents/praisonaiagents/agents/agents.py:837` - TypeScript: `src/praisonai-ts/src/agent/team.ts:93` (ctor `src/praisonai-ts/src/agent/team.ts:381`) - Python aliases: PraisonAIAgents, Agents - Counts: 23 python params: 20 exact, 3 camelCase, 0 alias, 0 flattened, 0 missing; 1 mismatches; 1 waived; 3 TS-only of 26 @@ -247,7 +247,7 @@ TS-only members: `previousResult`?, `options`? ### `AgentTeam.start` -- Python: `src/praisonai-agents/praisonaiagents/agents/agents.py:2484` +- Python: `src/praisonai-agents/praisonaiagents/agents/agents.py:2573` - TypeScript: `src/praisonai-ts/src/agent/team.ts:1125` - Counts: 3 python params: 2 exact, 1 camelCase, 0 alias, 0 flattened, 0 missing; 0 mismatches; 0 waived; 1 TS-only of 4 @@ -279,7 +279,7 @@ TS-only members: `condition`?, `contextPolicy`?, `maxContextTokens`?, `maxContex ### `LLM.__init__` -- Python: `src/praisonai-agents/praisonaiagents/llm/llm.py:514` +- Python: `src/praisonai-agents/praisonaiagents/llm/llm.py:673` - TypeScript: `src/praisonai-ts/src/llm/index.ts:236` (ctor `src/praisonai-ts/src/llm/index.ts:416`) - Counts: 25 python params: 8 exact, 16 camelCase, 1 alias, 0 flattened, 0 missing; 1 mismatches; 1 waived; 1 TS-only of 26 diff --git a/src/praisonai-ts/signature-parity.json b/src/praisonai-ts/signature-parity.json index 68096c439f..da1869528a 100644 --- a/src/praisonai-ts/signature-parity.json +++ b/src/praisonai-ts/signature-parity.json @@ -2714,7 +2714,7 @@ "PraisonAIAgents", "Agents" ], - "location": "src/praisonai-agents/praisonaiagents/agents/agents.py:804", + "location": "src/praisonai-agents/praisonaiagents/agents/agents.py:837", "resolved_class": "AgentTeam" }, "ts_only": [ @@ -2852,7 +2852,7 @@ } ], "python": { - "location": "src/praisonai-agents/praisonaiagents/agents/agents.py:2484", + "location": "src/praisonai-agents/praisonaiagents/agents/agents.py:2573", "resolved_class": "AgentTeam" }, "ts_only": [ @@ -4659,7 +4659,7 @@ } ], "python": { - "location": "src/praisonai-agents/praisonaiagents/llm/llm.py:514", + "location": "src/praisonai-agents/praisonaiagents/llm/llm.py:673", "resolved_class": "LLM" }, "ts_only": [ diff --git a/src/praisonai/praisonai/_dev/parity/signatures/inventory-baseline.json b/src/praisonai/praisonai/_dev/parity/signatures/inventory-baseline.json index 1e07f62ebb..7369339698 100644 --- a/src/praisonai/praisonai/_dev/parity/signatures/inventory-baseline.json +++ b/src/praisonai/praisonai/_dev/parity/signatures/inventory-baseline.json @@ -122,6 +122,8 @@ "AgentTeam.set_state", "AgentTeam.spawn_sub_agent", "AgentTeam.start_for_each", + "AgentTeam.to_graph", + "AgentTeam.to_mermaid", "AgentTeam.update_plan_step_status", "AgentTeam.update_state", "AgentTeam.wait_for_completions", From e5983ca96461adf2f8b999db77a5d72436177373 Mon Sep 17 00:00:00 2001 From: MervinPraison <454862+MervinPraison@users.noreply.github.com> Date: Mon, 26 Oct 2026 13:18:22 +0000 Subject: [PATCH 11/12] fix(replay): keep distinct route labels when merging a shared graph edge Co-Authored-By: Claude Opus 5.5 --- .../praisonaiagents/workflows/diagram.py | 9 ++++++--- .../tests/unit/workflows/test_graph.py | 15 +++++++++++++++ 2 files changed, 21 insertions(+), 3 deletions(-) diff --git a/src/praisonai-agents/praisonaiagents/workflows/diagram.py b/src/praisonai-agents/praisonaiagents/workflows/diagram.py index a3b00744aa..09b6c36aac 100644 --- a/src/praisonai-agents/praisonaiagents/workflows/diagram.py +++ b/src/praisonai-agents/praisonaiagents/workflows/diagram.py @@ -114,6 +114,8 @@ def __init__(self, kind: str, name: Optional[str], process: str) -> None: # time to stamp node_id on step events. self.ids: Dict[int, str] = {} self._seen: Dict[str, int] = {} + # edge id -> distinct route labels, joined only for display. + self._labels: Dict[str, List[str]] = {} def unique(self, base: str) -> str: """``base`` the first time, then ``base#2``, ``base#3`` ...""" @@ -159,12 +161,13 @@ def edge( "kind": kind, "label": label, } + self._labels[edge_id] = [label] if label else [] elif label: # Two conditions to the same target share one edge; keep both labels. - existing = self.edges[edge_id] - labels = (existing["label"] or "").split(" / ") if existing["label"] else [] + labels = self._labels[edge_id] if label not in labels: - existing["label"] = " / ".join(labels + [label]) + labels.append(label) + self.edges[edge_id]["label"] = " / ".join(labels) if draw: self.ops.append(("edge", source, target, label)) diff --git a/src/praisonai-agents/tests/unit/workflows/test_graph.py b/src/praisonai-agents/tests/unit/workflows/test_graph.py index 2f47eda604..e68a55a89d 100644 --- a/src/praisonai-agents/tests/unit/workflows/test_graph.py +++ b/src/praisonai-agents/tests/unit/workflows/test_graph.py @@ -125,6 +125,21 @@ def test_team_graph_keeps_when_routes_and_shared_target_labels(): assert edges["task:revise->task:publish:routes"]["label"] == "approved / forced" +def test_shared_target_keeps_a_label_that_looks_like_a_merged_one(): + from praisonaiagents import Agent, Task + from praisonaiagents.agents.agents import AgentTeam + + a = Agent(name="A", instructions="a") + tasks = [ + Task(name="revise", description="r", agent=a, is_start=True, + condition={"approved / forced": ["publish"], "approved": ["publish"]}), + Task(name="publish", description="p", agent=a), + ] + graph = to_graph(AgentTeam(agents=[a], tasks=tasks, process="workflow")) + edges = {e["id"]: e for e in graph["edges"]} + assert edges["task:revise->task:publish:routes"]["label"] == "approved / forced / approved" + + def test_repeat_with_a_step_list_draws_each_member(): from praisonaiagents.workflows.workflows import AgentFlow, Repeat From c763f770145aa4ad70b5fe8180b4fc6f4e6d1d29 Mon Sep 17 00:00:00 2001 From: MervinPraison <454862+MervinPraison@users.noreply.github.com> Date: Mon, 26 Oct 2026 13:48:11 +0000 Subject: [PATCH 12/12] docs(replay): note that the recording is the run boundary for run_id Co-Authored-By: Claude Opus 5.5 --- src/praisonai-agents/praisonaiagents/trace/context_events.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/src/praisonai-agents/praisonaiagents/trace/context_events.py b/src/praisonai-agents/praisonaiagents/trace/context_events.py index 967fe5b3bb..7e4d7433ad 100644 --- a/src/praisonai-agents/praisonaiagents/trace/context_events.py +++ b/src/praisonai-agents/praisonaiagents/trace/context_events.py @@ -145,6 +145,10 @@ def _run_scope(run_id: Optional[str] = None) -> Generator[str, None, None]: Nested teams and flows inherit the OUTER run id: when a run id is already in scope it is kept and ``run_id`` is ignored. A fresh uuid is used when no id is given. Yields the effective run id. + + The recording is the run boundary: inside an outer scope (for example + ``praisonai.replay.record()``) every team or flow start shares one run id; + outside one, a team run is stamped with its own ``AgentTeam.run_id``. """ outer = _run_id_var.get() if outer is not None: