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 f785f08288..3d7622b2ab 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._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) + 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) @@ -2441,6 +2515,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 @@ -2519,8 +2608,27 @@ 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._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( + 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/llm/llm.py b/src/praisonai-agents/praisonaiagents/llm/llm.py index 1ae66b4420..5661ecd5fb 100644 --- a/src/praisonai-agents/praisonaiagents/llm/llm.py +++ b/src/praisonai-agents/praisonaiagents/llm/llm.py @@ -43,6 +43,165 @@ 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) + # Chat Completions sends ``messages``; the Responses API sends ``input``. + messages = params.get("messages", params.get("input")) + # 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, + )) + 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)) + # 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: + """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 + # 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 + 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"], + )) + + +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.""" + + __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 +1895,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 +1934,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: @@ -7636,6 +7815,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. @@ -7649,6 +7854,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. @@ -7778,7 +7997,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, @@ -7931,7 +8163,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/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/trace/context_events.py b/src/praisonai-agents/praisonaiagents/trace/context_events.py index 4215a2b2e0..7e4d7433ad 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,130 @@ 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. + + 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: + 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 +301,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 +373,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 +392,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 +417,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 +434,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 +456,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 +554,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 +594,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 +623,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 +786,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 +803,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 +853,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 +901,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 +1001,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 +1030,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 +1043,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 +1056,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 +1102,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 +1114,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 +1147,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 +1168,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, diff --git a/src/praisonai-agents/praisonaiagents/workflows/diagram.py b/src/praisonai-agents/praisonaiagents/workflows/diagram.py index af8e211763..09b6c36aac 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,179 @@ 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] = {} + # 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`` ...""" + 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, + } + self._labels[edge_id] = [label] if label else [] + elif label: + # Two conditions to the same target share one edge; keep both labels. + labels = self._labels[edge_id] + if label not in labels: + labels.append(label) + self.edges[edge_id]["label"] = " / ".join(labels) + 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 +252,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 +264,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 +281,96 @@ 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 = 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) + 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 +380,300 @@ 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 + # 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]): + 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..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 @@ -1498,6 +1596,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``. @@ -1601,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 = "", @@ -1821,6 +1956,7 @@ def _run_impl( continue # Normalize single step + raw_step = step step = self._normalize_single_step(step, i) # Create context for handlers @@ -1839,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: @@ -1854,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: @@ -1904,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}") @@ -2196,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 @@ -2212,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: @@ -2984,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. " @@ -3036,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( @@ -3423,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, @@ -3453,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") @@ -3704,6 +3868,7 @@ def _merge_branch_variables( winners[key] = idx all_variables[key] = value + @_replay_pattern("parallel") def _execute_parallel( self, parallel_step: Parallel, @@ -3956,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, @@ -4342,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 ): @@ -4425,6 +4592,7 @@ def _execute_discussion( "stop": discussion_stopped, } + @_replay_pattern("repeat") def _execute_repeat( self, repeat_step: Repeat, @@ -4499,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, @@ -4530,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" @@ -4583,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, 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 diff --git a/src/praisonai-agents/tests/unit/test_lazy_imports.py b/src/praisonai-agents/tests/unit/test_lazy_imports.py index 66fc17cf04..7a96faa6b6 100644 --- a/src/praisonai-agents/tests/unit/test_lazy_imports.py +++ b/src/praisonai-agents/tests/unit/test_lazy_imports.py @@ -278,3 +278,58 @@ 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_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 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 new file mode 100644 index 0000000000..722e6f1774 --- /dev/null +++ b/src/praisonai-agents/tests/unit/trace/test_run_events.py @@ -0,0 +1,823 @@ +"""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_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) + 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_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) + 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 == [] + + +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" 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..e68a55a89d --- /dev/null +++ b/src/praisonai-agents/tests/unit/workflows/test_graph.py @@ -0,0 +1,160 @@ +"""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_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_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 + + 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) 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",