Compare commits

...
Author SHA1 Message Date
MervinPraisonandClaude Opus 5.5 f37513d3b8 test: set the project root for headless YAML output_file tests
#5770 contains task output_file writes to the project root (cwd or
PRAISONAI_PROJECT_ROOT), matching Agent.output_file. These tests write under
tmp_path, which is outside the cwd, so main's test-core (root) failed with
'output_file not written'. Point the project root at tmp_path.

Co-Authored-By: Claude Opus 5.5 (1M context) <[email protected]>
2026-10-26 21:18:14 +00:00
Mervin Praison a5f35a7a9e Merge pull request #6016 from MervinPraison/fix/replay-rerun-prefix-variables
fix(replay): restore handler step variables on re-run; reject steps inserted before the target
2026-10-26 19:46:09 +01:00
MervinPraison 597868883f docs: auto-update feature parity trackers [skip ci] 2026-10-26 18:42:35 +00:00
56c34fd055 fix(ts): send maxOutputTokens so the AI SDK output-token cap applies (#6013)
* fix(ts): send maxOutputTokens to the AI SDK so the output-token cap applies

AI SDK 5 renamed generateText/streamText's `maxTokens` to `maxOutputTokens`
and ignores the old key. The package pins ai@^6, so AISDKBackend (and
src/ai/generate-text.ts) passed a key the SDK drops: neither
SAFE_DEFAULTS.maxOutputTokens (4096) nor a caller's maxTokens reached the
model on any AI SDK-backed provider (Anthropic, Google, ...).

Send maxOutputTokens, and keep maxTokens for the ai@4 peer range.
Found with SoreqAI (soreqai.com).

Co-Authored-By: Claude Opus 5.5 <[email protected]>

* test(ts): cover maxOutputTokens on all five changed AI SDK calls

Per review: generateText alone left four calls untested. Adds
AISDKBackend.streamText (stream drained so the SDK call runs) and
generateObject, plus the generateText/streamText helpers in
src/ai/generate-text.ts, each checking the default cap and a caller override.

Co-Authored-By: Claude Opus 5.5 <[email protected]>

* fix(ts): send maxOutputTokens from the generateObject/streamObject helpers too

Per CodeRabbit: src/ai/generate-object.ts passed only the old maxTokens key
to sdk.generateObject and sdk.streamObject, the same dropped-cap bug. Adds
maxOutputTokens there and a test for each helper.

Co-Authored-By: Claude Opus 5.5 <[email protected]>

---------

Co-authored-by: Shaked Yacoby <[email protected]>
Co-authored-by: Claude Opus 5.5 <[email protected]>
2026-10-26 19:42:01 +01:00
Swapnil Yadav 13d02e41d1 fix(context-agent): write saved artifacts as UTF-8 (#5989)
* fix(context-agent): write saved artifacts as UTF-8

ContextAgent saved its PRP, analysis summary, and results/interaction
files with the platform default encoding, while other writers in the same
file already pass encoding="utf-8". Writing a PRP or summary that contains
non-ASCII (any non-English content) raised UnicodeEncodeError on Windows,
and the JSON results escaped non-ASCII to \uXXXX. Open every ContextAgent
output file as UTF-8 and dump JSON with ensure_ascii=False so the saved
artifacts are readable and portable, matching the rest of the module.

* test(context-agent): keep the regression test's files inside tmp_path

* test(context-agent): cover all result JSON writers for UTF-8

Populate non-ASCII in every result field and assert each saved file
(comprehensive_analysis, pattern_library, validation_framework,
all_agent_interactions) round-trips as readable UTF-8, not just
comprehensive_analysis.json.
2026-10-26 19:41:56 +01:00
Swapnil Yadav e4b1088f68 fix(memory): use timezone-aware UTC timestamps in the learn store (#5988)
* fix(memory): use timezone-aware UTC timestamps in the learn store

The learn store built timestamps with the deprecated datetime.utcnow(),
which returns a naive datetime. Retention pruning subtracted a naive now
from timestamps parsed back from storage, and datetime.utcnow().timestamp()
computed the wrong epoch by treating the naive UTC value as local time.

Use datetime.now(timezone.utc) everywhere, and when parsing a stored
timestamp that predates this change (naive, no offset), treat it as UTC so
pruning keeps working on existing stores instead of raising "can't subtract
offset-naive and offset-aware datetimes".

* test(memory): add docstrings to learn-store timestamp tests

* fix(memory): normalize timestamps before sorting in the learn store

list_all() and the excess-entry pruning sorted raw ISO strings, so a
stored timestamp with a non-UTC offset could order against a UTC value in
the wrong direction. Add a shared sort key that parses each timestamp,
treats naive values as UTC, and buckets unparseable ones last. Covers the
updated_at ordering and the last_used/created_at pruning tie-breakers.

Also add an end-to-end prune() test and a list ordering test for mixed
offsets.

* fix(memory): sort missing/malformed timestamps as oldest in the learn store

Bucketing unparseable values last flipped prune()'s eviction order: an
entry never used (last_used=None, which parses as empty) sorted after used
ones, so it was evicted last instead of first. Bucket missing/malformed
values as the oldest instant instead. prune() (ascending) then evicts them
first, matching the previous raw-string behaviour, and list_all()
(descending) still lists them last. Add a regression test with two
use_count=0 entries where only one was ever used.
2026-10-26 19:41:49 +01:00
51accf43c6 fix: window doom-loop per-tool count and harden download_file SSRF/size cap (#5959)
* fix: window doom-loop per-tool count and route download_file through shared SSRF guard (fixes #5957)

Co-Authored-By: Claude Opus 5.5 <[email protected]>

* fix: keep download_file SSRF guard independent of ALLOW_LOCAL_CRAWL (#5957)

Co-Authored-By: Claude Opus 5.5 <[email protected]>

* fix: download_file writes to temp file so a rejected download keeps the existing destination (#5957)

Co-Authored-By: Claude Opus 5.5 <[email protected]>

---------

Co-authored-by: MervinPraison <[email protected]>
Co-authored-by: Claude Opus 5.5 <[email protected]>
2026-10-26 19:41:43 +01:00
64d3924985 fix: detect cancel_token support by signature instead of replaying on TypeError (fixes #5884) (#5885)
Co-authored-by: MervinPraison <[email protected]>
Co-authored-by: Claude Opus 5.5 <[email protected]>
2026-10-26 19:41:37 +01:00
abc3d72197 fix(security): contain Task and Workflow step output_file writes in project root (#5770)
* fix(security): contain Task and Workflow step output_file writes in project root

AgentTeam.save_output_to_file and both workflow step writers opened
output_file without the resolve_within_root check Agent.output_file
already applies, so a relative or {{input}}-templated path could write
anywhere on disk. Route all three through a shared write_output_file
helper that rejects paths outside the project root (PRAISONAI_PROJECT_ROOT
or cwd), honours create_directory and writes UTF-8.

Co-Authored-By: Claude Sonnet 5.5 <[email protected]>
Signed-off-by: Stanley Shen <[email protected]>

* fix: preserve output_file permissions on overwrite

atomic_write_text replaces the destination with a mkstemp file (0600),
which silently made shared reports owner-only. Use plain open() like
Agent.output_file so existing modes and umask are kept.

Co-Authored-By: Claude Opus 5.5 <[email protected]>

* fix: replace existing output_file instead of truncating it

Truncating in place writes through a hardlink to a file outside the
project root. Create new files exclusively so umask applies, and replace
existing ones atomically, restoring their previous mode.

Co-Authored-By: Claude Opus 5.5 <[email protected]>
Signed-off-by: Stanley Shen <[email protected]>

* fix: overwrite output_file in place unless it is hardlinked

Replacing the file through a temp file needed write access to the parent
directory and changed the mode through the destination path afterwards.
Truncate the existing file in place when it has a single link, and only
unlink and recreate it (with the old mode) when it is hardlinked.

Co-Authored-By: Claude Opus 5.5 <[email protected]>
Signed-off-by: Stanley Shen <[email protected]>

* fix: swap in hardlinked output_file replacement after writing it

Unlinking first lost the old report if creating or writing the new one
failed, and the replacement's mode was masked by umask. Write a temp file
in the same directory, set the old mode on its descriptor, then
os.replace it over the hardlinked name.

Co-Authored-By: Claude Opus 5.5 <[email protected]>
Signed-off-by: Stanley Shen <[email protected]>

* fix: keep group on hardlinked output_file replacement or drop group bits

Co-Authored-By: Claude Opus 5.5 <[email protected]>

---------

Signed-off-by: Stanley Shen <[email protected]>
Co-authored-by: Claude Sonnet 5.5 <[email protected]>
Co-authored-by: MervinPraison <[email protected]>
2026-10-26 19:41:31 +01:00
MervinPraisonandClaude Opus 5.5 aef2f8ecee fix(replay): snapshot nested-step variables only when the step really runs, not when skipped or cached
Co-Authored-By: Claude Opus 5.5 <[email protected]>
2026-10-26 17:38:25 +00:00
praisonai-triage-agent[bot] a28549c977 Merge pull request #6019 from MervinPraison/claude/issue-6018-20261026
fix: remove dead praisonaiagents/eval/utils.py (duplicate of BaseEvaluator._save_result)
2026-10-26 17:30:15 +00:00
MervinPraisonandClaude Opus 5.5 a62ae64398 fix: remove unreferenced eval/utils.py duplicating BaseEvaluator._save_result (fixes #6018)
Co-Authored-By: Claude Opus 5.5 <[email protected]>
2026-10-26 17:17:07 +00:00
MervinPraisonandClaude Opus 5.5 3e73cfd68b fix(replay): copy nested-step variables per value so one uncopyable value does not disable change detection
Before, one variable that could not be deep-copied (a lock or a client, for
example) switched the whole snapshot to identity comparison, so a list another
variable's step changed in place was no longer recorded on step_end. Use the
existing _isolate_scope, which shares only the uncopyable values.

Co-Authored-By: Claude Opus 5.5 <[email protected]>
2026-10-26 16:38:18 +00:00
MervinPraison 629bc88a9c docs: auto-update feature parity trackers [skip ci] 2026-10-26 16:22:43 +00:00
praisonai-triage-agent[bot] bc3ed37841 Merge pull request #5981 from MervinPraison/claude/issue-5980-20261025
fix: re-runnable AgentTeam.start/astart, non-mutating RAG overrides, SmartRetriever filter/error contracts
2026-10-26 16:22:00 +00:00
MervinPraison b6a8ab1d46 docs: auto-update api.md [skip ci] 2026-10-26 16:13:16 +00:00
praisonai-triage-agent[bot] 783a940d11 Merge pull request #6017 from MervinPraison/feat/replay-show
feat(replay): watch() / replay_app(): a Run button that draws each run live, and embedding
2026-10-26 16:11:15 +00:00
MervinPraisonandClaude Opus 5.5 8eed25c3c2 fix: keep propagated-variable ownership across start_for_each so later runs pick up current team variables
Co-Authored-By: Claude Opus 5.5 <[email protected]>
2026-10-26 16:03:32 +00:00
MervinPraisonandClaude Opus 5.5 693a3337b2 fix: keep caller-defined Task.variables on AgentTeam rerun, retain injected content across resume, deterministic race tests
Co-Authored-By: Claude Opus 5.5 <[email protected]>
2026-10-26 15:58:49 +00:00
MervinPraisonandClaude Opus 5.5 fcf4bebc15 fix: re-runnable AgentTeam.start/astart, non-mutating RAG.retrieve overrides, and SmartRetriever filter/error contracts (fixes #5980)
- AgentTeam: a repeated start()/astart() resets per-task run state and the
  content the previous run injected, instead of skipping completed tasks and
  returning stale output. restore_session_state() still resumes.
- RAG.retrieve: resolve top_k/rerank per call instead of writing them to the
  shared RAGConfig (aretrieve runs concurrently in an executor).
- SmartRetriever: over-fetch before post-filtering, raise ValueError on an
  invalid path_filter regex, return store errors from _search instead of
  storing them on self.
- SubQuestionEngine: unwrap RetrievalResult.chunks from SmartRetriever.

Co-Authored-By: Claude Opus 5.5 <[email protected]>
2026-10-26 15:58:49 +00:00
MervinPraison 6f09f5e02f docs: auto-update feature parity trackers [skip ci] 2026-10-26 15:57:45 +00:00
Wu Shuwen f74706d573 fix: check failed workflow dependencies before dispatch (#5936) 2026-10-26 16:57:13 +01:00
Wu Shuwen 2f93c7bf82 fix: keep synchronous input guardrails off the async event loop (#5932)
* fix: offload synchronous input guardrails in achat

* perf: retain disabled input guardrail fast path

* perf: avoid thread dispatch for output-only guardrails
2026-10-26 16:57:07 +01:00
Wu Shuwen b5713a302d fix: avoid default memory-tool import for supplied adapters (#5896)
* fix: use supplied memory adapter without importing default tool

* test: focus memory adapter checks on capability and input ownership
2026-10-26 16:57:01 +01:00
Wu Shuwen 87a3baea82 perf: defer handoff import when no conversion is needed (#5894)
* perf: defer handoff import when no conversion is needed

* test: document handoff import regression scenarios
2026-10-26 16:56:56 +01:00
Wu Shuwen f50f64ef93 fix: populate existing after-LLM hook token usage (#5890) 2026-10-26 16:56:50 +01:00
Wu Shuwen 942c5dbac0 fix: reject undefined YAML route targets (#5888) 2026-10-26 16:56:44 +01:00
Wu Shuwen 53afc70905 fix: return retained completions from synchronous default wait (#5887)
* fix: return retained completions from synchronous default wait

* fix: preserve explicit empty completion target selection
2026-10-26 16:56:38 +01:00
Wu Shuwen dabc3f5593 fix: align asynchronous moderation default model (#5886) 2026-10-26 16:56:32 +01:00
MervinPraisonandClaude Opus 5.5 3bee2f57de fix(replay): restore team task context between runs, await async targets, bracket IPv6 URL
Co-Authored-By: Claude Opus 5.5 <[email protected]>
2026-10-26 15:53:29 +00:00
MervinPraisonandClaude Opus 5.5 06c97d3d7a test(replay): prove a recorded skip stays skipped on re-run (replaces a test that only checked a file existed)
Co-Authored-By: Claude Opus 5.5 (1M context) <[email protected]>
2026-10-26 15:52:48 +00:00
MervinPraisonandClaude Opus 5.5 e3a4319b50 fix(replay): keep full outputs out of cached step_end; check nested targets; compare nested variables by value
- core: a step-cache hit traced the cached variables, which include the step's
  output variable (its full output), bypassing preview-only recording. Trace
  only what the handler set, as the cold run does.
- core: nested-step variables are diffed against a deep copy (recording only),
  so a list or dict changed in place counts as set; identity missed it and a
  re-run restored the stale value.
- wrapper: the inserted-step check now covers a target nested in a pattern
  (e.g. If): every top-level step before its container must have run.

Addresses the merge-gate findings on #6016 (workflows.py:2044, rerun.py:316,
workflows.py:3143).

Co-Authored-By: Claude Opus 5.5 (1M context) <[email protected]>
2026-10-26 15:51:27 +00:00
MervinPraisonandClaude Opus 5.5 faba52b940 feat(replay): watch() / replay_app(): a Run button that draws each run live, and embedding
- praisonai.replay.watch(target): serves a page with a Run button; each click
  records the Agent/AgentTeam/AgentFlow/YAML/callable under a new run id and the
  page draws it live (map: not reached / running / done). A team's tasks start
  afresh each run (AgentTeam otherwise skips completed tasks).
- praisonai.replay.replay_app(target=None): the same page as an ASGI app to mount
  in your own FastAPI/Starlette app; /replay/?run=<id> shows one run (iframe).
- CLI: praisonai replay watch [agents.yaml] ("show" was taken).
- examples/replay: watch_flow.py (no key), watch_team.py, embed_fastapi.py.

Needs aiui with create_replay_app (MervinPraison/PraisonAIUI#317); older aiui
gets a clear "pip install -U praisonai[ui]" error and the tests skip.

Co-Authored-By: Claude Opus 5.5 (1M context) <[email protected]>
2026-10-26 15:47:17 +00:00
MervinPraisonandClaude Opus 5.5 91641a906d fix(replay): record nested-step variables; reject redacted restores and unnamed inserted steps
- core: steps run via _execute_single_step_internal (patterns, hierarchical)
  now record the variables they set on step_end, like the sequential path.
- wrapper: plan_rerun refuses to restore a variable the recording redacted,
  instead of feeding the literal [REDACTED] marker to later steps.
- wrapper: the inserted-step check matches top-level steps by graph node id,
  so an unnamed pattern (e.g. If) before the target is rejected too.

Co-Authored-By: Claude Opus 5.5 <[email protected]>
2026-10-26 15:14:29 +00:00
MervinPraison 930d6d9837 docs: auto-update api.md [skip ci] 2026-10-26 14:27:30 +00:00
MervinPraisonandClaude Opus 5.5 746cb1fcdf fix(replay): restore handler step variables on re-run; reject steps inserted before the target
- core: step_end records the variables a step set (redacted), so a re-run can
  restore them; before this, step_end carried none and the wrapper's restore
  path had nothing to read from a real recording.
- wrapper: plan_rerun records which steps started before the target; seed_flow
  refuses a flow that has an earlier step the recording never ran, instead of
  silently feeding the target the wrong prefix.

Addresses review threads on #6011 (rerun.py:240, rerun.py:297).

Co-Authored-By: Claude Opus 5.5 (1M context) <[email protected]>
2026-10-26 14:27:05 +00:00
Mervin Praison 05d38d1705 Merge pull request #6010 from MervinPraison/feat/replay-core
feat(replay): complete run events (run/parent ids, tasks, steps, hand-offs, approvals, LiteLLM) and graph export for Replay
2026-10-26 15:26:56 +01:00
adbc6d7b3d fix: route Telegram/Discord/Slack group gating through core admission primitive (#5911)
* test: cover cross-channel group admission parity (#5908)

Adapter routing landed on main via #5962 (_group_policy); this keeps the
parity tests, retargeted at decide_group_admission.

Co-Authored-By: Claude Opus 5.5 <[email protected]>

* fix: cap litellm <1.106 to stay compatible with mcp<2 pin

Co-Authored-By: Claude Opus 5.5 <[email protected]>

* fix: sync praisonai-code uv.lock with litellm<1.106 cap

Co-Authored-By: Claude Opus 5.5 <[email protected]>

---------

Co-authored-by: MervinPraison <[email protected]>
Co-authored-by: Claude Opus 5.5 <[email protected]>
2026-10-26 14:53:18 +01:00
MervinPraisonandClaude Opus 5.5 c763f77014 docs(replay): note that the recording is the run boundary for run_id
Co-Authored-By: Claude Opus 5.5 <[email protected]>
2026-10-26 13:48:11 +00:00
MervinPraison a0516b47db docs: auto-update api.md [skip ci] 2026-10-26 13:44:43 +00:00
Mervin Praison 1d65000a30 Merge pull request #6011 from MervinPraison/feat/replay-wrapper
feat(replay): praisonai replay view/export/rerun, record(), OTel and self-contained HTML export
2026-10-26 14:44:17 +01:00
MervinPraisonandClaude Opus 5.5 e5983ca964 fix(replay): keep distinct route labels when merging a shared graph edge
Co-Authored-By: Claude Opus 5.5 <[email protected]>
2026-10-26 13:18:22 +00:00
MervinPraisonandClaude Opus 5.5 d5b87bb052 chore(parity): record AgentTeam.to_graph/to_mermaid as Python-only
The signature gate flags new Python public methods without a TypeScript
counterpart. Replay's graph export is Python-only for now (the TS SDK has no
Replay viewer or run-event schema 1.1 yet), so the two methods are recorded in
the inventory baseline; the regenerated parity reports only move line numbers.

Co-Authored-By: Claude Opus 5.5 (1M context) <[email protected]>
2026-10-26 13:10:45 +00:00
Mervin PraisonandClaude Opus 5.5 f060fcfbdd fix(browser): assert start_automation on send_json; run test_server.py in CI (#6012)
#5818 moved start_automation delivery from send_text to send_json, but
test_ws_no_origin_cli_start_session still asserted send_text, and a
not-awaited check on send_text became vacuous. CI never ran test_server.py
(only test_protocol_compliance.py), so the break went unseen.

Co-authored-by: Claude Opus 5.5 (1M context) <[email protected]>
2026-10-26 13:52:09 +01:00
MervinPraisonandClaude Opus 5.5 3250b28893 fix(replay): export filename ignores unsafe trace run ids; rerun rejects reordered flows
Co-Authored-By: Claude Opus 5.5 <[email protected]>
2026-10-26 12:42:02 +00:00
MervinPraisonandClaude Opus 5.5 e36fa35ef5 fix(replay): strip ANSI codes before checking view --help flags
Rich forces colour under GITHUB_ACTIONS and splits '--live' with escape
codes, so the plain substring check failed on CI only.

Co-Authored-By: Claude Opus 5.5 <[email protected]>
2026-10-26 12:38:55 +00:00
MervinPraisonandClaude Opus 5.5 7f92faff5e fix(replay): address review findings on graph export and LLM events
- llm: snapshot the messages list in llm_request; record the model that
  served the response (failover) for name and cost
- diagram: draw when/then_task/else_task team routes; keep both labels
  when two conditions share a target; expand Repeat(step=[...]) members
- tests: replace wall-clock import budget with a deterministic check;
  add regression tests for each fix

Co-Authored-By: Claude Opus 5.5 <[email protected]>
2026-10-26 12:28:29 +00:00
MervinPraisonandClaude Opus 5.5 be0b44ff35 fix(replay): rerun keeps step variables, rejects shortened answers and mismatched flows
Co-Authored-By: Claude Opus 5.5 <[email protected]>
2026-10-26 12:26:32 +00:00
MervinPraisonandClaude Opus 5.5 c90f145745 fix(replay): HTML export runs aiui's renderer instead of the fallback table
aiui's replay-core.js and replay.js are ES modules (export statements). Inlined
as classic <script> elements they threw a SyntaxError, so an exported run only
ever showed the built-in fallback table. Inline them as module scripts.

Co-Authored-By: Claude Opus 5.5 (1M context) <[email protected]>
2026-10-26 12:05:37 +00:00
MervinPraisonandClaude Opus 5.5 2337d08a01 fix(replay): rerun accepts the Path that record() returns
plan_rerun only treated str as a trace path, so passing rec.path (a
pathlib.Path) straight from record() crashed with 'PosixPath is not iterable'.

Co-Authored-By: Claude Opus 5.5 (1M context) <[email protected]>
2026-10-26 12:05:37 +00:00
MervinPraisonandClaude Opus 5.5 d73f4d4b8e fix(replay): a team YAML's map uses the ids its run emits
`replay flow team.yaml` built the static graph from the YAML flow conversion,
which gives step: node ids, while running the same YAML builds an AgentTeam
whose events carry task:/agent: ids, so live events never landed on the map.
Team YAMLs now go through the PraisonAI adapter's own builders (no LLM call);
flow YAMLs and older SDKs keep the previous path.

Co-Authored-By: Claude Opus 5.5 (1M context) <[email protected]>
2026-10-26 12:05:37 +00:00
MervinPraisonandClaude Opus 5.5 e6f6ba9d57 fix(replay): record() stamps its run_id on every event
record() passed run_id only to emitters that took it as an argument, and the
schema 1.1 emitter stamps it through a run scope instead, so events outside a
team or flow (session start/end, a single agent's hand-off) had no run_id.
record() now enters the SDK run scope for the whole recording; nested teams
and flows inherit it. Older SDKs without run scopes are unaffected.

Co-Authored-By: Claude Opus 5.5 (1M context) <[email protected]>
2026-10-26 12:05:37 +00:00
MervinPraisonandClaude Opus 5.5 5dfac0c01f test(replay): real agentic record -> export test (gated)
RUN_REAL_KEY_TESTS=1 PRAISONAI_ALLOW_NETWORK=1 with OPENAI_API_KEY: one
gpt-4o-mini call under record(live=True), then replay export --format
html and --format otel on the recorded trace.

Co-Authored-By: Claude Opus 5.5 (1M context) <[email protected]>
2026-10-26 12:05:37 +00:00
MervinPraisonandClaude Opus 5.5 d49a32f713 feat(replay): rerun a recorded run from a node (P4)
replay rerun RUN --from NODE_ID [--file YAML]; Python:
praisonai.replay.rerun.rerun(flow_or_team, run, node_id).

- plan_rerun reads the recording: outputs of every task/step that
  finished before NODE started (completed/cached). A preview cut at 500
  chars is recovered from the span's last llm_response, otherwise the
  rerun refuses rather than feed the next node a truncated input. 1.0
  traces (no node_id) and unknown nodes fail with a clear message.
- AgentFlow: TraceStepCache implements StepCacheProtocol and serves the
  recorded outputs through the flow's own cache-hit path (variables,
  previous_output, callbacks), identifying the asking step by the
  "running" status the flow sets just before the lookup; set() is a
  no-op so NODE onwards always executes.
- AgentTeam: the team's checkpoint restore (_restore_serialised_task_state)
  marks earlier tasks completed with their recorded output; tasks from
  NODE on are reset to "not started".
- YAML runs: a zero-cost ContextVar hook (replay/_hook.py) seeds the
  team/flow AgentsGenerator builds; the CLI checks the YAML kind first so
  a mismatch is refused before any LLM call.
- The new run records forked_from = {run_id, node_id} on session_start
  and on its first workflow_start (sink wrapper, no core change needed).

Co-Authored-By: Claude Opus 5.5 (1M context) <[email protected]>
2026-10-26 12:05:37 +00:00
MervinPraisonandClaude Opus 5.5 27672202e2 feat(replay): replay view / export CLI, mermaid flow and YAML static graph
- replay flow <run|yaml> --format ascii|json|mermaid: a YAML file gives
  the static graph (diagram.to_graph / graph_to_mermaid when the SDK has
  them, else AgentFlow.to_mermaid()); a run gives the sequence diagram.
  Existing ascii/json output for runs is unchanged.
- replay view [RUN|JSONL|YAML] [--live] [--port 8090] [--host 127.0.0.1]
  launches aiui with the new praisonai/ui_replay preset (pages=["replay"])
  and opens http://127.0.0.1:<port>/#/replay?run=<id>. A YAML file is run
  recorded live in a background thread while serving; a JSONL path
  serves its own directory (PRAISONAI_TRACES_DIR for the server). Warns
  on a non-loopback host, fails with the pip install "praisonai[ui]"
  hint when aiui is missing, and stops the server on exit/SIGTERM.
- replay export RUN --format html|json|mermaid|otel [-o FILE]
  [--no-redact].

Co-Authored-By: Claude Opus 5.5 (1M context) <[email protected]>
2026-10-26 12:05:37 +00:00
MervinPraisonandClaude Opus 5.5 cf241acdf1 feat(replay): export runs to JSON, Mermaid, OTLP-JSON and one HTML file
praisonai.replay.export, stdlib only:
- to_json: {schema, run_id, graph, events}.
- to_mermaid: sequenceDiagram with lanes You / Manager / one per agent;
  messages are task boundaries, hand-offs/delegations and approvals
  (1.0 traces fall back to agent start/end).
- to_otlp_json: offline OTLP/JSON with OTel GenAI semconv
  (invoke_workflow / invoke_agent / chat / execute_tool, INTERNAL
  task/step/approval spans), traceId = sha256(run_id)[:32],
  spanId = sha256(run_id:start_seq)[:16], parents from parent_seq or
  nesting; scope marked gen_ai.semconv.stability=development.
- to_html: self-contained page, no CDN. Inlines aiui's standalone
  replay-core.js + replay.js when aiui ships them (detected at runtime),
  always includes a plain table timeline; embedded JSON escapes every
  "<" so event data cannot close the script.
- Everything is redacted by default (praisonaiagents.trace.redact key
  rules + scrub_pii_text), keeping numeric token counters; redact=False
  opts out.

Co-Authored-By: Claude Opus 5.5 (1M context) <[email protected]>
2026-10-26 12:05:37 +00:00
MervinPraisonandClaude Opus 5.5 4629a75276 feat(replay): record() context manager and env-aware traces dir
- Traces live in PRAISONAI_TRACES_DIR, else $PRAISON_HOME/traces (the
  root trace/last_run.py already uses), else ~/.praison/traces.
- praisonai.replay.record(run_id=None, *, live=False, meta=None,
  full_content=None) owns writer + emitter + set_context_emitter and the
  session_start/session_end pair; live=True flushes every event
  (buffer_size=1). It yields .run_id / .path.
- The three duplicated setup blocks (legacy praison_ai.py via the
  wrapper bridge, workflow run, recipe core) now use it. Recipe early
  returns (not found / missing deps / policy denied / dry run) used to
  leave the trace open and the emitter set; they now always close.
- save_replay: true in agents.yaml / workflow YAML records like --save
  (added to YAMLConfig, schema artefact regenerated). YAML runs store
  yaml_path in session_start so a rerun can find the definition.
- events.py: tolerant raw-dict reader (partial last line, unknown event
  types, schema 1.0 and 1.1) plus start/end span pairing.

Co-Authored-By: Claude Opus 5.5 (1M context) <[email protected]>
2026-10-26 12:05:37 +00:00
MervinPraisonandClaude Opus 5.5 12d9ee0744 fix(replay): label an unnamed team run by its agents
AgentTeam.name defaults to None, so workflow_start carried no name and viewers
fell back to the first agent's name ('Researcher') as the run's title. An
unnamed team is now labelled by its agents ('Researcher + Writer'); a named
team keeps its name.

Co-Authored-By: Claude Opus 5.5 (1M context) <[email protected]>
2026-10-26 12:05:36 +00:00
MervinPraisonandClaude Opus 5.5 82680b0970 test(replay): real agentic test for team, LLM and hand-off events
A two-agent AgentTeam on gpt-4o-mini (one agent on the OpenAI-client path, one
on the LiteLLM path) plus a real Triage -> Billing hand-off, recorded with a
ContextListSink. Asserts workflow/task/agent/llm events, the hand-off, the
team graph, run ids and resolvable parent links. Gated on RUN_REAL_KEY_TESTS=1,
PRAISONAI_ALLOW_NETWORK=1 and a real OPENAI_API_KEY; about five small calls.

Co-Authored-By: Claude Opus 5.5 (1M context) <[email protected]>
2026-10-26 12:05:36 +00:00
MervinPraisonandClaude Opus 5.5 da00df8e69 feat(llm): record the Responses API calls on the LiteLLM path too
OpenAI models reached through LiteLLM (e.g. llm="openai/gpt-4o-mini") use
litellm.responses, not litellm.completion, so the real agentic run showed the
Writer agent with no LLM events. The four Responses API entry points (sync,
async, streamed sync, streamed async) are now thin traced wrappers around the
unchanged bodies, recording input/output tokens, text and cost.

Co-Authored-By: Claude Opus 5.5 (1M context) <[email protected]>
2026-10-26 12:05:36 +00:00
MervinPraisonandClaude Opus 5.5 fea0c82680 test(replay): cover run events, spans, branches and every emission site
Schema round-trip and 1.0 back-compat, unknown event types, parent_seq and
run_id inheritance, 8 concurrent Parallel branches (branch_id, parent_seq,
unique sequence numbers), flow step statuses and chosen routes, team task
completed/skipped/failed (sync and async), the manager's delegate hand-off,
human-review and tool-call approval pairs, every hand-off path emitting once,
the LiteLLM path (mocked: sync, async, streaming, errors, no double count),
and a micro-benchmark that the disabled emitter adds next to nothing.

Co-Authored-By: Claude Opus 5.5 (1M context) <[email protected]>
2026-10-26 12:05:36 +00:00
MervinPraisonandClaude Opus 5.5 9ba8c09bd9 feat(llm): record llm_request/llm_response on the LiteLLM path
Only the agent's OpenAI-client path recorded model calls, so a team or flow
running on any LiteLLM model showed no LLM activity in a replay. The two
LiteLLM choke points (_completion_with_retry and its async twin) now emit
the pair with model, duration_ms, prompt/completion tokens and cost when
litellm knows the price. Streams are wrapped in a pass-through that records
the response when the stream ends; with tracing off the stream is untouched.

A call the agent already recorded (its OpenAI-client dispatch can reach the
LiteLLM adapter) is not counted twice: the agent's llm_request leaves a
context marker the LLM class honours inside the same span. chat_mixin.py is
not modified.

Co-Authored-By: Claude Opus 5.5 (1M context) <[email protected]>
2026-10-26 12:05:36 +00:00
MervinPraisonandClaude Opus 5.5 08ca3aca18 feat(replay): record hand-offs, team tasks, flow steps, manager choices and approvals
- agent/handoff.py: every place a hand-off actually happens (programmatic,
  async, LLM tool, typed sync/async) emits one agent_handoff(kind="handoff")
  with the context policy and redacted args.
- agents/agents.py: task_start/task_end (completed, failed after retries,
  skipped by should_run) on the sync and async paths in every process mode;
  workflow_start/workflow_end with the team graph inside _run_scope(run_id);
  approval_requested/approval_decided around human review.
- workflows/workflows.py: step_start/step_end for top-level and nested steps
  (completed/failed/skipped/cached) via helpers that wrap the existing
  on_step_start/on_step_complete calls; Route and If record the branch taken;
  pattern executors are wrapped once so every nesting level is covered;
  workflow_start/workflow_end with the flow graph and a fresh run id.
- process/process.py: the hierarchical manager's choice is an
  agent_handoff from "Manager" with kind="delegate", the task node id and a
  task-status summary (sync and async).
- approval/registry.py: approval_requested/approval_decided around the real
  backend call only (not the fast paths), args redacted.

All of it is skipped when the context emitter is disabled; the graph builder
is imported only for a traced run.

Co-Authored-By: Claude Opus 5.5 (1M context) <[email protected]>
2026-10-26 12:05:36 +00:00
MervinPraisonandClaude Opus 5.5 662edacc9e feat(workflows): graph data for agents, teams and flows
One builder walks a definition and produces praisonai.graph/1 data with ids
derived from names and positions only (agent:, task:, step:, tool:,
human:you, manager:Manager, start:begin, end:done), so the ids are the same
in every process and are what replay events carry as node_id. Covers task
context/next_tasks/conditions, the hierarchical manager, hand-offs, tools,
human review and Route/Parallel/Loop/Repeat/If/Include/Discussion.

AgentFlow.to_mermaid() output is byte-identical (the builder records the
legacy draw order). Adds to_graph(), graph_to_mermaid(), AgentFlow.to_graph()
and AgentTeam.to_graph()/to_mermaid(). Stdlib only and never imported by
import praisonaiagents; golden JSON and a two-subprocess stability check pin
the shape.

Co-Authored-By: Claude Opus 5.5 (1M context) <[email protected]>
2026-10-26 12:05:36 +00:00
MervinPraisonandClaude Opus 5.5 12a04358b3 feat(trace): replay schema 1.1 with run_id, parent_seq and node_id
Context events gain three optional fields (omitted from to_dict when unset)
so a replay can tell which run, which enclosing span and which graph node an
event belongs to. 1.0 traces still load, and from_dict keeps an unknown
event_type as its raw string (with .value) instead of relabelling it.

Branch id and the open-span stack move into contextvars keyed by emitter, so
concurrent Parallel/Loop branches no longer clobber each other; sequence
numbers come from itertools.count under a lock so append order matches
sequence order. Adds workflow/task/step/approval start/end events, a kind on
agent_handoff, and private _run_scope/_workflow_scope helpers. Every new
method is a no-op when the emitter is disabled.

Co-Authored-By: Claude Opus 5.5 (1M context) <[email protected]>
2026-10-26 12:05:36 +00:00
92 changed files with 10414 additions and 964 deletions

No files matched your search

+1 -1
View File
@@ -256,7 +256,7 @@ jobs:
"
/tmp/praisonai-browser-standalone/bin/praisonai-browser --help
bash "$GITHUB_WORKSPACE/scripts/check_c11_browser_imports.sh"
python -m pytest src/praisonai-browser/tests/browser/test_protocol_compliance.py -q -m "not network"
python -m pytest src/praisonai-browser/tests/browser/test_protocol_compliance.py src/praisonai-browser/tests/browser/test_server.py -q -m "not network"
- name: Standalone praisonai-mcp smoke (C12)
run: |
+7
View File
@@ -717,6 +717,7 @@ Methods:
* <code title="class AgentFlow">AgentFlow.<a href="./src/praisonai-agents/praisonaiagents/workflows/workflows.py">state</a>()</code>
* <code title="class AgentFlow">AgentFlow.<a href="./src/praisonai-agents/praisonaiagents/workflows/workflows.py">stream</a>() -> bool</code>
* <code title="class AgentFlow">AgentFlow.<a href="./src/praisonai-agents/praisonaiagents/workflows/workflows.py">to_dict</a>() -> Dict[str, Any]</code>
* <code title="class AgentFlow">AgentFlow.<a href="./src/praisonai-agents/praisonaiagents/workflows/workflows.py">to_graph</a>() -> Dict[str, Any]</code>
* <code title="class AgentFlow">AgentFlow.<a href="./src/praisonai-agents/praisonaiagents/workflows/workflows.py">to_mermaid</a>() -> str</code>
* <code title="class AgentFlow">AgentFlow.<a href="./src/praisonai-agents/praisonaiagents/workflows/workflows.py">validate_variables</a>() -> None</code>
* <code title="class AgentFlow">AgentFlow.<a href="./src/praisonai-agents/praisonaiagents/workflows/workflows.py">verbose</a>() -> bool</code>
@@ -770,6 +771,8 @@ Methods:
* <code title="class AgentTeam">AgentTeam.<a href="./src/praisonai-agents/praisonaiagents/agents/agents.py">start</a>(content = None, return_dict = False, output = None, **kwargs)</code>
* <code title="class AgentTeam">AgentTeam.<a href="./src/praisonai-agents/praisonaiagents/agents/agents.py">start_for_each</a>(inputs, **kwargs)</code>
* <code title="class AgentTeam">AgentTeam.<a href="./src/praisonai-agents/praisonaiagents/agents/agents.py">stream_emitter</a>()</code>
* <code title="class AgentTeam">AgentTeam.<a href="./src/praisonai-agents/praisonaiagents/agents/agents.py">to_graph</a>() -> Dict[str, Any]</code>
* <code title="class AgentTeam">AgentTeam.<a href="./src/praisonai-agents/praisonaiagents/agents/agents.py">to_mermaid</a>() -> str</code>
* <code title="class AgentTeam">AgentTeam.<a href="./src/praisonai-agents/praisonaiagents/agents/agents.py">todo_list</a>()</code>
* <code title="class AgentTeam">AgentTeam.<a href="./src/praisonai-agents/praisonaiagents/agents/agents.py">update_plan_step_status</a>(step_id: str, status: str) -> bool</code>
* <code title="class AgentTeam">AgentTeam.<a href="./src/praisonai-agents/praisonaiagents/agents/agents.py">update_state</a>(updates: Dict) -> None</code>
@@ -898,9 +901,13 @@ Methods:
* <code title="cli">praisonai replay context <a href="./src/praisonai/praisonai/cli/commands/replay.py">--help</a></code>
* <code title="cli">praisonai replay dashboard <a href="./src/praisonai/praisonai/cli/commands/replay.py">--help</a></code>
* <code title="cli">praisonai replay delete <a href="./src/praisonai/praisonai/cli/commands/replay.py">--help</a></code>
* <code title="cli">praisonai replay export <a href="./src/praisonai/praisonai/cli/commands/replay.py">--help</a></code>
* <code title="cli">praisonai replay flow <a href="./src/praisonai/praisonai/cli/commands/replay.py">--help</a></code>
* <code title="cli">praisonai replay list <a href="./src/praisonai/praisonai/cli/commands/replay.py">--help</a></code>
* <code title="cli">praisonai replay rerun <a href="./src/praisonai/praisonai/cli/commands/replay.py">--help</a></code>
* <code title="cli">praisonai replay show <a href="./src/praisonai/praisonai/cli/commands/replay.py">--help</a></code>
* <code title="cli">praisonai replay view <a href="./src/praisonai/praisonai/cli/commands/replay.py">--help</a></code>
* <code title="cli">praisonai replay watch <a href="./src/praisonai/praisonai/cli/commands/replay.py">--help</a></code>
* <code title="cli">praisonai run <a href="./src/praisonai/praisonai/cli/main.py">--help</a></code>
* <code title="cli">praisonai schedule add <a href="./src/praisonai/praisonai/cli/commands/schedule.py">--help</a></code>
* <code title="cli">praisonai schedule blueprint <a href="./src/praisonai/praisonai/cli/commands/schedule.py">--help</a></code>
+16
View File
@@ -0,0 +1,16 @@
# PraisonAI Replay: watch runs live
| Example | What it shows |
|---|---|
| `watch_flow.py` | Click **Run** and watch each step light up on the map; **Play** goes through it again. No API key. |
| `watch_team.py` | The same for two real agents (needs `OPENAI_API_KEY`). |
| `embed_fastapi.py` | Your own app starts a run and shows its progress in an `<iframe>` (`app.mount("/replay", replay_app())`). |
The same from the command line, for a YAML file:
```bash
praisonai replay watch agents.yaml
```
Install: `pip install "praisonai[ui]"`. Keep the page on localhost or behind your app's
login: anyone who can reach it can read the runs, and with a Run button, start one.
+60
View File
@@ -0,0 +1,60 @@
"""Show agent progress inside your own web app.
Your app starts the run and records it; the Replay view, mounted at /replay,
draws it live in an <iframe>. No dashboard, no second server.
pip install "praisonai[ui]" fastapi uvicorn
uvicorn embed_fastapi:app --port 8000 # open http://127.0.0.1:8000/
"""
import threading
import time
import uuid
from fastapi import FastAPI
from fastapi.responses import HTMLResponse
from praisonaiagents import AgentFlow
from praisonaiagents.workflows import StepResult
from praisonai.replay import record, replay_app
def research(ctx):
time.sleep(1.2)
return StepResult(output="notes")
def write(ctx):
time.sleep(1.2)
return StepResult(output="article")
flow = AgentFlow(name="Article", steps=[research, write])
app = FastAPI()
app.mount("/replay", replay_app()) # read-only view: /replay/?run=<id> shows one run
@app.post("/generate")
def generate():
run_id = f"run-{uuid.uuid4().hex[:12]}"
def job():
with record(run_id, live=True): # live: the view follows it as it is written
flow.start("solar sails")
threading.Thread(target=job, daemon=True).start()
return {"run_id": run_id}
@app.get("/", response_class=HTMLResponse)
def home():
return """<!doctype html><title>My app</title>
<h1>My app</h1>
<button id="go">Generate article</button>
<iframe id="progress" title="Progress" style="width:100%;height:640px;border:1px solid #ccc"></iframe>
<script>
go.onclick = async () => {
const { run_id } = await (await fetch('/generate', { method: 'POST' })).json();
progress.src = '/replay/?run=' + run_id;
};
</script>"""
+40
View File
@@ -0,0 +1,40 @@
"""Watch a flow run, step by step: click Run and each step lights up on the map.
No API key needed (the steps are plain Python that take a moment each).
pip install "praisonai[ui]"
python watch_flow.py # opens http://127.0.0.1:8090/
Map: dashed = not reached, blue = running now, green = done.
After a run, press Play to go through it again one step at a time.
"""
import time
from praisonaiagents import AgentFlow
from praisonaiagents.workflows import StepResult
from praisonai.replay import watch
def research(ctx):
time.sleep(1.5)
return StepResult(output=f"3 sources about {ctx.input}")
def outline(ctx):
time.sleep(1.5)
return StepResult(output="intro, body, close")
def draft(ctx):
time.sleep(1.5)
return StepResult(output="first draft")
def review(ctx):
time.sleep(1.5)
return StepResult(output="approved")
flow = AgentFlow(name="Write an article", steps=[research, outline, draft, review])
watch(flow, input="solar sails")
+18
View File
@@ -0,0 +1,18 @@
"""Watch two agents work together, live (calls the LLM: needs OPENAI_API_KEY).
python watch_team.py # opens http://127.0.0.1:8090/, press Run
"""
from praisonaiagents import Agent, AgentTeam, Task
from praisonai.replay import watch
researcher = Agent(name="Researcher", instructions="Give 3 short facts.", llm="gpt-4o-mini")
writer = Agent(name="Writer", instructions="Write one short paragraph from the facts.", llm="gpt-4o-mini")
team = AgentTeam(
agents=[researcher, writer],
tasks=[
Task(name="research", description="3 facts about solar sails", agent=researcher, expected_output="3 facts"),
Task(name="write", description="One paragraph using the facts", agent=writer, expected_output="a paragraph"),
],
)
watch(team)
@@ -7091,6 +7091,9 @@ Answer:"""
def _process_handoffs(self):
"""Process handoffs and convert them to tools that can be used by the agent."""
if not self.handoffs and not (isinstance(self.tools, list) and self.tools):
return
# Import here to avoid circular imports
from .handoff import Handoff
@@ -2336,6 +2336,7 @@ Your Goal: {self.goal}"""
messages=messages,
response=str(final_response),
model=self.llm if isinstance(self.llm, str) else str(self.llm),
tokens_used=_prompt_tokens + _completion_tokens,
latency_ms=(time.time() - start_time) * 1000
)
self._hook_runner.execute_sync(HookEvent.AFTER_LLM, after_llm_input)
@@ -3037,6 +3038,7 @@ Your Goal: {self.goal}"""
messages=messages,
response=str(final_response),
model=self.llm if isinstance(self.llm, str) else str(self.llm),
tokens_used=_prompt_tokens + _completion_tokens,
latency_ms=(time.time() - _achat_start_time) * 1000
)
await self._hook_runner.execute(HookEvent.AFTER_LLM, after_llm_input)
@@ -4240,11 +4242,17 @@ Output MUST be JSON with 'reflection' and 'satisfactory'.
except Exception as e:
logger.warning(f"Steering check failed, continuing without steering: {e}")
# Input-side guardrail validation (parity with sync chat()). No-op
# unless the guardrail exposes validate_input; runs before any
# backend/LLM dispatch so a rejected prompt never reaches the model.
# Input-side validators may perform blocking LLM calls. Keep them off
# the event loop, before any backend/LLM dispatch, just like tool guards.
if hasattr(self, '_validate_input_with_guardrail'):
_in_ok, _in_prompt, _in_err = self._validate_input_with_guardrail(prompt)
_input_guard = self._guardrail_fn
if (_input_guard is None or not hasattr(_input_guard, "validate_input")
or getattr(_input_guard, "_praison_output_only", False)):
_in_ok, _in_prompt, _in_err = self._validate_input_with_guardrail(prompt)
else:
_in_ok, _in_prompt, _in_err = await asyncio.to_thread(
self._validate_input_with_guardrail, prompt
)
if not _in_ok:
logging.warning(f"Agent {self.name}: input blocked by guardrail: {_in_err}")
return None
@@ -400,8 +400,8 @@ This report contains all agent interactions and outputs from a complete ContextA
response_filename = f"{self.interaction_counter:03d}_{safe_agent_name}_{timestamp.replace(':', '_')}.json"
response_path = os.path.join(self.output_dir, "agent_responses", response_filename)
with open(response_path, "w") as f:
json.dump(interaction_data, f, indent=2)
with open(response_path, "w", encoding="utf-8") as f:
json.dump(interaction_data, f, indent=2, ensure_ascii=False)
# Save individual markdown files - debug only
md_filename = f"{self.interaction_counter:03d}_{safe_agent_name}_{timestamp.replace(':', '_')}"
@@ -1201,20 +1201,20 @@ detailed, actionable implementation steps."""
def _save_context_engineering_results(self):
"""Save all Context Engineering results following PRD methodology."""
# Save comprehensive analysis
with open(os.path.join(self.output_dir, "comprehensive_analysis.json"), "w") as f:
json.dump(self.codebase_analysis, f, indent=2)
with open(os.path.join(self.output_dir, "comprehensive_analysis.json"), "w", encoding="utf-8") as f:
json.dump(self.codebase_analysis, f, indent=2, ensure_ascii=False)
# Save pattern library
with open(os.path.join(self.output_dir, "pattern_library.json"), "w") as f:
json.dump(self.pattern_library, f, indent=2)
with open(os.path.join(self.output_dir, "pattern_library.json"), "w", encoding="utf-8") as f:
json.dump(self.pattern_library, f, indent=2, ensure_ascii=False)
# Save validation framework
with open(os.path.join(self.output_dir, "validation_framework.json"), "w") as f:
json.dump(self.validation_framework, f, indent=2)
with open(os.path.join(self.output_dir, "validation_framework.json"), "w", encoding="utf-8") as f:
json.dump(self.validation_framework, f, indent=2, ensure_ascii=False)
# Save all agent interactions summary
with open(os.path.join(self.output_dir, "all_agent_interactions.json"), "w") as f:
json.dump(self.agent_interactions, f, indent=2)
with open(os.path.join(self.output_dir, "all_agent_interactions.json"), "w", encoding="utf-8") as f:
json.dump(self.agent_interactions, f, indent=2, ensure_ascii=False)
print(f"💾 Context Engineering results saved to: {self.output_dir}/")
print(f"💾 Agent interactions ({len(self.agent_interactions)} total) saved to: agent_responses/")
@@ -1267,7 +1267,7 @@ Product Requirements Prompts following PRD template methodology.
Every agent interaction has been saved for full audit trail and reproducibility.
"""
with open(os.path.join(self.output_dir, "analysis_summary.md"), "w") as f:
with open(os.path.join(self.output_dir, "analysis_summary.md"), "w", encoding="utf-8") as f:
f.write(summary)
print(f"📋 Analysis summary saved following PRD methodology")
@@ -1308,7 +1308,7 @@ Every agent interaction has been saved for full audit trail and reproducibility.
safe_filename = feature_request.lower().replace(' ', '_').replace('-', '_')[:50]
prp_path = os.path.join(self.output_dir, "PRPs", f"{safe_filename}.md")
with open(prp_path, "w") as f:
with open(prp_path, "w", encoding="utf-8") as f:
f.write(prp)
print(f"📄 Comprehensive PRP following PRD template saved to: {prp_path}")
@@ -630,6 +630,27 @@ class Handoff:
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 @@ class Handoff:
# 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 @@ class Handoff:
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 @@ class Handoff:
# 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 @@ class TypedHandoff(Handoff, Generic[T]):
# 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 @@ class TypedHandoff(Handoff, Generic[T]):
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)
@@ -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 @@ class AgentTeam(SpawnAnnounceProtocol):
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
@@ -1684,10 +1729,21 @@ class AgentTeam(SpawnAnnounceProtocol):
except Exception as e:
logger.error(f"Error in on_task_start callback: {e}")
# Apply global variables to task if not already set. Use a shallow copy so a
# task mutating its variables doesn't leak back into the shared AgentTeam state.
self._propagate_variables(task, task_id)
def _propagate_variables(self, task, task_id):
"""Apply global variables to a task that has none of its own.
Uses a shallow copy so a task mutating its variables doesn't leak back
into the shared AgentTeam state. The copy is remembered so a later fresh
``start()`` can drop it (and re-propagate current values) without
touching caller-defined ``Task.variables``.
"""
if self.variables and not getattr(task, 'variables', None):
task.variables = self.variables.copy()
if not hasattr(self, "_propagated_variables"):
self._propagated_variables = {}
self._propagated_variables[task_id] = task.variables
def _run_task_complete_hook(self, task, task_output):
"""Run the on_task_complete hook. Shared by run_task and arun_task."""
@@ -1714,8 +1770,7 @@ class AgentTeam(SpawnAnnounceProtocol):
except Exception as e:
logger.error(f"Error in on_task_start callback: {e}")
if self.variables and not getattr(task, 'variables', None):
task.variables = self.variables.copy()
self._propagate_variables(task, task_id)
async def _arun_task_complete_hook(self, task, task_output):
"""Async-aware on_task_complete hook for arun_task (see _arun_task_start_hook)."""
@@ -1740,6 +1795,18 @@ class AgentTeam(SpawnAnnounceProtocol):
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 +1814,7 @@ class AgentTeam(SpawnAnnounceProtocol):
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)
@@ -1927,18 +1994,20 @@ class AgentTeam(SpawnAnnounceProtocol):
if self._depends_on_pending(task, tasks_to_run):
await self._gather_with_isolation([c for _, c in tasks_to_run])
tasks_to_run = []
# A just-flushed prerequisite may now be failed; don't
# queue a dependent that would run with missing context
# unless it opted into graceful degradation.
if self._deps_failed(task_id) and not self._should_continue_on_dep_failure(task_id):
self.tasks[task_id].status = "failed"
continue
# The predecessor may have run synchronously (no batch),
# or may just have finished in the batch flushed above.
if self._deps_failed(task_id) and not self._should_continue_on_dep_failure(task_id):
task.status = "failed"
continue
tasks_to_run.append((task_id, self.arun_task(task_id)))
else:
# If we encounter a sync task, we must wait for the previous async tasks to finish.
if tasks_to_run:
await self._gather_with_isolation([c for _, c in tasks_to_run])
tasks_to_run = []
if self._deps_failed(task_id) and not self._should_continue_on_dep_failure(task_id):
task.status = "failed"
continue
# Run sync task in an executor to avoid blocking the event loop
# Use copy_context_to_callable to propagate contextvars (needed for trace emission)
@@ -2201,7 +2270,16 @@ class AgentTeam(SpawnAnnounceProtocol):
# 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)
self._reset_if_previous_run_finished()
try:
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
finally:
self._previous_run_finished = True
async def _astart_impl(self, content=None, return_dict=False, **kwargs):
# Track execution via telemetry
@@ -2215,6 +2293,7 @@ class AgentTeam(SpawnAnnounceProtocol):
if not task.context:
task.context = []
task.context.append(content)
self._injected_content.append((task, content))
if self.planning:
await self._arun_with_planning()
@@ -2253,13 +2332,11 @@ class AgentTeam(SpawnAnnounceProtocol):
def save_output_to_file(self, task, task_output):
if task.output_file:
try:
if task.create_directory:
dir_name = os.path.dirname(task.output_file)
if dir_name:
os.makedirs(dir_name, exist_ok=True)
with open(task.output_file, "w") as f:
f.write(str(task_output))
if self.verbose >= 1:
from ..utils.output_file import write_output_file
saved = write_output_file(
task.output_file, str(task_output), create_dirs=task.create_directory
)
if saved and self.verbose >= 1:
logger.info(f"Task output saved to {task.output_file}")
except Exception as e:
display_error(f"Error saving task output to file: {e}")
@@ -2318,6 +2395,18 @@ class AgentTeam(SpawnAnnounceProtocol):
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 +2414,7 @@ class AgentTeam(SpawnAnnounceProtocol):
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)
@@ -2417,6 +2506,9 @@ class AgentTeam(SpawnAnnounceProtocol):
if self.process == "workflow":
for task_id in process.workflow():
if self._deps_failed(task_id) and not self._should_continue_on_dep_failure(task_id):
self.tasks[task_id].status = "failed"
continue
self.run_task(task_id)
elif self.process == "sequential":
for task_id in process.sequential():
@@ -2441,6 +2533,21 @@ class AgentTeam(SpawnAnnounceProtocol):
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 +2626,31 @@ class AgentTeam(SpawnAnnounceProtocol):
# 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)
self._reset_if_previous_run_finished()
try:
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
finally:
self._previous_run_finished = True
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
@@ -2551,6 +2681,7 @@ class AgentTeam(SpawnAnnounceProtocol):
if not task.context:
task.context = []
task.context.append(content)
self._injected_content.append((task, content))
# ─────────────────────────────────────────────────────────────
# Verbose Mode: Show Rich panels for multi-agent workflow
@@ -2746,14 +2877,19 @@ class AgentTeam(SpawnAnnounceProtocol):
Captures the fields that ``run_task``/``arun_task`` mutate during a run
(``status``, ``result``, ``retry_count``) plus caller-defined
``variables``, keeping the public Task API backward-compatible.
``propagated`` records whether those variables were copied from the
team's ``variables`` so a restore keeps that ownership tracking.
"""
snapshot = {}
propagated = getattr(self, "_propagated_variables", {})
for task_id, task in getattr(self, "tasks", {}).items():
variables = getattr(task, "variables", None)
snapshot[task_id] = {
"status": getattr(task, "status", None),
"result": getattr(task, "result", None),
"retry_count": getattr(task, "retry_count", None),
"variables": getattr(task, "variables", None),
"variables": variables,
"propagated": variables is not None and propagated.get(task_id) is variables,
}
return snapshot
@@ -2772,8 +2908,41 @@ class AgentTeam(SpawnAnnounceProtocol):
task.retry_count = 0
task.variables = {}
def _reset_if_previous_run_finished(self):
"""Make a repeated ``start()``/``astart()`` a fresh run.
``run_task``/``arun_task`` skip tasks whose status is "completed", so
without this a second ``start()`` on the same team executed nothing and
returned the previous run's output. Resetting only once a previous run
has finished keeps ``restore_session_state()`` + ``start()`` resuming,
and also drops the ``content`` earlier runs appended to task context.
Caller-defined ``Task.variables`` are kept; only copies propagated from
the team's ``variables`` are cleared so current values re-propagate.
"""
injected = getattr(self, "_injected_content", [])
if getattr(self, "_previous_run_finished", False):
propagated = getattr(self, "_propagated_variables", {})
for task_id, task in self.tasks.items():
task.status = "not started"
task.result = None
if hasattr(task, "retry_count"):
task.retry_count = 0
if task_id in propagated and getattr(task, "variables", None) is propagated[task_id]:
task.variables = {}
propagated.clear()
for task, item in injected:
try:
task.context.remove(item)
except (AttributeError, ValueError):
pass
injected = []
# A resumed run keeps the tracked entries so the next fresh run removes them.
self._previous_run_finished = False
self._injected_content = injected
def _restore_task_state(self, snapshot):
"""Restore per-task run state captured by ``_snapshot_task_state``."""
propagated = {}
for task_id, state in snapshot.items():
task = self.tasks.get(task_id)
if task is None:
@@ -2783,6 +2952,9 @@ class AgentTeam(SpawnAnnounceProtocol):
if state["retry_count"] is not None:
task.retry_count = state["retry_count"]
task.variables = state["variables"]
if state.get("propagated"):
propagated[task_id] = task.variables
self._propagated_variables = propagated
def _build_batch_result(self, batch_id, items):
"""Aggregate per-item results into a batch summary dict.
@@ -3131,6 +3303,9 @@ class AgentTeam(SpawnAnnounceProtocol):
if state.get("variables"):
task.variables = state["variables"]
restored += 1
if restored:
# A restored checkpoint is meant to be resumed, not reset.
self._previous_run_finished = False
return restored
@staticmethod
@@ -4461,7 +4636,10 @@ class AgentTeam(SpawnAnnounceProtocol):
# Read target agents inside lock to avoid race condition
with self._spawn_lock:
target_agents = agent_ids or list(self._spawned_agents.keys())
if not target_agents:
# Return retained history only for the default target selection.
return list(self._completion_events) if agent_ids is None else []
completed_agents = set()
while True:
@@ -689,29 +689,41 @@ class ApprovalRegistry:
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
)
# 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:
return stale
# 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
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 @@ class ApprovalRegistry:
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")
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
# 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
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]:
@@ -1,147 +0,0 @@
"""
Utility functions for PraisonAI Agents evaluation framework.
"""
import json
from praisonaiagents._logging import get_logger
from pathlib import Path
from typing import Any, Dict, Optional, Union
logger = get_logger(__name__)
def save_result_to_file(
result: Any,
file_path: str,
eval_id: Optional[str] = None,
name: Optional[str] = None
) -> bool:
"""
Save evaluation result to a file.
Args:
result: The evaluation result to save
file_path: Path to save the result (supports {name}, {eval_id} placeholders)
eval_id: Optional eval ID for placeholder substitution
name: Optional name for placeholder substitution
Returns:
True if save was successful, False otherwise
"""
try:
formatted_path = file_path.format(name=name or "", eval_id=eval_id or "")
path = Path(formatted_path)
path.parent.mkdir(parents=True, exist_ok=True)
if hasattr(result, 'to_dict'):
data = result.to_dict()
elif hasattr(result, '__dict__'):
data = result.__dict__
else:
data = {"result": str(result)}
path.write_text(json.dumps(data, indent=2, default=str))
logger.debug(f"Saved result to: {formatted_path}")
return True
except Exception as e:
logger.warning(f"Failed to save result: {e}")
return False
def load_result_from_file(file_path: str) -> Optional[Dict[str, Any]]:
"""
Load evaluation result from a file.
Args:
file_path: Path to the result file
Returns:
Dictionary containing the result data, or None if load failed
"""
try:
path = Path(file_path)
if not path.exists():
logger.warning(f"Result file not found: {file_path}")
return None
data = json.loads(path.read_text())
return data
except Exception as e:
logger.warning(f"Failed to load result: {e}")
return None
def format_score(score: float, max_score: float = 10.0) -> str:
"""
Format a score for display.
Args:
score: The score value
max_score: Maximum possible score
Returns:
Formatted score string
"""
return f"{score:.1f}/{max_score:.0f}"
def format_percentage(value: float) -> str:
"""
Format a value as a percentage.
Args:
value: Value between 0 and 1
Returns:
Formatted percentage string
"""
return f"{value * 100:.1f}%"
def format_duration(seconds: float) -> str:
"""
Format a duration in seconds for display.
Args:
seconds: Duration in seconds
Returns:
Formatted duration string
"""
if seconds < 0.001:
return f"{seconds * 1000000:.0f}μs"
elif seconds < 1:
return f"{seconds * 1000:.2f}ms"
elif seconds < 60:
return f"{seconds:.2f}s"
else:
minutes = int(seconds // 60)
secs = seconds % 60
return f"{minutes}m {secs:.1f}s"
def format_memory(mb: float) -> str:
"""
Format memory usage for display.
Args:
mb: Memory in megabytes
Returns:
Formatted memory string
"""
if mb < 1:
return f"{mb * 1024:.1f} KB"
elif mb < 1024:
return f"{mb:.2f} MB"
else:
return f"{mb / 1024:.2f} GB"
def calculate_pass_rate(passed: int, total: int) -> float:
"""
Calculate pass rate.
Args:
passed: Number of passed evaluations
total: Total number of evaluations
Returns:
Pass rate as a float between 0 and 1
"""
if total == 0:
return 0.0
return passed / total
@@ -281,6 +281,8 @@ class SubQuestionEngine:
# Retrieve for each sub-question
try:
results = retriever.retrieve(sub_q, top_k=3)
# SmartRetriever returns a RetrievalResult, not a list.
results = getattr(results, "chunks", results)
for r in results:
if hasattr(r, 'text'):
all_contexts.append(r.text)
+267 -22
View File
@@ -43,6 +43,165 @@ _TOOL_ARGUMENTS_PARSE_FAILED = object()
# 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 @@ Respond with ONLY a valid JSON tool call in this format:
# 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 @@ Respond with ONLY a valid JSON tool call in this format:
"""
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:
@@ -1837,10 +2016,10 @@ Respond with ONLY a valid JSON tool call in this format:
return None
if self._claude_memory_tool is None:
from ..tools.claude_memory_tool import ClaudeMemoryTool
if isinstance(self.claude_memory, bool):
# Create default memory tool
# Only the default implementation needs this dependency.
from ..tools.claude_memory_tool import ClaudeMemoryTool
self._claude_memory_tool = ClaudeMemoryTool()
else:
# Use provided memory tool instance
@@ -7636,6 +7815,32 @@ Output MUST be JSON with 'reflection' and 'satisfactory'.
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 @@ Output MUST be JSON with 'reflection' and 'satisfactory'.
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 @@ Output MUST be JSON with 'reflection' and 'satisfactory'.
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 @@ Output MUST be JSON with 'reflection' and 'satisfactory'.
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,
@@ -13,7 +13,7 @@ Each store handles a specific type of learning data:
from abc import ABC, abstractmethod
from dataclasses import dataclass, field
from datetime import datetime
from datetime import datetime, timezone
import json
from pathlib import Path
from typing import Any, Dict, List, Optional, TYPE_CHECKING
@@ -31,8 +31,8 @@ class LearnEntry:
id: str
content: str
metadata: Dict[str, Any] = field(default_factory=dict)
created_at: str = field(default_factory=lambda: datetime.utcnow().isoformat())
updated_at: str = field(default_factory=lambda: datetime.utcnow().isoformat())
created_at: str = field(default_factory=lambda: datetime.now(timezone.utc).isoformat())
updated_at: str = field(default_factory=lambda: datetime.now(timezone.utc).isoformat())
use_count: int = 0
last_used: Optional[str] = None
@@ -53,8 +53,8 @@ class LearnEntry:
id=data["id"],
content=data["content"],
metadata=data.get("metadata", {}),
created_at=data.get("created_at", datetime.utcnow().isoformat()),
updated_at=data.get("updated_at", datetime.utcnow().isoformat()),
created_at=data.get("created_at", datetime.now(timezone.utc).isoformat()),
updated_at=data.get("updated_at", datetime.now(timezone.utc).isoformat()),
use_count=data.get("use_count", 0),
last_used=data.get("last_used"),
)
@@ -154,13 +154,13 @@ class BaseStore(ABC):
# Update metadata if provided, but don't create duplicate
if metadata:
existing_entry.metadata.update(metadata)
existing_entry.updated_at = datetime.utcnow().isoformat()
existing_entry.updated_at = datetime.now(timezone.utc).isoformat()
self._save()
self._was_updated = True
return existing_entry
# No duplicate found, create new entry
entry_id = f"{self.store_name}_{len(self._entries)}_{datetime.utcnow().timestamp()}"
entry_id = f"{self.store_name}_{len(self._entries)}_{datetime.now(timezone.utc).timestamp()}"
entry = LearnEntry(
id=entry_id,
content=content,
@@ -183,7 +183,7 @@ class BaseStore(ABC):
"""
if not entries:
return
now = datetime.utcnow().isoformat()
now = datetime.now(timezone.utc).isoformat()
for entry in entries:
entry.use_count += 1
entry.last_used = now
@@ -229,7 +229,7 @@ class BaseStore(ABC):
"""List all entries."""
refreshed = self._refresh_for_retrieval()
entries = list(self._entries.values())
entries.sort(key=lambda x: x.updated_at, reverse=True)
entries.sort(key=lambda x: self._timestamp_sort_key(x.updated_at), reverse=True)
entries = entries[:limit]
if refreshed:
self._touch(entries)
@@ -244,7 +244,7 @@ class BaseStore(ABC):
entry.content = content
if metadata:
entry.metadata.update(metadata)
entry.updated_at = datetime.utcnow().isoformat()
entry.updated_at = datetime.now(timezone.utc).isoformat()
self._save()
self._was_updated = True
return entry
@@ -308,9 +308,29 @@ class BaseStore(ABC):
ref_dt = datetime.fromisoformat(reference)
except (ValueError, TypeError):
return False
if ref_dt.tzinfo is None:
ref_dt = ref_dt.replace(tzinfo=timezone.utc)
age_days = (now - ref_dt).total_seconds() / 86400.0
return age_days > self.retention_days
@staticmethod
def _timestamp_sort_key(value: str) -> tuple:
"""Comparable key that orders ISO timestamps chronologically in UTC.
Naive values are treated as UTC (matching how older stores were written).
Missing or unparseable values sort as the oldest instant, so prune()
evicts them first (an entry never used sorts before used ones, as it did
under the previous raw-string sort) and list_all() lists them last. A
datetime is never compared against a string, so mixed data never raises.
"""
try:
parsed = datetime.fromisoformat(value)
except (ValueError, TypeError):
return (-1, value)
if parsed.tzinfo is None:
parsed = parsed.replace(tzinfo=timezone.utc)
return (0, parsed.astimezone(timezone.utc))
def prune(self, save: bool = True, protect_id: Optional[str] = None) -> int:
"""Archive stale and least-used/oldest excess entries.
@@ -331,7 +351,7 @@ class BaseStore(ABC):
if self.max_entries <= 0 and self.retention_days <= 0:
return 0
now = datetime.utcnow()
now = datetime.now(timezone.utc)
evicted: List[LearnEntry] = []
# 1. Staleness: archive entries beyond the retention window
@@ -349,7 +369,11 @@ class BaseStore(ABC):
]
remaining = sorted(
candidates,
key=lambda e: (e.use_count, e.last_used or "", e.created_at),
key=lambda e: (
e.use_count,
self._timestamp_sort_key(e.last_used or ""),
self._timestamp_sort_key(e.created_at),
),
)
overflow = len(self._entries) - self.max_entries
for entry in remaining[:overflow]:
@@ -5,6 +5,7 @@ Detects and prevents agents from getting stuck in repetitive loops.
"""
from praisonaiagents._logging import get_logger
import threading
import time
from collections import defaultdict
from dataclasses import dataclass, field
@@ -93,8 +94,13 @@ class DoomLoopDetector:
self.max_tool_calls = max_tool_calls
self._records: List[ToolCallRecord] = []
# Lifetime per-tool totals, reported by get_stats() only. Loop
# decisions use the windowed _records so a slow, healthy tool is
# never permanently denied.
self._call_counts: Dict[str, int] = defaultdict(int)
self._session_id: Optional[str] = None
# check_doom_loop is reached from concurrent tool execution.
self._lock = threading.RLock()
def set_session(self, session_id: str):
"""Set the current session ID."""
@@ -125,19 +131,18 @@ class DoomLoopDetector:
Returns:
The recorded ToolCallRecord
"""
self._cleanup_old_records()
args_hash = self._hash_arguments(arguments or {})
record = ToolCallRecord(
tool_name=tool_name,
arguments_hash=args_hash,
session_id=self._session_id,
)
self._records.append(record)
self._call_counts[tool_name] += 1
with self._lock:
self._cleanup_old_records()
record = ToolCallRecord(
tool_name=tool_name,
arguments_hash=args_hash,
session_id=self._session_id,
)
self._records.append(record)
self._call_counts[tool_name] += 1
return record
def check(
@@ -155,13 +160,15 @@ class DoomLoopDetector:
Returns:
DoomLoopResult indicating if a loop is detected
"""
self._cleanup_old_records()
args_hash = self._hash_arguments(arguments or {})
with self._lock:
self._cleanup_old_records()
records = list(self._records)
# Count identical calls in window
identical_count = sum(
1 for r in self._records
1 for r in records
if r.tool_name == tool_name and r.arguments_hash == args_hash
)
@@ -176,7 +183,7 @@ class DoomLoopDetector:
)
# Check for excessive tool calls
total_calls = len(self._records)
total_calls = len(records)
if total_calls >= self.max_tool_calls:
return DoomLoopResult(
is_loop=True,
@@ -185,8 +192,8 @@ class DoomLoopDetector:
recommendation="Reduce tool usage or break task into smaller steps",
)
# Check for tool-specific excessive calls
tool_count = self._call_counts.get(tool_name, 0)
# Check for tool-specific excessive calls (within the window)
tool_count = sum(1 for r in records if r.tool_name == tool_name)
if tool_count >= self.max_tool_calls // 2:
return DoomLoopResult(
is_loop=True,
@@ -219,29 +226,32 @@ class DoomLoopDetector:
Returns:
DoomLoopResult indicating if a loop is detected
"""
result = self.check(tool_name, arguments)
if not result.is_loop:
self.record(tool_name, arguments)
with self._lock:
result = self.check(tool_name, arguments)
if not result.is_loop:
self.record(tool_name, arguments)
return result
def get_stats(self) -> Dict[str, Any]:
"""Get statistics about tool calls."""
self._cleanup_old_records()
return {
"total_calls": len(self._records),
"unique_tools": len(set(r.tool_name for r in self._records)),
"call_counts": dict(self._call_counts),
"window_seconds": self.window_seconds,
"loop_threshold": self.loop_threshold,
}
with self._lock:
self._cleanup_old_records()
return {
"total_calls": len(self._records),
"unique_tools": len(set(r.tool_name for r in self._records)),
"call_counts": dict(self._call_counts),
"window_seconds": self.window_seconds,
"loop_threshold": self.loop_threshold,
}
def reset(self):
"""Reset all records."""
self._records.clear()
self._call_counts.clear()
with self._lock:
self._records.clear()
self._call_counts.clear()
def reset_tool(self, tool_name: str):
"""Reset records for a specific tool."""
self._records = [r for r in self._records if r.tool_name != tool_name]
self._call_counts.pop(tool_name, None)
with self._lock:
self._records = [r for r in self._records if r.tool_name != tool_name]
self._call_counts.pop(tool_name, None)
@@ -80,6 +80,31 @@ class Process:
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 @@ Provide a JSON with the structure:
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 @@ Provide a JSON with the structure:
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
@@ -154,6 +154,7 @@ class RAG:
user_id: Optional[str] = None,
agent_id: Optional[str] = None,
run_id: Optional[str] = None,
top_k: Optional[int] = None,
**kwargs,
) -> List[Dict[str, Any]]:
"""
@@ -193,8 +194,8 @@ class RAG:
filtered.append(r)
results = filtered
# Limit to top_k
results = results[:self.config.top_k]
# Limit to top_k (per-call value, never written back to shared config)
results = results[:self.config.top_k if top_k is None else top_k]
return results
@@ -202,6 +203,7 @@ class RAG:
self,
query: str,
results: List[Dict[str, Any]],
rerank: Optional[bool] = None,
) -> List[Dict[str, Any]]:
"""
Rerank results if reranker is configured.
@@ -213,7 +215,9 @@ class RAG:
Returns:
Reranked results
"""
if not self.config.rerank or not self.reranker:
if rerank is None:
rerank = self.config.rerank
if not rerank or not self.reranker:
return results
try:
@@ -354,54 +358,47 @@ class RAG:
"""
start_time = time.time()
# Apply overrides
original_top_k = self.config.top_k
original_rerank = self.config.rerank
# Resolve per-call overrides without mutating the shared config:
# aretrieve() runs this in an executor, so concurrent calls on one RAG
# would otherwise see (and permanently "restore") each other's values.
eff_top_k = self.config.top_k if top_k is None else top_k
eff_rerank = self.config.rerank if rerank is None else rerank
if top_k is not None:
self.config.top_k = top_k
if rerank is not None:
self.config.rerank = rerank
# 1. Retrieve
results = self._retrieve(
query=query,
user_id=user_id,
agent_id=agent_id,
run_id=run_id,
filters=filters,
top_k=eff_top_k,
**kwargs,
)
try:
# 1. Retrieve
results = self._retrieve(
query=query,
user_id=user_id,
agent_id=agent_id,
run_id=run_id,
filters=filters,
**kwargs,
)
# 2. Rerank (optional)
results = self._rerank(query, results)
# 3. Build context
context, used_results = self._build_context(results)
# 4. Format citations
citations = self._format_citations(used_results)
elapsed = time.time() - start_time
return ContextPack(
context=context,
citations=citations,
query=query,
metadata={
"elapsed_seconds": elapsed,
"num_results": len(results),
"num_citations": len(citations),
"retrieval_strategy": self.config.retrieval_strategy.value,
"reranked": self.config.rerank,
"top_k": self.config.top_k,
},
)
finally:
# Restore original config
self.config.top_k = original_top_k
self.config.rerank = original_rerank
# 2. Rerank (optional)
results = self._rerank(query, results, rerank=eff_rerank)
# 3. Build context
context, used_results = self._build_context(results)
# 4. Format citations
citations = self._format_citations(used_results)
elapsed = time.time() - start_time
return ContextPack(
context=context,
citations=citations,
query=query,
metadata={
"elapsed_seconds": elapsed,
"num_results": len(results),
"num_citations": len(citations),
"retrieval_strategy": self.config.retrieval_strategy.value,
"reranked": eff_rerank,
"top_k": eff_top_k,
},
)
async def aretrieve(
self,
@@ -131,26 +131,44 @@ class SmartRetriever:
if not self._knowledge:
return result
has_filters = bool(include_glob or exclude_glob or path_filter or metadata_filters)
if path_filter:
import re
try:
re.compile(path_filter)
except re.error as e:
# Fail closed: silently dropping the filter would return every
# chunk while still reporting the result as filtered.
raise ValueError(f"Invalid path_filter regex {path_filter!r}: {e}") from e
# Filters are applied after the store truncates, so over-fetch when any
# is set (same rule as Memory's metadata_filter search) -- otherwise a
# filter whose matches sit just below the top_k cut returns nothing.
fetch_k = top_k * 3 if strategy in ("reranked", "hybrid") else top_k
if has_filters:
fetch_k = max(fetch_k, top_k * 10)
# Get initial results from knowledge store
search_results = self._search(
search_results, error = self._search(
query=query,
top_k=top_k * 3 if strategy in ("reranked", "hybrid") else top_k,
top_k=fetch_k,
user_id=user_id,
agent_id=agent_id,
)
# Surface a store failure on the result. Without this the caller sees an
# empty chunk list and cannot tell whether the knowledge base had no
# match or was never reached.
if getattr(self, "_last_error", None):
# match or was never reached. The error is returned rather than stored
# on self so concurrent calls on one retriever cannot see each other's.
if error:
result.retrieval_failed = True
result.error = self._last_error
result.metadata["retrieval_error"] = self._last_error
result.error = error
result.metadata["retrieval_error"] = error
chunks = self._normalize_results(search_results)
result.total_found = len(chunks)
# Apply filters
if include_glob or exclude_glob or path_filter or metadata_filters:
if has_filters:
chunks = self._apply_filters(
chunks,
include_glob=include_glob,
@@ -176,23 +194,22 @@ class SmartRetriever:
top_k: int,
user_id: Optional[str] = None,
agent_id: Optional[str] = None,
) -> Any:
"""Perform search on knowledge store."""
self._last_error = None
) -> tuple:
"""Perform search on knowledge store. Returns ``(results, error)``."""
try:
return self._knowledge.search(
query,
user_id=user_id,
agent_id=agent_id,
limit=top_k,
)
), None
except Exception as e:
# Record the failure rather than returning a bare [] that is
# indistinguishable from "nothing matched". Knowledge.search itself
# re-raises, so reaching here means the store really is broken --
# misconfigured, unreachable, or refusing auth -- and the agent is
# about to answer with no context at all.
self._last_error = f"{type(e).__name__}: {e}"
error = f"{type(e).__name__}: {e}"
# Do not log the raw query: it may carry personal data, credentials
# or proprietary text, and this line lands at ERROR in application
# logs. The exception text is enough to diagnose the store; the
@@ -200,9 +217,9 @@ class SmartRetriever:
logger.error(
"Knowledge retrieval failed (query length %d): %s. "
"Returning no chunks; the answer will not use the knowledge base.",
len(query or ""), self._last_error,
len(query or ""), error,
)
return []
return [], error
def _normalize_results(self, results: Any) -> List[Dict[str, Any]]:
"""Normalize search results to list of dicts."""
@@ -275,12 +292,8 @@ class SmartRetriever:
continue
# Check path filter regex
if path_filter:
try:
if not re.search(path_filter, filepath):
continue
except re.error:
pass
if path_filter and not re.search(path_filter, filepath):
continue
# Check metadata filters
if metadata_filters:
@@ -29,6 +29,10 @@ DEFAULT_MAX_LINES = 2000
# cannot dominate the window.
DEFAULT_MAX_LINE_CHARS = 2000
# Default byte cap for ``download_file`` so an endless or oversized response
# cannot fill the disk. Override with PRAISONAI_MAX_DOWNLOAD_BYTES.
_DEFAULT_MAX_DOWNLOAD_BYTES = 100 * 1024 * 1024
def _render_with_line_numbers(lines, start_index, max_line_chars):
"""Render ``lines`` with a right-aligned 1-based line-number gutter.
@@ -485,58 +489,63 @@ class FileTools:
# Validate destination path
safe_path = self._validate_path(destination)
# Validate URL to prevent SSRF
# Validate URL to prevent SSRF (shared guard used by all HTTP tools)
from urllib.parse import urlparse
import ipaddress
import socket
from .url_safety import is_safe_http_url
parsed = urlparse(url)
if parsed.scheme not in ("http", "https"):
return {
"success": False, "path": "", "size": 0,
"error": f"Unsupported URL scheme: {parsed.scheme}. Only http/https allowed."
}
hostname = parsed.hostname or ""
# Resolve hostname and check for private/reserved IPs
try:
resolved = socket.getaddrinfo(hostname, parsed.port or 443)
for family, _, _, _, sockaddr in resolved:
addr = ipaddress.ip_address(sockaddr[0])
if addr.is_private or addr.is_loopback or addr.is_reserved or addr.is_link_local:
return {
"success": False, "path": "", "size": 0,
"error": f"Access to private/internal addresses is blocked: {hostname}"
}
except (socket.gaierror, ValueError):
pass # Let httpx handle DNS failures naturally
# Block cloud metadata endpoints
_blocked_hosts = {"169.254.169.254", "metadata.google.internal", "metadata.google.com"}
if hostname in _blocked_hosts:
# allow_local=False: download_file never had a local-access opt-in,
# so ALLOW_LOCAL_CRAWL must not bypass this check.
if not is_safe_http_url(url, allow_local=False):
return {
"success": False, "path": "", "size": 0,
"error": f"Access to metadata service is blocked: {hostname}"
"error": f"Access to private/internal addresses is blocked: {parsed.hostname or ''}"
}
max_bytes = int(os.environ.get("PRAISONAI_MAX_DOWNLOAD_BYTES", _DEFAULT_MAX_DOWNLOAD_BYTES))
# Create directory if needed
os.makedirs(os.path.dirname(safe_path) or ".", exist_ok=True)
# Download file
with httpx.stream("GET", url, timeout=timeout, follow_redirects=False) as response:
response.raise_for_status()
total_size = int(response.headers.get("content-length", 0))
downloaded = 0
with open(safe_path, "wb") as f:
dest_dir = os.path.dirname(safe_path) or "."
os.makedirs(dest_dir, exist_ok=True)
# Download into a temp file next to the destination and only
# replace the destination on success, so a failed or oversized
# download never clobbers an existing file.
import tempfile
fd, tmp_path = tempfile.mkstemp(dir=dest_dir, suffix=".part")
try:
with os.fdopen(fd, "wb") as f, \
httpx.stream("GET", url, timeout=timeout, follow_redirects=False) as response:
response.raise_for_status()
total_size = int(response.headers.get("content-length", 0))
downloaded = 0
for chunk in response.iter_bytes(chunk_size=8192):
f.write(chunk)
downloaded += len(chunk)
if downloaded > max_bytes:
break
f.write(chunk)
if progress_callback and total_size:
progress_callback(downloaded, total_size)
if downloaded > max_bytes:
return {
"success": False, "path": "", "size": 0,
"error": f"Download exceeds {max_bytes} bytes"
}
os.replace(tmp_path, safe_path)
finally:
if os.path.exists(tmp_path):
os.unlink(tmp_path)
return {
"success": True,
"path": safe_path,
@@ -89,7 +89,10 @@ def is_safe_http_url(
continue
return False
if (
ip.is_private
# is_global also catches shared/CGNAT space (100.64.0.0/10,
# e.g. Alibaba Cloud metadata), which is_private misses.
not ip.is_global
or ip.is_private
or ip.is_link_local
or ip.is_multicast
or ip.is_unspecified
@@ -9,10 +9,12 @@ Zero Performance Impact:
- 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 @@ from typing import Any, Dict, Generator, List, Optional, Protocol, runtime_check
import json
CONTEXT_SCHEMA_VERSION = "1.0"
CONTEXT_SCHEMA_VERSION = "1.1"
# =============================================================================
@@ -37,6 +39,28 @@ _context_emitter: contextvars.ContextVar["ContextTraceEmitter"] = contextvars.Co
# 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 @@ class ContextEvent:
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 @@ class 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 @@ class 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 @@ class ContextListSink:
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 @@ class ContextTraceEmitter:
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
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) -> None:
"""Internal emit with enabled check and sequence assignment.
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 @@ class ContextTraceEmitter:
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 @@ class ContextTraceEmitter:
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 @@ class ContextTraceEmitter:
"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 @@ class ContextTraceEmitter:
"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 @@ class ContextTraceEmitter:
"""
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 @@ class ContextTraceEmitter:
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 @@ class ContextTraceEmitter:
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 @@ class ContextTraceEmitter:
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 @@ class ContextTraceEmitter:
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 @@ class ContextTraceEmitter:
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 @@ class ContextTraceEmitter:
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,217 @@ class ContextTraceEmitter:
"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:<name or index>``)."""
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,
variables: Optional[Dict[str, Any]] = None,
) -> None:
"""A flow step ends: ``completed``, ``failed``, ``skipped`` or ``cached``.
``variables`` are the workflow variables the step itself set (a handler's
``StepResult.variables``); recorded redacted so a re-run that reuses the
step can restore them.
"""
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
if variables:
data["variables"] = self._redact_dict(dict(variables))
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,
@@ -0,0 +1,62 @@
import os
import stat
import tempfile
from praisonaiagents._logging import get_logger
from ..tools.path_safety import resolve_within_root
logger = get_logger(__name__)
def _open_nofollow(path, flags, mode=0o666):
return os.open(path, flags | getattr(os, "O_NOFOLLOW", 0), mode)
def write_output_file(path: str, content: str, create_dirs: bool = True) -> str | None:
"""Write *content* to *path* if it resolves inside the project root.
The root is ``PRAISONAI_PROJECT_ROOT`` or the cwd, matching ``Agent.output_file``.
Returns the resolved path, or ``None`` when the path escapes the root.
"""
project_root = os.environ.get("PRAISONAI_PROJECT_ROOT") or os.getcwd()
file_path = resolve_within_root(path, project_root)
if file_path is None:
logger.warning(
"Output file %r is outside project root %r; skipping save",
path,
project_root,
)
return None
if create_dirs:
os.makedirs(os.path.dirname(file_path), exist_ok=True)
with open(file_path, "a", encoding="utf-8", opener=_open_nofollow) as f:
st = os.fstat(f.fileno())
if st.st_nlink == 1:
f.truncate(0)
f.write(content)
return file_path
# A hardlink may share its inode with a file outside the root, so swap in
# a new file under the same name instead of truncating the shared one.
fd, tmp = tempfile.mkstemp(
dir=os.path.dirname(file_path),
prefix=os.path.basename(file_path) + ".",
suffix=".tmp",
)
try:
mode = stat.S_IMODE(st.st_mode)
if hasattr(os, "fchown"):
try:
os.fchown(fd, -1, st.st_gid)
except OSError:
# Can't keep the original group; don't grant its bits to ours.
mode &= ~0o070
if hasattr(os, "fchmod"):
os.fchmod(fd, mode)
with os.fdopen(fd, "w", encoding="utf-8") as f:
f.write(content)
os.replace(tmp, file_path)
except BaseException:
os.unlink(tmp)
raise
return file_path
@@ -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:<name>``, ``task:<name or index>``, ``step:<name>``
(duplicates become ``#2``, ``#3`` in definition order),
``tool:<agent>/<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 @@ GitHub and the docs render it inline.
...
"""
from typing import Any, List, Optional, Tuple
from typing import Any, Dict, List, Optional, Tuple
__all__ = ["to_mermaid", "flow_to_mermaid"]
__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"
_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 node(self, shape: str) -> str:
node_id = f"n{self._n}"
self._n += 1
self.lines.append(f" {node_id}{shape}")
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,
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 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 _branch(steps: Any, b: "_Builder") -> Optional[Tuple[str, str]]:
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:<name>`` 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 _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 = _flow_step(step, g, prev, None)
if rendered:
prev = rendered[1]
end = g.node(END_ID, "end", "Delivered", shape='([" end "])')
if prev:
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:<name>`` (or ``task:<index>`` 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``."""
b = _Builder()
start = b.node('([" start "])' if not name else f'(["{_escape(name)}"])')
prev: Optional[str] = start
for step in (steps or []):
rendered = _render(step, b, prev)
if rendered:
prev = rendered[1]
end = b.node('([" end "])')
if prev:
b.edge(prev, end)
return "\n".join(b.lines)
return _build_flow(steps, name).ops_to_mermaid()
def flow_to_mermaid(flow: Any) -> str:
@@ -725,6 +725,115 @@ 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 _differs(a: Any, b: Any) -> bool:
"""``a != b`` that treats a failing or ambiguous comparison as "changed"."""
try:
return bool(a != b)
except Exception:
return True
def _snapshot_scope(span, variables):
"""Before a recorded leaf step really runs (not skipped or cached), keep a
per-value deep copy of the scope on its span, so a list or dict the step
changes in place still counts as set by it; only an uncopyable value is
shared by reference."""
if span is not None:
span["before"] = _isolate_scope(variables)
def _replay_step_end(span, status, output=None, variables=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"],
variables=variables,
)
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 +1573,26 @@ class AgentFlow:
"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 +1614,15 @@ class AgentFlow:
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 +1726,35 @@ class AgentFlow:
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,
variables: Optional[Dict[str, 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, variables)
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 +1975,7 @@ class AgentFlow:
continue
# Normalize single step
raw_step = step
step = self._normalize_single_step(step, i)
# Create context for handlers
@@ -1839,11 +1994,7 @@ class AgentFlow:
# 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 +2005,9 @@ class AgentFlow:
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 +2053,16 @@ class AgentFlow:
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}")
# The cache also holds the step's output variable (its full
# output): trace only what the handler set, like a cold run.
_out_var = step.output_variable or f"{step.name}_output"
self._notify_step_complete(
_replay, step.name, StepResult(
output=previous_output or "",
stop_workflow=bool(_cached.get("stop")),
), "cached", previous_output,
variables={k: v for k, v in (_cached.get("variables") or {}).items() if k != _out_var},
)
if _cached.get("stop"):
if verbose:
print(f"🛑 Workflow stopped at: {step.name}")
@@ -1933,6 +2082,8 @@ class AgentFlow:
validation_feedback = None
guardrail_failed = False
cached_variable_updates = {} if _cache_key is not None else None
# Variables this step set (StepResult.variables), for replay's step_end.
step_variable_updates: Dict[str, Any] = {}
while retry_count <= max_retries:
step_error = None
@@ -1955,6 +2106,7 @@ class AgentFlow:
# accepted attempt that wrote {a} leaves both
# keys in the live run, so both must be cached.
all_variables.update(result.variables)
step_variable_updates.update(result.variables)
if cached_variable_updates is not None:
cached_variable_updates.update(result.variables)
else:
@@ -2196,14 +2348,10 @@ class AgentFlow:
})
# 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,25 +2360,21 @@ class AgentFlow:
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,
variables=step_variable_updates,
)
# Handle output_file - save output to file
if hasattr(step, 'output_file') and step.output_file and output:
try:
import os
from ..utils.output_file import write_output_file
output_path = step.output_file
# Substitute variables in path
for key, value in all_variables.items():
output_path = output_path.replace(f"{{{{{key}}}}}", str(value))
# Create directory if needed
os.makedirs(os.path.dirname(output_path) if os.path.dirname(output_path) else ".", exist_ok=True)
with open(output_path, "w") as f:
f.write(str(output))
if verbose:
if write_output_file(output_path, str(output)) and verbose:
print(f"📁 Saved output to: {output_path}")
except Exception as e:
logger.error(f"Failed to save output to file: {e}")
@@ -2984,10 +3128,57 @@ Create a brief execution plan (2-3 sentences) describing how to best accomplish
# 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"
)
# Variables this step set, minus its output variable, for replay's step_end.
# A cache hit leaves all_variables untouched, so it is its own baseline.
before = span.pop("before", all_variables)
step_vars = {}
if is_dict:
step_vars = {
k: v for k, v in (result.get("variables") or {}).items()
if k not in before or _differs(before[k], v)
}
step_vars.pop(getattr(step, "output_variable", None) or f"{result.get('step')}_output", None)
_replay_step_end(span, status, result.get("output") if is_dict else None, step_vars)
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. "
@@ -3014,6 +3205,7 @@ Create a brief execution plan (2-3 sentences) describing how to best accomplish
}
cache = getattr(self, "_step_cache", None)
if cache is None:
_snapshot_scope(_replay, all_variables)
return self._execute_single_step_uncached(
step, previous_output, input, all_variables, model, verbose, index,
stream=stream, depth=depth,
@@ -3036,8 +3228,11 @@ Create a brief execution plan (2-3 sentences) describing how to best accomplish
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)
_snapshot_scope(_replay, all_variables)
result = self._execute_single_step_uncached(
step, previous_output, input, all_variables, model, verbose, index,
stream=stream, depth=depth,
@@ -3410,19 +3605,18 @@ Create a brief execution plan (2-3 sentences) describing how to best accomplish
# Handle output_file - save output to file
if hasattr(step, 'output_file') and step.output_file and output and not step_error:
try:
from ..utils.output_file import write_output_file
output_path = step.output_file
for key, value in all_variables.items():
output_path = output_path.replace(f"{{{{{key}}}}}", str(value))
os.makedirs(os.path.dirname(output_path) if os.path.dirname(output_path) else ".", exist_ok=True)
with open(output_path, "w") as f:
f.write(str(output))
if verbose:
if write_output_file(output_path, str(output)) and verbose:
print(f"📁 Saved output to: {output_path}")
except Exception as e:
logger.error(f"Failed to save output to file: {e}")
return output, step_error
@_replay_pattern("route")
def _execute_route(
self,
route_step: Route,
@@ -3453,6 +3647,9 @@ Create a brief execution plan (2-3 sentences) describing how to best accomplish
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 +3901,7 @@ CONCISE SUMMARY:"""
winners[key] = idx
all_variables[key] = value
@_replay_pattern("parallel")
def _execute_parallel(
self,
parallel_step: Parallel,
@@ -3956,6 +4154,7 @@ CONCISE SUMMARY:"""
return {"steps": results, "output": combined_output, "variables": all_variables, "stop": parallel_stopped}
@_replay_pattern("loop")
def _execute_loop(
self,
loop_step: Loop,
@@ -4342,6 +4541,7 @@ CONCISE SUMMARY:"""
# 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 +4625,7 @@ CONCISE SUMMARY:"""
"stop": discussion_stopped,
}
@_replay_pattern("repeat")
def _execute_repeat(
self,
repeat_step: Repeat,
@@ -4499,6 +4700,7 @@ CONCISE SUMMARY:"""
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 +4732,7 @@ CONCISE SUMMARY:"""
# 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 +4786,7 @@ CONCISE SUMMARY:"""
from praisonaiagents.conditions.evaluator import evaluate_condition
return evaluate_condition(condition, variables, previous_output)
@_replay_pattern("include")
def _execute_include(
self,
include_step: Include,
@@ -1077,11 +1077,15 @@ class YAMLWorkflowParser:
# Build routes dictionary with resolved agents
routes = {}
for route_key, agent_ids in route_config.items():
if isinstance(agent_ids, list):
routes[route_key] = [self._agents[aid] for aid in agent_ids if aid in self._agents]
else:
if agent_ids in self._agents:
routes[route_key] = [self._agents[agent_ids]]
if not isinstance(agent_ids, list):
agent_ids = [agent_ids]
try:
routes[route_key] = [self._agents[aid] for aid in agent_ids]
except KeyError as exc:
raise ValueError(
f"Agent '{exc.args[0]}' in route '{route_key}' "
"not defined in agents section"
) from exc
return route(routes)
@@ -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
@@ -1274,12 +1274,13 @@ class TestTaskTools:
class TestMigratedFeatures:
"""Test features migrated from process='workflow'."""
def test_step_output_file(self):
def test_step_output_file(self, monkeypatch):
"""Test Task with output_file parameter."""
import tempfile
import os
with tempfile.TemporaryDirectory() as tmpdir:
monkeypatch.setenv("PRAISONAI_PROJECT_ROOT", tmpdir)
output_path = os.path.join(tmpdir, "output.txt")
def generate_content(ctx: WorkflowContext) -> StepResult:
@@ -0,0 +1,27 @@
from praisonaiagents.agent.context_agent import ContextAgent
NON_ASCII = "日本語 café ☕"
def test_context_engineering_results_written_as_readable_utf8(tmp_path, monkeypatch):
"""Every saved Context Engineering result file keeps non-ASCII content as readable UTF-8."""
# ContextAgent defaults output to .praisonai/prp under the cwd, so run inside tmp.
monkeypatch.chdir(tmp_path)
agent = ContextAgent(llm="gpt-4o-mini", auto_analyze=False)
agent.output_dir = str(tmp_path)
agent.codebase_analysis = {"summary": NON_ASCII}
agent.pattern_library = {"summary": NON_ASCII}
agent.validation_framework = {"summary": NON_ASCII}
agent.agent_interactions = [{"response": NON_ASCII}]
agent._save_context_engineering_results()
for filename in (
"comprehensive_analysis.json",
"pattern_library.json",
"validation_framework.json",
"all_agent_interactions.json",
):
saved = (tmp_path / filename).read_text(encoding="utf-8")
assert NON_ASCII in saved, filename
assert "\\u" not in saved, filename
@@ -0,0 +1,35 @@
"""Retained completed events are available to both wait entry points."""
import asyncio
from praisonaiagents import AgentTeam
def test_sync_default_returns_already_completed_events():
with AgentTeam(agents=[]) as team:
team.announce_completion("finished-a", "task-a", "result-a")
team.announce_completion("finished-b", "task-b", "result-b")
events = team.wait_for_completions(timeout=0.2)
assert [event.agent_id for event in events] == ["finished-a", "finished-b"]
def test_explicit_empty_ids_do_not_return_retained_events():
"""An explicit empty target set retains its previous empty result."""
with AgentTeam(agents=[]) as team:
team.announce_completion("finished-a", "task-a", "result-a")
assert team.wait_for_completions(timeout=0.2, agent_ids=[]) == []
def test_explicit_ids_filter_retained_events():
with AgentTeam(agents=[]) as team:
team.announce_completion("finished-a", "task-a", "result-a")
team.announce_completion("finished-b", "task-b", "result-b")
events = team.wait_for_completions(timeout=0.2, agent_ids=["finished-b"])
assert [event.agent_id for event in events] == ["finished-b"]
def test_async_default_returns_already_completed_events():
with AgentTeam(agents=[]) as team:
team.announce_completion("finished-a", "task-a", "result-a")
team.announce_completion("finished-b", "task-b", "result-b")
events = asyncio.run(team.await_for_completions(timeout=0.2))
assert [event.agent_id for event in events] == ["finished-a", "finished-b"]
@@ -0,0 +1,49 @@
"""Workflow dependents must honor failure policy after synchronous predecessors."""
import asyncio
import pytest
from praisonaiagents import Agent, AgentTeam, Task
class ScriptedAgent(Agent):
"""Exercise production task scheduling without provider requests."""
def __init__(self, name, reply):
super().__init__(name=name, instructions="test", llm="gpt-4o-mini")
self.reply = reply
self.calls = 0
def chat(self, prompt, **kwargs):
"""Record a real dispatch of this task's agent."""
self.calls += 1
return self.reply
async def achat(self, prompt, **kwargs):
"""Record asynchronous dispatch through the same fixture."""
return self.chat(prompt, **kwargs)
@pytest.mark.parametrize("entry", ["start", "astart"])
@pytest.mark.parametrize("upstream_async", [False, True])
@pytest.mark.parametrize("downstream_async", [False, True])
@pytest.mark.parametrize("policy", ["stop", "continue", "skip"])
def test_workflow_failed_predecessor_policy(entry, upstream_async, downstream_async, policy):
"""Default dependents do not execute; explicit degraded-mode ones still do."""
upstream = ScriptedAgent("upstream", None)
downstream = ScriptedAgent("downstream", "done")
first = Task(name="first", description="first", agent=upstream,
is_start=True, next_tasks=["second"], async_execution=upstream_async, max_retries=1)
flags = {"on_error": "continue"} if policy == "continue" else {}
if policy == "skip":
flags["skip_on_failure"] = True
second = Task(name="second", description="second", agent=downstream,
async_execution=downstream_async, max_retries=1, **flags)
team = AgentTeam(agents=[upstream, downstream], tasks=[first, second], process="workflow")
if entry == "astart":
asyncio.run(team.astart())
else:
team.start()
assert first.status == "failed"
assert upstream.calls == 1
assert downstream.calls == (0 if policy == "stop" else 1)
@@ -1,222 +0,0 @@
"""Unit tests for evaluation utility functions."""
import json
import tempfile
from pathlib import Path
from unittest.mock import MagicMock
from praisonaiagents.eval.utils import (
save_result_to_file,
load_result_from_file,
format_score,
format_percentage,
format_duration,
format_memory,
calculate_pass_rate,
)
class TestSaveResultToFile:
"""Tests for save_result_to_file function."""
def test_save_with_to_dict(self):
"""Test saving result with to_dict method."""
with tempfile.TemporaryDirectory() as tmpdir:
path = f"{tmpdir}/result.json"
result = MagicMock()
result.to_dict.return_value = {"score": 8.0, "passed": True}
success = save_result_to_file(result, path)
assert success is True
assert Path(path).exists()
with open(path) as f:
data = json.load(f)
assert data["score"] == 8.0
def test_save_with_dict_attr(self):
"""Test saving result with __dict__ attribute."""
with tempfile.TemporaryDirectory() as tmpdir:
path = f"{tmpdir}/result.json"
class SimpleResult:
def __init__(self):
self.score = 9.0
self.name = "test"
result = SimpleResult()
success = save_result_to_file(result, path)
assert success is True
with open(path) as f:
data = json.load(f)
assert data["score"] == 9.0
def test_save_with_placeholders(self):
"""Test saving with path placeholders."""
with tempfile.TemporaryDirectory() as tmpdir:
path = f"{tmpdir}/{{name}}_{{eval_id}}.json"
result = MagicMock()
result.to_dict.return_value = {"score": 8.0}
success = save_result_to_file(
result, path,
eval_id="abc123",
name="test_eval"
)
assert success is True
expected_path = f"{tmpdir}/test_eval_abc123.json"
assert Path(expected_path).exists()
def test_save_creates_directories(self):
"""Test that save creates parent directories."""
with tempfile.TemporaryDirectory() as tmpdir:
path = f"{tmpdir}/nested/dir/result.json"
result = MagicMock()
result.to_dict.return_value = {"score": 8.0}
success = save_result_to_file(result, path)
assert success is True
assert Path(path).exists()
def test_save_failure_returns_false(self):
"""Test that save failure returns False."""
result = MagicMock()
result.to_dict.side_effect = Exception("Test error")
success = save_result_to_file(result, "/invalid/path/result.json")
assert success is False
class TestLoadResultFromFile:
"""Tests for load_result_from_file function."""
def test_load_existing_file(self):
"""Test loading existing result file."""
with tempfile.TemporaryDirectory() as tmpdir:
path = f"{tmpdir}/result.json"
data = {"score": 8.0, "passed": True}
with open(path, 'w') as f:
json.dump(data, f)
loaded = load_result_from_file(path)
assert loaded == data
def test_load_nonexistent_file(self):
"""Test loading nonexistent file returns None."""
result = load_result_from_file("/nonexistent/path/result.json")
assert result is None
def test_load_invalid_json(self):
"""Test loading invalid JSON returns None."""
with tempfile.TemporaryDirectory() as tmpdir:
path = f"{tmpdir}/invalid.json"
with open(path, 'w') as f:
f.write("not valid json")
result = load_result_from_file(path)
assert result is None
class TestFormatScore:
"""Tests for format_score function."""
def test_format_default_max(self):
"""Test formatting with default max score."""
assert format_score(8.5) == "8.5/10"
def test_format_custom_max(self):
"""Test formatting with custom max score."""
assert format_score(85.0, 100.0) == "85.0/100"
def test_format_zero(self):
"""Test formatting zero score."""
assert format_score(0.0) == "0.0/10"
class TestFormatPercentage:
"""Tests for format_percentage function."""
def test_format_full(self):
"""Test formatting 100%."""
assert format_percentage(1.0) == "100.0%"
def test_format_half(self):
"""Test formatting 50%."""
assert format_percentage(0.5) == "50.0%"
def test_format_zero(self):
"""Test formatting 0%."""
assert format_percentage(0.0) == "0.0%"
def test_format_decimal(self):
"""Test formatting decimal percentage."""
assert format_percentage(0.333) == "33.3%"
class TestFormatDuration:
"""Tests for format_duration function."""
def test_format_microseconds(self):
"""Test formatting microseconds."""
result = format_duration(0.0001)
assert "μs" in result
def test_format_milliseconds(self):
"""Test formatting milliseconds."""
result = format_duration(0.5)
assert "ms" in result
def test_format_seconds(self):
"""Test formatting seconds."""
result = format_duration(5.5)
assert "s" in result
assert "m" not in result
def test_format_minutes(self):
"""Test formatting minutes."""
result = format_duration(90.0)
assert "m" in result
class TestFormatMemory:
"""Tests for format_memory function."""
def test_format_kilobytes(self):
"""Test formatting kilobytes."""
result = format_memory(0.5)
assert "KB" in result
def test_format_megabytes(self):
"""Test formatting megabytes."""
result = format_memory(100.0)
assert "MB" in result
def test_format_gigabytes(self):
"""Test formatting gigabytes."""
result = format_memory(2048.0)
assert "GB" in result
class TestCalculatePassRate:
"""Tests for calculate_pass_rate function."""
def test_all_passed(self):
"""Test 100% pass rate."""
assert calculate_pass_rate(10, 10) == 1.0
def test_none_passed(self):
"""Test 0% pass rate."""
assert calculate_pass_rate(0, 10) == 0.0
def test_half_passed(self):
"""Test 50% pass rate."""
assert calculate_pass_rate(5, 10) == 0.5
def test_zero_total(self):
"""Test with zero total."""
assert calculate_pass_rate(0, 0) == 0.0
@@ -0,0 +1,53 @@
"""AFTER_LLM reports the usage counted by the existing completion path."""
import asyncio
from types import SimpleNamespace
from unittest.mock import AsyncMock, Mock
from openai.types.chat import ChatCompletion
import pytest
from praisonaiagents import Agent
from praisonaiagents.config.feature_configs import HooksConfig
from praisonaiagents.hooks import HookEvent, HookRegistry, HookResult
@pytest.mark.parametrize("asynchronous", [False, True])
@pytest.mark.parametrize("with_usage", [False, True])
def test_after_llm_receives_current_completion_usage(asynchronous, with_usage):
usage = {"prompt_tokens": 11, "completion_tokens": 7, "total_tokens": 18}
completion = ChatCompletion(
id="chatcmpl-fixture", created=1, object="chat.completion",
model="gpt-4o-mini",
choices=[{
"index": 0, "finish_reason": "stop",
"message": {"role": "assistant", "content": "Fixture greeting."},
}],
usage=usage if with_usage else None,
)
observed = []
registry = HookRegistry()
def after_llm(event):
observed.append(event.tokens_used)
return HookResult.allow()
registry.register_function(HookEvent.AFTER_LLM, after_llm)
agent = Agent(
name="usage-fixture", instructions="Greet.", model="gpt-4o-mini",
api_key="synthetic-offline-fixture",
hooks=HooksConfig(registry=registry),
output={"verbose": False, "stream": False},
reflection=False, memory=False, rules=False, retry=False,
)
agent._unified_dispatcher = SimpleNamespace(
chat_completion=Mock(return_value=completion),
achat_completion=AsyncMock(return_value=completion),
)
if asynchronous:
result = asyncio.run(agent.achat("Hello", stream=False))
agent._unified_dispatcher.achat_completion.assert_awaited_once()
else:
result = agent.chat("Hello", stream=False)
agent._unified_dispatcher.chat_completion.assert_called_once()
assert result == "Fixture greeting."
assert observed == [18 if with_usage else 0]
@@ -0,0 +1,94 @@
"""A supplied memory adapter must not require the default implementation."""
import copy
import sys
import types
import pytest
pytest.importorskip("litellm", reason="optional LLM provider dependency")
from praisonaiagents.llm import LLM
class ProvidedMemory:
"""Side-effect-free implementation of the supplied memory adapter surface."""
def get_tool_definition(self):
"""Return the request tool definition without constructing a backend."""
return {"type": "memory_20250818", "name": "memory"}
def get_beta_header(self):
"""Return a synthetic beta marker for local parameter checks."""
return "synthetic-memory-fixture-beta"
def execute(self, **kwargs):
"""Provide the adapter execution surface without I/O."""
return "fixture memory result"
def make_llm(memory, **kwargs):
"""Construct the real params producer with an explicit synthetic key."""
return LLM(
model="anthropic/claude-sonnet-4-20250514",
api_key="synthetic-offline-fixture",
base_url="http://127.0.0.1:1",
claude_memory=memory,
**kwargs,
)
def test_supplied_memory_is_used_without_default_module(monkeypatch):
"""A supplied instance is reused even if the default cannot be imported."""
monkeypatch.setitem(sys.modules, "praisonaiagents.tools.claude_memory_tool", None)
supplied = ProvidedMemory()
llm = make_llm(supplied)
assert llm._get_claude_memory_tool() is supplied
assert llm._get_claude_memory_tool() is supplied
@pytest.mark.parametrize("tools, headers", [
(None, None),
([], {}),
([{"type": "function", "function": {"name": "fixture", "parameters": {"type": "object", "properties": {}}}}], {"x-fixture": "retained"}),
])
def test_real_params_builder_accepts_supplied_memory_without_default_module(monkeypatch, tools, headers):
"""Supplied memory is included without ordering assumptions or input mutation."""
monkeypatch.setitem(sys.modules, "praisonaiagents.tools.claude_memory_tool", None)
supplied = ProvidedMemory()
llm = make_llm(supplied, web_fetch=True)
original_tools, original_headers = copy.deepcopy(tools), copy.deepcopy(headers)
first = llm._build_completion_params(tools=tools, extra_headers=headers)
second = llm._build_completion_params(tools=tools, extra_headers=headers)
assert supplied.get_tool_definition() in first["tools"]
assert first["tools"] == second["tools"]
assert supplied.get_beta_header() in first["extra_headers"]["anthropic-beta"]
assert tools == original_tools
assert headers == original_headers
@pytest.mark.parametrize("memory", [None, False])
def test_disabled_memory_still_needs_no_implementation(monkeypatch, memory):
"""Disabled memory never imports an implementation."""
monkeypatch.setitem(sys.modules, "praisonaiagents.tools.claude_memory_tool", None)
assert make_llm(memory)._get_claude_memory_tool() is None
def test_boolean_opt_in_still_constructs_and_caches_default(monkeypatch):
"""Boolean opt-in retains default factory construction and caching."""
calls = []
class DefaultMemory(ProvidedMemory):
"""Optional implementation stand-in that records construction."""
def __init__(self):
"""Record one default factory invocation."""
calls.append(self)
module = types.ModuleType("praisonaiagents.tools.claude_memory_tool")
module.ClaudeMemoryTool = DefaultMemory
monkeypatch.setitem(sys.modules, module.__name__, module)
llm = make_llm(True)
tool = llm._get_claude_memory_tool()
assert isinstance(tool, DefaultMemory)
assert llm._get_claude_memory_tool() is tool
assert calls == [tool]
@@ -0,0 +1,80 @@
from datetime import datetime, timezone
from praisonaiagents.memory.learn.stores import InsightStore, LearnEntry
def test_new_entry_timestamps_are_timezone_aware():
"""New learn entries carry timezone-aware created_at and updated_at."""
entry = LearnEntry(id="e1", content="hello")
assert datetime.fromisoformat(entry.created_at).tzinfo is not None
assert datetime.fromisoformat(entry.updated_at).tzinfo is not None
def test_prune_handles_legacy_naive_timestamps(tmp_path):
"""Pruning treats legacy naive timestamps as UTC instead of raising on the subtraction."""
store = InsightStore(store_path=str(tmp_path / "store.json"), retention_days=1)
now = datetime.now(timezone.utc)
# An entry persisted before the timezone-aware change: naive isoformat, no offset.
legacy = LearnEntry(
id="old",
content="old",
created_at="2020-01-01T00:00:00",
updated_at="2020-01-01T00:00:00",
)
# Must not raise "can't subtract offset-naive and offset-aware datetimes".
assert store._is_stale(legacy, now) is True
fresh = LearnEntry(id="new", content="new")
assert store._is_stale(fresh, now) is False
def test_prune_evicts_stale_legacy_entry_end_to_end(tmp_path):
"""prune() archives a stale legacy (naive-timestamp) entry without raising."""
store = InsightStore(store_path=str(tmp_path / "store.json"), retention_days=1)
store._entries["old"] = LearnEntry(
id="old",
content="old",
created_at="2020-01-01T00:00:00",
updated_at="2020-01-01T00:00:00",
)
assert store.prune(save=False) == 1
assert "old" not in store._entries
def test_list_all_orders_by_instant_not_raw_string(tmp_path):
"""list_all() orders by the actual UTC instant, even when offsets differ.
00:30+01:00 is an earlier instant than 00:00+00:00, but a raw-string sort
would order them the other way round.
"""
store = InsightStore(store_path=str(tmp_path / "store.json"))
store._entries = {
"earlier": LearnEntry(id="earlier", content="e", updated_at="2024-01-01T00:30:00+01:00"),
"later": LearnEntry(id="later", content="l", updated_at="2024-01-01T00:00:00+00:00"),
}
store._save() # persist so list_all's refresh reloads these entries
ordered = [e.id for e in store.list_all()] # reverse=True: most recent first
assert ordered == ["later", "earlier"]
def test_prune_evicts_never_used_entry_before_used_one(tmp_path):
"""At equal use_count, a never-used entry (last_used=None) is evicted first."""
store = InsightStore(store_path=str(tmp_path / "store.json"), max_entries=1)
store._entries = {
"never": LearnEntry(
id="never", content="n", use_count=0, last_used=None,
created_at="2024-01-01T00:00:00+00:00",
),
"used": LearnEntry(
id="used", content="u", use_count=0,
last_used="2024-06-01T00:00:00+00:00",
created_at="2024-01-01T00:00:00+00:00",
),
}
store.prune(save=False)
assert "never" not in store._entries
assert "used" in store._entries
@@ -0,0 +1,62 @@
"""Input validation must not run blocking validators on the achat event loop."""
import asyncio
import threading
import pytest
from praisonaiagents import Agent
class RecordingInputGuardrail:
"""Reject before provider dispatch while recording the execution thread."""
def __init__(self):
self.thread_id = None
self.content = None
def validate_input(self, content, **kwargs):
self.thread_id = threading.get_ident()
self.content = content
return False, "blocked by test validator"
def validate_output(self, content, **kwargs):
return True, content
def test_achat_offloads_sync_input_guardrail():
"""Public achat invokes a synchronous validator outside its event-loop thread."""
guardrail = RecordingInputGuardrail()
agent = Agent(name="T", instructions="test", llm="gpt-4o-mini", guardrails=guardrail)
async def run():
loop_thread = threading.get_ident()
await agent.achat("blocked prompt")
assert guardrail.content == "blocked prompt"
assert guardrail.thread_id is not None
assert guardrail.thread_id != loop_thread
asyncio.run(run())
@pytest.mark.parametrize("guardrails", [None, "Check output quality", lambda output: (True, output)])
def test_achat_disabled_guardrail_keeps_fast_path(monkeypatch, guardrails):
"""Absent, callable and output-only guards need no worker-thread dispatch."""
agent = Agent(name="T", instructions="test", llm="gpt-4o-mini", guardrails=guardrails)
guardrail = RecordingInputGuardrail()
def record(prompt):
guardrail.validate_input(prompt)
return False, None, "stop before provider dispatch"
monkeypatch.setattr(agent, "_validate_input_with_guardrail", record)
asyncio.run(agent.achat("blocked prompt"))
assert guardrail.thread_id == threading.get_ident()
def test_chat_keeps_sync_input_guardrail_on_caller_thread():
"""The synchronous public API retains its existing execution context."""
guardrail = RecordingInputGuardrail()
agent = Agent(name="T", instructions="test", llm="gpt-4o-mini", guardrails=guardrail)
agent.chat("blocked prompt")
assert guardrail.content == "blocked prompt"
assert guardrail.thread_id == threading.get_ident()
@@ -0,0 +1,61 @@
"""Disabled handoffs do not load their implementation; enabled paths still convert."""
import sys
import pytest
from praisonaiagents import Agent
from praisonaiagents.agent.handoff import Handoff
def make_agent(name, **kwargs):
"""Construct a minimal offline agent for handoff conversion checks."""
return Agent(
name=name,
instructions="Return a short answer.",
api_key="synthetic-offline-fixture",
reflection=False,
memory=False,
knowledge=False,
**kwargs,
)
@pytest.mark.parametrize("tools", [None, []])
def test_disabled_handoffs_do_not_require_handoff_module(monkeypatch, tools):
"""Empty tool/handoff configuration must not load handoff support."""
# A blocked import makes an unnecessary local import observable even when
# another test has already loaded this module in the same interpreter.
monkeypatch.setitem(sys.modules, "praisonaiagents.agent.handoff", None)
agent = make_agent("no-handoff", tools=tools)
assert agent.tools == []
assert agent.handoffs == []
@pytest.mark.parametrize("entry", ["tools", "handoffs"])
def test_handoff_objects_still_convert_to_bound_tools(entry):
"""Both supported Handoff entry points retain their bound source object."""
target = make_agent("target")
source = Handoff(target)
agent = make_agent("parent", **{entry: [source]})
assert len(agent.tools) == 1
assert callable(agent.tools[0])
assert agent.tools[0]._praison_handoff_source is source
def test_direct_agent_handoff_still_converts_to_bound_tool():
"""A direct Agent handoff still creates a tool for that target."""
target = make_agent("target")
agent = make_agent("parent", handoffs=[target])
assert len(agent.tools) == 1
assert callable(agent.tools[0])
assert agent.tools[0]._praison_handoff_source.agent is target
def test_regular_tool_is_unchanged():
"""Ordinary callable tools remain unchanged by handoff conversion."""
def echo(text: str) -> str:
"""Return the input text."""
return text
agent = make_agent("regular-tool", tools=[echo])
assert agent.tools == [echo]
@@ -0,0 +1,139 @@
"""Regression tests for issue #5957 (doom-loop window, download_file SSRF/size cap)."""
import socket
import time
from unittest.mock import MagicMock, patch
import pytest
from praisonaiagents.approval import reset_yaml_approved_tools, set_yaml_approved_tools
from praisonaiagents.permissions.doom_loop import DoomLoopDetector
from praisonaiagents.tools import url_safety
def test_doom_loop_tool_count_is_windowed():
detector = DoomLoopDetector(window_seconds=0.05, max_tool_calls=10)
# One call per window with fresh arguments is not a loop.
for i in range(8):
result = detector.record_and_check("read_file", {"path": f"f{i}"})
assert not result.is_loop, result.reason
time.sleep(0.06)
# Lifetime totals are still reported for stats.
assert detector.get_stats()["call_counts"]["read_file"] == 8
def test_doom_loop_tool_count_still_trips_within_window():
detector = DoomLoopDetector(window_seconds=60, max_tool_calls=10)
for i in range(5):
detector.record("read_file", {"path": f"f{i}"})
result = detector.check("read_file", {"path": "new"})
assert result.is_loop
assert "excessively" in result.reason
def _addrinfo(ip):
return [(socket.AF_INET, socket.SOCK_STREAM, 0, "", (ip, 0))]
@pytest.mark.parametrize("ip", ["100.100.100.200", "100.64.0.1"])
def test_url_safety_blocks_shared_address_space(ip, monkeypatch):
monkeypatch.delenv("ALLOW_LOCAL_CRAWL", raising=False)
with patch.object(url_safety.socket, "getaddrinfo", return_value=_addrinfo(ip)):
assert not url_safety.is_safe_http_url(f"http://{ip}/latest/meta-data/")
def test_url_safety_allows_public_address(monkeypatch):
monkeypatch.delenv("ALLOW_LOCAL_CRAWL", raising=False)
with patch.object(url_safety.socket, "getaddrinfo", return_value=_addrinfo("93.184.216.34")):
assert url_safety.is_safe_http_url("https://example.com/")
@pytest.fixture
def file_tools(tmp_path, monkeypatch):
pytest.importorskip("httpx")
from praisonaiagents.tools.file_tools import FileTools
monkeypatch.chdir(tmp_path)
monkeypatch.delenv("ALLOW_LOCAL_CRAWL", raising=False)
tools = FileTools()
monkeypatch.setattr(tools, "_require_workspace_access", lambda write=False: None)
monkeypatch.setattr(tools, "_validate_path", lambda p: str(tmp_path / p))
token = set_yaml_approved_tools(["download_file"])
try:
yield tools
finally:
reset_yaml_approved_tools(token)
@pytest.mark.parametrize("url", [
"http://localhost./x",
"http://169.254.169.254./x",
"http://0177.0.0.1/x",
])
def test_download_file_uses_shared_ssrf_guard(file_tools, url):
with patch("httpx.stream") as stream:
result = file_tools.download_file(url, "out.bin")
assert not result["success"]
stream.assert_not_called()
def test_download_file_ignores_allow_local_crawl(file_tools, monkeypatch):
monkeypatch.setenv("ALLOW_LOCAL_CRAWL", "true")
with patch("httpx.stream") as stream:
result = file_tools.download_file("http://127.0.0.1/x", "out.bin")
assert not result["success"]
stream.assert_not_called()
def test_download_file_fails_closed_on_unresolvable_host(file_tools):
with patch.object(url_safety.socket, "getaddrinfo", side_effect=socket.gaierror), \
patch("httpx.stream") as stream:
result = file_tools.download_file("http://does-not-resolve.invalid/x", "out.bin")
assert not result["success"]
stream.assert_not_called()
def test_download_file_enforces_byte_cap(file_tools, tmp_path, monkeypatch):
monkeypatch.setenv("PRAISONAI_MAX_DOWNLOAD_BYTES", "10000")
response = MagicMock()
response.headers = {}
response.iter_bytes.return_value = iter([b"x" * 8192] * 10)
ctx = MagicMock()
ctx.__enter__.return_value = response
with patch.object(url_safety.socket, "getaddrinfo", return_value=_addrinfo("93.184.216.34")), \
patch("httpx.stream", return_value=ctx):
result = file_tools.download_file("https://example.com/big", "big.bin")
assert not result["success"]
assert "exceeds" in result["error"]
assert not (tmp_path / "big.bin").exists()
def _stream_ctx(chunks):
response = MagicMock()
response.headers = {}
response.iter_bytes.return_value = iter(chunks)
ctx = MagicMock()
ctx.__enter__.return_value = response
return ctx
def test_download_file_cap_keeps_existing_file(file_tools, tmp_path, monkeypatch):
monkeypatch.setenv("PRAISONAI_MAX_DOWNLOAD_BYTES", "10000")
(tmp_path / "keep.bin").write_bytes(b"original")
with patch.object(url_safety.socket, "getaddrinfo", return_value=_addrinfo("93.184.216.34")), \
patch("httpx.stream", return_value=_stream_ctx([b"x" * 8192] * 10)):
result = file_tools.download_file("https://example.com/big", "keep.bin")
assert not result["success"]
assert (tmp_path / "keep.bin").read_bytes() == b"original"
assert sorted(p.name for p in tmp_path.iterdir()) == ["keep.bin"]
def test_download_file_success_replaces_destination(file_tools, tmp_path):
(tmp_path / "out.bin").write_bytes(b"old")
with patch.object(url_safety.socket, "getaddrinfo", return_value=_addrinfo("93.184.216.34")), \
patch("httpx.stream", return_value=_stream_ctx([b"new", b"data"])):
result = file_tools.download_file("https://example.com/f", "out.bin")
assert result["success"]
assert result["size"] == 7
assert (tmp_path / "out.bin").read_bytes() == b"newdata"
assert sorted(p.name for p in tmp_path.iterdir()) == ["out.bin"]
@@ -278,3 +278,58 @@ class TestTelemetryConfig:
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 == "[]"
@@ -0,0 +1,173 @@
"""Task and Workflow step output_file writes must stay inside the project root."""
import asyncio
import os
import pytest
from praisonaiagents import (
Agent,
AgentTeam,
StepResult,
Task,
Workflow,
WorkflowContext,
)
@pytest.fixture
def root(tmp_path, monkeypatch):
work = tmp_path / "work"
work.mkdir()
monkeypatch.chdir(work)
monkeypatch.delenv("PRAISONAI_PROJECT_ROOT", raising=False)
return work
def _team(output_file, create_directory=True):
agent = Agent(name="a", instructions="test")
task = Task(
description="d",
expected_output="e",
agent=agent,
output_file=output_file,
create_directory=create_directory,
)
return AgentTeam(agents=[agent], tasks=[task]), task
def _workflow(output_file):
def generate(ctx: WorkflowContext) -> StepResult:
return StepResult(output="report")
return Workflow(steps=[Task(name="gen", handler=generate, output_file=output_file)])
def test_task_output_file_traversal_blocked(root):
team, task = _team("../escaped/out.txt")
team.save_output_to_file(task, "PWNED")
assert not (root.parent / "escaped").exists()
def test_task_output_file_symlink_escape_blocked(root, tmp_path):
outside = tmp_path / "outside"
outside.mkdir()
os.symlink(outside, root / "link")
team, task = _team("link/out.txt")
team.save_output_to_file(task, "PWNED")
assert not (outside / "out.txt").exists()
def test_task_output_file_inside_root(root):
team, task = _team("sub/out.txt")
team.save_output_to_file(task, "héllo 世界")
assert (root / "sub" / "out.txt").read_text(encoding="utf-8") == "héllo 世界"
team, task = _team(str(root / "abs.txt"))
team.save_output_to_file(task, "ok")
assert (root / "abs.txt").read_text() == "ok"
@pytest.mark.skipif(os.name == "nt", reason="POSIX permission bits")
def test_task_output_file_overwrite_preserves_permissions(root):
target = root / "shared.txt"
target.write_text("old")
os.chmod(target, 0o644)
team, task = _team("shared.txt")
team.save_output_to_file(task, "new")
assert target.read_text() == "new"
assert target.stat().st_mode & 0o777 == 0o644
@pytest.mark.skipif(os.name == "nt", reason="POSIX hardlinks")
def test_task_output_file_hardlink_does_not_write_outside(root, tmp_path):
outside = tmp_path / "outside.txt"
outside.write_text("secret")
os.link(outside, root / "linked.txt")
team, task = _team("linked.txt")
team.save_output_to_file(task, "PWNED")
assert (root / "linked.txt").read_text() == "PWNED"
assert outside.read_text() == "secret"
@pytest.mark.skipif(os.name == "nt", reason="POSIX hardlinks")
def test_task_output_file_hardlink_keeps_mode_under_umask(root, tmp_path):
outside = tmp_path / "outside.txt"
outside.write_text("secret")
os.chmod(outside, 0o664)
os.link(outside, root / "linked.txt")
team, task = _team("linked.txt")
old_umask = os.umask(0o077)
try:
team.save_output_to_file(task, "new")
finally:
os.umask(old_umask)
assert (root / "linked.txt").stat().st_mode & 0o777 == 0o664
assert outside.read_text() == "secret"
@pytest.mark.skipif(os.name == "nt" or os.geteuid() == 0, reason="POSIX permission bits")
def test_task_output_file_overwrite_in_readonly_directory(root):
(root / "ro").mkdir()
target = root / "ro" / "out.txt"
target.write_text("old")
os.chmod(root / "ro", 0o555)
try:
team, task = _team("ro/out.txt")
team.save_output_to_file(task, "new")
finally:
os.chmod(root / "ro", 0o755)
assert target.read_text() == "new"
def test_task_output_file_respects_create_directory(root):
team, task = _team("missing/out.txt", create_directory=False)
team.save_output_to_file(task, "ok")
assert not (root / "missing").exists()
def test_task_output_file_project_root_override(root, tmp_path, monkeypatch):
monkeypatch.setenv("PRAISONAI_PROJECT_ROOT", str(tmp_path))
team, task = _team(str(tmp_path / "shared" / "out.txt"))
team.save_output_to_file(task, "ok")
assert (tmp_path / "shared" / "out.txt").read_text() == "ok"
def test_workflow_output_file_input_traversal_blocked(root):
_workflow("reports/{{input}}.md").start("../../x")
assert not (root.parent / "x.md").exists()
_workflow("reports/{{input}}.md").start("ok")
assert (root / "reports" / "ok.md").read_text() == "report"
def test_workflow_async_output_file_input_traversal_blocked(root):
asyncio.run(_workflow("reports/{{input}}.md").arun("../../x"))
assert not (root.parent / "x.md").exists()
asyncio.run(_workflow("reports/{{input}}.md").arun("ok"))
assert (root / "reports" / "ok.md").read_text() == "report"
@pytest.mark.skipif(not os.environ.get("OPENAI_API_KEY"), reason="OPENAI_API_KEY not set")
def test_task_output_file_traversal_blocked_real_run(root):
team, _ = _team("../escaped.txt")
team.start()
assert not (root.parent / "escaped.txt").exists()
@pytest.mark.skipif(os.name == "nt", reason="POSIX hardlinks")
def test_task_output_file_hardlink_group_kept_or_bits_dropped(root, tmp_path, monkeypatch):
outside = tmp_path / "outside.txt"
outside.write_text("secret")
os.chmod(outside, 0o660)
os.link(outside, root / "linked.txt")
def deny(*args):
raise PermissionError
monkeypatch.setattr(os, "fchown", deny)
team, task = _team("linked.txt")
team.save_output_to_file(task, "new")
assert (root / "linked.txt").read_text() == "new"
assert (root / "linked.txt").stat().st_mode & 0o777 == 0o600
assert outside.read_text() == "secret"
@@ -0,0 +1,179 @@
"""Regression tests for #5980.
- A second AgentTeam.start()/astart() must re-execute tasks, not return stale output.
- RAG.retrieve(top_k=...) must not mutate the shared RAGConfig.
- SmartRetriever must over-fetch before post-filtering, fail closed on a bad
path_filter regex, and not share store errors across calls via instance state.
"""
import asyncio
import threading
import pytest
from praisonaiagents import Agent, AgentTeam, Task
from praisonaiagents.knowledge.query_engine import SubQuestionEngine
from praisonaiagents.rag.models import RAGConfig
from praisonaiagents.rag.pipeline import RAG
from praisonaiagents.rag.retriever import SmartRetriever
def _team():
calls = []
agent = Agent(name="a", instructions="x", llm="gpt-4o-mini")
def chat(prompt, **kw):
calls.append(prompt)
return f"answer#{len(calls)}"
async def achat(prompt, **kw):
return chat(prompt)
agent.chat = chat
agent.achat = achat
task = Task(name="t1", description="Summarise the ticket", expected_output="e", agent=agent)
return AgentTeam(agents=[agent], tasks=[task], process="sequential"), task, calls
def test_second_start_reexecutes_and_drops_previous_content():
team, task, calls = _team()
assert team.start("ticket A", output="silent") == "answer#1"
assert team.start("ticket B", output="silent") == "answer#2"
assert len(calls) == 2
assert task.context == ["ticket B"]
def test_second_astart_reexecutes():
team, _, calls = _team()
async def run():
return await team.astart(), await team.astart()
assert asyncio.run(run()) == ("answer#1", "answer#2")
assert len(calls) == 2
def test_restored_checkpoint_is_resumed_not_reset():
team, task, calls = _team()
team.start(output="silent")
snapshot = team._serialisable_task_state()
assert team._restore_serialised_task_state(snapshot) == 1
team.start(output="silent")
assert len(calls) == 1 # completed task restored -> skipped
assert task.status == "completed"
def test_rerun_keeps_caller_defined_task_variables():
team, task, calls = _team()
task.description = "Summarise {{ticket}}"
task.variables = {"ticket": "T-1"}
team.start(output="silent")
team.start(output="silent")
assert task.variables == {"ticket": "T-1"}
assert all("T-1" in str(p) and "{{ticket}}" not in str(p) for p in calls)
def test_rerun_repropagates_current_team_variables():
team, task, calls = _team()
task.description = "Summarise {{ticket}}"
team.variables = {"ticket": "T-1"}
team.start(output="silent")
team.variables = {"ticket": "T-2"}
team.start(output="silent")
assert "T-2" in str(calls[1])
def test_team_variables_repropagate_after_batch():
team, task, calls = _team()
task.description = "Summarise {{ticket}}"
team.variables = {"ticket": "T-1"}
team.start(output="silent")
team.start_for_each([{"ticket": "X"}])
team.variables = {"ticket": "T-2"}
team.start(output="silent")
assert "T-2" in str(calls[-1])
def test_fresh_run_after_resume_drops_earlier_content():
team, task, calls = _team()
team.start("ticket A", output="silent")
team._restore_serialised_task_state(team._serialisable_task_state())
team.start(output="silent") # resume: completed task skipped
team.start("ticket B", output="silent")
assert task.context == ["ticket B"]
def test_rag_retrieve_override_does_not_mutate_config():
# Hold the top_k=1 call inside search() until a default call has finished,
# so the default call is guaranteed to run while the override is active.
override_in_search = threading.Event()
default_done = threading.Event()
class Knowledge:
def search(self, query, **kw):
if query == "override":
override_in_search.set()
default_done.wait(2)
return [{"text": f"d{i}", "score": 1.0, "metadata": {}} for i in range(20)]
rag = RAG(knowledge=Knowledge(), config=RAGConfig(top_k=5))
out = {}
ta = threading.Thread(target=lambda: out.setdefault("a", rag.retrieve("override", top_k=1)))
ta.start()
assert override_in_search.wait(2)
out["b"] = rag.retrieve("default")
default_done.set()
ta.join()
assert rag.config.top_k == 5
assert out["a"].metadata["top_k"] == 1
assert out["b"].metadata["top_k"] == 5
class _MixedStore:
def search(self, query, limit=5, **kw):
docs = [{"text": f"p{i}", "metadata": {"filename": f"p{i}.pdf"}, "score": 1 - i / 100}
for i in range(10)]
docs += [{"text": f"m{i}", "metadata": {"filename": f"m{i}.md"}, "score": 0.5}
for i in range(3)]
return docs[:limit]
def test_filters_over_fetch_before_truncation():
result = SmartRetriever(_MixedStore()).retrieve("q", top_k=5, include_glob=["*.md"])
assert [c["text"] for c in result.chunks] == ["m0", "m1", "m2"]
def test_invalid_path_filter_fails_closed():
with pytest.raises(ValueError):
SmartRetriever(_MixedStore()).retrieve("q", path_filter="([unclosed")
def test_store_error_is_not_shared_across_concurrent_calls():
# Hold the good call inside search() until the failing call has fully
# recorded its error, the interleaving that leaked it via instance state.
good_in_search = threading.Event()
bad_done = threading.Event()
class Store:
def search(self, query, **kw):
if query == "bad":
raise RuntimeError("down")
good_in_search.set()
bad_done.wait(2)
return [{"text": "ok", "metadata": {}, "score": 1.0}]
retriever = SmartRetriever(Store())
out = {}
tg = threading.Thread(target=lambda: out.setdefault("good", retriever.retrieve("good")))
tg.start()
assert good_in_search.wait(2)
out["bad"] = retriever.retrieve("bad")
bad_done.set()
tg.join()
assert out["bad"].retrieval_failed is True
assert out["good"].retrieval_failed is False
def test_sub_question_engine_accepts_smart_retriever():
result = SubQuestionEngine().query("what is m?", retriever=SmartRetriever(_MixedStore()))
assert result.sources
@@ -0,0 +1,939 @@
"""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"
class TestStepVariables:
"""step_end records the variables a step set, so a re-run can restore them."""
def test_step_end_records_handler_variables_redacted(self, recorder):
from praisonaiagents.workflows import StepResult
from praisonaiagents.workflows.workflows import AgentFlow
def research(ctx):
return StepResult(output="notes", variables={"customer_id": 42, "api_key": "sk-abcdefghijklmnop1234"})
def draft(ctx):
return StepResult(output="draft")
AgentFlow(name="v", steps=[research, draft]).run("go")
ends = {e.node_id: e.data for e in _of(recorder, "step_end")}
assert ends["step:research"]["variables"]["customer_id"] == 42
assert "sk-abcdefghijklmnop1234" not in str(ends["step:research"]["variables"])
assert "variables" not in ends["step:draft"]
def test_nested_step_end_records_handler_variables(self, recorder):
"""Steps run through _execute_single_step_internal (patterns, hierarchical) record them too."""
from praisonaiagents.workflows import StepResult
from praisonaiagents.workflows.workflows import AgentFlow, If
def research(ctx):
return StepResult(output="notes", variables={"customer_id": 42})
AgentFlow(name="v", steps=[If(condition="1 > 0", then_steps=[research])]).run("go")
ends = {e.node_id: e.data for e in _of(recorder, "step_end")}
variables = next(d["variables"] for n, d in ends.items() if n.endswith("research"))
assert variables == {"customer_id": 42}
def test_cache_hit_step_end_leaves_out_the_output_variable(self, recorder):
"""The step cache stores the output variable (the full output); a cached
step_end must trace only what the handler set, like the cold run."""
from praisonaiagents.workflows import StepResult
from praisonaiagents.workflows.workflows import AgentFlow
def research(ctx):
return StepResult(output="the full step output " * 50, variables={"customer_id": 42})
flow = AgentFlow(name="v", steps=[research], cache=True)
flow.run("go")
flow.run("go") # cache hit
ends = [e.data for e in _of(recorder, "step_end")]
assert [d["status"] for d in ends] == ["completed", "cached"]
for d in ends:
assert d["variables"] == {"customer_id": 42}
assert "the full step output" not in str(d["variables"])
def test_nested_step_records_a_variable_changed_in_place(self, recorder):
"""A list the step appends to (same object) still counts as set by it."""
from praisonaiagents.workflows import StepResult
from praisonaiagents.workflows.workflows import AgentFlow, If
def seed(ctx):
return StepResult(output="s", variables={"items": ["a"]})
def add(ctx):
ctx.variables["items"].append("b")
return StepResult(output="m")
AgentFlow(name="v", steps=[seed, If(condition="1 > 0", then_steps=[add])]).run("go")
ends = {e.node_id: e.data for e in _of(recorder, "step_end")}
variables = next(d.get("variables") for n, d in ends.items() if n.endswith("add"))
assert variables == {"items": ["a", "b"]}
def test_uncopyable_variable_does_not_hide_an_in_place_change(self, recorder):
"""One variable that cannot be deep-copied must not stop the others being compared by value."""
import threading
from praisonaiagents.workflows import StepResult
from praisonaiagents.workflows.workflows import AgentFlow, If
def seed(ctx):
return StepResult(output="s", variables={"items": ["a"]})
def add(ctx):
ctx.variables["items"].append("b")
return StepResult(output="m")
flow = AgentFlow(name="v", steps=[seed, If(condition="1 > 0", then_steps=[add])],
variables={"lock": threading.Lock()})
flow.run("go")
ends = {e.node_id: e.data for e in _of(recorder, "step_end")}
variables = next(d.get("variables") for n, d in ends.items() if n.endswith("add"))
assert variables == {"items": ["a", "b"]}
def test_skipped_or_cached_nested_step_does_not_copy_the_scope(self, recorder, monkeypatch):
"""Only a nested step that really runs snapshots the variables; a cache
hit still records what the handler set."""
from praisonaiagents.workflows import StepResult, Task
from praisonaiagents.workflows import workflows as wf
copies = []
real = wf._isolate_scope
monkeypatch.setattr(wf, "_isolate_scope", lambda v: copies.append(1) or real(v))
def research(ctx):
return StepResult(output="notes", variables={"customer_id": 42})
def never(ctx):
return StepResult(output="x")
steps = [wf.If(condition="1 > 0", then_steps=[
research, Task(name="never", handler=never, should_run=lambda ctx: False),
])]
flow = wf.AgentFlow(name="v", steps=steps, variables={"items": list(range(1000))}, cache=True)
flow.run("go")
assert len(copies) == 1 # research ran; the skipped step copied nothing
copies.clear()
flow.run("go") # research is a cache hit
assert copies == []
cached = [e.data for e in _of(recorder, "step_end") if e.data["status"] == "cached"]
assert [d.get("variables") for d in cached] == [{"customer_id": 42}]
@@ -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,
}
@@ -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"
}
@@ -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"
}
@@ -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"
}
@@ -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"
}
@@ -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)
@@ -0,0 +1,28 @@
"""Route references must not silently discard undefined agents."""
import pytest
from praisonaiagents.workflows.yaml_parser import YAMLWorkflowParser
PREFIX = """name: route-validation
agents:
known:
name: Known
instructions: Greet.
llm: openai/fixture-model
steps:
"""
@pytest.mark.parametrize("target", ["missing", "[known, missing]"])
def test_route_rejects_undefined_target(target):
source = PREFIX + f" - route:\n chosen: {target}\n"
with pytest.raises(ValueError, match="Agent 'missing'.*route 'chosen'"):
YAMLWorkflowParser().parse_string(source)
@pytest.mark.parametrize("target", ["known", "[known]"])
def test_route_keeps_defined_targets(target):
source = PREFIX + f" - route:\n chosen: {target}\n"
flow = YAMLWorkflowParser().parse_string(source)
assert [agent.name for agent in flow.steps[0].routes["chosen"]] == ["Known"]
@@ -9,6 +9,7 @@ from __future__ import annotations
import asyncio
import hashlib
import inspect
import logging
import math
import os
@@ -4636,17 +4637,29 @@ class WebSocketGateway:
except Exception:
pass
async def _call_async(fn: Any) -> Any:
def _kwargs_for(fn: Any) -> Dict[str, Any]:
# Issue #5884: decide the calling convention *before* invoking, so
# an application ``TypeError`` raised inside the entry point is
# never mistaken for a missing ``cancel_token`` and replayed.
if not _kw:
return _kw
try:
result = await fn(content, **_kw)
except TypeError:
if not _kw:
raise
# Entry point predates cancel_token: fall back to the shared
# attribute for this turn (best-effort, non-isolated).
if hasattr(agent, "interrupt_controller"):
agent.interrupt_controller = interrupt
result = await fn(content)
params = inspect.signature(fn).parameters.values()
except (TypeError, ValueError):
return _kw
if any(
p.name == "cancel_token" or p.kind is inspect.Parameter.VAR_KEYWORD
for p in params
):
return _kw
# Entry point predates cancel_token: fall back to the shared
# attribute for this turn (best-effort, non-isolated).
if hasattr(agent, "interrupt_controller"):
agent.interrupt_controller = interrupt
return {}
async def _call_async(fn: Any) -> Any:
result = await fn(content, **_kwargs_for(fn))
# No ``await`` between the turn completing and the snapshot: the
# single-threaded loop cannot run another turn in between.
_snapshot()
@@ -4660,14 +4673,7 @@ class WebSocketGateway:
loop = asyncio.get_running_loop()
def _call_sync() -> Any:
try:
result = agent.chat(content, **_kw)
except TypeError:
if not _kw:
raise
if hasattr(agent, "interrupt_controller"):
agent.interrupt_controller = interrupt
result = agent.chat(content)
result = agent.chat(content, **_kwargs_for(agent.chat))
# Snapshot in *this* worker thread, immediately after the turn
# produced its result, before the thread returns to the pool and a
# concurrent turn can overwrite the agent's shared metrics.
@@ -0,0 +1,80 @@
#!/usr/bin/env python3
"""Issue #5884: ``_dispatch_agent_turn`` must not replay a turn whose entry
point raised an internal ``TypeError``; legacy entry points without
``cancel_token`` are detected from the signature before invoking."""
import sys
from pathlib import Path
import pytest
REPO_ROOT = Path(__file__).resolve().parents[5]
sys.path.insert(0, str(REPO_ROOT / "src" / "praisonai-agents"))
from praisonai_bot.gateway.server import WebSocketGateway
_dispatch = WebSocketGateway._dispatch_agent_turn
_TOKEN = object()
class _FailingSync:
interrupt_controller = None
def __init__(self):
self.calls = []
def chat(self, content, cancel_token=None):
self.calls.append(cancel_token)
raise TypeError("application failure")
class _FailingAsync(_FailingSync):
async def achat(self, content, cancel_token=None):
self.calls.append(cancel_token)
raise TypeError("application failure")
class _LegacySync:
interrupt_controller = None
def __init__(self):
self.calls = 0
def chat(self, content):
self.calls += 1
return f"sync:{content}"
class _LegacyAsync(_LegacySync):
async def achat(self, content):
self.calls += 1
return f"async:{content}"
class _KwargsAsync:
async def achat(self, content, **kwargs):
return kwargs.get("cancel_token")
@pytest.mark.asyncio
@pytest.mark.parametrize("cls", [_FailingSync, _FailingAsync])
async def test_internal_type_error_is_not_replayed(cls):
agent = cls()
with pytest.raises(TypeError, match="application failure"):
await _dispatch(agent, "hi", interrupt=_TOKEN)
assert agent.calls == [_TOKEN]
assert agent.interrupt_controller is None
@pytest.mark.asyncio
@pytest.mark.parametrize("cls,prefix", [(_LegacySync, "sync"), (_LegacyAsync, "async")])
async def test_legacy_signature_uses_shared_controller(cls, prefix):
agent = cls()
assert await _dispatch(agent, "hi", interrupt=_TOKEN) == f"{prefix}:hi"
assert agent.calls == 1
assert agent.interrupt_controller is _TOKEN
@pytest.mark.asyncio
async def test_var_keyword_receives_cancel_token():
assert await _dispatch(_KwargsAsync(), "hi", interrupt=_TOKEN) is _TOKEN
@@ -0,0 +1,53 @@
"""Cross-channel group admission parity via the core primitive (Issue #5908)."""
from types import SimpleNamespace
import pytest
from praisonaiagents.bots.admission import MentionFacts
from praisonai_bot.bots._group_policy import decide_group_admission
def _admit(chat_type="group", explicit=False, reply_to_bot=False, is_command=False, **cfg):
return decide_group_admission(
SimpleNamespace(**cfg),
chat_type=chat_type,
mention=MentionFacts(explicit=explicit, reply_to_bot=reply_to_bot),
is_command=is_command,
)
def test_dm_always_admitted():
assert _admit(chat_type="dm", group_policy="command_only").admit
assert _admit(chat_type="im", group_policy="command_only").admit
def test_mention_only_admits_explicit_and_reply_to_bot():
cfg = dict(group_policy="mention_only", mention_required=True)
assert _admit(explicit=True, **cfg).admit
assert _admit(reply_to_bot=True, **cfg).admit
decision = _admit(**cfg)
assert not decision.admit
assert decision.reason_code == "group_mention_only"
def test_command_only_drops_mentions():
# Gateway sets mention_required=False for command_only; it must not open
# the gate to all group traffic.
cfg = dict(group_policy="command_only", mention_required=False)
assert not _admit(explicit=True, **cfg).admit
assert _admit(is_command=True, **cfg).admit
def test_observe_flags_passive_record():
decision = _admit(group_policy="observe", mention_required=False)
assert not decision.admit and decision.observe
def test_legacy_mention_required_false_responds_all():
assert _admit(group_policy="mention_only", mention_required=False).admit
@pytest.mark.parametrize("policy", [None, "bogus"])
def test_unset_or_unknown_policy_fails_safe(policy):
assert not _admit(group_policy=policy).admit
@@ -553,7 +553,7 @@ class TestRegisterRole:
{"type": "start_session", "goal": "Test"}, conn
)
assert response["code"] == "NO_EXTENSION"
conn.websocket.send_text.assert_not_awaited()
conn.websocket.send_json.assert_not_awaited()
@pytest.mark.asyncio
async def test_ws_no_origin_cli_start_session(self):
@@ -571,8 +571,8 @@ class TestRegisterRole:
{"type": "start_session", "goal": "Test"}, cli
)
assert response["status"] == "running"
ext.websocket.send_text.assert_awaited_once()
cli.websocket.send_text.assert_not_awaited()
ext.websocket.send_json.assert_awaited_once()
cli.websocket.send_json.assert_not_awaited()
if server._sessions:
server._sessions.close()
@@ -889,29 +889,31 @@ class PraisonAI:
except Exception:
pass # Continue even if flow display fails
# Initialize trace variables for cleanup
trace_writer = None
trace_emitter = None
trace_emitter_token = None
# Get save flag for replay trace
save_replay = getattr(args, 'save', False)
# Initialize replay trace writer if --save flag is set
# Replay trace: --save, or ``save_replay: true`` in the YAML. Recording
# lives in the wrapper (praisonai.replay.record); reach it through the
# bridge so a standalone praisonai-code install still runs, untraced.
import contextlib
import uuid
from praisonai_code._wrapper_bridge import optional_wrapper_attr
run_id = f"run-{uuid.uuid4().hex[:12]}"
save_replay = getattr(args, 'save', False)
if not save_replay:
yaml_wants_replay = optional_wrapper_attr('praisonai.replay.yaml_run', 'yaml_wants_replay')
if yaml_wants_replay is not None:
save_replay = yaml_wants_replay(self.agent_file, self.agent_yaml)
replay_stack = contextlib.ExitStack()
recording = None
if save_replay:
try:
from praisonai_code._wrapper_bridge import import_wrapper_module
_mod = import_wrapper_module('praisonai.replay')
ContextTraceWriter = getattr(_mod, 'ContextTraceWriter')
from praisonaiagents.trace.context_events import ContextTraceEmitter, set_context_emitter
trace_writer = ContextTraceWriter(session_id=run_id)
trace_emitter = ContextTraceEmitter(sink=trace_writer, session_id=run_id)
# Set as global emitter so agents can access it
trace_emitter_token = set_context_emitter(trace_emitter)
trace_emitter.session_start({"agents_file": self.agent_file, "run_id": run_id})
record = optional_wrapper_attr('praisonai.replay', 'record')
if record is None:
raise ImportError("praisonai.replay is not installed")
meta = {"agents_file": self.agent_file, "run_id": run_id}
if self.agent_file and os.path.isfile(self.agent_file):
meta["yaml_path"] = os.path.abspath(self.agent_file)
recording = replay_stack.enter_context(record(run_id, meta=meta))
print(f"[cyan]📝 Replay trace enabled: {run_id}[/cyan]")
except ImportError as e:
import logging
@@ -919,8 +921,8 @@ class PraisonAI:
except Exception as e:
import logging
logging.warning(f"Failed to initialize trace writer: {e}")
try:
with replay_stack:
AgentsGenerator = _get_agents_generator()
# Extract CLI configuration for YAML CLI parity
cli_config = self._extract_cli_config_for_yaml()
@@ -934,29 +936,9 @@ class PraisonAI:
) as agents_generator:
result = agents_generator.generate_crew_and_kickoff()
print(result)
# Close trace writer on success
if trace_emitter:
trace_emitter.session_end()
print(f"[cyan]📝 Replay trace saved: {run_id}[/cyan]")
if trace_writer:
trace_writer.close()
# Reset global emitter
if trace_emitter_token:
from praisonaiagents.trace.context_events import reset_context_emitter
reset_context_emitter(trace_emitter_token)
return result
except Exception as e:
# Cleanup trace on error
if trace_emitter:
trace_emitter.session_end()
if trace_writer:
trace_writer.close()
if trace_emitter_token:
from praisonaiagents.trace.context_events import reset_context_emitter
reset_context_emitter(trace_emitter_token)
raise
if recording is not None:
print(f"[cyan]📝 Replay trace saved: {run_id}[/cyan]")
return result
def parse_args(self):
"""
+1 -1
View File
@@ -16,7 +16,7 @@ dependencies = [
"textual>=0.47.0",
"PyYAML>=6.0",
"python-dotenv>=0.19.0",
"litellm>=1.83.14,<2",
"litellm>=1.83.14,<1.106", # 1.106+ imports mcp>=2.2 types; we pin mcp<2
"mcp>=1.20.0,<2",
"pydantic>=2.0.0",
"toml>=0.10.2",
+3 -3
View File
@@ -1389,7 +1389,7 @@ dependencies = [
[package.metadata]
requires-dist = [
{ name = "click", specifier = ">=8.4.2" },
{ name = "litellm", specifier = ">=1.83.14,<2" },
{ name = "litellm", specifier = ">=1.83.14,<1.106" },
{ name = "mcp", specifier = ">=1.20.0,<2" },
{ name = "praisonaiagents", directory = "../praisonai-agents" },
{ name = "pydantic", specifier = ">=2.0.0" },
@@ -1432,8 +1432,8 @@ requires-dist = [
{ name = "fastapi", marker = "extra == 'api'", specifier = ">=0.115.0" },
{ name = "fastapi", marker = "extra == 'mcp'", specifier = ">=0.115.0" },
{ name = "fastapi", marker = "extra == 'os'", specifier = ">=0.115.0" },
{ name = "litellm", marker = "extra == 'llm'", specifier = ">=1.83.14" },
{ name = "litellm", marker = "extra == 'memory'", specifier = ">=1.83.14" },
{ name = "litellm", marker = "extra == 'llm'", specifier = ">=1.83.14,<1.106" },
{ name = "litellm", marker = "extra == 'memory'", specifier = ">=1.83.14,<1.106" },
{ name = "markitdown", extras = ["all"], marker = "extra == 'knowledge'", specifier = ">=0.1.0" },
{ name = "mcp", marker = "extra == 'mcp'", specifier = ">=1.20.0,<2" },
{ name = "mcp", marker = "extra == 'parallel-search'", specifier = ">=1.30.0,<2" },
+4 -4
View File
@@ -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
@@ -219,7 +219,7 @@ TS-only members: `previousResult`?, `onToken`?, `signal`?, `onEvent`?, `options`
### `Agent.chat`
- Python: `src/praisonai-agents/praisonaiagents/agent/chat_mixin.py:3369`
- Python: `src/praisonai-agents/praisonaiagents/agent/chat_mixin.py:3371`
- TypeScript: `src/praisonai-ts/src/agent/simple.ts:3592`
- Counts: 17 python params: 7 exact, 9 camelCase, 1 alias, 0 flattened, 0 missing; 0 mismatches; 0 waived; 2 TS-only of 19
@@ -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:2591`
- 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
+4 -4
View File
@@ -1797,7 +1797,7 @@
}
],
"python": {
"location": "src/praisonai-agents/praisonaiagents/agent/chat_mixin.py:3369",
"location": "src/praisonai-agents/praisonaiagents/agent/chat_mixin.py:3371",
"resolved_class": "ChatMixin"
},
"ts_only": [
@@ -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:2591",
"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": [
@@ -181,6 +181,7 @@ export async function generateObject<T = unknown>(
prompt: options.prompt,
messages: options.messages,
system: options.system,
maxOutputTokens: options.maxTokens,
maxTokens: options.maxTokens,
temperature: options.temperature,
topP: options.topP,
@@ -239,6 +240,7 @@ export async function streamObject<T = unknown>(
prompt: options.prompt,
messages: options.messages,
system: options.system,
maxOutputTokens: options.maxTokens,
maxTokens: options.maxTokens,
temperature: options.temperature,
topP: options.topP,
+2
View File
@@ -257,6 +257,7 @@ export async function generateText(options: GenerateTextOptions): Promise<Genera
system: options.system,
tools: options.tools,
toolChoice: options.toolChoice,
maxOutputTokens: options.maxTokens,
maxTokens: options.maxTokens,
temperature: options.temperature,
topP: options.topP,
@@ -331,6 +332,7 @@ export async function streamText(options: StreamTextOptions): Promise<StreamText
system: options.system,
tools: options.tools,
toolChoice: options.toolChoice,
maxOutputTokens: options.maxTokens,
maxTokens: options.maxTokens,
temperature: options.temperature,
topP: options.topP,
@@ -151,6 +151,9 @@ export class AISDKBackend implements LLMProvider {
const callOptions: Record<string, unknown> = {
model: this.model,
messages: prompt,
// AI SDK >= 5 reads `maxOutputTokens` (it renamed `maxTokens`); `maxTokens`
// stays for the ai@4 peer range. Without it the cap never reached the model.
maxOutputTokens: options.maxTokens || this.config.maxOutputTokens,
maxTokens: options.maxTokens || this.config.maxOutputTokens,
temperature: options.temperature,
topP: options.topP,
@@ -253,6 +256,9 @@ export class AISDKBackend implements LLMProvider {
const callOptions: Record<string, unknown> = {
model: this.model,
messages: prompt,
// AI SDK >= 5 reads `maxOutputTokens` (it renamed `maxTokens`); `maxTokens`
// stays for the ai@4 peer range. Without it the cap never reached the model.
maxOutputTokens: options.maxTokens || this.config.maxOutputTokens,
maxTokens: options.maxTokens || this.config.maxOutputTokens,
temperature: options.temperature,
topP: options.topP,
@@ -372,6 +378,9 @@ export class AISDKBackend implements LLMProvider {
model: this.model,
messages: prompt,
schema: options.schema,
// AI SDK >= 5 reads `maxOutputTokens` (it renamed `maxTokens`); `maxTokens`
// stays for the ai@4 peer range. Without it the cap never reached the model.
maxOutputTokens: options.maxTokens || this.config.maxOutputTokens,
maxTokens: options.maxTokens || this.config.maxOutputTokens,
temperature: options.temperature,
maxRetries: options.maxRetries || this.config.maxRetries,
@@ -0,0 +1,112 @@
/**
* AI SDK >= 5 renamed `maxTokens` to `maxOutputTokens` and ignores the old key.
* Every call into the SDK must send `maxOutputTokens`, or neither SAFE_DEFAULTS'
* cap nor a caller's `maxTokens` ever reaches the model. One test per changed
* call: AISDKBackend.generateText / streamText / generateObject, the
* generateText / streamText helpers in src/ai/generate-text.ts, and the
* generateObject / streamObject helpers in src/ai/generate-object.ts.
*/
const calls: Record<string, any[]> = { generateText: [], streamText: [], generateObject: [], streamObject: [] };
jest.mock('ai', () => ({
generateText: jest.fn(async (opts: any) => {
calls.generateText.push(opts);
return { text: 'ok', toolCalls: [], usage: {}, finishReason: 'stop' };
}),
streamText: jest.fn((opts: any) => {
calls.streamText.push(opts);
return {
textStream: (async function* () { yield 'ok'; })(),
toolCalls: Promise.resolve([]),
usage: undefined,
finishReason: 'stop',
};
}),
generateObject: jest.fn(async (opts: any) => {
calls.generateObject.push(opts);
return { object: { ok: true }, usage: undefined };
}),
streamObject: jest.fn((opts: any) => {
calls.streamObject.push(opts);
return {};
}),
}));
// generate-object.ts resolves a model string through the provider package.
jest.mock('@ai-sdk/openai', () => ({ openai: (id: string) => ({ id }) }), { virtual: true });
import { AISDKBackend } from '../../../../src/llm/providers/ai-sdk/backend';
import { SAFE_DEFAULTS } from '../../../../src/llm/providers/ai-sdk/types';
import { generateText, streamText } from '../../../../src/ai/generate-text';
import { generateObject, streamObject } from '../../../../src/ai/generate-object';
const messages = [{ role: 'user', content: 'hi' }];
function backend(config = {}) {
const b = new AISDKBackend('anthropic/claude-sonnet-4-5', config);
// A model already set makes ensureInitialized() return early: no provider, no key.
(b as any).model = {};
return b;
}
async function drain(stream: AsyncIterable<unknown>) {
for await (const _ of stream) { /* run the SDK call */ }
}
describe('AISDKBackend sends the output-token cap as maxOutputTokens', () => {
beforeEach(() => { for (const k of Object.keys(calls)) calls[k].length = 0; });
it('generateText: the safe default cap', async () => {
await backend().generateText({ messages } as any);
expect(calls.generateText[0].maxOutputTokens).toBe(SAFE_DEFAULTS.maxOutputTokens);
});
it("generateText: the caller's maxTokens, and maxTokens kept for ai@4", async () => {
await backend().generateText({ messages, maxTokens: 256 } as any);
expect(calls.generateText[0].maxOutputTokens).toBe(256);
expect(calls.generateText[0].maxTokens).toBe(256);
});
it('streamText: the safe default cap and a caller override', async () => {
await drain(await backend().streamText({ messages } as any));
await drain(await backend().streamText({ messages, maxTokens: 128 } as any));
expect(calls.streamText[0].maxOutputTokens).toBe(SAFE_DEFAULTS.maxOutputTokens);
expect(calls.streamText[1].maxOutputTokens).toBe(128);
});
it('generateObject: the safe default cap and a caller override', async () => {
await backend().generateObject({ messages, schema: {} } as any);
await backend().generateObject({ messages, schema: {}, maxTokens: 64 } as any);
expect(calls.generateObject[0].maxOutputTokens).toBe(SAFE_DEFAULTS.maxOutputTokens);
expect(calls.generateObject[1].maxOutputTokens).toBe(64);
});
});
describe('ai/generate-text helpers send maxTokens as maxOutputTokens', () => {
beforeEach(() => { for (const k of Object.keys(calls)) calls[k].length = 0; });
it('generateText', async () => {
await generateText({ model: {} as any, prompt: 'hi', maxTokens: 300 } as any);
expect(calls.generateText[0].maxOutputTokens).toBe(300);
});
it('streamText', async () => {
await streamText({ model: {} as any, prompt: 'hi', maxTokens: 400 } as any);
expect(calls.streamText[0].maxOutputTokens).toBe(400);
});
});
describe('ai/generate-object helpers send maxTokens as maxOutputTokens', () => {
beforeEach(() => { for (const k of Object.keys(calls)) calls[k].length = 0; });
it('generateObject', async () => {
await generateObject({ model: 'gpt-4o', schema: {} as any, prompt: 'hi', maxTokens: 500 } as any);
expect(calls.generateObject[0].maxOutputTokens).toBe(500);
});
it('streamObject', async () => {
await streamObject({ model: 'gpt-4o', schema: {} as any, prompt: 'hi', maxTokens: 600 } as any);
expect(calls.streamObject[0].maxOutputTokens).toBe(600);
});
});
@@ -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",
@@ -1980,6 +1980,10 @@ class AgentsGenerator:
parser = YAMLWorkflowParser()
workflow = parser.parse_string(yaml_content)
# ``praisonai replay rerun``: reuse recorded step outputs (no-op otherwise).
from .replay._hook import apply_rerun_seed
apply_rerun_seed(workflow)
# Get input: 'input' is canonical, 'topic' is alias for backward compatibility
input_data = config.get('input', config.get('topic', ''))
@@ -95,7 +95,7 @@ def moderate(
async def amoderate(
input: Union[str, List[str]],
model: str = "text-moderation-latest",
model: str = "omni-moderation-latest",
timeout: float = 600.0,
api_key: Optional[str] = None,
api_base: Optional[str] = None,
+412 -3
View File
@@ -4,6 +4,11 @@ Replay command group for PraisonAI CLI.
Provides context replay functionality:
- replay list: List available traces
- replay context <session_id>: Interactive replay of a trace
- replay flow <run|yaml> --format ascii|json|mermaid: Agent flow / static graph
- replay view [run|yaml] [--live]: Open PraisonAI Replay (aiui) on a run
- replay watch [yaml]: A page with a Run button that draws each run live
- replay export <run> --format html|json|mermaid|otel: Export a run
- replay rerun <run> --from NODE: Re-run from a node, reusing earlier outputs
"""
from typing import Optional
@@ -465,19 +470,42 @@ def replay_delete(
def replay_flow(
session_id: str = typer.Argument(
...,
help="Session ID or path to trace file",
help="Session ID, path to trace file, or a YAML definition (static graph)",
),
format: str = typer.Option(
"ascii",
"--format",
"-f",
help="Output format: ascii, json",
help="Output format: ascii, json, mermaid",
),
):
"""Visualize agent flow from a context trace."""
"""Visualize agent flow from a context trace, or a YAML definition before it runs."""
from ..output.console import get_output_controller
output = get_output_controller()
if format not in ("ascii", "json", "mermaid"):
output.print_error(f"Unknown format: {format} (use ascii, json or mermaid)")
raise typer.Exit(2)
from praisonai.replay.yaml_run import is_yaml_path
if is_yaml_path(session_id):
_print_static_flow(output, session_id, format)
return
if format == "mermaid":
from praisonai.replay.events import load_events
from praisonai.replay.export import to_mermaid
try:
raw_events = load_events(session_id)
except FileNotFoundError:
output.print_error(f"Trace not found: {session_id}")
raise typer.Exit(1)
# Plain echo: Rich would treat mermaid's [brackets] as markup.
typer.echo(to_mermaid(raw_events), nl=False)
return
try:
from praisonai.replay import ContextTraceReader
@@ -507,6 +535,33 @@ def replay_flow(
_print_ascii_flow(output, flow_data, session_id)
def _print_static_flow(output, yaml_file: str, format: str) -> None:
"""Static graph of a YAML definition (nothing is executed)."""
from praisonai.replay.yaml_run import static_graph, static_mermaid
try:
if format == "mermaid":
typer.echo(static_mermaid(yaml_file), nl=False)
return
graph, _flow = static_graph(yaml_file)
except Exception as e:
output.print_error(f"Could not read {yaml_file}: {e}")
raise typer.Exit(1)
if format == "json":
output.print_json(graph)
return
typer.echo(f"\n STATIC GRAPH: {yaml_file} ({graph.get('kind')}, {graph.get('process')})\n")
for node in graph.get("nodes", []):
typer.echo(f" [{node.get('kind')}] {node.get('label')} ({node.get('id')})")
typer.echo("")
for edge in graph.get("edges", []):
label = f" [{edge['label']}]" if edge.get("label") else ""
typer.echo(f" {edge.get('source')} --{edge.get('kind')}--> {edge.get('target')}{label}")
typer.echo("")
def _build_flow_from_events(events) -> dict:
"""Build flow structure from context events."""
agents = {}
@@ -828,3 +883,357 @@ def replay_cleanup(
output.print_success(f"Deleted {deleted} old traces.")
else:
output.print_info("No traces needed cleanup.")
# ---------------------------------------------------------------------------
# PraisonAI Replay: view / export / rerun
# ---------------------------------------------------------------------------
_AIUI_INSTALL_HINT = 'pip install "praisonai[ui]"'
_LOOPBACK_HOSTS = {"127.0.0.1", "localhost", "::1"}
_EXPORT_FORMATS = ("html", "json", "mermaid", "otel")
def _aiui_installed() -> bool:
import importlib.util
try:
return importlib.util.find_spec("praisonaiui") is not None
except (ImportError, ValueError):
return False
def _replay_app_path():
"""The bundled ``praisonai/ui_replay/default_app.py`` preset."""
from pathlib import Path
import praisonai
return Path(praisonai.__file__).resolve().parent / "ui_replay" / "default_app.py"
def _spawn_aiui(app_path, host: str, port: int, env: dict):
"""Start ``aiui run`` for the replay app (falls back to ``python -m praisonaiui.cli``)."""
import shutil
import subprocess
import sys
args = ["run", str(app_path), "--port", str(port), "--host", host]
exe = shutil.which("aiui")
cmd = [exe, *args] if exe else [sys.executable, "-m", "praisonaiui.cli", *args]
return subprocess.Popen(cmd, env=env)
def _open_browser(url: str, host: str, port: int, timeout: float = 20.0) -> None:
"""Open ``url`` once the server accepts connections (in a daemon thread)."""
import socket
import threading
import time
import webbrowser
probe_host = "127.0.0.1" if host in ("0.0.0.0", "::", "") else host
def _wait_then_open():
deadline = time.time() + timeout
while time.time() < deadline:
try:
with socket.create_connection((probe_host, port), timeout=0.5):
break
except OSError:
time.sleep(0.25)
webbrowser.open(url)
threading.Thread(target=_wait_then_open, name="replay-open-browser", daemon=True).start()
def _replay_url(host: str, port: int, run_id: Optional[str], live: bool) -> str:
url_host = "127.0.0.1" if host in ("0.0.0.0", "::", "") else host
if ":" in url_host and not url_host.startswith("["):
url_host = f"[{url_host}]"
url = f"http://{url_host}:{port}/#/replay"
params = []
if run_id:
params.append(f"run={run_id}")
if live:
params.append("live=1")
return url + ("?" + "&".join(params) if params else "")
def _start_yaml_run(yaml_file: str, run_id: str):
"""Run a YAML definition recorded live in a background thread."""
import threading
from praisonai.replay.yaml_run import run_yaml_recorded
def _target():
try:
run_yaml_recorded(yaml_file, run_id=run_id, live=True)
typer.echo(f"Run finished: {run_id}")
except Exception as e: # the viewer keeps serving the partial trace
typer.echo(f"Run failed: {e}", err=True)
thread = threading.Thread(target=_target, name="replay-yaml-run", daemon=True)
thread.start()
return thread
@app.command("view")
def replay_view(
target: Optional[str] = typer.Argument(
None,
help="Run id, trace .jsonl path, or a YAML file to run and watch (default: recent runs)",
),
live: bool = typer.Option(
False,
"--live",
help="Follow the run as it is written (implied for a YAML file)",
),
port: int = typer.Option(8090, "--port", "-p", help="Port for the replay UI"),
host: str = typer.Option("127.0.0.1", "--host", help="Host to bind to"),
):
"""Open PraisonAI Replay (timeline + map) for a run in the browser.
Examples:
praisonai replay view
praisonai replay view run-1a2b3c4d5e6f
praisonai replay view agents.yaml # runs it, recorded, and follows it live
"""
import os
from pathlib import Path
from praisonai.replay.events import SAFE_RUN_ID, resolve_trace_path
from praisonai.replay.recorder import new_run_id
from praisonai.replay.storage import resolve_traces_dir
from praisonai.replay.yaml_run import is_yaml_path
if not _aiui_installed():
typer.echo("ERROR: PraisonAI Replay needs PraisonAI UI (aiui), which is not installed.", err=True)
typer.echo(f"Install with:\n {_AIUI_INSTALL_HINT}", err=True)
raise typer.Exit(1)
if host not in _LOOPBACK_HOSTS:
typer.echo(
f"WARNING: serving replay on {host} makes your traces (prompts, outputs, tool "
"arguments) reachable from the network. Use --host 127.0.0.1 to keep them local.",
err=True,
)
traces_dir = resolve_traces_dir()
run_id: Optional[str] = None
yaml_file: Optional[str] = None
if target and is_yaml_path(target):
yaml_file = target
run_id = new_run_id()
live = True
elif target:
path = resolve_trace_path(target)
if str(target).endswith(".jsonl") or "/" in str(target):
traces_dir = path.parent
run_id = path.stem
if not SAFE_RUN_ID.match(run_id):
typer.echo(f"ERROR: not a valid run id: {run_id}", err=True)
raise typer.Exit(2)
if not path.exists() and not live:
typer.echo(f"ERROR: Trace not found: {target}", err=True)
typer.echo("Use 'praisonai replay list' to see available traces (or --live to wait for it).", err=True)
raise typer.Exit(1)
app_path = _replay_app_path()
if not app_path.exists():
typer.echo(f"ERROR: bundled replay app missing: {app_path}", err=True)
raise typer.Exit(1)
env = os.environ.copy()
env["PRAISONAI_TRACES_DIR"] = str(Path(traces_dir))
url = _replay_url(host, port, run_id, live)
if yaml_file:
typer.echo(f"Running {yaml_file} as {run_id} (recording live to {Path(traces_dir) / (run_id + '.jsonl')})")
_start_yaml_run(yaml_file, run_id)
typer.echo(f"PraisonAI Replay: {url}")
proc = _spawn_aiui(app_path, host, port, env)
_open_browser(url, host, port)
code = _wait_for_server(proc)
if code:
raise typer.Exit(code)
def _wait_for_server(proc) -> int:
"""Wait for the aiui server; never leave it running after this CLI exits."""
import signal
import threading
previous = None
if threading.current_thread() is threading.main_thread():
def _on_term(signum, _frame):
raise SystemExit(128 + signum)
previous = signal.signal(signal.SIGTERM, _on_term)
try:
return proc.wait()
except KeyboardInterrupt:
return 0
finally:
if proc.poll() is None:
proc.terminate()
try:
proc.wait(timeout=5)
except Exception:
proc.kill()
if previous is not None:
signal.signal(signal.SIGTERM, previous)
def _load_raw_events(output, run: str):
from praisonai.replay.events import load_events
try:
events = load_events(run)
except FileNotFoundError:
output.print_error(f"Trace not found: {run}")
output.print_info("Use 'praisonai replay list' to see available traces.")
raise typer.Exit(1)
if not events:
output.print_error(f"No events in trace: {run}")
raise typer.Exit(1)
return events
@app.command("watch")
def replay_watch(
yaml_file: Optional[str] = typer.Argument(None, help="agents.yaml to run from the page (default: just view recorded runs)"),
port: int = typer.Option(8090, "--port", "-p", help="Port for the page"),
host: str = typer.Option("127.0.0.1", "--host", help="Host to bind to"),
no_browser: bool = typer.Option(False, "--no-browser", help="Do not open the browser"),
):
"""Open a page with a Run button: each run is recorded and drawn live, step by step.
Examples:
praisonai replay watch agents.yaml
praisonai replay watch # recorded runs only
"""
from praisonai.replay.live import watch
from praisonai.replay.yaml_run import is_yaml_path
if not _aiui_installed():
typer.echo(f"ERROR: PraisonAI Replay needs PraisonAI UI (aiui). Install with:\n {_AIUI_INSTALL_HINT}", err=True)
raise typer.Exit(1)
if yaml_file and not is_yaml_path(yaml_file):
typer.echo(f"ERROR: not a YAML file: {yaml_file}", err=True)
raise typer.Exit(2)
try:
watch(yaml_file, host=host, port=port, open_browser=not no_browser)
except ImportError as e:
typer.echo(f"ERROR: {e}", err=True)
raise typer.Exit(1)
@app.command("export")
def replay_export(
run: str = typer.Argument(..., help="Run id or path to trace file"),
format: str = typer.Option("html", "--format", "-f", help="html, json, mermaid or otel (OTLP-JSON)"),
output_file: Optional[str] = typer.Option(
None, "--output", "-o", help="Write to FILE (default: stdout; html: ./<run>.html)"
),
no_redact: bool = typer.Option(
False, "--no-redact", help="Export as recorded, without removing secrets/PII"
),
):
"""Export a run: a self-contained HTML replay, JSON, Mermaid or OTLP-JSON.
Everything is redacted by default; --no-redact exports the trace as recorded.
"""
import json
from pathlib import Path
from ..output.console import get_output_controller
from praisonai.replay import export as exporter
from praisonai.replay.events import SAFE_RUN_ID, run_id_of
output = get_output_controller()
if format not in _EXPORT_FORMATS:
output.print_error(f"Unknown format: {format} (use {', '.join(_EXPORT_FORMATS)})")
raise typer.Exit(2)
events = _load_raw_events(output, run)
redact = not no_redact
if format == "html":
text = exporter.to_html(events, redact=redact)
if not output_file:
# The run id comes from the trace: never let it pick a path outside ./
name = run_id_of(events, Path(run).stem)
if not SAFE_RUN_ID.match(name):
name = Path(run).stem if SAFE_RUN_ID.match(Path(run).stem) else "replay"
output_file = f"{name}.html"
elif format == "json":
text = exporter.to_json(events, redact=redact)
elif format == "mermaid":
text = exporter.to_mermaid(events, redact=redact)
else:
text = json.dumps(exporter.to_otlp_json(events, redact=redact), indent=2)
if output_file:
Path(output_file).write_text(text, encoding="utf-8")
typer.echo(f"Wrote {format} export to {output_file}")
else:
typer.echo(text, nl=not text.endswith("\n"))
@app.command("rerun")
def replay_rerun(
run: str = typer.Argument(..., help="Run id or path to the recorded trace"),
from_node: str = typer.Option(..., "--from", help="Node id to re-run from (e.g. task:write, step:draft)"),
file: Optional[str] = typer.Option(
None, "--file", help="YAML definition (default: the one recorded with the run)"
),
):
"""Re-run a recorded run from NODE: earlier nodes reuse their recorded output.
The new run is recorded with workflow_start.data.forked_from = {run_id, node_id}.
"""
from pathlib import Path
from ..output.console import get_output_controller
from praisonai.replay.events import session_meta
from praisonai.replay.recorder import new_run_id
from praisonai.replay.rerun import RerunError, plan_rerun
from praisonai.replay.yaml_run import run_yaml_recorded, yaml_kind
output = get_output_controller()
events = _load_raw_events(output, run)
try:
plan = plan_rerun(events, from_node)
except RerunError as e:
output.print_error(str(e))
raise typer.Exit(1)
meta = session_meta(events)
definition = file or meta.get("yaml_path") or meta.get("agents_file") or meta.get("workflow")
if not definition:
output.print_error("This run did not record its definition; pass it with --file agents.yaml")
raise typer.Exit(1)
if not Path(definition).expanduser().is_file():
output.print_error(f"Definition not found: {definition} (pass --file)")
raise typer.Exit(1)
if yaml_kind(definition) != plan.kind:
output.print_error(
f"{definition} defines a {yaml_kind(definition)} but {from_node} is a {plan.kind} node; "
"pass the definition that recorded this run with --file."
)
raise typer.Exit(1)
seeded = []
new_id = new_run_id()
reused = ", ".join(plan.outputs) or "nothing"
typer.echo(f"Re-running {plan.source_run_id} from {from_node} as {new_id} (reusing: {reused})")
try:
rec, result = run_yaml_recorded(definition, run_id=new_id, plan=plan, on_seeded=seeded.append)
except Exception as e:
output.print_error(f"Re-run failed: {e}")
raise typer.Exit(1)
if not seeded:
output.print_warning("The definition was run, but the recorded outputs could not be applied to it.")
typer.echo(str(result) if result is not None else "")
typer.echo(f"Recorded: {rec.run_id} ({rec.path})")
@@ -329,11 +329,16 @@ def _run_yaml_workflow(self, yaml_file: str, action_args: list, variables: dict
variables: Workflow variables
args: Parsed command line arguments
"""
# Initialize trace variables for cleanup
trace_writer = None
trace_emitter = None
trace_emitter_token = None
# Replay trace (--save): praisonai.replay.record() owns writer + emitter setup.
import contextlib
replay_stack = contextlib.ExitStack()
recording = None
def _finish_replay(failed: bool) -> None:
if recording is not None and failed:
recording.status = "failed"
replay_stack.close()
try:
from praisonaiagents.workflows import WorkflowManager
from rich import print
@@ -362,22 +367,22 @@ def _run_yaml_workflow(self, yaml_file: str, action_args: list, variables: dict
# Get save flag for replay trace
save_replay = '--save' in action_args or '-s' in action_args or (getattr(args, 'save', False) if args else False)
if not save_replay:
# ``save_replay: true`` in the YAML records every run, like --save.
from praisonai.replay.yaml_run import yaml_wants_replay
save_replay = yaml_wants_replay(yaml_file)
# Initialize replay trace writer if --save flag is set
run_id = f"run-{uuid.uuid4().hex[:12]}"
if save_replay:
try:
from praisonai_code._wrapper_bridge import import_wrapper_module
_mod = import_wrapper_module('praisonai.replay')
ContextTraceWriter = getattr(_mod, 'ContextTraceWriter')
from praisonaiagents.trace.context_events import ContextTraceEmitter, set_context_emitter
from pathlib import Path
trace_writer = ContextTraceWriter(session_id=run_id)
trace_emitter = ContextTraceEmitter(sink=trace_writer, session_id=run_id)
# Set as global emitter so agents can access it
trace_emitter_token = set_context_emitter(trace_emitter)
trace_emitter.session_start({"workflow": yaml_file, "run_id": run_id})
import os as _os
from praisonai.replay.recorder import record
meta = {"workflow": yaml_file, "run_id": run_id}
if _os.path.isfile(yaml_file):
meta["yaml_path"] = _os.path.abspath(yaml_file)
recording = replay_stack.enter_context(record(run_id, meta=meta))
print(f"[cyan]📝 Replay trace enabled: {run_id}[/cyan]")
except ImportError as e:
import logging
@@ -512,39 +517,19 @@ def _run_yaml_workflow(self, yaml_file: str, action_args: list, variables: dict
print(f"\n[red]❌ Workflow failed: {result.get('error', 'Unknown error')}[/red]")
exit_code = 1
# Close trace writer on completion
if trace_emitter:
trace_emitter.session_end()
# Close the replay trace on completion
_finish_replay(failed=exit_code != 0)
if recording is not None:
print(f"[cyan]📝 Replay trace saved: {run_id}[/cyan]")
if trace_writer:
trace_writer.close()
# Reset global emitter
if trace_emitter_token:
from praisonaiagents.trace.context_events import reset_context_emitter
reset_context_emitter(trace_emitter_token)
return exit_code
except FileNotFoundError:
# Cleanup trace on error
if trace_emitter:
trace_emitter.session_end()
if trace_writer:
trace_writer.close()
if trace_emitter_token:
from praisonaiagents.trace.context_events import reset_context_emitter
reset_context_emitter(trace_emitter_token)
_finish_replay(failed=True)
print(f"[red]ERROR: YAML file not found: {yaml_file}[/red]")
return 1
except Exception as e:
# Cleanup trace on error
if trace_emitter:
trace_emitter.session_end()
if trace_writer:
trace_writer.close()
if trace_emitter_token:
from praisonaiagents.trace.context_events import reset_context_emitter
reset_context_emitter(trace_emitter_token)
_finish_replay(failed=True)
print(f"[red]ERROR: YAML workflow failed: {e}[/red]")
import traceback
traceback.print_exc()
@@ -1443,6 +1443,19 @@
"default": null,
"description": "Task dependency declarations",
"title": "Dependencies"
},
"save_replay": {
"anyOf": [
{
"type": "boolean"
},
{
"type": "null"
}
],
"default": null,
"description": "Record a replay trace for each run (praisonai replay)",
"title": "Save Replay"
}
},
"title": "PraisonAI Agents Configuration",
+4 -1
View File
@@ -402,7 +402,10 @@ class YAMLConfig(BaseModel):
# Deployment / dependency declarations
deploy: Optional[Dict[str, Any]] = Field(default=None, description="Deployment configuration")
dependencies: Optional[Any] = Field(default=None, description="Task dependency declarations")
# PraisonAI Replay: record every run of this file (same as --save)
save_replay: Optional[bool] = Field(default=None, description="Record a replay trace for each run (praisonai replay)")
@model_validator(mode='after')
def validate_config_structure(self):
"""Validate overall configuration structure."""
@@ -992,6 +992,11 @@ class PraisonAIAdapter(BaseFrameworkAdapter):
else:
logger.info(f"No prior state for session: {resume_session}")
# ``praisonai replay rerun``: restore recorded tasks before NODE and
# reset the rest (no-op unless a re-run is in progress).
from praisonai.replay._hook import apply_rerun_seed
apply_rerun_seed(team)
# Bridge the team's aggregate per-step events onto the CLI
# structured output stream so `--output stream-json` on a YAML/team
# run emits the same per-agent NDJSON events as a single-agent run.
+28 -60
View File
@@ -222,19 +222,19 @@ def _run_impl(
"session_id": session_id,
}
# Initialize replay trace writer if --save flag is set
trace_writer = None
# Replay trace (--save): praisonai.replay.record() owns writer + emitter
# setup and always closes it, including on the early-return paths below.
import contextlib
replay_stack = contextlib.ExitStack()
recording = None
trace_emitter = None
trace_emitter_token = None # Token for resetting global emitter
if options.get("save_replay", False):
try:
from praisonai.replay import ContextTraceWriter
from praisonaiagents.trace.context_events import ContextTraceEmitter, set_context_emitter
trace_writer = ContextTraceWriter(session_id=run_id)
trace_emitter = ContextTraceEmitter(sink=trace_writer, session_id=run_id, full_content=True)
# Set as global emitter so agents can access it
trace_emitter_token = set_context_emitter(trace_emitter)
trace_emitter.session_start({"recipe": name, "run_id": run_id})
from praisonai.replay.recorder import record
recording = replay_stack.enter_context(
record(run_id, meta={"recipe": name, "run_id": run_id}, full_content=True)
)
trace_emitter = recording.emitter
print(f"📝 Replay trace enabled: {run_id}")
except ImportError as e:
import logging
@@ -242,12 +242,17 @@ def _run_impl(
except Exception as e:
import logging
logging.warning(f"Failed to initialize trace writer: {e}")
def _replay_failed() -> None:
if recording is not None:
recording.status = "failed"
try:
# Load template
recipe_config = _load_recipe(name, offline=options.get("offline", False))
if recipe_config is None:
_replay_failed()
return RecipeResult(
run_id=run_id,
recipe=name,
@@ -263,6 +268,7 @@ def _run_impl(
dep_result = _check_dependencies(recipe_config)
if not dep_result["all_satisfied"]:
missing = _format_missing_deps(dep_result)
_replay_failed()
return RecipeResult(
run_id=run_id,
recipe=recipe_config.name,
@@ -277,6 +283,7 @@ def _run_impl(
if not options.get("allow_dangerous_tools", False):
policy_error = _check_tool_policy(recipe_config, options)
if policy_error:
_replay_failed()
return RecipeResult(
run_id=run_id,
recipe=recipe_config.name,
@@ -289,11 +296,6 @@ def _run_impl(
# Dry run mode
if options.get("dry_run", False):
# Close trace writer for dry-run
if trace_emitter:
trace_emitter.session_end()
if trace_writer:
trace_writer.close()
return RecipeResult(
run_id=run_id,
recipe=recipe_config.name,
@@ -335,16 +337,10 @@ def _run_impl(
duration = time.time() - start_time
# Close trace writer on success
if trace_emitter:
trace_emitter.session_end()
if trace_writer:
trace_writer.close()
# Close the replay trace on success
replay_stack.close()
if recording is not None:
print(f"📝 Replay trace saved: {run_id}")
# Reset global emitter
if trace_emitter_token:
from praisonaiagents.trace.context_events import reset_context_emitter
reset_context_emitter(trace_emitter_token)
return RecipeResult(
run_id=run_id,
@@ -357,13 +353,7 @@ def _run_impl(
)
except RecipeNotFoundError as e:
if trace_emitter:
trace_emitter.session_end()
if trace_writer:
trace_writer.close()
if trace_emitter_token:
from praisonaiagents.trace.context_events import reset_context_emitter
reset_context_emitter(trace_emitter_token)
_replay_failed()
return RecipeResult(
run_id=run_id,
recipe=name,
@@ -374,13 +364,7 @@ def _run_impl(
trace=trace,
)
except RecipeDependencyError as e:
if trace_emitter:
trace_emitter.session_end()
if trace_writer:
trace_writer.close()
if trace_emitter_token:
from praisonaiagents.trace.context_events import reset_context_emitter
reset_context_emitter(trace_emitter_token)
_replay_failed()
return RecipeResult(
run_id=run_id,
recipe=e.recipe or name,
@@ -391,13 +375,7 @@ def _run_impl(
trace=trace,
)
except RecipePolicyError as e:
if trace_emitter:
trace_emitter.session_end()
if trace_writer:
trace_writer.close()
if trace_emitter_token:
from praisonaiagents.trace.context_events import reset_context_emitter
reset_context_emitter(trace_emitter_token)
_replay_failed()
return RecipeResult(
run_id=run_id,
recipe=e.recipe or name,
@@ -408,13 +386,7 @@ def _run_impl(
trace=trace,
)
except RecipeTimeoutError as e:
if trace_emitter:
trace_emitter.session_end()
if trace_writer:
trace_writer.close()
if trace_emitter_token:
from praisonaiagents.trace.context_events import reset_context_emitter
reset_context_emitter(trace_emitter_token)
_replay_failed()
return RecipeResult(
run_id=run_id,
recipe=name,
@@ -425,13 +397,7 @@ def _run_impl(
trace=trace,
)
except Exception as e:
if trace_emitter:
trace_emitter.session_end()
if trace_writer:
trace_writer.close()
if trace_emitter_token:
from praisonaiagents.trace.context_events import reset_context_emitter
reset_context_emitter(trace_emitter_token)
_replay_failed()
return RecipeResult(
run_id=run_id,
recipe=name,
@@ -441,6 +407,8 @@ def _run_impl(
metrics={"duration_sec": time.time() - start_time},
trace=trace,
)
finally:
replay_stack.close()
def run_stream(
@@ -20,9 +20,29 @@ Usage:
# Interactive replay
player = ReplayPlayer(reader)
player.run()
# Record a run (PraisonAI Replay) and export it
from praisonai.replay import record
from praisonai.replay.events import load_events
from praisonai.replay import export
with record(live=True) as rec:
agent.start("Say hello")
html = export.to_html(load_events(rec.path))
# Watch it live: a page with a Run button (or mount it in your own app)
from praisonai.replay import watch, replay_app
watch(flow) # opens the browser
app.mount("/replay", replay_app(flow)) # FastAPI / Starlette
"""
__all__ = [
"record",
"Recording",
"watch",
"replay_app",
"TraceStepCache",
"load_events",
"ContextTraceWriter",
"ContextTraceReader",
"ReplayPlayer",
@@ -40,6 +60,22 @@ __all__ = [
def __getattr__(name: str):
"""Lazy load module components."""
if name in ("record", "Recording"):
from . import recorder
return getattr(recorder, name)
if name in ("watch", "replay_app"):
from . import live
return getattr(live, name)
if name == "TraceStepCache":
from .rerun import TraceStepCache
return TraceStepCache
if name == "load_events":
from .events import load_events
return load_events
if name == "ContextTraceWriter":
from .writer import ContextTraceWriter
return ContextTraceWriter
+22
View File
@@ -0,0 +1,22 @@
"""Zero-cost hook the YAML runner calls on the team/flow it builds.
Kept apart from :mod:`praisonai.replay.rerun` (which imports the workflow
engine) so a normal run pays one ContextVar lookup and nothing else.
"""
import contextvars
from typing import Any
#: (RerunPlan, on_seeded) while ``praisonai.replay.rerun.rerun_scope`` is active.
ACTIVE_PLAN: contextvars.ContextVar = contextvars.ContextVar("praisonai_replay_rerun_plan", default=None)
def apply_rerun_seed(obj: Any) -> None:
"""Seed ``obj`` from the active re-run plan, if any (no-op otherwise)."""
active = ACTIVE_PLAN.get()
if active is None:
return
from .rerun import _seed_built
plan, on_seeded = active
_seed_built(obj, plan, on_seeded)
+202
View File
@@ -0,0 +1,202 @@
"""Raw trace-event helpers shared by export, rerun and the replay CLI.
Events are handled as plain dicts (one ``ContextEvent.to_dict()`` per JSONL
line) rather than ``ContextEvent`` objects, so that:
- event types this version of ``praisonaiagents`` does not know survive intact
(``ContextEvent.from_dict`` may relabel them);
- schema 1.1 fields (``run_id``, ``parent_seq``, ``node_id``) are read when
present and simply absent on 1.0 traces;
- a partially written last line (a live recording) is skipped, not fatal.
Stdlib only.
"""
from __future__ import annotations
import json
import re
from pathlib import Path
from typing import Any, Dict, Iterable, List, Optional, Tuple, Union
from .storage import get_trace_path, resolve_traces_dir
__all__ = [
"SAFE_RUN_ID",
"SPAN_PAIRS",
"as_dicts",
"etype",
"find_graph",
"load_events",
"pair_spans",
"resolve_trace_path",
"run_id_of",
"session_meta",
]
#: Run ids that are safe to use as file names and URL fragments (contract §5).
SAFE_RUN_ID = re.compile(r"^[A-Za-z0-9._-]+$")
#: start event type -> matching end event type.
SPAN_PAIRS: Dict[str, str] = {
"session_start": "session_end",
"workflow_start": "workflow_end",
"task_start": "task_end",
"step_start": "step_end",
"agent_start": "agent_end",
"llm_request": "llm_response",
"tool_call_start": "tool_call_end",
"approval_requested": "approval_decided",
}
_END_TO_START = {end: start for start, end in SPAN_PAIRS.items()}
EventLike = Union[Dict[str, Any], Any]
def etype(event: Dict[str, Any]) -> str:
"""The raw event type string (an unknown type is kept as-is)."""
data = event.get("data") or {}
raw = data.get("_raw_event_type") if isinstance(data, dict) else None
value = raw or event.get("event_type") or ""
return getattr(value, "value", value)
def as_dicts(events: Iterable[EventLike]) -> List[Dict[str, Any]]:
"""Normalise ``ContextEvent`` objects or dicts to plain dicts."""
out: List[Dict[str, Any]] = []
for ev in events or []:
if isinstance(ev, dict):
out.append(ev)
elif hasattr(ev, "to_dict"):
out.append(ev.to_dict())
return out
def resolve_trace_path(target: Union[str, Path], traces_dir: Optional[Path] = None) -> Path:
"""Map a run id, session id or JSONL path to the trace file path."""
text = str(target)
if text.endswith(".jsonl") or "/" in text or "\\" in text:
return Path(text).expanduser()
return get_trace_path(text, Path(traces_dir) if traces_dir else resolve_traces_dir())
def load_events(target: Union[str, Path], traces_dir: Optional[Path] = None) -> List[Dict[str, Any]]:
"""Read every complete event of a trace as dicts, in file order.
Blank lines and lines that are not valid JSON objects (typically the
still-being-written last line of a live trace) are skipped.
Raises:
FileNotFoundError: if the trace does not exist.
"""
path = resolve_trace_path(target, traces_dir)
if not path.exists():
raise FileNotFoundError(f"Trace not found: {target}")
events: List[Dict[str, Any]] = []
with open(path, "r", encoding="utf-8") as fh:
for line in fh:
line = line.strip()
if not line:
continue
try:
obj = json.loads(line)
except ValueError:
continue
if isinstance(obj, dict):
events.append(obj)
return events
def run_id_of(events: List[Dict[str, Any]], default: str = "run") -> str:
"""The run id of a trace: first ``run_id`` field, else ``session_id``."""
for ev in events:
if ev.get("run_id"):
return str(ev["run_id"])
for ev in events:
if ev.get("session_id"):
return str(ev["session_id"])
return default
def session_meta(events: List[Dict[str, Any]]) -> Dict[str, Any]:
"""The ``data`` of the first ``session_start`` event (recording metadata)."""
for ev in events:
if etype(ev) == "session_start":
data = ev.get("data")
return dict(data) if isinstance(data, dict) else {}
return {}
def find_graph(events: List[Dict[str, Any]]) -> Optional[Dict[str, Any]]:
"""The static graph embedded in the first ``workflow_start`` (schema 1.1)."""
for ev in events:
if etype(ev) == "workflow_start":
graph = (ev.get("data") or {}).get("graph")
if isinstance(graph, dict):
return graph
return None
def _seq(event: Dict[str, Any], index: int) -> int:
seq = event.get("sequence_num")
return seq if isinstance(seq, int) else index
def _same_subject(start: Dict[str, Any], end: Dict[str, Any]) -> bool:
"""Heuristic match for 1.0 traces that carry no ``parent_seq``."""
if start.get("agent_name") != end.get("agent_name") and end.get("agent_name"):
return False
s_data, e_data = start.get("data") or {}, end.get("data") or {}
for key in ("tool_name", "name"):
if s_data.get(key) and e_data.get(key) and s_data[key] != e_data[key]:
return False
if start.get("node_id") and end.get("node_id") and start["node_id"] != end["node_id"]:
return False
return True
def pair_spans(events: List[Dict[str, Any]]) -> Tuple[List[Dict[str, Any]], List[Dict[str, Any]]]:
"""Turn start/end event pairs into spans.
Returns ``(spans, points)``. Each span is
``{"start": ev, "end": ev|None, "seq": int, "parent": int|None, "type": str}``
where ``parent`` is the ``seq`` of the enclosing span. ``points`` are events
that are neither a start nor an end (messages, hand-offs, snapshots...).
Schema 1.1 traces are paired exactly through ``parent_seq``; 1.0 traces fall
back to a stack (innermost open span of the same kind and subject).
"""
spans: Dict[int, Dict[str, Any]] = {}
order: List[int] = []
points: List[Dict[str, Any]] = []
open_stack: List[int] = []
for index, ev in enumerate(events):
kind = etype(ev)
seq = _seq(ev, index)
if kind in SPAN_PAIRS:
parent = ev.get("parent_seq")
if not isinstance(parent, int):
parent = open_stack[-1] if open_stack else None
spans[seq] = {"start": ev, "end": None, "seq": seq, "parent": parent, "type": kind}
order.append(seq)
open_stack.append(seq)
elif kind in _END_TO_START:
start_kind = _END_TO_START[kind]
target = ev.get("parent_seq")
if not (isinstance(target, int) and target in spans and spans[target]["end"] is None):
target = None
for cand in reversed(open_stack):
span = spans[cand]
if span["type"] == start_kind and _same_subject(span["start"], ev):
target = cand
break
if target is None:
points.append(ev)
continue
spans[target]["end"] = ev
if target in open_stack:
open_stack.remove(target)
else:
points.append(ev)
return [spans[s] for s in order], points
+577
View File
@@ -0,0 +1,577 @@
"""Export a replay trace: JSON, Mermaid, OTLP-JSON (OTel GenAI) or one HTML file.
from praisonai.replay.events import load_events
from praisonai.replay import export
events = load_events("run-1a2b3c4d5e6f")
html = export.to_html(events) # self-contained, no CDN
otlp = export.to_otlp_json(events) # dict, OTLP/JSON trace payload
Every exporter redacts by default (secrets by key name, plus secret/PII
patterns in free text) using the helpers in ``praisonaiagents.trace.redact``.
Pass ``redact=False`` to export the trace as recorded.
Stdlib only: no OpenTelemetry dependency, nothing fetched at view time.
"""
from __future__ import annotations
import copy
import hashlib
import html as _html
import json
from pathlib import Path
from typing import Any, Dict, Iterable, List, Optional, Tuple
from .events import as_dicts, etype, find_graph, pair_spans, run_id_of
__all__ = [
"redact_events",
"to_json",
"to_mermaid",
"to_otlp_json",
"to_html",
"aiui_renderer_files",
"SEMCONV_STABILITY",
]
#: The OTel GenAI semantic conventions are still "Development" status.
SEMCONV_STABILITY = "development"
_REDACTED = "[REDACTED]"
# ---------------------------------------------------------------------------
# Redaction
# ---------------------------------------------------------------------------
def _redactors():
"""Key test + free-text scrubber, reused from ``praisonaiagents.trace.redact``."""
from praisonaiagents.trace import redact as _r
should_redact = getattr(_r, "_should_redact", None)
if should_redact is None: # pragma: no cover - older praisonaiagents
def should_redact(key: str) -> bool:
return _r.redact_dict({key: "x"}).get(key) == _r.REDACTED_VALUE
return should_redact, _r.scrub_pii_text
def _redact_value(value: Any, should_redact, scrub) -> Any:
if isinstance(value, str):
return scrub(value)
if isinstance(value, dict):
out = {}
for key, item in value.items():
# Numbers/bools under a "token"-ish key are counters, not secrets
# (``prompt_tokens``, ``max_tokens``); keep them readable.
if isinstance(key, str) and should_redact(key) and not isinstance(item, (int, float, bool, type(None))):
out[key] = _REDACTED
else:
out[key] = _redact_value(item, should_redact, scrub)
return out
if isinstance(value, list):
return [_redact_value(item, should_redact, scrub) for item in value]
return value
def redact_events(events: Iterable[Any]) -> List[Dict[str, Any]]:
"""Copies of ``events`` with secrets removed from ``data`` (envelope kept)."""
should_redact, scrub = _redactors()
out = []
for ev in as_dicts(events):
ev = copy.deepcopy(ev)
if isinstance(ev.get("data"), dict):
ev["data"] = _redact_value(ev["data"], should_redact, scrub)
out.append(ev)
return out
def _prepare(events: Iterable[Any], redact: bool) -> List[Dict[str, Any]]:
return redact_events(events) if redact else as_dicts(events)
# ---------------------------------------------------------------------------
# JSON
# ---------------------------------------------------------------------------
def to_json(events: Iterable[Any], *, redact: bool = True, indent: Optional[int] = 2) -> str:
"""``{"run_id", "events", "graph"}`` as a JSON string."""
evs = _prepare(events, redact)
payload = {
"schema": "praisonai.replay/1",
"run_id": run_id_of(evs),
"graph": find_graph(evs),
"events": evs,
}
return json.dumps(payload, indent=indent, default=str)
# ---------------------------------------------------------------------------
# Mermaid sequence diagram
# ---------------------------------------------------------------------------
_YOU = "You"
_MANAGER = "Manager"
def _mm_text(text: Any, limit: int = 60) -> str:
"""Mermaid-safe message text (``;`` ends a statement, ``#`` starts an entity)."""
s = " ".join(str(text or "").split())
s = s.replace(";", ",").replace("#", "no.").replace("%%", "%")
if len(s) > limit:
s = s[: limit - 1] + "…"
return s
def to_mermaid(events: Iterable[Any], *, redact: bool = True) -> str:
"""A Mermaid ``sequenceDiagram`` of who handed what to whom.
Lanes are "You", "Manager" (when a hierarchical manager appears) and one per
agent in order of first appearance. Messages are task boundaries,
hand-offs/delegations and approvals. A 1.0 trace without task events falls
back to agent start/end.
"""
evs = _prepare(events, redact)
has_tasks = any(etype(e) == "task_start" for e in evs)
agents: List[str] = []
manager = False
def lane(name: Optional[str]) -> str:
nonlocal manager
if not name:
return _YOU
if name == _MANAGER:
manager = True
return _MANAGER
if name not in agents:
agents.append(name)
return name
messages: List[Tuple[str, str, str, str]] = [] # (src, arrow, dst, text)
open_agents: Dict[str, int] = {}
task_starts: Dict[Any, Dict[str, Any]] = {} # seq / node_id -> task_start event
for ev in evs:
kind = etype(ev)
data = ev.get("data") or {}
who = ev.get("agent_name")
if kind == "task_start":
task_starts[ev.get("sequence_num")] = ev
if ev.get("node_id"):
task_starts[ev["node_id"]] = ev
dst = lane(data.get("agent") or who)
messages.append((_YOU, "->>", dst, f"task {data.get('name') or ev.get('node_id') or ''}"))
elif kind == "task_end":
# task_end carries only the outcome; name/agent live on its start.
start = task_starts.get(ev.get("parent_seq")) or task_starts.get(ev.get("node_id")) or {}
s_data = start.get("data") or {}
src = lane(data.get("agent") or s_data.get("agent") or who or start.get("agent_name"))
name = data.get("name") or s_data.get("name") or ev.get("node_id") or ""
messages.append((src, "-->>", _YOU, f"task {name} {data.get('status') or 'done'}"))
elif kind == "agent_handoff":
src = lane(data.get("from_agent") or who)
dst = lane(data.get("to_agent"))
label = data.get("kind") or "handoff"
reason = data.get("reason")
messages.append((src, "->>", dst, f"{label}: {reason}" if reason else label))
elif kind == "approval_requested":
src = lane(who)
messages.append((src, "->>", _YOU, f"approve {data.get('tool') or 'action'}?"))
elif kind == "approval_decided":
dst = lane(who)
verdict = "approved" if data.get("approved") else "denied"
messages.append((_YOU, "-->>", dst, verdict))
elif kind in ("agent_start", "agent_end") and who:
name = lane(who)
if has_tasks:
continue
if kind == "agent_start":
open_agents[name] = open_agents.get(name, 0) + 1
if open_agents[name] == 1:
messages.append((_YOU, "->>", name, "start"))
else:
open_agents[name] = max(0, open_agents.get(name, 0) - 1)
if open_agents[name] == 0:
messages.append((name, "-->>", _YOU, "done"))
ids: Dict[str, str] = {_YOU: _YOU, _MANAGER: _MANAGER}
lines = ["sequenceDiagram", f" participant {_YOU}"]
if manager:
lines.append(f" participant {_MANAGER}")
for i, name in enumerate(agents):
ids[name] = f"A{i}"
lines.append(f" participant A{i} as {_mm_text(name, 40)}")
for src, arrow, dst, text in messages:
lines.append(f" {ids[src]}{arrow}{ids[dst]}: {_mm_text(text)}")
return "\n".join(lines) + "\n"
# ---------------------------------------------------------------------------
# OTLP-JSON (OpenTelemetry GenAI semantic conventions)
# ---------------------------------------------------------------------------
_KIND_INTERNAL = 1
_KIND_CLIENT = 3
_STATUS_UNSET, _STATUS_OK, _STATUS_ERROR = 0, 1, 2
def _trace_id(run_id: str) -> str:
return hashlib.sha256(run_id.encode("utf-8")).hexdigest()[:32]
def _span_id(run_id: str, start_seq: int) -> str:
return hashlib.sha256(f"{run_id}:{start_seq}".encode("utf-8")).hexdigest()[:16]
def _attr(key: str, value: Any) -> Optional[Dict[str, Any]]:
if value is None or value == "":
return None
if isinstance(value, bool):
return {"key": key, "value": {"boolValue": value}}
if isinstance(value, int):
return {"key": key, "value": {"intValue": str(value)}}
if isinstance(value, float):
return {"key": key, "value": {"doubleValue": value}}
if isinstance(value, (list, tuple)) and all(isinstance(v, str) for v in value):
return {"key": key, "value": {"arrayValue": {"values": [{"stringValue": v} for v in value]}}}
if not isinstance(value, str):
value = json.dumps(value, default=str)
return {"key": key, "value": {"stringValue": value}}
def _nanos(ts: Any) -> str:
try:
return str(int(float(ts) * 1_000_000_000))
except (TypeError, ValueError):
return "0"
def _span_shape(span: Dict[str, Any]) -> Tuple[str, int, List[Tuple[str, Any]]]:
"""(name, kind, attributes) for one start/end pair."""
start, end = span["start"], span["end"] or {}
s_data = start.get("data") or {}
e_data = end.get("data") or {}
agent = start.get("agent_name")
kind = span["type"]
attrs: List[Tuple[str, Any]] = [
("gen_ai.conversation.id", start.get("session_id")),
("praisonai.node_id", start.get("node_id")),
("praisonai.event_type", kind),
]
if kind == "workflow_start":
name = s_data.get("name") or "workflow"
attrs += [
("gen_ai.operation.name", "invoke_workflow"),
("gen_ai.workflow.name", name),
("praisonai.workflow.kind", s_data.get("kind")),
("praisonai.workflow.process", s_data.get("process")),
("praisonai.forked_from", s_data.get("forked_from")),
]
return f"invoke_workflow {name}", _KIND_INTERNAL, attrs
if kind == "agent_start":
attrs += [("gen_ai.operation.name", "invoke_agent"), ("gen_ai.agent.name", agent)]
return f"invoke_agent {agent or ''}".strip(), _KIND_INTERNAL, attrs
if kind == "llm_request":
model = e_data.get("model") or s_data.get("model")
attrs += [
("gen_ai.operation.name", "chat"),
("gen_ai.agent.name", agent),
("gen_ai.request.model", model),
("gen_ai.response.model", e_data.get("model")),
("gen_ai.usage.input_tokens", end.get("prompt_tokens") or None),
("gen_ai.usage.output_tokens", end.get("completion_tokens") or None),
("gen_ai.response.finish_reasons", [e_data["finish_reason"]] if e_data.get("finish_reason") else None),
("praisonai.cost_usd", end.get("cost_usd") or None),
]
return f"chat {model or ''}".strip(), _KIND_CLIENT, attrs
if kind == "tool_call_start":
tool = s_data.get("tool_name") or e_data.get("tool_name")
attrs += [
("gen_ai.operation.name", "execute_tool"),
("gen_ai.tool.name", tool),
("gen_ai.agent.name", agent),
("gen_ai.tool.call.arguments", s_data.get("tool_args") or None),
]
return f"execute_tool {tool or ''}".strip(), _KIND_INTERNAL, attrs
if kind == "task_start":
attrs += [
("praisonai.task.name", s_data.get("name")),
("gen_ai.agent.name", s_data.get("agent") or agent),
("praisonai.task.status", e_data.get("status")),
("praisonai.task.retries", e_data.get("retries")),
]
return f"task {s_data.get('name') or ''}".strip(), _KIND_INTERNAL, attrs
if kind == "step_start":
attrs += [
("praisonai.step.name", s_data.get("name")),
("praisonai.step.pattern", s_data.get("pattern")),
("praisonai.step.status", e_data.get("status")),
("praisonai.step.route", e_data.get("route")),
]
return f"step {s_data.get('name') or ''}".strip(), _KIND_INTERNAL, attrs
if kind == "approval_requested":
attrs += [
("praisonai.approval.id", s_data.get("approval_id")),
("praisonai.approval.tool", s_data.get("tool")),
("praisonai.approval.risk", s_data.get("risk")),
("praisonai.approval.approved", e_data.get("approved")),
("praisonai.approval.approver", e_data.get("approver")),
]
return f"approval {s_data.get('tool') or ''}".strip(), _KIND_INTERNAL, attrs
# session_start / anything else
return "praisonai.session", _KIND_INTERNAL, attrs
def _span_status(span: Dict[str, Any]) -> Dict[str, Any]:
end = span["end"]
if end is None:
return {"code": _STATUS_UNSET}
data = end.get("data") or {}
failed = data.get("status") == "failed" or bool(data.get("error"))
if span["type"] == "approval_requested" and data.get("approved") is False:
return {"code": _STATUS_OK, "message": "denied"}
if failed:
return {"code": _STATUS_ERROR, "message": str(data.get("error") or "failed")[:200]}
return {"code": _STATUS_OK}
def to_otlp_json(events: Iterable[Any], *, redact: bool = True, service_name: str = "praisonai") -> Dict[str, Any]:
"""An OTLP/JSON ``ExportTraceServiceRequest`` for the trace.
Ids are deterministic so re-exporting a run gives the same trace:
``traceId = sha256(run_id)[:32]`` and ``spanId = sha256(f"{run_id}:{start_seq}")[:16]``;
``parentSpanId`` comes from ``parent_seq`` (schema 1.1) or nesting (1.0).
Attributes follow the OTel GenAI semantic conventions, which are still in
*Development* status (recorded as ``gen_ai.semconv.stability`` on the scope).
"""
evs = _prepare(events, redact)
run_id = run_id_of(evs)
trace_id = _trace_id(run_id)
spans, _points = pair_spans(evs)
last_ts = max((e.get("timestamp") or 0 for e in evs), default=0)
known = {span["seq"] for span in spans}
out_spans = []
for span in spans:
name, kind, attrs = _span_shape(span)
start, end = span["start"], span["end"]
record: Dict[str, Any] = {
"traceId": trace_id,
"spanId": _span_id(run_id, span["seq"]),
"name": name,
"kind": kind,
"startTimeUnixNano": _nanos(start.get("timestamp")),
"endTimeUnixNano": _nanos((end or {}).get("timestamp", last_ts)),
"attributes": [a for a in (_attr(k, v) for k, v in attrs) if a],
"status": _span_status(span),
}
if span["parent"] is not None and span["parent"] in known:
record["parentSpanId"] = _span_id(run_id, span["parent"])
out_spans.append(record)
try:
from praisonai.version import __version__ as version
except Exception: # pragma: no cover - version module is always present in the wrapper
version = ""
return {
"resourceSpans": [{
"resource": {"attributes": [
_attr("service.name", service_name),
_attr("praisonai.run_id", run_id),
]},
"scopeSpans": [{
"scope": {
"name": "praisonai.replay",
"version": version,
"attributes": [_attr("gen_ai.semconv.stability", SEMCONV_STABILITY)],
},
"spans": out_spans,
}],
}]
}
# ---------------------------------------------------------------------------
# HTML (self-contained)
# ---------------------------------------------------------------------------
_RENDERER_FILES = ("replay-core.js", "replay.js")
def aiui_renderer_files() -> Optional[List[Path]]:
"""The aiui standalone replay renderer, if aiui is installed and ships it."""
import importlib.util
try:
spec = importlib.util.find_spec("praisonaiui")
except (ImportError, ValueError):
return None
if spec is None or not spec.submodule_search_locations:
return None
for root in spec.submodule_search_locations:
views = Path(root) / "templates" / "frontend" / "plugins" / "views"
files = [views / name for name in _RENDERER_FILES]
if all(f.is_file() for f in files):
return files
return None
def _script_safe(text: str) -> str:
"""Make text safe inside an inline ``<script>`` element."""
return text.replace("</", "<\\/").replace("<!--", "<\\!--")
def _json_for_script(value: Any) -> str:
"""JSON that cannot end or re-open the surrounding ``<script>`` (every ``<`` escaped)."""
text = json.dumps(value, default=str)
return text.replace("<", "\\u003c").replace("\u2028", "\\u2028").replace("\u2029", "\\u2029")
_PAGE = """<!doctype html>
<html lang="en">
<head>
<meta charset="utf-8">
<meta name="viewport" content="width=device-width, initial-scale=1">
<meta name="generator" content="PraisonAI Replay">
<title>PraisonAI Replay: {title}</title>
<style>
:root {{ --bg:#ffffff; --fg:#1d1d1f; --muted:#5f6368; --line:#d9dce1; --agent:#8B0000; --tool:#0F7287; --human:#8A5A00; --manager:#3C4A5C; }}
@media (prefers-color-scheme: dark) {{ :root {{ --bg:#16181c; --fg:#ecedef; --muted:#a3a8b0; --line:#30343b; }} }}
body {{ margin:0; background:var(--bg); color:var(--fg); font:14px/1.45 system-ui,-apple-system,Segoe UI,Roboto,sans-serif; }}
header {{ padding:16px; border-bottom:1px solid var(--line); }}
h1 {{ font-size:18px; margin:0 0 4px; }}
.meta {{ color:var(--muted); }}
main {{ padding:16px; }}
table {{ border-collapse:collapse; width:100%; }}
th, td {{ text-align:left; vertical-align:top; padding:6px 8px; border-bottom:1px solid var(--line); }}
th {{ font-weight:600; color:var(--muted); }}
td.seq {{ font-variant-numeric:tabular-nums; color:var(--muted); width:3em; }}
.k {{ display:inline-block; padding:0 6px; border-radius:4px; border:1px solid currentColor; font-size:12px; }}
.k-agent {{ color:var(--agent); }} .k-tool {{ color:var(--tool); }} .k-human {{ color:var(--human); }} .k-manager {{ color:var(--manager); }}
td.indent0 {{ padding-left:8px; }}
details summary {{ cursor:pointer; color:var(--muted); }}
pre {{ white-space:pre-wrap; word-break:break-word; margin:4px 0 0; }}
:focus-visible {{ outline:2px solid var(--tool); outline-offset:2px; }}
@media (prefers-reduced-motion: reduce) {{ * {{ transition:none !important; animation:none !important; }} }}
</style>
</head>
<body>
<div id="replay-root">{fallback}</div>
<script id="replay-data">window.__REPLAY_DATA__ = {data};</script>
{scripts}
</body>
</html>
"""
def _lane_class(ev: Dict[str, Any]) -> Tuple[str, str]:
kind = etype(ev)
if kind.startswith("tool_call"):
return "k-tool", "tool"
if kind.startswith("approval"):
return "k-human", "you"
if ev.get("agent_name") == _MANAGER:
return "k-manager", "manager"
return "k-agent", "agent"
def _describe(ev: Dict[str, Any]) -> str:
kind = etype(ev)
data = ev.get("data") or {}
who = ev.get("agent_name") or ""
if kind == "agent_handoff":
return f"{data.get('from_agent')} → {data.get('to_agent')} ({data.get('kind') or 'handoff'})"
if kind.startswith("tool_call"):
return f"{who} · {data.get('tool_name') or ''}"
if kind in ("task_start", "task_end", "step_start", "step_end", "workflow_start", "workflow_end"):
bits = [str(data.get("name") or ev.get("node_id") or ""), str(data.get("status") or "")]
return " · ".join(b for b in bits if b)
if kind.startswith("approval"):
if kind == "approval_decided":
return "approved" if data.get("approved") else "denied"
return f"approve {data.get('tool') or ''}?"
return who
def _fallback_table(evs: List[Dict[str, Any]], run_id: str) -> str:
spans, _ = pair_spans(evs)
depth: Dict[int, int] = {}
indent_of: Dict[int, int] = {} # id(event) -> nesting level
for span in spans:
parent = span["parent"]
depth[span["seq"]] = depth.get(parent, -1) + 1 if parent is not None else 0
indent_of[id(span["start"])] = depth[span["seq"]]
if span["end"] is not None:
indent_of[id(span["end"])] = depth[span["seq"]]
total_cost = sum(float(e.get("cost_usd") or 0) for e in evs)
total_tokens = sum(int(e.get("prompt_tokens") or 0) + int(e.get("completion_tokens") or 0) for e in evs)
agents = sorted({e.get("agent_name") for e in evs if e.get("agent_name")})
t0 = min((e.get("timestamp") or 0 for e in evs), default=0)
rows = []
for i, ev in enumerate(evs):
seq = ev.get("sequence_num", i)
cls, label = _lane_class(ev)
offset = max(0.0, float(ev.get("timestamp") or t0) - t0)
indent = indent_of.get(id(ev))
if indent is None:
parent = ev.get("parent_seq")
indent = depth.get(parent, -1) + 1 if isinstance(parent, int) else 0
raw = _html.escape(json.dumps(ev.get("data") or {}, indent=2, default=str))
rows.append(
"<tr>"
f"<td class=\"seq\">{_html.escape(str(seq))}</td>"
f"<td>+{offset:.2f}s</td>"
f"<td><span class=\"k {cls}\">{label}</span></td>"
f"<td style=\"padding-left:{8 + 16 * min(indent, 8)}px\"><strong>{_html.escape(etype(ev))}</strong> "
f"{_html.escape(_describe(ev))}"
f"<details><summary>data</summary><pre>{raw}</pre></details></td>"
"</tr>"
)
return (
"<header><h1>PraisonAI Replay</h1>"
f"<div class=\"meta\">Run <code>{_html.escape(run_id)}</code> · {len(evs)} events · "
f"agents: {_html.escape(', '.join(agents) or '-')} · tokens: {total_tokens:,} · cost: ${total_cost:.6f}</div>"
"</header><main><table aria-label=\"Replay timeline\">"
"<thead><tr><th>#</th><th>Time</th><th>Lane</th><th>What happened</th></tr></thead>"
f"<tbody>{''.join(rows)}</tbody></table></main>"
)
def to_html(
events: Iterable[Any],
*,
redact: bool = True,
graph: Optional[Dict[str, Any]] = None,
title: Optional[str] = None,
) -> str:
"""One self-contained HTML file: the events, the graph and a renderer.
Uses aiui's standalone replay renderer (``replay-core.js`` + ``replay.js``)
when aiui is installed and ships it, inlined so the file works offline;
otherwise a built-in table timeline. Nothing is loaded from a CDN.
"""
evs = _prepare(events, redact)
run_id = run_id_of(evs)
graph = graph if graph is not None else find_graph(evs)
data = {"run_id": run_id, "events": evs, "graph": graph}
renderer = aiui_renderer_files()
if renderer:
# aiui's renderer files are ES modules (``export`` statements), so they
# must be inlined as module scripts; a classic <script> rejects them.
scripts = "\n".join(
f"<script type=\"module\" data-src=\"{_html.escape(p.name)}\">\n{_script_safe(p.read_text(encoding='utf-8'))}\n</script>"
for p in renderer
)
else:
scripts = ""
# The table is always present: it is the no-JS view and what a renderer
# replaces when it mounts into #replay-root.
fallback = _fallback_table(evs, run_id)
return _PAGE.format(
title=_html.escape(title or run_id),
fallback=fallback,
data=_json_for_script(data),
scripts=scripts,
)
+138
View File
@@ -0,0 +1,138 @@
"""Watch an agent, team or flow run live: a page with a Run button (PraisonAI Replay).
from praisonai.replay import watch
watch(flow) # opens the browser; each Run is recorded and drawn live
from praisonai.replay import replay_app
app.mount("/replay", replay_app(flow)) # inside your own FastAPI/Starlette app
The page is aiui's Replay view without the dashboard around it
(``praisonaiui.features.replay.create_replay_app``). Requires ``pip install "praisonai[ui]"``.
"""
from __future__ import annotations
import asyncio
import inspect
import threading
import webbrowser
from pathlib import Path
from typing import Any, Callable, Optional
_LOOPBACK_HOSTS = {"127.0.0.1", "localhost", "::1"}
_INSTALL_HINT = 'PraisonAI Replay needs PraisonAI UI: pip install -U "praisonai[ui]"'
def _runner(target: Any) -> Optional[Callable[[str, str], None]]:
"""``run(input, run_id)`` that records *target* under ``run_id``, live."""
if target is None:
return None
from .recorder import record
from .yaml_run import is_yaml_path, run_yaml_recorded
if isinstance(target, (str, Path)):
if not is_yaml_path(str(target)):
raise ValueError(f"not a YAML file: {target}")
path = str(target)
def run_yaml(_text: str, run_id: str) -> None:
run_yaml_recorded(path, run_id=run_id, live=True)
return run_yaml
start = getattr(target, "start", None)
call = start if callable(start) else target
if not callable(call):
raise TypeError("target must be an Agent, AgentTeam, AgentFlow, a YAML path or a callable(input)")
# An AgentTeam skips tasks it already completed and appends each input to
# every task's context; each Run starts them afresh with the original context.
reset = getattr(target, "_reset_task_state_for_item", None)
tasks = getattr(target, "tasks", None)
tasks = list(tasks.values()) if isinstance(tasks, dict) and callable(reset) else []
contexts = [(t, list(t.context) if isinstance(getattr(t, "context", None), list) else None) for t in tasks]
def run(text: str, run_id: str) -> None:
if callable(reset):
reset()
for task, context in contexts:
task.context = list(context) if context is not None else None
with record(run_id, live=True):
out = call(text)
if inspect.isawaitable(out):
asyncio.run(out)
return run
def _title(target: Any) -> str:
if isinstance(target, (str, Path)):
return f"PraisonAI Replay: {Path(target).name}"
name = getattr(target, "name", None) or getattr(target, "__name__", None)
return f"PraisonAI Replay: {name}" if isinstance(name, str) and name else "PraisonAI Replay"
def replay_app(target: Any = None, *, input: str = "", placeholder: str = "", title: Optional[str] = None):
"""The Replay page as an ASGI app; with *target*, it has a Run button.
Args:
target: An Agent, AgentTeam or AgentFlow (``.start(input)`` is called), a
YAML file (run as written; no input box), or any ``callable(input)``.
``None`` gives a read-only view of recorded runs.
input: Text the input box starts with.
placeholder: Placeholder for the input box.
title: Page title (default: from the target's name).
Each Run is recorded with :func:`praisonai.replay.record` under a new run id
and the page follows it live. One run at a time. The same *target* object is
reused: a team's tasks start afresh each run, an Agent keeps its chat history.
Open ``<mount>/`` for the run list, or ``<mount>/?run=<run_id>`` for one run
(handy in an ``<iframe>``).
"""
try:
from praisonaiui.features.replay import create_replay_app
except ImportError as e: # aiui missing or older than the embeddable app
raise ImportError(_INSTALL_HINT) from e
is_yaml = isinstance(target, (str, Path))
return create_replay_app(
_runner(target),
title=title or _title(target),
input=None if is_yaml else input,
placeholder=placeholder,
)
def watch(
target: Any = None,
*,
input: str = "",
placeholder: str = "",
title: Optional[str] = None,
host: str = "127.0.0.1",
port: int = 8090,
open_browser: bool = True,
) -> None:
"""Serve :func:`replay_app` and open it in the browser (blocks until Ctrl+C).
Example::
from praisonai.replay import watch
watch(flow, input="Write about solar sails")
"""
import uvicorn
app = replay_app(target, input=input, placeholder=placeholder, title=title)
if host not in _LOOPBACK_HOSTS:
print(
f"WARNING: serving on {host} lets anyone who can reach it read your runs"
+ (" and start new ones." if target is not None else ".")
)
browser_host = "127.0.0.1" if host in ("0.0.0.0", "::") else host
if ":" in browser_host and not browser_host.startswith("["):
browser_host = f"[{browser_host}]"
url = f"http://{browser_host}:{port}/"
print(f"PraisonAI Replay: {url} (Ctrl+C to stop)")
if open_browser:
threading.Timer(1.0, webbrowser.open, [url]).start()
uvicorn.run(app, host=host, port=port, log_level="warning")
+164
View File
@@ -0,0 +1,164 @@
"""Record a run to a replay trace: ``with praisonai.replay.record(): ...``.
One place owns the writer + emitter + context-var setup that the CLI, the
workflow runner and recipes used to duplicate::
from praisonai.replay import record
with record(live=True) as rec:
Agent(name="a", instructions="...").start("Say hello")
print(rec.run_id, rec.path)
Everything the emitter records is redacted by default (the emitter's own
``redact=True``). ``live=True`` writes every event to disk as it happens
(``buffer_size=1``) so ``praisonai replay view --live`` can follow the run.
"""
from __future__ import annotations
import contextlib
import inspect
import time
import uuid
from dataclasses import dataclass, field
from pathlib import Path
from typing import Any, Dict, Iterator, Optional
__all__ = ["Recording", "record", "new_run_id"]
#: Events buffered before a flush when not recording live (the writer's default).
BATCH_BUFFER_SIZE = 100
def new_run_id() -> str:
"""A fresh run id in the format the CLI has always used (``run-<12 hex>``)."""
return f"run-{uuid.uuid4().hex[:12]}"
@dataclass
class Recording:
"""Handle yielded by :func:`record`."""
run_id: str
path: Path
emitter: Any = field(repr=False, default=None)
writer: Any = field(repr=False, default=None)
live: bool = False
#: Set to ``"failed"`` when the caller handled a failure itself, so
#: ``session_end`` still reports it (an escaping exception does this too).
status: Optional[str] = None
class _ForkTagSink:
"""Sink wrapper that stamps ``forked_from`` on the run's first ``workflow_start``."""
def __init__(self, inner: Any, forked_from: Dict[str, Any]):
self._inner = inner
self._forked_from = dict(forked_from)
self._tagged = False
def emit(self, event: Any) -> None:
if not self._tagged:
kind = getattr(event, "event_type", None)
kind = getattr(kind, "value", kind)
if kind == "workflow_start":
data = dict(getattr(event, "data", None) or {})
data.setdefault("forked_from", dict(self._forked_from))
try:
event.data = data
except Exception:
pass
self._tagged = True
self._inner.emit(event)
def flush(self) -> None:
self._inner.flush()
def close(self) -> None:
self._inner.close()
def _make_emitter(sink: Any, run_id: str, full_content: bool):
from praisonaiagents.trace.context_events import ContextTraceEmitter
kwargs: Dict[str, Any] = {"sink": sink, "session_id": run_id, "full_content": full_content}
# Schema 1.1 emitters stamp ``run_id`` on every event; older ones do not
# take it. The session id equals the run id either way.
try:
if "run_id" in inspect.signature(ContextTraceEmitter.__init__).parameters:
kwargs["run_id"] = run_id
except (TypeError, ValueError):
pass
return ContextTraceEmitter(**kwargs)
def _enter_run_scope(run_id: str):
"""Enter the SDK's run scope when it has one (schema 1.1); else ``None``."""
try:
from praisonaiagents.trace.context_events import _run_scope
except ImportError:
return None
scope = _run_scope(run_id)
scope.__enter__()
return scope
@contextlib.contextmanager
def record(
run_id: Optional[str] = None,
*,
live: bool = False,
meta: Optional[Dict[str, Any]] = None,
full_content: Optional[bool] = None,
_forked_from: Optional[Dict[str, Any]] = None,
) -> Iterator[Recording]:
"""Record everything run inside the block to ``<traces dir>/<run_id>.jsonl``.
Args:
run_id: Trace id and file name (default: ``run-<12 hex>``).
live: Flush every event immediately so a viewer can follow the run.
meta: Extra ``session_start`` data (e.g. ``{"yaml_path": ...}``, which
``praisonai replay rerun`` uses to find the definition again).
full_content: Keep tool results / LLM responses untruncated.
Yields:
:class:`Recording` with ``run_id`` and ``path``.
"""
from praisonaiagents.trace.context_events import reset_context_emitter, set_context_emitter
from .writer import ContextTraceWriter
run_id = run_id or new_run_id()
writer = ContextTraceWriter(
session_id=run_id, buffer_size=1 if live else BATCH_BUFFER_SIZE
)
sink: Any = _ForkTagSink(writer, _forked_from) if _forked_from else writer
emitter = _make_emitter(sink, run_id, bool(full_content))
token = set_context_emitter(emitter)
# Schema 1.1 SDKs stamp ``run_id`` through a run scope; nested teams and
# flows inherit it, so every event in the recording shares this id.
run_scope = _enter_run_scope(run_id)
started = time.time()
rec = Recording(run_id=run_id, path=writer.path, emitter=emitter, writer=writer, live=live)
start_meta: Dict[str, Any] = dict(meta or {})
start_meta.setdefault("run_id", run_id)
if _forked_from:
start_meta.setdefault("forked_from", dict(_forked_from))
status = "completed"
try:
emitter.session_start(start_meta)
yield rec
except BaseException:
status = "failed"
raise
finally:
try:
emitter.session_end(
{"status": rec.status or status, "duration_ms": round((time.time() - started) * 1000, 1)}
)
finally:
writer.close()
if run_scope is not None:
run_scope.__exit__(None, None, None)
reset_context_emitter(token)
+485
View File
@@ -0,0 +1,485 @@
"""Re-run a recorded run from one node: ``praisonai replay rerun RUN --from NODE``.
Everything before NODE is served from the recording instead of being executed
again; NODE and everything after it runs fresh, and the new run is recorded
with ``workflow_start.data.forked_from = {run_id, node_id}``.
- **AgentFlow** -- :class:`TraceStepCache`, a ``StepCacheProtocol`` seeded with
the recorded ``step_end`` outputs. The flow's own cache-hit path replays them,
so variables, ``{{previous_output}}`` and callbacks behave as on a real hit.
- **AgentTeam** -- the team's existing task-state restore: earlier tasks come
back ``completed`` with their recorded output (which the process layer then
skips), tasks from NODE on are reset to ``not started``.
Python::
from praisonai.replay.rerun import rerun
rec, result = rerun(flow, "run-1a2b3c4d5e6f", "step:draft", input="AI")
Needs a schema 1.1 trace (events carry ``node_id``). A 1.0 trace, an unknown
node or an output that was only recorded as a truncated preview fail with
:class:`RerunError` rather than silently re-running with the wrong inputs.
"""
from __future__ import annotations
import contextlib
import os
import re
import threading
from collections import OrderedDict
from typing import Any, Callable, Dict, Iterator, List, Optional, Union
from praisonaiagents.trace.redact import REDACTED_VALUE
from praisonaiagents.workflows.step_cache import StepCacheProtocol
from .events import as_dicts, etype, load_events, pair_spans, run_id_of
__all__ = [
"RerunError",
"RerunPlan",
"TraceStepCache",
"plan_rerun",
"rerun",
"rerun_scope",
"seed_flow",
"seed_team",
]
#: Previews are truncated to this many characters by the emitter (contract §2).
PREVIEW_LIMIT = 500
_REUSABLE = {"completed", "cached"}
#: The emitter's marker for an ``llm_response`` cut to its head and tail.
_SHORTENED = re.compile(r"\n\.\.\.\[[\d,]+ chars, showing first/last portions\]\.\.\.\n")
class RerunError(RuntimeError):
"""The recording cannot support a re-run from the requested node."""
class RerunPlan:
"""What to reuse: recorded outputs (in order) for every node before ``node_id``."""
def __init__(
self,
source_run_id: str,
node_id: str,
kind: str,
outputs: "OrderedDict[str, Any]",
variables: Optional[Dict[str, Dict[str, Any]]] = None,
started_before: Optional[List[str]] = None,
):
self.source_run_id = source_run_id
self.node_id = node_id
self.kind = kind # "flow" | "team"
self.outputs = outputs
# Variable updates a step's handler made, when its step_end recorded them.
self.variables = variables or {}
# Every node of the target's kind that STARTED before it, whatever its
# status (skipped and failed steps too): what the recorded run executed.
self.started_before = list(started_before or [])
@property
def forked_from(self) -> Dict[str, str]:
return {"run_id": self.source_run_id, "node_id": self.node_id}
def __repr__(self) -> str:
return f"RerunPlan(from={self.node_id!r}, reuse={list(self.outputs)!r})"
# ---------------------------------------------------------------------------
# Reading the recording
# ---------------------------------------------------------------------------
def _looks_truncated(preview: str) -> bool:
return len(preview) >= PREVIEW_LIMIT - 3 or preview.endswith(("…", "..."))
def _recover_output(span: Dict[str, Any], events: List[Dict[str, Any]]) -> Any:
"""The full output of a finished task/step, or RerunError if only a stub exists."""
end_data = span["end"].get("data") or {}
if "output" in end_data and end_data["output"] is not None:
return end_data["output"]
preview = end_data.get("output_preview")
if preview is None:
raise RerunError(
f"{span['start'].get('node_id')} has no recorded output (no output/output_preview "
"on its end event); re-record the run with this version of PraisonAI."
)
preview = str(preview)
if not _looks_truncated(preview):
return preview
# The preview was cut at 500 chars: the agent's last LLM answer inside the
# span is the untruncated text when it starts the same way.
head = preview.rstrip("…").rstrip(".")[:100]
lo, hi = span["seq"], span["end"].get("sequence_num", float("inf"))
for ev in reversed(events):
seq = ev.get("sequence_num")
if not isinstance(seq, int) or not (lo < seq < hi) or etype(ev) != "llm_response":
continue
content = (ev.get("data") or {}).get("response_content")
if isinstance(content, str) and content.startswith(head):
if _SHORTENED.search(content):
break # the emitter kept only its head and tail: not the full output
return content
raise RerunError(
f"Only a truncated preview of {span['start'].get('node_id')}'s output was recorded, so "
"re-running after it would feed the next node the wrong input. Re-record with "
"full content (praisonai.replay.record(full_content=True)) or re-run from an earlier node."
)
def plan_rerun(events: Union[str, List[Any]], node_id: str) -> RerunPlan:
"""Work out which recorded outputs to reuse for a re-run from ``node_id``."""
evs = load_events(events) if isinstance(events, (str, os.PathLike)) else as_dicts(events)
if not evs:
raise RerunError("The trace is empty.")
if not any(ev.get("node_id") for ev in evs):
raise RerunError(
"This trace has no node ids (schema 1.0). Re-running from a node needs a run "
"recorded with PraisonAI Replay schema 1.1 or later; record it again."
)
spans, _ = pair_spans(evs)
nodes = [s for s in spans if s["type"] in ("task_start", "step_start") and s["start"].get("node_id")]
target = next((s for s in nodes if s["start"]["node_id"] == node_id), None)
if target is None:
known = ", ".join(dict.fromkeys(s["start"]["node_id"] for s in nodes)) or "none"
raise RerunError(f"Node {node_id!r} is not in this trace. Re-runnable nodes: {known}")
kind = "team" if target["type"] == "task_start" else "flow"
for ev in evs:
if etype(ev) == "workflow_start":
recorded = (ev.get("data") or {}).get("kind")
if recorded in ("team", "flow"):
kind = recorded
break
outputs: "OrderedDict[str, Any]" = OrderedDict()
variables: Dict[str, Dict[str, Any]] = {}
started_before = [
s["start"]["node_id"] for s in nodes
if s["type"] == target["type"] and s["seq"] < target["seq"]
]
for span in nodes:
if span["type"] != target["type"] or span["end"] is None:
continue
end_seq = span["end"].get("sequence_num")
# Reuse only what had FINISHED before NODE started: that excludes NODE,
# everything after it, and containers that enclose it.
if not isinstance(end_seq, int) or end_seq >= target["seq"]:
continue
status = (span["end"].get("data") or {}).get("status", "completed")
if status not in _REUSABLE:
continue
outputs[span["start"]["node_id"]] = _recover_output(span, evs)
recorded_vars = (span["end"].get("data") or {}).get("variables")
if isinstance(recorded_vars, dict):
# A redacted value cannot be restored: the marker would replace the
# real value in the re-run's variables.
hidden = [k for k, v in recorded_vars.items() if REDACTED_VALUE in repr(v)]
if hidden:
raise RerunError(
f"{span['start']['node_id']} set {', '.join(map(repr, hidden))}, which the "
"recording redacted, so it cannot be restored; re-run from that node or earlier."
)
variables[span["start"]["node_id"]] = recorded_vars
return RerunPlan(run_id_of(evs), node_id, kind, outputs, variables, started_before)
# ---------------------------------------------------------------------------
# AgentFlow: a step cache seeded from the recording
# ---------------------------------------------------------------------------
def _step_name(node_id: str) -> str:
"""``step:draft#2`` -> ``draft`` (duplicate suffixes are served in order)."""
name = node_id.split(":", 1)[1] if ":" in node_id else node_id
head, sep, tail = name.rpartition("#")
return head if sep and tail.isdigit() else name
class TraceStepCache(StepCacheProtocol):
"""``StepCacheProtocol`` that serves recorded outputs for the steps before NODE.
It never caches anything new: a miss means "run this step", and ``set`` is
a no-op so later steps always execute. Bind it to the flow
(:func:`seed_flow` does) so it can tell which step is asking -- the flow
marks the step ``running`` just before it consults the cache.
``hits`` lists the step names served from the recording. ``variables``
(node id -> dict) restores the variable updates a handler made, when the
recording kept them.
"""
def __init__(
self,
outputs: "OrderedDict[str, Any]",
variables: Optional[Dict[str, Dict[str, Any]]] = None,
):
variables = variables or {}
self._queues: Dict[str, List[Any]] = {}
for node_id, output in outputs.items():
self._queues.setdefault(_step_name(node_id), []).append((output, variables.get(node_id) or {}))
self._flow: Any = None
self._lock = threading.Lock()
self.hits: List[str] = []
self.misses = 0
def bind(self, flow: Any) -> "TraceStepCache":
self._flow = flow
return self
def _current_step(self) -> Optional[str]:
statuses = getattr(self._flow, "step_statuses", None) or {}
running = [name for name, status in statuses.items() if status == "running" and self._queues.get(name)]
# Two candidates (parallel branches) are ambiguous: run them fresh
# rather than risk serving one step another's output.
return running[0] if len(running) == 1 else None
def _step_obj(self, name: str) -> Any:
for step in getattr(self._flow, "steps", None) or []:
if getattr(step, "name", None) == name or getattr(step, "__name__", None) == name:
return step
return None
def get(self, key: str) -> Optional[Dict[str, Any]]:
with self._lock:
name = self._current_step()
if name is None:
self.misses += 1
return None
output, recorded_vars = self._queues[name].pop(0)
self.hits.append(name)
step = self._step_obj(name)
var = getattr(step, "output_variable", None) or f"{name}_output"
return {
"output": output,
"variables": {**recorded_vars, var: output},
"stop": False,
"step_record": {"step": name, "output": output, "status": "completed", "retries": 0},
}
def set(self, key: str, value: Dict[str, Any]) -> None:
return None
def _flow_step_names(steps: Any) -> set:
"""Names of every step in a flow, including those nested in patterns."""
names: set = set()
stack = list(steps or [])
while stack:
step = stack.pop()
for attr in ("name", "__name__", "_yaml_step_name"): # YAML agent steps run as their step name
value = getattr(step, attr, None)
if isinstance(value, str):
names.add(value)
routes = getattr(step, "routes", None)
if isinstance(routes, dict):
for branch in routes.values():
stack.extend(branch if isinstance(branch, list) else [branch])
for attr in ("steps", "default", "then_steps", "else_steps"):
value = getattr(step, attr, None)
if isinstance(value, list):
stack.extend(value)
inner = getattr(step, "step", None)
if inner is not None:
stack.append(inner)
return names
def _check_flow_order(
steps: List[Any],
recorded: List[str],
started_before: Optional[List[str]] = None,
node_ids: Optional[Dict[int, str]] = None,
) -> None:
"""Reject a sequential flow whose top-level steps are not in the recorded order.
Only checked when the flow runs strictly in sequence: ``next_steps`` /
``branch_condition`` jumps or a repeated node (``#2``) can revisit steps.
Also rejects a step inserted before the target that never ran in the
recording: it would run fresh and replace the reused output the target
expects as its input. Steps are matched by graph node id (``node_ids``,
id(step) -> node id), so unnamed patterns such as ``If`` are covered too.
A target nested in a pattern is checked up to that pattern: the last
top-level step the recording started before the target.
"""
if any(getattr(s, "next_steps", None) or getattr(s, "branch_condition", None) for s in steps):
return
if any("#" in node_id for node_id in recorded):
return
position: Dict[str, int] = {}
for index, step in enumerate(steps):
for attr in ("name", "__name__", "_yaml_step_name"):
value = getattr(step, attr, None)
if isinstance(value, str):
position.setdefault(value, index)
if started_before is not None and node_ids is not None and recorded:
ran = set(started_before)
top = [node_ids.get(id(step)) for step in steps]
if recorded[-1] in top:
anchor = top.index(recorded[-1])
else: # nested target: its container is the last top-level step that started
anchor = max((i for i, node_id in enumerate(top) if node_id in ran), default=-1)
if anchor > 0:
for step, node_id in zip(steps, top[:anchor]):
if node_id not in ran:
raise RerunError(
f"Step {node_id or type(step).__name__!r} comes before {recorded[-1]!r} in "
"this flow but did not run in the recording; re-run against the flow that recorded it."
)
last = -1
for node_id in recorded: # nested steps have no top-level position: skipped
index = position.get(_step_name(node_id))
if index is None:
continue
if index < last:
raise RerunError(
f"Recorded node {node_id!r} runs in a different order in this flow; "
"re-run against the flow that recorded it."
)
last = index
def seed_flow(flow: Any, plan: RerunPlan) -> TraceStepCache:
"""Install a :class:`TraceStepCache` for ``plan`` on an AgentFlow."""
from praisonaiagents.workflows.diagram import _flow_graph
from praisonaiagents.workflows.step_cache import resolve_step_cache
# A different or edited flow would silently run unmatched steps fresh.
names = _flow_step_names(getattr(flow, "steps", None))
recorded = [*plan.outputs, plan.node_id]
for node_id in recorded:
if _step_name(node_id) not in names:
raise RerunError(f"Recorded node {node_id!r} does not match any step in this flow.")
_check_flow_order(
getattr(flow, "steps", None) or [], recorded, plan.started_before, _flow_graph(flow).ids,
)
cache = TraceStepCache(plan.outputs, plan.variables).bind(flow)
flow.cache = cache
# AgentFlow resolves ``cache=`` once at construction; resolve the new one
# the same way so the run uses it.
flow._step_cache = resolve_step_cache(cache)
return cache
# ---------------------------------------------------------------------------
# AgentTeam: restore finished tasks, reset the rest
# ---------------------------------------------------------------------------
def _task_key(team: Any, node_id: str) -> Any:
ref = node_id.split(":", 1)[1] if ":" in node_id else node_id
tasks = getattr(team, "tasks", {}) or {}
for key, task in tasks.items():
if getattr(task, "name", None) == ref:
return key
for index, key in enumerate(tasks):
if str(index) == ref or str(key) == ref:
return key
return None
def seed_team(team: Any, plan: RerunPlan) -> List[Any]:
"""Mark the tasks before NODE completed (with their recorded output), reset the rest.
Returns the task keys restored from the recording.
"""
tasks = getattr(team, "tasks", None)
if not isinstance(tasks, dict):
raise RerunError("Expected an AgentTeam (an object with a tasks dict).")
snapshot: Dict[str, Dict[str, Any]] = {}
restored: List[Any] = []
for node_id, output in plan.outputs.items():
key = _task_key(team, node_id)
if key is None:
raise RerunError(f"Recorded node {node_id!r} does not match any task in this team.")
snapshot[str(key)] = {"status": "completed", "result": output}
restored.append(key)
if _task_key(team, plan.node_id) is None:
raise RerunError(f"Node {plan.node_id!r} does not match any task in this team.")
for key, task in tasks.items():
if key in restored:
continue
task.status = "not started"
task.result = None
if hasattr(task, "retry_count"):
task.retry_count = 0
restore = getattr(team, "_restore_serialised_task_state", None)
if callable(restore):
restore(snapshot) # the team's own checkpoint-restore path (wraps text as TaskOutput)
else: # pragma: no cover - older praisonaiagents
for key in restored:
tasks[key].status = "completed"
tasks[key].result = snapshot[str(key)]["result"]
return restored
# ---------------------------------------------------------------------------
# Seeding objects built elsewhere (YAML runs go through AgentsGenerator)
# ---------------------------------------------------------------------------
def _is_team(obj: Any) -> bool:
return isinstance(getattr(obj, "tasks", None), dict) and hasattr(obj, "run_all_tasks")
def _is_flow(obj: Any) -> bool:
return hasattr(obj, "step_statuses") and hasattr(obj, "steps")
def _seed_built(obj: Any, plan: RerunPlan, on_seeded: Callable[[Any], None]) -> None:
"""Seed a team/flow the YAML runner just built (called by ``_hook``)."""
if plan.kind == "flow" and _is_flow(obj):
on_seeded(seed_flow(obj, plan))
elif plan.kind == "team" and _is_team(obj):
on_seeded(seed_team(obj, plan))
@contextlib.contextmanager
def rerun_scope(plan: RerunPlan, on_seeded: Optional[Callable[[Any], None]] = None) -> Iterator[RerunPlan]:
"""While active, the next team/flow built by the YAML runner is seeded from ``plan``."""
from . import _hook
token = _hook.ACTIVE_PLAN.set((plan, on_seeded or (lambda _x: None)))
try:
yield plan
finally:
_hook.ACTIVE_PLAN.reset(token)
# ---------------------------------------------------------------------------
# Python entry point
# ---------------------------------------------------------------------------
def rerun(
target: Any,
source: Union[str, List[Any]],
node_id: str,
*,
input: Any = None,
run_id: Optional[str] = None,
live: bool = False,
meta: Optional[Dict[str, Any]] = None,
):
"""Re-run ``target`` (an AgentFlow or AgentTeam) from ``node_id`` of ``source``.
Returns ``(recording, result)``; the new trace is at ``recording.path``.
"""
from .recorder import record
plan = plan_rerun(source, node_id)
if _is_flow(target):
if plan.kind != "flow":
raise RerunError(f"{node_id!r} is a team task; pass the AgentTeam that recorded it.")
seed_flow(target, plan)
elif _is_team(target):
if plan.kind != "team":
raise RerunError(f"{node_id!r} is a flow step; pass the AgentFlow that recorded it.")
seed_team(target, plan)
else:
raise RerunError(f"Cannot re-run a {type(target).__name__}; expected an AgentFlow or AgentTeam.")
start_meta = dict(meta or {})
with record(run_id, live=live, meta=start_meta, _forked_from=plan.forked_from) as rec:
result = target.start(input) if input is not None else target.start()
return rec, result
+24 -1
View File
@@ -15,6 +15,29 @@ from praisonaiagents.storage.base import list_json_sessions, cleanup_old_session
DEFAULT_TRACES_DIR = "~/.praison/traces"
#: Overrides the traces directory outright (shared with the aiui replay view).
TRACES_DIR_ENV = "PRAISONAI_TRACES_DIR"
#: State root shared with ``praisonaiagents.trace.last_run`` (``$PRAISON_HOME/traces``).
PRAISON_HOME_ENV = "PRAISON_HOME"
def resolve_traces_dir() -> Path:
"""Where traces live, without creating anything.
Precedence (read live on every call so tests and sandboxes can redirect it):
1. ``PRAISONAI_TRACES_DIR``
2. ``$PRAISON_HOME/traces`` (the same root ``trace/last_run.py`` uses)
3. ``~/.praison/traces``
"""
override = os.environ.get(TRACES_DIR_ENV)
if override:
return Path(os.path.expanduser(override))
home = os.environ.get(PRAISON_HOME_ENV)
if home:
return Path(os.path.expanduser(home)) / "traces"
return Path(os.path.expanduser(DEFAULT_TRACES_DIR))
def get_traces_dir() -> Path:
"""
@@ -23,7 +46,7 @@ def get_traces_dir() -> Path:
Returns:
Path to traces directory (creates if not exists)
"""
traces_dir = Path(os.path.expanduser(DEFAULT_TRACES_DIR))
traces_dir = resolve_traces_dir()
traces_dir.mkdir(parents=True, exist_ok=True)
return traces_dir
+218
View File
@@ -0,0 +1,218 @@
"""YAML definitions for replay: their static graph, and running them recorded.
- :func:`static_graph` / :func:`static_mermaid` picture an ``agents.yaml`` or
workflow YAML *before* it runs (``praisonai replay flow file.yaml``).
- :func:`run_yaml_recorded` runs one under :func:`praisonai.replay.record`,
storing the YAML path in ``session_start`` so ``replay rerun`` can find the
definition again (``praisonai replay view file.yaml``).
"""
from __future__ import annotations
import contextlib
from pathlib import Path
from typing import Any, Callable, Dict, Optional, Tuple, Union
__all__ = [
"is_yaml_path",
"load_flow",
"static_graph",
"static_mermaid",
"run_yaml_recorded",
"yaml_kind",
"yaml_meta",
"yaml_wants_replay",
]
GRAPH_SCHEMA = "praisonai.graph/1"
def is_yaml_path(target: Union[str, Path, None]) -> bool:
"""True for an existing ``.yaml``/``.yml`` file."""
if not target:
return False
path = Path(str(target)).expanduser()
return path.suffix.lower() in (".yaml", ".yml") and path.is_file()
def yaml_meta(yaml_path: Union[str, Path]) -> Dict[str, Any]:
"""``session_start`` metadata that lets ``replay rerun`` find the definition."""
resolved = str(Path(yaml_path).expanduser().resolve())
return {"yaml_path": resolved, "agents_file": resolved}
def yaml_wants_replay(yaml_path: Optional[Union[str, Path]] = None, yaml_text: Optional[str] = None) -> bool:
"""True when the YAML sets ``save_replay: true`` at the top level."""
try:
import yaml
if yaml_text is None:
if not yaml_path or not Path(yaml_path).is_file():
return False
yaml_text = Path(yaml_path).read_text(encoding="utf-8")
data = yaml.safe_load(yaml_text)
except Exception:
return False
return isinstance(data, dict) and data.get("save_replay") is True
def yaml_kind(yaml_path: Union[str, Path]) -> str:
"""``"flow"`` when ``praisonai file.yaml`` would run it as a workflow, else ``"team"``.
Same rule as ``AgentsGenerator._is_workflow_yaml``.
"""
import yaml
data = yaml.safe_load(Path(yaml_path).expanduser().read_text(encoding="utf-8")) or {}
if not isinstance(data, dict):
return "team"
is_flow = (
data.get("process", "sequential") == "workflow"
or ("steps" in data and "workflow" in data)
or data.get("type") in {"job", "hybrid"}
)
return "flow" if is_flow else "team"
def load_flow(yaml_path: Union[str, Path]) -> Any:
"""Parse a YAML definition into an AgentFlow with the SDK's own parser (no run)."""
from praisonaiagents.workflows.yaml_parser import YAMLWorkflowParser
return YAMLWorkflowParser().parse_file(Path(yaml_path).expanduser())
def _core_diagram(name: str) -> Optional[Callable[..., Any]]:
try:
from praisonaiagents.workflows import diagram
except Exception:
return None
fn = getattr(diagram, name, None)
return fn if callable(fn) else None
def _fallback_graph(flow: Any) -> Dict[str, Any]:
"""A minimal ``praisonai.graph/1`` (linear steps) for SDKs without ``to_graph``."""
nodes = [{"id": "start:begin", "kind": "start", "label": "Brief", "parent": None, "meta": {}}]
edges = []
seen: Dict[str, int] = {}
prev = "start:begin"
for step in getattr(flow, "steps", None) or []:
name = getattr(step, "name", None) or getattr(step, "__name__", None) or type(step).__name__
seen[name] = seen.get(name, 0) + 1
node_id = f"step:{name}" if seen[name] == 1 else f"step:{name}#{seen[name]}"
nodes.append({"id": node_id, "kind": "step", "label": str(name), "parent": None, "meta": {}})
edges.append({"id": f"{prev}->{node_id}:next", "source": prev, "target": node_id, "kind": "next", "label": None})
prev = node_id
nodes.append({"id": "end:done", "kind": "end", "label": "Delivered", "parent": None, "meta": {}})
edges.append({"id": f"{prev}->end:done:next", "source": prev, "target": "end:done", "kind": "next", "label": None})
return {
"schema": GRAPH_SCHEMA,
"kind": "flow",
"name": getattr(flow, "name", None) or "workflow",
"process": "workflow",
"nodes": nodes,
"edges": edges,
}
def static_graph(yaml_path: Union[str, Path]) -> Tuple[Dict[str, Any], Any]:
"""``(graph, flow)`` for a YAML definition, without running it.
Uses ``praisonaiagents.workflows.diagram.to_graph`` when this SDK has it
(schema ``praisonai.graph/1``); otherwise a linear fallback graph.
"""
flow = load_flow(yaml_path)
to_graph = _core_diagram("to_graph")
if to_graph is not None:
# A team YAML runs as an AgentTeam (task:/agent: node ids), so build the
# same team the run builds; the flow view would give step: ids that
# never match the run's events.
if yaml_kind(yaml_path) == "team":
team = _build_team_from_yaml(yaml_path)
if team is not None:
try:
graph = to_graph(team)
if isinstance(graph, dict):
return graph, flow
except Exception:
pass
try:
graph = to_graph(flow)
if isinstance(graph, dict):
return graph, flow
except Exception:
pass
return _fallback_graph(flow), flow
def _build_team_from_yaml(yaml_path: Union[str, Path]) -> Any:
"""The ``AgentTeam`` a ``praisonai file.yaml`` run would build, without running it.
Uses the PraisonAI adapter's own builders so task and agent ids match the run.
Returns ``None`` when the YAML can't be built offline (e.g. missing tools).
"""
import yaml
try:
from praisonai.framework_adapters.praisonai_adapter import PraisonAIAdapter
from praisonai_code.llm.config import build_config_list
except Exception:
return None
try:
config = yaml.safe_load(Path(yaml_path).expanduser().read_text(encoding="utf-8")) or {}
adapter = PraisonAIAdapter()
model = adapter._pick_model(build_config_list())
agents, tasks = adapter._build_agents_and_tasks(
config, config.get("topic", ""), {}, None, None, model
)
return adapter._build_team(config, agents, tasks, model)
except Exception:
return None
def static_mermaid(yaml_path: Union[str, Path]) -> str:
"""Mermaid for a YAML definition: ``graph_to_mermaid`` when available, else ``AgentFlow.to_mermaid()``."""
graph, flow = static_graph(yaml_path)
graph_to_mermaid = _core_diagram("graph_to_mermaid")
if graph_to_mermaid is not None and graph.get("schema") == GRAPH_SCHEMA:
try:
return graph_to_mermaid(graph)
except Exception:
pass
return flow.to_mermaid()
def run_yaml_recorded(
yaml_path: Union[str, Path],
*,
run_id: Optional[str] = None,
live: bool = False,
plan: Any = None,
on_seeded: Optional[Callable[[Any], None]] = None,
) -> Tuple[Any, Any]:
"""Run a YAML definition the way ``praisonai file.yaml`` does, recorded.
Args:
plan: a :class:`praisonai.replay.rerun.RerunPlan` to re-run from a node.
on_seeded: called with what the plan restored (TraceStepCache or task keys).
Returns:
``(recording, result)``.
"""
from praisonai.agents_generator import AgentsGenerator
from praisonai_code.llm.config import build_config_list
from .recorder import record
path = Path(yaml_path).expanduser().resolve()
forked = plan.forked_from if plan is not None else None
with record(run_id, live=live, meta=yaml_meta(path), _forked_from=forked) as rec:
scope: Any = contextlib.nullcontext()
if plan is not None:
from .rerun import rerun_scope
scope = rerun_scope(plan, on_seeded)
with scope:
with AgentsGenerator(str(path), "", build_config_list()) as generator:
result = generator.generate_crew_and_kickoff()
return rec, result
@@ -0,0 +1 @@
"""PraisonAI Replay — bundled default app for ``praisonai replay view``."""
@@ -0,0 +1,23 @@
"""PraisonAI Replay — aiui app that opens straight to the replay page.
Loaded by ``praisonai replay view`` -> ``aiui run default_app.py``. The replay
page reads traces from ``PRAISONAI_TRACES_DIR`` (set by the CLI), else
``$PRAISON_HOME/traces``, else ``~/.praison/traces``.
Requires: pip install "praisonai[ui]"
Run:
praisonai replay view # recent runs
praisonai replay view RUN_ID # one run
praisonai replay view agents.yaml --live
"""
from praisonai.integration.host_app import UIPreset, build_ui_app
app = build_ui_app(UIPreset(
title="PraisonAI Replay",
pages=["replay"],
sidebar=False,
theme={"preset": "blue", "dark_mode": False, "radius": "lg"},
welcome="PraisonAI Replay: pick a run to see what happened.",
))
+161
View File
@@ -0,0 +1,161 @@
"""Fixtures for PraisonAI Replay tests: an isolated traces dir and contract-format traces."""
import json
from pathlib import Path
import pytest
@pytest.fixture
def traces_dir(tmp_path, monkeypatch):
"""Point every trace read/write at a temp dir (never the developer's ~/.praison)."""
d = tmp_path / "traces"
monkeypatch.setenv("PRAISONAI_TRACES_DIR", str(d))
monkeypatch.delenv("PRAISON_HOME", raising=False)
return d
def write_trace(directory: Path, run_id: str, events, partial_tail: str = "") -> Path:
directory.mkdir(parents=True, exist_ok=True)
path = directory / f"{run_id}.jsonl"
with open(path, "w", encoding="utf-8") as fh:
for ev in events:
fh.write(json.dumps(ev) + "\n")
if partial_tail:
fh.write(partial_tail) # a live trace's half-written last line
return path
def _ev(seq, event_type, *, agent=None, data=None, parent=None, node=None, run_id="run-team01", t0=1_700_000_000.0, **extra):
ev = {
"schema_version": "1.1",
"event_type": event_type,
"timestamp": t0 + seq * 0.5,
"session_id": run_id,
"agent_name": agent,
"sequence_num": seq,
"messages_count": 0,
"tokens_used": 0,
"tokens_budget": 0,
"data": data or {},
"prompt_tokens": 0,
"completion_tokens": 0,
"cost_usd": 0.0,
"run_id": run_id,
}
if parent is not None:
ev["parent_seq"] = parent
if node is not None:
ev["node_id"] = node
ev.update(extra)
return ev
GRAPH = {
"schema": "praisonai.graph/1", "kind": "team", "name": "Blog team", "process": "sequential",
"nodes": [
{"id": "start:begin", "kind": "start", "label": "Brief", "parent": None, "meta": {}},
{"id": "task:research", "kind": "task", "label": "research", "parent": None, "meta": {}},
{"id": "task:write", "kind": "task", "label": "write", "parent": None, "meta": {}},
{"id": "end:done", "kind": "end", "label": "Delivered", "parent": None, "meta": {}},
],
"edges": [
{"id": "start:begin->task:research:next", "source": "start:begin", "target": "task:research", "kind": "next", "label": None},
{"id": "task:research->task:write:next", "source": "task:research", "target": "task:write", "kind": "next", "label": None},
{"id": "task:write->end:done:next", "source": "task:write", "target": "end:done", "kind": "next", "label": None},
],
}
def team_events(run_id="run-team01", yaml_path=None):
"""A schema 1.1 sequential team run: research -> (handoff) -> write, with an approval."""
meta = {"run_id": run_id}
if yaml_path:
meta["yaml_path"] = str(yaml_path)
e = lambda *a, **k: _ev(*a, run_id=run_id, **k) # noqa: E731
return [
e(0, "session_start", data=meta),
e(1, "workflow_start", parent=0, data={"kind": "team", "name": "Blog team", "process": "sequential", "graph": GRAPH}),
e(2, "task_start", parent=1, node="task:research",
data={"name": "research", "description_preview": "Research AI", "agent": "Researcher"}),
e(3, "agent_start", agent="Researcher", parent=2, node="agent:Researcher"),
e(4, "llm_request", agent="Researcher", parent=3, data={"model": "gpt-4o-mini"}),
e(5, "llm_response", agent="Researcher", parent=4, prompt_tokens=10, completion_tokens=5, cost_usd=0.0001,
data={"model": "gpt-4o-mini", "duration_ms": 12.0, "finish_reason": "stop",
"response_content": "Findings: AI is big"}),
e(6, "tool_call_start", agent="Researcher", parent=3, node="tool:Researcher/web_search",
data={"tool_name": "web_search", "tool_args": {"query": "ai", "api_key": "sk-abcdefghijklmnop1234"}}),
e(7, "tool_call_end", agent="Researcher", parent=6, node="tool:Researcher/web_search",
data={"tool_name": "web_search", "result": "contact [email protected]", "duration_ms": 3.0}),
e(8, "agent_end", agent="Researcher", parent=3, node="agent:Researcher"),
e(9, "task_end", parent=2, node="task:research",
data={"status": "completed", "retries": 0, "duration_ms": 3500, "output_preview": "Findings: AI is big"}),
e(10, "agent_handoff", agent="Researcher", parent=1,
data={"from_agent": "Researcher", "to_agent": "Writer", "reason": "draft it", "context_passed": {}, "kind": "handoff"}),
e(11, "task_start", parent=1, node="task:write",
data={"name": "write", "description_preview": "Write post", "agent": "Writer"}),
e(12, "approval_requested", agent="Writer", parent=11, node="tool:Writer/publish",
data={"approval_id": "ap1", "tool": "publish", "args": {}, "risk": "high", "kind": "tool_call"}),
e(13, "approval_decided", agent="Writer", parent=12, node="tool:Writer/publish",
data={"approval_id": "ap1", "approved": True, "reason": "ok", "approver": "you", "modified": False}),
e(14, "task_end", parent=11, node="task:write",
data={"status": "completed", "retries": 0, "duration_ms": 1000, "output_preview": "Post"}),
e(15, "workflow_end", parent=1, data={"status": "completed", "duration_ms": 7000}),
e(16, "session_end", parent=0, data={"status": "completed"}),
]
def legacy_events(run_id="run-legacy1"):
"""A schema 1.0 trace (no run_id / parent_seq / node_id), as written before Replay."""
base = dict(session_id=run_id, messages_count=0, tokens_used=0, tokens_budget=0,
prompt_tokens=0, completion_tokens=0, cost_usd=0.0, schema_version="1.0")
rows = [
("session_start", None, {"agents_file": "agents.yaml", "run_id": run_id}),
("agent_start", "Researcher", {}),
("llm_request", "Researcher", {"model": "gpt-4o-mini"}),
("llm_response", "Researcher", {"response_tokens": 5, "duration_ms": 10.0, "finish_reason": "stop"}),
("agent_end", "Researcher", {}),
("agent_handoff", "Researcher", {"from_agent": "Researcher", "to_agent": "Writer", "reason": "next", "context_passed": {}}),
("agent_start", "Writer", {}),
("agent_end", "Writer", {}),
("session_end", None, {}),
]
return [
dict(base, event_type=t, agent_name=a, data=d, sequence_num=i, timestamp=1_600_000_000.0 + i)
for i, (t, a, d) in enumerate(rows)
]
def flow_events(run_id="run-flow01", yaml_path=None, step_names=("research", "draft", "polish"), outputs=None):
"""A schema 1.1 AgentFlow run of handler steps."""
outputs = outputs or {n: f"{n} output" for n in step_names}
meta = {"run_id": run_id}
if yaml_path:
meta["yaml_path"] = str(yaml_path)
e = lambda *a, **k: _ev(*a, run_id=run_id, **k) # noqa: E731
events = [e(0, "session_start", data=meta), e(1, "workflow_start", parent=0, data={"kind": "flow", "name": "f", "process": "workflow"})]
seq = 2
for name in step_names:
start = seq
events.append(e(seq, "step_start", parent=1, node=f"step:{name}", data={"name": name}))
seq += 1
events.append(e(seq, "step_end", parent=start, node=f"step:{name}",
data={"status": "completed", "duration_ms": 1, "output_preview": outputs[name]}))
seq += 1
events.append(e(seq, "workflow_end", parent=1, data={"status": "completed", "duration_ms": 9}))
events.append(e(seq + 1, "session_end", parent=0, data={"status": "completed"}))
return events
@pytest.fixture
def fx():
"""Trace builders: ``fx.team_events()``, ``fx.legacy_events()``, ``fx.flow_events()``, ``fx.write_trace()``."""
import types
return types.SimpleNamespace(
team_events=team_events,
legacy_events=legacy_events,
flow_events=flow_events,
write_trace=write_trace,
GRAPH=GRAPH,
)
@@ -0,0 +1,379 @@
"""``praisonai replay`` CLI: flow (mermaid / YAML), export, view, and 1.0 traces still readable."""
import json
import re
import textwrap
import pytest
from typer.testing import CliRunner
def _invoke(*args):
from praisonai.cli.commands.replay import app
return CliRunner().invoke(app, list(args))
FLOW_YAML = textwrap.dedent("""\
name: blog-flow
process: workflow
input: AI
agents:
researcher:
role: Researcher
goal: Research
instructions: You research.
writer:
role: Writer
goal: Write
instructions: You write.
steps:
- name: research
agent: researcher
action: "Research {{input}}"
- name: draft
agent: writer
action: "Write about {{previous_output}}"
""")
@pytest.fixture
def flow_yaml(tmp_path, monkeypatch):
monkeypatch.setenv("OPENAI_API_KEY", "sk-test-not-used")
y = tmp_path / "flow.yaml"
y.write_text(FLOW_YAML)
return y
# ---------------------------------------------------------------------------
# flow
# ---------------------------------------------------------------------------
def test_flow_mermaid_on_yaml_is_the_static_graph(flow_yaml, monkeypatch):
from praisonai.replay import yaml_run
result = _invoke("flow", str(flow_yaml), "--format", "mermaid")
assert result.exit_code == 0, result.output
expected = yaml_run.static_mermaid(flow_yaml)
assert result.output == expected
assert result.output.startswith(("graph", "flowchart"))
assert "researcher" in result.output and "writer" in result.output
def test_flow_mermaid_on_yaml_falls_back_without_core_to_graph(flow_yaml, monkeypatch):
"""Without ``to_graph`` in the SDK, the existing AgentFlow.to_mermaid() is used."""
from praisonai.replay import yaml_run
monkeypatch.setattr(yaml_run, "_core_diagram", lambda name: None)
result = _invoke("flow", str(flow_yaml), "--format", "mermaid")
assert result.exit_code == 0, result.output
assert result.output == yaml_run.load_flow(flow_yaml).to_mermaid()
def test_flow_json_on_yaml_is_a_graph(flow_yaml):
result = _invoke("flow", str(flow_yaml), "--format", "json")
assert result.exit_code == 0, result.output
graph = json.loads(result.output)
assert graph["schema"] == "praisonai.graph/1"
ids = [n["id"] for n in graph["nodes"]]
assert "start:begin" in ids and "end:done" in ids
def test_flow_mermaid_on_a_run_is_a_sequence_diagram(traces_dir, fx):
fx.write_trace(traces_dir, "run-team01", fx.team_events())
result = _invoke("flow", "run-team01", "--format", "mermaid")
assert result.exit_code == 0, result.output
assert result.output.startswith("sequenceDiagram\n")
assert "A0->>A1: handoff: draft it" in result.output
def test_flow_rejects_unknown_format(traces_dir, fx):
fx.write_trace(traces_dir, "run-team01", fx.team_events())
assert _invoke("flow", "run-team01", "--format", "dot").exit_code == 2
# ---------------------------------------------------------------------------
# 1.0 traces keep working with show / flow
# ---------------------------------------------------------------------------
def test_legacy_trace_show_and_flow(traces_dir, fx):
fx.write_trace(traces_dir, "run-legacy1", fx.legacy_events())
shown = _invoke("show", "run-legacy1", "--json")
assert shown.exit_code == 0, shown.output
assert json.loads(shown.output)["total_events"] == 9
flow = _invoke("flow", "run-legacy1")
assert flow.exit_code == 0, flow.output
assert "Researcher" in flow.output and "Writer" in flow.output
flow_json = _invoke("flow", "run-legacy1", "--format", "json")
assert json.loads(flow_json.output)["agent_order"] == ["Researcher", "Writer"]
mermaid = _invoke("flow", "run-legacy1", "--format", "mermaid")
assert mermaid.exit_code == 0 and "A0->>A1: handoff: next" in mermaid.output
def test_schema_1_1_trace_show_and_flow(traces_dir, fx):
"""New event types must not break the existing readers."""
fx.write_trace(traces_dir, "run-team01", fx.team_events())
assert _invoke("show", "run-team01").exit_code == 0
flow = _invoke("flow", "run-team01", "--format", "json")
assert flow.exit_code == 0
assert set(json.loads(flow.output)["agent_order"]) == {"Researcher"}
# ---------------------------------------------------------------------------
# export
# ---------------------------------------------------------------------------
def test_export_html_default_file(traces_dir, fx, tmp_path, monkeypatch):
from praisonai.replay import export
monkeypatch.setattr(export, "aiui_renderer_files", lambda: None)
fx.write_trace(traces_dir, "run-team01", fx.team_events())
monkeypatch.chdir(tmp_path)
result = _invoke("export", "run-team01")
assert result.exit_code == 0, result.output
html = (tmp_path / "run-team01.html").read_text()
assert "window.__REPLAY_DATA__" in html
assert "sk-abcdefghijklmnop1234" not in html
def test_export_html_default_file_ignores_unsafe_trace_run_id(traces_dir, fx, tmp_path, monkeypatch):
from praisonai.replay import export
monkeypatch.setattr(export, "aiui_renderer_files", lambda: None)
evs = fx.team_events()
for ev in evs:
ev["run_id"] = "../escaped"
fx.write_trace(traces_dir, "run-team01", evs)
work = tmp_path / "work"
work.mkdir()
monkeypatch.chdir(work)
result = _invoke("export", "run-team01")
assert result.exit_code == 0, result.output
assert (work / "run-team01.html").exists()
assert not (tmp_path / "escaped.html").exists()
@pytest.mark.parametrize("fmt", ["json", "mermaid", "otel"])
def test_export_formats_to_file(traces_dir, fx, tmp_path, fmt):
fx.write_trace(traces_dir, "run-team01", fx.team_events())
out = tmp_path / f"out.{fmt}"
result = _invoke("export", "run-team01", "--format", fmt, "-o", str(out))
assert result.exit_code == 0, result.output
text = out.read_text()
assert "sk-abcdefghijklmnop1234" not in text
if fmt == "json":
assert json.loads(text)["run_id"] == "run-team01"
elif fmt == "otel":
assert json.loads(text)["resourceSpans"][0]["scopeSpans"][0]["spans"]
else:
assert text.startswith("sequenceDiagram")
def test_export_no_redact(traces_dir, fx):
fx.write_trace(traces_dir, "run-team01", fx.team_events())
result = _invoke("export", "run-team01", "--format", "json", "--no-redact")
assert result.exit_code == 0
assert "sk-abcdefghijklmnop1234" in result.output
def test_export_missing_run(traces_dir):
result = _invoke("export", "run-nope", "--format", "json")
assert result.exit_code == 1
# ---------------------------------------------------------------------------
# view (aiui is never actually started here)
# ---------------------------------------------------------------------------
class _FakeProc:
def wait(self, timeout=None):
return 0
def poll(self):
return 0
def terminate(self):
pass
@pytest.fixture
def fake_aiui(monkeypatch):
from praisonai.cli.commands import replay as replay_cli
calls = {}
def spawn(app_path, host, port, env):
calls["spawn"] = {"app": app_path, "host": host, "port": port, "env": env}
return _FakeProc()
monkeypatch.setattr(replay_cli, "_aiui_installed", lambda: True)
monkeypatch.setattr(replay_cli, "_spawn_aiui", spawn)
monkeypatch.setattr(replay_cli, "_open_browser", lambda url, host, port: calls.setdefault("url", url))
monkeypatch.setattr(replay_cli, "_start_yaml_run", lambda y, run_id: calls.setdefault("yaml_run", (y, run_id)))
return calls
def test_view_help():
result = _invoke("view", "--help")
assert result.exit_code == 0
# Rich forces colour on CI (GITHUB_ACTIONS), which splits flags with ANSI codes.
output = re.sub(r"\x1b\[[0-9;]*m", "", result.output)
for flag in ("--live", "--port", "--host"):
assert flag in output
def test_view_run_opens_replay_page(traces_dir, fx, fake_aiui):
fx.write_trace(traces_dir, "run-team01", fx.team_events())
result = _invoke("view", "run-team01")
assert result.exit_code == 0, result.output
assert fake_aiui["url"] == "http://127.0.0.1:8090/#/replay?run=run-team01"
spawn = fake_aiui["spawn"]
assert spawn["app"].name == "default_app.py" and spawn["app"].parent.name == "ui_replay"
assert spawn["app"].exists()
assert (spawn["host"], spawn["port"]) == ("127.0.0.1", 8090)
assert spawn["env"]["PRAISONAI_TRACES_DIR"] == str(traces_dir)
assert "WARNING" not in result.output
def test_view_jsonl_path_serves_its_directory(tmp_path, fx, fake_aiui, monkeypatch):
monkeypatch.delenv("PRAISONAI_TRACES_DIR", raising=False)
path = fx.write_trace(tmp_path / "elsewhere", "run-team01", fx.team_events())
result = _invoke("view", str(path), "--port", "9001")
assert result.exit_code == 0, result.output
assert fake_aiui["spawn"]["env"]["PRAISONAI_TRACES_DIR"] == str(tmp_path / "elsewhere")
assert fake_aiui["url"] == "http://127.0.0.1:9001/#/replay?run=run-team01"
def test_view_yaml_runs_it_recorded_live(traces_dir, flow_yaml, fake_aiui):
result = _invoke("view", str(flow_yaml))
assert result.exit_code == 0, result.output
yaml_file, run_id = fake_aiui["yaml_run"]
assert yaml_file == str(flow_yaml) and run_id.startswith("run-")
assert fake_aiui["url"] == f"http://127.0.0.1:8090/#/replay?run={run_id}&live=1"
def test_view_non_loopback_host_warns(traces_dir, fx, fake_aiui):
fx.write_trace(traces_dir, "run-team01", fx.team_events())
result = _invoke("view", "run-team01", "--host", "0.0.0.0")
assert result.exit_code == 0
assert "WARNING" in result.output and "0.0.0.0" in result.output
assert fake_aiui["url"].startswith("http://127.0.0.1:8090/")
def test_view_missing_run_fails(traces_dir, fake_aiui):
result = _invoke("view", "run-nope")
assert result.exit_code == 1
assert "spawn" not in fake_aiui
def test_view_live_allows_a_run_that_has_not_started(traces_dir, fake_aiui):
result = _invoke("view", "run-soon", "--live")
assert result.exit_code == 0
assert fake_aiui["url"].endswith("?run=run-soon&live=1")
def test_view_without_aiui_gives_install_hint(traces_dir, monkeypatch):
from praisonai.cli.commands import replay as replay_cli
spawned = []
monkeypatch.setattr(replay_cli, "_aiui_installed", lambda: False)
monkeypatch.setattr(replay_cli, "_spawn_aiui", lambda *a: spawned.append(a))
result = _invoke("view")
assert result.exit_code == 1
assert 'pip install "praisonai[ui]"' in result.output
assert spawned == []
def test_replay_preset_targets_the_replay_page():
import ast
from praisonai.cli.commands.replay import _replay_app_path
tree = ast.parse(_replay_app_path().read_text())
pages = [
kw.value for node in ast.walk(tree) if isinstance(node, ast.Call)
for kw in node.keywords if kw.arg == "pages"
]
assert pages and ast.literal_eval(pages[0]) == ["replay"]
def test_view_stops_the_server_when_the_cli_exits():
from praisonai.cli.commands.replay import _wait_for_server
class Running:
terminated = False
def wait(self, timeout=None):
if timeout is None:
raise KeyboardInterrupt
return 0
def poll(self):
return None if not self.terminated else 0
def terminate(self):
self.terminated = True
proc = Running()
assert _wait_for_server(proc) == 0
assert proc.terminated
def test_team_yaml_static_graph_uses_the_run_ids(tmp_path):
"""A team YAML's map is built from the same AgentTeam a run builds (task:/agent: ids)."""
import pytest
from praisonai.replay import yaml_run
if yaml_run._core_diagram("to_graph") is None:
pytest.skip("SDK has no to_graph (pre schema 1.1)")
path = tmp_path / "team.yaml"
path.write_text(
"framework: praisonai\n"
"topic: t\n"
"roles:\n"
" researcher:\n role: Researcher\n goal: g\n backstory: b\n"
" tasks:\n research:\n description: d\n expected_output: e\n"
" writer:\n role: Writer\n goal: g\n backstory: b\n"
" tasks:\n write:\n description: d\n expected_output: e\n",
encoding="utf-8",
)
graph, _ = yaml_run.static_graph(path)
ids = {n["id"] for n in graph["nodes"]}
assert {"agent:Researcher", "agent:Writer", "task:0", "task:1"} <= ids
assert not any(i.startswith("step:") for i in ids)
# ---------------------------------------------------------------------------
# watch: a page with a Run button (the server is never started here)
# ---------------------------------------------------------------------------
def test_watch_needs_aiui(monkeypatch):
from praisonai.cli.commands import replay as replay_cli
monkeypatch.setattr(replay_cli, "_aiui_installed", lambda: False)
result = _invoke("watch", "agents.yaml")
assert result.exit_code == 1 and "praisonai[ui]" in result.output
def test_watch_rejects_a_non_yaml_file(monkeypatch):
from praisonai.cli.commands import replay as replay_cli
monkeypatch.setattr(replay_cli, "_aiui_installed", lambda: True)
result = _invoke("watch", "notes.txt")
assert result.exit_code == 2 and "not a YAML file" in result.output
def test_watch_serves_the_yaml_with_a_run_button(monkeypatch):
from praisonai.cli.commands import replay as replay_cli
from praisonai.replay import live
calls = []
monkeypatch.setattr(replay_cli, "_aiui_installed", lambda: True)
monkeypatch.setattr(live, "watch", lambda target, **kw: calls.append((target, kw)))
result = _invoke("watch", "agents.yaml", "--port", "8123", "--no-browser")
assert result.exit_code == 0, result.output
assert calls == [("agents.yaml", {"host": "127.0.0.1", "port": 8123, "open_browser": False})]
@@ -0,0 +1,259 @@
"""PraisonAI Replay exports: JSON, Mermaid (golden), OTLP-JSON and self-contained HTML."""
import hashlib
import json
import re
GOLDEN_MERMAID = """sequenceDiagram
participant You
participant A0 as Researcher
participant A1 as Writer
You->>A0: task research
A0-->>You: task research completed
A0->>A1: handoff: draft it
You->>A1: task write
A1->>You: approve publish?
You-->>A1: approved
A1-->>You: task write completed
"""
def test_load_events_skips_partial_last_line(traces_dir, fx):
from praisonai.replay.events import load_events
fx.write_trace(traces_dir, "run-team01", fx.team_events(), partial_tail='{"event_type": "task_st')
events = load_events("run-team01")
assert len(events) == len(fx.team_events())
assert events[-1]["event_type"] == "session_end"
def test_load_events_keeps_unknown_event_types(traces_dir, fx):
from praisonai.replay.events import load_events
evs = fx.team_events() + [{"event_type": "future_thing", "sequence_num": 99, "data": {}}]
path = fx.write_trace(traces_dir, "run-x", evs)
assert load_events(str(path))[-1]["event_type"] == "future_thing"
def test_mermaid_golden(fx):
from praisonai.replay.export import to_mermaid
assert to_mermaid(fx.team_events()) == GOLDEN_MERMAID
def test_mermaid_manager_lane_and_delegation(fx):
from praisonai.replay.export import to_mermaid
evs = [
{"event_type": "agent_handoff", "agent_name": "Manager", "sequence_num": 0,
"data": {"from_agent": "Manager", "to_agent": "Writer", "reason": "write; now #1", "kind": "delegate"}},
]
out = to_mermaid(evs)
assert " participant Manager\n" in out
# ';' and '#' would break the diagram
assert " Manager->>A0: delegate: write, now no.1\n" in out
def test_mermaid_legacy_trace_uses_agent_start_end(fx):
from praisonai.replay.export import to_mermaid
out = to_mermaid(fx.legacy_events())
assert out.splitlines() == [
"sequenceDiagram",
" participant You",
" participant A0 as Researcher",
" participant A1 as Writer",
" You->>A0: start",
" A0-->>You: done",
" A0->>A1: handoff: next",
" You->>A1: start",
" A1-->>You: done",
]
def _spans_by_name(otlp):
spans = otlp["resourceSpans"][0]["scopeSpans"][0]["spans"]
return spans, {s["name"]: s for s in spans}
def _attrs(span):
out = {}
for a in span["attributes"]:
v = a["value"]
out[a["key"]] = next(iter(v.values()))
return out
def test_otlp_ids_parents_and_attributes(fx):
from praisonai.replay.export import to_otlp_json
otlp = to_otlp_json(fx.team_events())
scope = otlp["resourceSpans"][0]["scopeSpans"][0]["scope"]
assert scope["name"] == "praisonai.replay"
assert {"key": "gen_ai.semconv.stability", "value": {"stringValue": "development"}} in scope["attributes"]
spans, by_name = _spans_by_name(otlp)
trace_id = hashlib.sha256(b"run-team01").hexdigest()[:32]
assert {s["traceId"] for s in spans} == {trace_id}
def sid(seq):
return hashlib.sha256(f"run-team01:{seq}".encode()).hexdigest()[:16]
wf = by_name["invoke_workflow Blog team"]
assert wf["spanId"] == sid(1) and wf["parentSpanId"] == sid(0)
assert _attrs(wf)["gen_ai.operation.name"] == "invoke_workflow"
task = by_name["task research"]
assert task["spanId"] == sid(2) and task["parentSpanId"] == sid(1) and task["kind"] == 1
agent = by_name["invoke_agent Researcher"]
assert agent["parentSpanId"] == sid(2)
assert _attrs(agent)["gen_ai.agent.name"] == "Researcher"
chat = by_name["chat gpt-4o-mini"]
assert chat["kind"] == 3 and chat["parentSpanId"] == sid(3)
a = _attrs(chat)
assert a["gen_ai.operation.name"] == "chat"
assert a["gen_ai.request.model"] == "gpt-4o-mini"
assert a["gen_ai.usage.input_tokens"] == "10" and a["gen_ai.usage.output_tokens"] == "5"
tool = by_name["execute_tool web_search"]
assert tool["parentSpanId"] == sid(3)
assert _attrs(tool)["gen_ai.tool.name"] == "web_search"
# Redacted by default: the api_key in the tool arguments never leaves.
assert "sk-abcdefghijklmnop1234" not in json.dumps(otlp)
approval = by_name["approval publish"]
assert approval["parentSpanId"] == sid(11)
assert _attrs(approval)["praisonai.approval.approved"] is True
assert int(chat["endTimeUnixNano"]) > int(chat["startTimeUnixNano"])
assert all(s["status"]["code"] == 1 for s in spans)
def test_otlp_legacy_trace_pairs_by_nesting(fx):
from praisonai.replay.export import to_otlp_json
spans, by_name = _spans_by_name(to_otlp_json(fx.legacy_events()))
session = by_name["praisonai.session"]
assert "parentSpanId" not in session
agent_spans = [s for s in spans if s["name"].startswith("invoke_agent")]
assert len(agent_spans) == 2
assert all(s["parentSpanId"] == session["spanId"] for s in agent_spans)
chat = by_name["chat gpt-4o-mini"]
assert chat["parentSpanId"] == by_name["invoke_agent Researcher"]["spanId"]
def test_otlp_unfinished_span_is_unset(fx):
from praisonai.replay.export import to_otlp_json
evs = fx.team_events()[:5] # live trace cut mid-run
spans, by_name = _spans_by_name(to_otlp_json(evs))
assert by_name["invoke_agent Researcher"]["status"] == {"code": 0}
_EXTERNAL_SRC = re.compile(r"<(script|link|img|iframe)[^>]+(src|href)\s*=\s*[\"']?(https?:)?//", re.I)
def test_html_is_self_contained_and_redacted(fx, monkeypatch):
from praisonai.replay import export
monkeypatch.setattr(export, "aiui_renderer_files", lambda: None)
html = export.to_html(fx.team_events())
assert html.startswith("<!doctype html>")
assert not _EXTERNAL_SRC.search(html)
assert "@import" not in html
assert "window.__REPLAY_DATA__" in html
assert "sk-abcdefghijklmnop1234" not in html
assert "[email protected]" not in html
assert "Researcher" in html and "task_start" in html
data = json.loads(re.search(r"window.__REPLAY_DATA__ = (.*?);</script>", html, re.S).group(1))
assert data["run_id"] == "run-team01"
assert data["graph"]["schema"] == "praisonai.graph/1"
assert len(data["events"]) == len(fx.team_events())
def test_html_no_redact_keeps_content(fx, monkeypatch):
from praisonai.replay import export
monkeypatch.setattr(export, "aiui_renderer_files", lambda: None)
assert "[email protected]" in export.to_html(fx.team_events(), redact=False)
def test_html_inlines_aiui_renderer_when_available(fx, monkeypatch, tmp_path):
from praisonai.replay import export
core = tmp_path / "replay-core.js"
view = tmp_path / "replay.js"
core.write_text("window.ReplayCore = {}; // </script> must not end the tag")
view.write_text("document.getElementById('replay-root').dataset.mounted = '1';")
monkeypatch.setattr(export, "aiui_renderer_files", lambda: [core, view])
html = export.to_html(fx.team_events())
assert "window.ReplayCore = {};" in html and "dataset.mounted" in html
assert html.index("window.__REPLAY_DATA__") < html.index("window.ReplayCore")
assert "// </script>" not in html # escaped so the inline script survives
assert not _EXTERNAL_SRC.search(html)
def test_html_escapes_script_breakers_in_event_data(fx, monkeypatch):
from praisonai.replay import export
monkeypatch.setattr(export, "aiui_renderer_files", lambda: None)
evs = fx.team_events()
evs[9]["data"]["output_preview"] = "</script><script>alert(1)</script>"
html = export.to_html(evs, redact=False)
# Only "</script" can end an inline script; it must never appear unescaped
# except as the real closing tags (one per <script>).
assert "</script><script>alert(1)" not in html
assert html.count("</script>") == html.count("<script")
def test_aiui_renderer_detection_requires_both_files(tmp_path, monkeypatch):
import importlib.util
from types import SimpleNamespace
from praisonai.replay import export
views = tmp_path / "pkg" / "templates" / "frontend" / "plugins" / "views"
views.mkdir(parents=True)
(views / "replay-core.js").write_text("x")
fake = SimpleNamespace(submodule_search_locations=[str(tmp_path / "pkg")])
monkeypatch.setattr(importlib.util, "find_spec", lambda name: fake if name == "praisonaiui" else None)
assert export.aiui_renderer_files() is None
(views / "replay.js").write_text("y")
assert [p.name for p in export.aiui_renderer_files()] == ["replay-core.js", "replay.js"]
def test_to_json_round_trip(fx):
from praisonai.replay.export import to_json
payload = json.loads(to_json(fx.team_events()))
assert payload["run_id"] == "run-team01"
assert payload["graph"]["name"] == "Blog team"
assert payload["events"][6]["data"]["tool_args"]["api_key"] == "[REDACTED]"
assert payload["events"][5]["prompt_tokens"] == 10 # envelope counters untouched
raw = json.loads(to_json(fx.team_events(), redact=False))
assert raw["events"][6]["data"]["tool_args"]["api_key"] == "sk-abcdefghijklmnop1234"
def test_redaction_keeps_numeric_token_counters():
from praisonai.replay.export import redact_events
out = redact_events([{"event_type": "llm_response", "data": {"max_tokens": 50, "token": "abc", "auth": {"x": 1}}}])
assert out[0]["data"] == {"max_tokens": 50, "token": "[REDACTED]", "auth": "[REDACTED]"}
def test_html_inlines_aiui_renderer_as_module_scripts(tmp_path, monkeypatch):
"""aiui's renderer files use ``export``; inlined as classic scripts they never run."""
from praisonai.replay import export
core = tmp_path / "replay-core.js"
view = tmp_path / "replay.js"
core.write_text("export const X = 1;\n", encoding="utf-8")
view.write_text("export function render() {}\n", encoding="utf-8")
monkeypatch.setattr(export, "aiui_renderer_files", lambda: [core, view])
events = [{"event_type": "session_start", "session_id": "s", "sequence_num": 0, "timestamp": 1.0, "data": {}}]
html = export.to_html(events)
assert html.count('<script type="module" data-src=') == 2
assert '<script data-src=' not in html
@@ -0,0 +1,173 @@
"""praisonai.replay.replay_app / watch: a page with a Run button that records each run live."""
import json
import sys
import threading
import pytest
def _flow(seen=None):
from praisonaiagents import AgentFlow
from praisonaiagents.workflows import StepResult
def research(ctx):
if seen is not None:
seen.append(ctx.input)
return StepResult(output="notes")
def write(ctx):
return StepResult(output="article")
return AgentFlow(name="Article", steps=[research, write])
def _events(path):
return [json.loads(line) for line in path.read_text().splitlines() if line.strip()]
def test_runner_records_a_flow_under_the_given_run_id(traces_dir):
from praisonai.replay.live import _runner
seen = []
_runner(_flow(seen))("solar sails", "run-live00000001")
path = traces_dir / "run-live00000001.jsonl"
types = [e["event_type"] for e in _events(path)]
assert seen == ["solar sails"]
assert types[0] == "session_start" and types[-1] == "session_end"
assert types.count("step_end") == 2
def test_runner_resets_a_team_before_each_run(traces_dir):
from praisonai.replay.live import _runner
calls = []
class Team:
def _reset_task_state_for_item(self):
calls.append("reset")
def start(self, text):
calls.append(("start", text))
run = _runner(Team())
run("a", "run-team-a")
run("b", "run-team-b")
assert calls == ["reset", ("start", "a"), "reset", ("start", "b")]
def test_runner_does_not_carry_earlier_inputs_into_later_team_runs(traces_dir):
from praisonai.replay.live import _runner
class Task:
def __init__(self, context):
self.context = context
seen = []
class Team:
tasks = {1: Task(["given"]), 2: Task(None)}
def _reset_task_state_for_item(self):
pass
def start(self, text): # like AgentTeam.start: appends the input to each task's context
for t in self.tasks.values():
t.context = (t.context or []) + [text]
seen.append([t.context for t in self.tasks.values()])
run = _runner(Team())
run("solar sails", "run-ctx-a")
run("comets", "run-ctx-b")
assert seen[1] == [["given", "comets"], ["comets"]]
def test_runner_awaits_an_async_callable(traces_dir):
from praisonai.replay.live import _runner
got = []
async def target(text):
got.append(text)
_runner(target)("hi", "run-async")
assert got == ["hi"]
def test_runner_accepts_a_plain_callable(traces_dir):
from praisonai.replay.live import _runner
got = []
_runner(lambda text: got.append(text))("hi", "run-fn")
assert got == ["hi"] and (traces_dir / "run-fn.jsonl").is_file()
def test_runner_runs_a_yaml_file_as_written(monkeypatch, tmp_path):
from praisonai.replay import live, yaml_run
calls = []
monkeypatch.setattr(yaml_run, "run_yaml_recorded", lambda path, **kw: calls.append((path, kw)))
yml = tmp_path / "agents.yaml"
yml.write_text("roles: {}\n")
live._runner(str(yml))("ignored", "run-yaml")
assert calls == [(str(yml), {"run_id": "run-yaml", "live": True})]
def test_runner_rejects_bad_targets():
from praisonai.replay.live import _runner
assert _runner(None) is None
with pytest.raises(ValueError, match="not a YAML file"):
_runner("notes.txt")
with pytest.raises(TypeError, match="Agent, AgentTeam, AgentFlow"):
_runner(42)
def test_replay_app_explains_how_to_install_aiui(monkeypatch):
from praisonai.replay import replay_app
monkeypatch.setitem(sys.modules, "praisonaiui.features.replay", None)
with pytest.raises(ImportError, match=r'pip install -U "praisonai\[ui\]"'):
replay_app()
def _aiui_embed():
mod = pytest.importorskip("praisonaiui.features.replay")
if not hasattr(mod, "create_replay_app"):
pytest.skip("installed aiui predates create_replay_app")
return mod
def test_run_button_records_and_serves_the_run_end_to_end(traces_dir):
"""Mounted like app.mount('/replay', replay_app(flow)): Run, then read the run back."""
_aiui_embed()
from starlette.applications import Starlette
from starlette.routing import Mount
from starlette.testclient import TestClient
from praisonai.replay import replay_app
traces_dir.mkdir(parents=True, exist_ok=True)
seen = []
client = TestClient(Starlette(routes=[Mount("/replay", app=replay_app(_flow(seen), input="solar sails"))]))
page = client.get("/replay/").text
assert "PraisonAI Replay: Article" in page and '"runner"' in page
run_id = client.post("/replay/api/replay/run", json={"input": "comets"}).json()["run_id"]
for _ in range(200):
data = client.get(f"/replay/api/replay/runs/{run_id}/events?wait=1").json()
if data["run"] and data["events"][-1]["event_type"] == "session_end":
break
threading.Event().wait(0.02) # time.sleep may be stubbed by the suite
assert seen == ["comets"]
assert data["run"]["status"] == "completed"
assert [e.get("node_id") for e in data["events"] if e["event_type"] == "step_end"] == ["step:research", "step:write"]
def test_yaml_target_has_no_input_box():
_aiui_embed()
from starlette.testclient import TestClient
from praisonai.replay import replay_app
page = TestClient(replay_app("agents.yaml")).get("/").text
assert '"ask":false' in page and "agents.yaml" in page
@@ -0,0 +1,87 @@
"""Real agentic test (AGENTS.md §9.4): record a live Agent run, export it to HTML.
Gated: RUN_REAL_KEY_TESTS=1 PRAISONAI_ALLOW_NETWORK=1 and OPENAI_API_KEY. One
cheap gpt-4o-mini call.
"""
import json
import os
import pytest
pytestmark = [
pytest.mark.skipif(
os.environ.get("RUN_REAL_KEY_TESTS") != "1" or not os.environ.get("OPENAI_API_KEY"),
reason="Real key tests disabled. Set RUN_REAL_KEY_TESTS=1 PRAISONAI_ALLOW_NETWORK=1 and OPENAI_API_KEY.",
),
pytest.mark.network,
pytest.mark.provider_openai,
]
def test_record_live_agent_then_export_html(traces_dir, tmp_path):
from typer.testing import CliRunner
from praisonaiagents import Agent
from praisonai.cli.commands.replay import app
from praisonai.replay import load_events, record
with record(live=True) as rec:
agent = Agent(name="Greeter", instructions="You are concise.", llm="gpt-4o-mini")
answer = agent.start("Say hello in one sentence")
print("\nAgent answer:", answer)
assert answer and isinstance(str(answer), str)
events = load_events(rec.path)
types = [e["event_type"] for e in events]
print("Recorded", rec.run_id, "->", rec.path)
print("Event types:", types)
assert types[0] == "session_start" and types[-1] == "session_end"
assert len(events) > 2, "the agent run itself must have emitted events"
out = tmp_path / "replay.html"
result = CliRunner().invoke(app, ["export", rec.run_id, "--format", "html", "-o", str(out)])
print(result.output)
assert result.exit_code == 0, result.output
html = out.read_text(encoding="utf-8")
assert "window.__REPLAY_DATA__" in html and "Greeter" in html
print("HTML export:", out, f"{len(html):,} bytes")
otlp = CliRunner().invoke(app, ["export", rec.run_id, "--format", "otel"])
spans = json.loads(otlp.output)["resourceSpans"][0]["scopeSpans"][0]["spans"]
print("OTLP spans:", [s["name"] for s in spans])
assert any(s["name"].startswith("invoke_agent") for s in spans)
def test_run_button_runs_a_real_team_and_draws_it(traces_dir):
"""replay_app(team): one Run click runs both agents, recorded live; a second Run runs them again."""
mod = pytest.importorskip("praisonaiui.features.replay")
if not hasattr(mod, "create_replay_app"):
pytest.skip("installed aiui predates create_replay_app")
import threading
from starlette.testclient import TestClient
from praisonaiagents import Agent, AgentTeam, Task
from praisonai.replay import replay_app
researcher = Agent(name="Researcher", instructions="Give 2 short facts.", llm="gpt-4o-mini")
writer = Agent(name="Writer", instructions="Write one sentence from the facts.", llm="gpt-4o-mini")
team = AgentTeam(agents=[researcher, writer], tasks=[
Task(name="research", description="2 facts about solar sails", agent=researcher, expected_output="2 facts"),
Task(name="write", description="One sentence using the facts", agent=writer, expected_output="a sentence"),
])
traces_dir.mkdir(parents=True, exist_ok=True)
client = TestClient(replay_app(team))
for attempt in (1, 2):
run_id = client.post("/api/replay/run", json={"input": ""}).json()["run_id"]
for _ in range(600):
data = client.get(f"/api/replay/runs/{run_id}/events?wait=1").json()
if data["run"] and data["events"] and data["events"][-1]["event_type"] == "session_end":
break
threading.Event().wait(0.1) # time.sleep may be stubbed by the suite
types = [e["event_type"] for e in data["events"]]
print(f"\nRun {attempt} ({run_id}):", types)
assert types.count("llm_response") >= 2, "both agents called the LLM"
assert [e["node_id"] for e in data["events"] if e["event_type"] == "task_end"] == ["task:research", "task:write"]
@@ -0,0 +1,599 @@
"""``praisonai replay rerun``: reuse recorded outputs before NODE, run NODE onwards fresh."""
import json
import textwrap
import pytest
@pytest.fixture
def fake_llm(monkeypatch):
"""Replace every Agent LLM call with a counted, deterministic answer."""
from praisonaiagents import Agent
calls = []
def fake_chat(self, prompt, *args, **kwargs):
calls.append((self.name, str(prompt)))
return f"{self.name} fresh output"
async def fake_achat(self, prompt, *args, **kwargs):
return fake_chat(self, prompt)
monkeypatch.setattr(Agent, "chat", fake_chat)
monkeypatch.setattr(Agent, "achat", fake_achat)
monkeypatch.setenv("OPENAI_API_KEY", "sk-test-not-used")
return calls
def _lines(path):
return [json.loads(l) for l in open(path, encoding="utf-8") if l.strip()]
# ---------------------------------------------------------------------------
# plan_rerun: what the recording can support
# ---------------------------------------------------------------------------
def test_plan_rejects_schema_1_0_trace(fx):
from praisonai.replay.rerun import RerunError, plan_rerun
with pytest.raises(RerunError, match="schema 1.0"):
plan_rerun(fx.legacy_events(), "task:write")
def test_plan_rejects_unknown_node_and_lists_known(fx):
from praisonai.replay.rerun import RerunError, plan_rerun
with pytest.raises(RerunError, match="task:research, task:write"):
plan_rerun(fx.team_events(), "task:nope")
def test_plan_reuses_only_nodes_finished_before_target(fx):
from praisonai.replay.rerun import plan_rerun
plan = plan_rerun(fx.flow_events(), "step:polish")
assert plan.kind == "flow"
assert dict(plan.outputs) == {"step:research": "research output", "step:draft": "draft output"}
assert plan.forked_from == {"run_id": "run-flow01", "node_id": "step:polish"}
assert dict(plan_rerun(fx.flow_events(), "step:research").outputs) == {}
def test_plan_recovers_long_output_from_llm_response(fx):
from praisonai.replay.rerun import plan_rerun
full = "Findings: " + "x" * 900
evs = fx.team_events()
evs[5]["data"]["response_content"] = full
evs[9]["data"]["output_preview"] = full[:500]
assert plan_rerun(evs, "task:write").outputs["task:research"] == full
def test_plan_refuses_truncated_output_it_cannot_recover(fx):
from praisonai.replay.rerun import RerunError, plan_rerun
evs = fx.team_events()
evs[5]["data"]["response_content"] = "something else"
evs[9]["data"]["output_preview"] = "y" * 500
with pytest.raises(RerunError, match="truncated preview"):
plan_rerun(evs, "task:write")
def test_plan_refuses_llm_response_the_emitter_shortened(fx):
from praisonai.replay.rerun import RerunError, plan_rerun
full = "Findings: " + "x" * 900
evs = fx.team_events()
evs[5]["data"]["response_content"] = full[:600] + "\n...[60,000 chars, showing first/last portions]...\n" + "tail"
evs[9]["data"]["output_preview"] = full[:500]
with pytest.raises(RerunError, match="truncated preview"):
plan_rerun(evs, "task:write")
# ---------------------------------------------------------------------------
# AgentFlow: TraceStepCache
# ---------------------------------------------------------------------------
def _counting_flow():
from praisonaiagents import AgentFlow
from praisonaiagents.workflows import StepResult
calls = []
def research(ctx):
calls.append("research")
return StepResult(output="research FRESH")
def draft(ctx):
calls.append("draft")
return StepResult(output=f"draft of [{ctx.previous_result}]")
def polish(ctx):
calls.append("polish")
return StepResult(output=f"polished [{ctx.previous_result}]")
return AgentFlow(steps=[research, draft, polish]), calls
def test_trace_step_cache_implements_step_cache_protocol():
from praisonaiagents.workflows.step_cache import StepCacheProtocol, resolve_step_cache
from praisonai.replay.rerun import TraceStepCache
cache = TraceStepCache({})
assert isinstance(cache, StepCacheProtocol)
assert resolve_step_cache(cache) is cache
assert cache.get("k") is None and cache.set("k", {}) is None
def test_flow_rerun_reuses_earlier_steps_and_runs_later_ones(traces_dir, fx):
from praisonai.replay.rerun import rerun
src = fx.write_trace(traces_dir, "run-flow01", fx.flow_events())
flow, calls = _counting_flow()
rec, result = rerun(flow, str(src), "step:draft", input="topic")
assert calls == ["draft", "polish"] # research came from the recording
assert result["output"] == "polished [draft of [research output]]"
assert flow._step_cache.hits == ["research"]
events = _lines(rec.path)
assert rec.run_id != "run-flow01"
assert events[0]["data"]["forked_from"] == {"run_id": "run-flow01", "node_id": "step:draft"}
for ev in events:
if ev["event_type"] == "workflow_start": # schema 1.1 core
assert ev["data"]["forked_from"] == {"run_id": "run-flow01", "node_id": "step:draft"}
def test_flow_rerun_from_last_step(traces_dir, fx):
from praisonai.replay.rerun import rerun
src = fx.write_trace(traces_dir, "run-flow01", fx.flow_events())
flow, calls = _counting_flow()
_rec, result = rerun(flow, str(src), "step:polish", input="topic")
assert calls == ["polish"]
assert result["output"] == "polished [draft output]"
def test_rerun_rejects_wrong_object_kind(traces_dir, fx):
from praisonai.replay.rerun import RerunError, rerun
src = fx.write_trace(traces_dir, "run-team01", fx.team_events())
flow, _calls = _counting_flow()
with pytest.raises(RerunError, match="team task"):
rerun(flow, str(src), "task:write")
# ---------------------------------------------------------------------------
# AgentTeam: restore finished tasks, reset the rest
# ---------------------------------------------------------------------------
def _team():
from praisonaiagents import Agent, AgentTeam, Task
r = Agent(name="Researcher", instructions="research")
w = Agent(name="Writer", instructions="write")
t1 = Task(name="research", description="Research AI", expected_output="notes", agent=r)
t2 = Task(name="write", description="Write post", expected_output="post", agent=w)
return AgentTeam(agents=[r, w], tasks=[t1, t2])
def test_team_rerun_skips_earlier_tasks(traces_dir, fx, fake_llm):
from praisonai.replay.rerun import rerun
src = fx.write_trace(traces_dir, "run-team01", fx.team_events())
team = _team()
rec, result = rerun(team, str(src), "task:write")
assert [name for name, _ in fake_llm] == ["Writer"] # Researcher not re-run
assert "Findings: AI is big" in fake_llm[0][1] # recorded output fed forward
assert "Writer fresh output" in str(result)
assert team.tasks[0].status == "completed"
assert _lines(rec.path)[0]["data"]["forked_from"] == {"run_id": "run-team01", "node_id": "task:write"}
def test_team_rerun_from_first_task_runs_everything(traces_dir, fx, fake_llm):
from praisonai.replay.rerun import rerun
src = fx.write_trace(traces_dir, "run-team01", fx.team_events())
rerun(_team(), str(src), "task:research")
assert [name for name, _ in fake_llm] == ["Researcher", "Writer"]
def test_seed_team_resets_tasks_from_node(fx):
from praisonai.replay.rerun import plan_rerun, seed_team
team = _team()
for task in team.tasks.values():
task.status = "completed"
task.result = "stale"
restored = seed_team(team, plan_rerun(fx.team_events(), "task:write"))
assert restored == [0]
assert team.tasks[0].status == "completed"
assert getattr(team.tasks[0].result, "raw", team.tasks[0].result) == "Findings: AI is big"
assert team.tasks[1].status == "not started" and team.tasks[1].result is None
# ---------------------------------------------------------------------------
# CLI: praisonai replay rerun (YAML definitions through the real runner)
# ---------------------------------------------------------------------------
FLOW_YAML = textwrap.dedent("""\
name: blog-flow
process: workflow
workflow:
default_llm: gpt-4o-mini
input: AI agents
agents:
researcher:
role: Researcher
goal: Research topics
instructions: You research.
writer:
role: Writer
goal: Write posts
instructions: You write.
steps:
- name: research
agent: researcher
action: "Research {{input}}"
- name: draft
agent: writer
action: "Write about {{previous_output}}"
""")
TEAM_YAML = textwrap.dedent("""\
framework: praisonai
topic: AI agents
roles:
researcher:
role: Researcher
goal: Research topics
backstory: You research.
tasks:
research:
description: Research {topic}
expected_output: notes
writer:
role: Writer
goal: Write posts
backstory: You write.
tasks:
write:
description: Write a post
expected_output: post
""")
def _invoke(*args):
from typer.testing import CliRunner
from praisonai.cli.commands.replay import app
return CliRunner().invoke(app, list(args))
def test_cli_rerun_flow_yaml(traces_dir, tmp_path, fx, fake_llm):
y = tmp_path / "flow.yaml"
y.write_text(FLOW_YAML)
fx.write_trace(traces_dir, "run-flow01", fx.flow_events(
yaml_path=y, step_names=("research", "draft"),
outputs={"research": "RECORDED NOTES", "draft": "old post"},
))
result = _invoke("rerun", "run-flow01", "--from", "step:draft")
assert result.exit_code == 0, result.output
assert [name for name, _ in fake_llm] == ["writer"] # only the draft step's agent ran
assert "RECORDED NOTES" in fake_llm[0][1]
assert "reusing: step:research" in result.output
new_ids = [p.stem for p in traces_dir.glob("*.jsonl") if p.stem != "run-flow01"]
assert len(new_ids) == 1
first = _lines(traces_dir / f"{new_ids[0]}.jsonl")[0]
assert first["data"]["forked_from"] == {"run_id": "run-flow01", "node_id": "step:draft"}
assert first["data"]["yaml_path"] == str(y.resolve())
def test_cli_rerun_team_yaml_with_file_option(traces_dir, tmp_path, fx, fake_llm):
y = tmp_path / "team.yaml"
y.write_text(TEAM_YAML)
evs = fx.team_events() # recorded without a yaml_path -> --file is required
for ev in evs: # YAML team tasks have no names: node ids are positional
if ev.get("node_id") == "task:research":
ev["node_id"] = "task:0"
elif ev.get("node_id") == "task:write":
ev["node_id"] = "task:1"
fx.write_trace(traces_dir, "run-team01", evs)
missing = _invoke("rerun", "run-team01", "--from", "task:1")
assert missing.exit_code == 1 and "--file" in missing.output
result = _invoke("rerun", "run-team01", "--from", "task:1", "--file", str(y))
assert result.exit_code == 0, result.output
assert [name for name, _ in fake_llm] == ["Writer"]
assert "Findings: AI is big" in fake_llm[0][1]
def test_cli_rerun_rejects_kind_mismatch(traces_dir, tmp_path, fx, fake_llm):
y = tmp_path / "team.yaml"
y.write_text(TEAM_YAML)
fx.write_trace(traces_dir, "run-flow01", fx.flow_events(yaml_path=y))
result = _invoke("rerun", "run-flow01", "--from", "step:draft")
assert result.exit_code == 1
assert "defines a team" in result.output
assert fake_llm == [] # refused before spending anything
def test_cli_rerun_legacy_trace_fails_clearly(traces_dir, fx):
fx.write_trace(traces_dir, "run-legacy1", fx.legacy_events())
result = _invoke("rerun", "run-legacy1", "--from", "task:write")
assert result.exit_code == 1
assert "schema 1.0" in result.output
def test_plan_rerun_accepts_a_path_object(tmp_path):
"""record() hands back a pathlib.Path; rerun must take it as-is."""
import json
from pathlib import Path
from praisonai.replay.rerun import plan_rerun
events = [
{"event_type": "session_start", "session_id": "s", "sequence_num": 0, "timestamp": 1.0, "data": {}, "run_id": "r"},
{"event_type": "workflow_start", "session_id": "s", "sequence_num": 1, "timestamp": 1.1, "data": {"kind": "flow"}, "run_id": "r"},
{"event_type": "step_start", "session_id": "s", "sequence_num": 2, "timestamp": 1.2, "data": {"name": "a"}, "node_id": "step:a", "run_id": "r"},
{"event_type": "step_end", "session_id": "s", "sequence_num": 3, "timestamp": 1.3, "parent_seq": 2, "data": {"status": "completed", "output": "A"}, "node_id": "step:a", "run_id": "r"},
{"event_type": "step_start", "session_id": "s", "sequence_num": 4, "timestamp": 1.4, "data": {"name": "b"}, "node_id": "step:b", "run_id": "r"},
]
path = tmp_path / "r.jsonl"
path.write_text("\n".join(json.dumps(e) for e in events) + "\n", encoding="utf-8")
assert plan_rerun(Path(path), "step:b") is not None
def test_flow_rerun_restores_recorded_step_variables(traces_dir, fx):
from praisonaiagents import AgentFlow
from praisonaiagents.workflows import StepResult
from praisonai.replay.rerun import rerun
seen = []
def research(ctx):
return StepResult(output="fresh", variables={"customer_id": 0})
def draft(ctx):
seen.append(ctx.variables.get("customer_id"))
return StepResult(output="draft")
def polish(ctx):
return StepResult(output="polish")
evs = fx.flow_events()
for ev in evs:
if ev["event_type"] == "step_end" and ev.get("node_id") == "step:research":
ev["data"]["variables"] = {"customer_id": 42}
src = fx.write_trace(traces_dir, "run-flow01", evs)
rerun(AgentFlow(steps=[research, draft, polish]), str(src), "step:draft", input="topic")
assert seen == [42]
def test_flow_rerun_rejects_flow_missing_recorded_steps(traces_dir, fx):
from praisonaiagents import AgentFlow
from praisonaiagents.workflows import StepResult
from praisonai.replay.rerun import RerunError, rerun
def other(ctx):
return StepResult(output="x")
src = fx.write_trace(traces_dir, "run-flow01", fx.flow_events())
with pytest.raises(RerunError, match="does not match any step"):
rerun(AgentFlow(steps=[other]), str(src), "step:draft", input="topic")
def test_flow_rerun_rejects_reordered_flow(traces_dir, fx):
from praisonaiagents import AgentFlow
from praisonaiagents.workflows import StepResult
from praisonai.replay.rerun import RerunError, rerun
ran = []
def research(ctx):
ran.append("research")
return StepResult(output="r")
def draft(ctx):
ran.append("draft")
return StepResult(output="d")
def polish(ctx):
ran.append("polish")
return StepResult(output="p")
src = fx.write_trace(traces_dir, "run-flow01", fx.flow_events())
with pytest.raises(RerunError, match="different order"):
rerun(AgentFlow(steps=[draft, research, polish]), str(src), "step:draft", input="topic")
assert ran == []
def test_flow_rerun_rejects_step_inserted_before_target(traces_dir, fx):
"""research -> extra -> draft: 'extra' never ran, so draft would get its output, not research's."""
from praisonaiagents import AgentFlow
from praisonaiagents.workflows import StepResult
from praisonai.replay.rerun import RerunError, rerun
ran = []
def research(ctx):
ran.append("research")
return StepResult(output="r")
def extra(ctx):
ran.append("extra")
return StepResult(output="e")
def draft(ctx):
ran.append("draft")
return StepResult(output="d")
def polish(ctx):
ran.append("polish")
return StepResult(output="p")
src = fx.write_trace(traces_dir, "run-flow01", fx.flow_events())
with pytest.raises(RerunError, match="did not run in the recording"):
rerun(AgentFlow(steps=[research, extra, draft, polish]), str(src), "step:draft", input="topic")
assert ran == []
def test_recorded_handler_variables_survive_a_rerun(traces_dir):
"""End to end: core records StepResult.variables on step_end; rerun restores them."""
import json
from praisonaiagents import AgentFlow
from praisonaiagents.workflows import StepResult
from praisonai.replay import record
from praisonai.replay.rerun import rerun
def make(customer_id, seen, calls):
def research(ctx):
calls.append("research")
return StepResult(output="notes", variables={"customer_id": customer_id})
def draft(ctx):
seen.append(ctx.variables.get("customer_id"))
return StepResult(output="draft")
return AgentFlow(steps=[research, draft])
with record(full_content=True) as rec:
make(42, [], []).start("topic")
ends = [json.loads(line) for line in open(rec.path) if '"step_end"' in line]
assert any((e["data"].get("variables") or {}).get("customer_id") == 42 for e in ends)
seen, calls = [], []
rerun(make(0, seen, calls), rec.path, "step:draft", input="topic")
assert calls == [] # research reused from the recording
assert seen == [42] # with the variable it set, not the fresh run's 0
def test_flow_rerun_rejects_unnamed_pattern_inserted_before_target(traces_dir, fx):
"""An If (no name) inserted before the target is caught by its graph node id."""
from praisonaiagents import AgentFlow
from praisonaiagents.workflows import StepResult
from praisonaiagents.workflows.workflows import If
from praisonai.replay.rerun import RerunError, rerun
ran = []
def step(name):
def handler(ctx):
ran.append(name)
return StepResult(output=name)
handler.__name__ = name
return handler
src = fx.write_trace(traces_dir, "run-flow01", fx.flow_events())
flow = AgentFlow(steps=[step("research"), If(condition="1 > 0", then_steps=[step("extra")]),
step("draft"), step("polish")])
with pytest.raises(RerunError, match="did not run in the recording"):
rerun(flow, str(src), "step:draft", input="topic")
assert ran == []
def test_flow_rerun_rejects_redacted_variables(traces_dir, fx):
"""A redacted variable would be restored as the literal marker: refuse instead."""
from praisonaiagents import AgentFlow
from praisonaiagents.workflows import StepResult
from praisonai.replay.rerun import RerunError, rerun
evs = fx.flow_events()
for ev in evs:
if ev["event_type"] == "step_end" and ev.get("node_id") == "step:research":
ev["data"]["variables"] = {"api_key": "[REDACTED]", "customer_id": 42}
src = fx.write_trace(traces_dir, "run-flow01", evs)
def research(ctx):
return StepResult(output="r")
def draft(ctx):
return StepResult(output="d")
def polish(ctx):
return StepResult(output="p")
with pytest.raises(RerunError, match="'api_key'.*redacted"):
rerun(AgentFlow(steps=[research, draft, polish]), str(src), "step:draft", input="topic")
def _nested_flow(ran, extra=False):
from praisonaiagents import AgentFlow
from praisonaiagents.workflows import StepResult
from praisonaiagents.workflows.workflows import If
def step(name):
def handler(ctx):
ran.append(name)
return StepResult(output=name)
handler.__name__ = name
return handler
steps = [step("research")] + ([step("extra")] if extra else [])
return AgentFlow(steps=steps + [If(condition="1 > 0", then_steps=[step("draft")]), step("polish")])
def test_flow_rerun_of_a_nested_target_rejects_a_step_inserted_before_its_container(traces_dir):
"""research -> If[draft] recorded; research -> extra -> If[draft] would feed draft extra's output."""
from praisonai.replay import record
from praisonai.replay.rerun import RerunError, rerun
ran = []
with record(full_content=True) as rec:
_nested_flow(ran).start("x")
ran.clear()
with pytest.raises(RerunError, match="did not run in the recording"):
rerun(_nested_flow(ran, extra=True), rec.path, "step:draft", input="x")
assert ran == []
rerun(_nested_flow(ran), rec.path, "step:draft", input="x") # the same flow still re-runs
assert ran == ["draft", "polish"]
def test_flow_rerun_keeps_a_recorded_skip_skipped(traces_dir):
"""A step skipped by its should_run is re-checked on re-run from the restored
state: the same condition skips it again and the target gets the same input."""
from praisonaiagents import AgentFlow, Task
from praisonaiagents.workflows import StepResult
from praisonai.replay import record
from praisonai.replay.rerun import rerun
ran, inputs = [], []
def make():
def research(ctx):
ran.append("research")
return StepResult(output="notes", variables={"mode": "brief"})
def extra(ctx):
ran.append("extra")
return StepResult(output="EXTRA")
def draft(ctx):
ran.append("draft")
inputs.append(ctx.previous_result)
return StepResult(output="draft")
return AgentFlow(steps=[
research,
Task(name="extra", handler=extra, should_run=lambda ctx: ctx.variables.get("mode") == "long"),
draft,
])
with record(full_content=True) as rec:
make().start("x")
assert ran == ["research", "draft"]
ran.clear()
rerun(make(), rec.path, "step:draft", input="x")
assert ran == ["draft"] # research reused, extra still skipped
assert inputs[0] == inputs[1] == "notes"
@@ -0,0 +1,344 @@
"""PraisonAI Replay: trace location precedence and ``praisonai.replay.record()``."""
import json
from pathlib import Path
import pytest
# ---------------------------------------------------------------------------
# storage: PRAISONAI_TRACES_DIR > $PRAISON_HOME/traces > ~/.praison/traces
# ---------------------------------------------------------------------------
def test_traces_dir_env_wins_over_praison_home(tmp_path, monkeypatch):
from praisonai.replay.storage import get_traces_dir
monkeypatch.setenv("PRAISON_HOME", str(tmp_path / "home"))
monkeypatch.setenv("PRAISONAI_TRACES_DIR", str(tmp_path / "explicit"))
assert get_traces_dir() == tmp_path / "explicit"
assert (tmp_path / "explicit").is_dir()
def test_traces_dir_uses_praison_home(tmp_path, monkeypatch):
from praisonai.replay.storage import get_trace_path, get_traces_dir
monkeypatch.delenv("PRAISONAI_TRACES_DIR", raising=False)
monkeypatch.setenv("PRAISON_HOME", str(tmp_path / "home"))
assert get_traces_dir() == tmp_path / "home" / "traces"
assert get_trace_path("abc") == tmp_path / "home" / "traces" / "abc.jsonl"
def test_traces_dir_default_is_home_dot_praison(tmp_path, monkeypatch):
from praisonai.replay.storage import resolve_traces_dir
monkeypatch.delenv("PRAISONAI_TRACES_DIR", raising=False)
monkeypatch.delenv("PRAISON_HOME", raising=False)
monkeypatch.setenv("HOME", str(tmp_path))
assert resolve_traces_dir() == tmp_path / ".praison" / "traces"
def test_praison_home_matches_last_run_pointer_root(tmp_path, monkeypatch):
"""Replay traces and the last-run pointer share one state root."""
from praisonaiagents.trace.last_run import last_run_pointer_path
from praisonai.replay.storage import resolve_traces_dir
monkeypatch.delenv("PRAISONAI_TRACES_DIR", raising=False)
monkeypatch.setenv("PRAISON_HOME", str(tmp_path))
assert resolve_traces_dir().parent == last_run_pointer_path().parent
def test_writer_and_list_follow_env(traces_dir):
from praisonai.replay import ContextTraceWriter, list_traces
w = ContextTraceWriter(session_id="envrun")
assert w.path == traces_dir / "envrun.jsonl"
w.close()
w.path.write_text('{"event_type": "session_start"}\n')
assert [t.session_id for t in list_traces()] == ["envrun"]
# ---------------------------------------------------------------------------
# record()
# ---------------------------------------------------------------------------
def _lines(path):
return [json.loads(l) for l in Path(path).read_text().splitlines() if l.strip()]
def test_record_writes_session_start_and_end(traces_dir):
from praisonaiagents.trace.context_events import get_context_emitter
from praisonai.replay import record
with record(meta={"yaml_path": "/x/agents.yaml"}) as rec:
assert rec.run_id.startswith("run-")
assert rec.path == traces_dir / f"{rec.run_id}.jsonl"
emitter = get_context_emitter()
assert emitter.enabled and emitter.session_id == rec.run_id
emitter.agent_start("A")
emitter.agent_end("A")
assert not get_context_emitter().enabled # reset afterwards
events = _lines(rec.path)
assert [e["event_type"] for e in events] == ["session_start", "agent_start", "agent_end", "session_end"]
assert events[0]["data"] == {"yaml_path": "/x/agents.yaml", "run_id": rec.run_id}
assert events[-1]["data"]["status"] == "completed"
def test_record_explicit_run_id(traces_dir):
from praisonai.replay import record
with record("my-run") as rec:
pass
assert rec.run_id == "my-run" and rec.path.name == "my-run.jsonl"
def test_record_live_flushes_every_event(traces_dir):
from praisonai.replay import record
with record(live=True) as rec:
assert len(_lines(rec.path)) == 1 # session_start already on disk
rec.emitter.agent_start("A")
assert len(_lines(rec.path)) == 2
def test_record_batched_by_default(traces_dir):
from praisonai.replay import record
with record() as rec:
rec.emitter.agent_start("A")
assert not rec.path.exists() or _lines(rec.path) == []
assert len(_lines(rec.path)) == 3
def test_record_marks_failure_and_resets_on_exception(traces_dir):
from praisonaiagents.trace.context_events import get_context_emitter
from praisonai.replay import record
with pytest.raises(ValueError):
with record() as rec:
raise ValueError("boom")
assert _lines(rec.path)[-1]["data"]["status"] == "failed"
assert not get_context_emitter().enabled
def test_record_redacts_by_default(traces_dir):
from praisonai.replay import record
with record() as rec:
rec.emitter.tool_call_start("A", "search", {"q": "x", "api_key": "sk-secret-value-123456"})
text = rec.path.read_text()
assert "sk-secret-value-123456" not in text and "[REDACTED]" in text
def test_record_forked_from_is_stamped_on_workflow_start(traces_dir):
from praisonaiagents.trace.context_events import ContextEvent
from praisonai.replay.recorder import record
with record(_forked_from={"run_id": "run-a", "node_id": "task:write"}) as rec:
ev = ContextEvent.from_dict({"event_type": "session_start", "session_id": rec.run_id})
ev.event_type = "workflow_start" # what a 1.1 emitter emits
ev.data = {"kind": "team"}
rec.emitter._emit(ev)
events = _lines(rec.path)
assert events[0]["data"]["forked_from"] == {"run_id": "run-a", "node_id": "task:write"}
assert events[1]["data"] == {"kind": "team", "forked_from": {"run_id": "run-a", "node_id": "task:write"}}
# ---------------------------------------------------------------------------
# The three callers that used to duplicate the setup still record
# ---------------------------------------------------------------------------
def test_recipe_save_replay_closes_trace_on_early_failure(traces_dir, monkeypatch):
"""A recipe that is not found used to leave the trace open and the emitter set."""
from praisonaiagents.trace.context_events import get_context_emitter
from praisonai.recipe import core
monkeypatch.setattr(core, "_load_recipe", lambda name, offline=False: None)
result = core._run_impl("nope", options={"save_replay": True, "trace_name": "recipe-run"})
assert result.error and "not found" in result.error
assert not get_context_emitter().enabled
events = _lines(traces_dir / "recipe-run.jsonl")
assert events[0]["event_type"] == "session_start"
assert events[0]["data"] == {"recipe": "nope", "run_id": "recipe-run"}
assert events[-1]["event_type"] == "session_end"
assert events[-1]["data"]["status"] == "failed"
def test_recipe_save_replay_success_passes_emitter(traces_dir, monkeypatch):
from types import SimpleNamespace
from praisonaiagents.trace.context_events import get_context_emitter
from praisonai.recipe import core
cfg = SimpleNamespace(name="r", version="1.0", defaults={})
seen = {}
def fake_execute(recipe_config, merged, session_id, options, trace_emitter=None):
seen["emitter"] = trace_emitter
seen["active"] = get_context_emitter()
trace_emitter.agent_start("RecipeAgent")
return {"ok": True}
monkeypatch.setattr(core, "_load_recipe", lambda name, offline=False: cfg)
monkeypatch.setattr(core, "_check_dependencies", lambda c: {"all_satisfied": True})
monkeypatch.setattr(core, "_check_tool_policy", lambda c, o: None)
monkeypatch.setattr(core, "_execute_recipe", fake_execute)
result = core._run_impl("r", options={"save_replay": True, "trace_name": "recipe-ok"})
assert result.output == {"ok": True}
assert seen["emitter"] is seen["active"] and seen["emitter"].enabled
events = _lines(traces_dir / "recipe-ok.jsonl")
assert [e["event_type"] for e in events] == ["session_start", "agent_start", "session_end"]
assert events[-1]["data"]["status"] == "completed"
assert not get_context_emitter().enabled
class _FakeWorkflow:
name = "wf"
steps = []
planning = False
reasoning = False
framework = "praisonai"
default_input = "hello"
verbose = False
def __init__(self):
self.variables = {}
self.seen_emitter = None
def start(self, _input):
from praisonaiagents.trace.context_events import get_context_emitter
self.seen_emitter = get_context_emitter()
self.seen_emitter.agent_start("Step")
return {"status": "completed", "output": "done"}
def _run_workflow_command(monkeypatch, yaml_file, action_args):
from praisonaiagents.workflows import WorkflowManager
from praisonai.cli.legacy import workflow_commands
wf = _FakeWorkflow()
monkeypatch.setattr(WorkflowManager, "load_yaml", lambda self, path, tool_registry=None: wf)
code = workflow_commands._run_yaml_workflow(None, str(yaml_file), action_args)
return code, wf
def test_workflow_run_save_records(traces_dir, tmp_path, monkeypatch):
from praisonaiagents.trace.context_events import get_context_emitter
y = tmp_path / "wf.yaml"
y.write_text("name: wf\nsteps: []\n")
code, wf = _run_workflow_command(monkeypatch, y, ["--save"])
assert code == 0
assert wf.seen_emitter.enabled
assert not get_context_emitter().enabled
events = _lines(traces_dir / f"{wf.seen_emitter.session_id}.jsonl")
assert events[0]["data"]["workflow"] == str(y)
assert events[0]["data"]["yaml_path"] == str(y.resolve())
assert [e["event_type"] for e in events] == ["session_start", "agent_start", "session_end"]
def test_workflow_run_yaml_save_replay_key_records(traces_dir, tmp_path, monkeypatch):
y = tmp_path / "wf.yaml"
y.write_text("name: wf\nsave_replay: true\nsteps: []\n")
code, wf = _run_workflow_command(monkeypatch, y, [])
assert code == 0 and wf.seen_emitter.enabled
assert (traces_dir / f"{wf.seen_emitter.session_id}.jsonl").exists()
def test_workflow_run_without_save_does_not_record(traces_dir, tmp_path, monkeypatch):
y = tmp_path / "wf.yaml"
y.write_text("name: wf\nsteps: []\n")
code, wf = _run_workflow_command(monkeypatch, y, [])
assert code == 0 and not wf.seen_emitter.enabled
assert not traces_dir.exists() or not list(traces_dir.iterdir())
def _drive_legacy_main(monkeypatch, tmp_path, argv, yaml_text):
import sys
pa = pytest.importorskip("praisonai_code.cli.legacy.praison_ai")
from praisonaiagents.trace.context_events import get_context_emitter
seen = {}
class FakeGenerator:
def __init__(self, *a, **k):
pass
def __enter__(self):
return self
def __exit__(self, *exc):
return False
def generate_crew_and_kickoff(self):
seen["emitter"] = get_context_emitter()
return "RESULT"
(tmp_path / "agents.yaml").write_text(yaml_text)
monkeypatch.chdir(tmp_path)
monkeypatch.setenv("OPENAI_API_KEY", "sk-test-not-used")
monkeypatch.setattr(pa, "_get_agents_generator", lambda: FakeGenerator)
monkeypatch.setattr(sys, "argv", argv)
monkeypatch.delenv("PYTEST_CURRENT_TEST", raising=False)
result = pa.PraisonAI(agent_file="agents.yaml").main()
return result, seen.get("emitter")
_AGENTS_YAML = (
"framework: praisonai\nroles:\n a:\n role: A\n goal: g\n backstory: b\n"
" tasks:\n t:\n description: d\n expected_output: e\n"
)
def test_agents_yaml_save_flag_records(traces_dir, tmp_path, monkeypatch):
from praisonaiagents.trace.context_events import get_context_emitter
result, emitter = _drive_legacy_main(monkeypatch, tmp_path, ["praisonai", "agents.yaml", "--save"], _AGENTS_YAML)
assert result == "RESULT"
assert emitter is not None and emitter.enabled
assert not get_context_emitter().enabled
events = _lines(traces_dir / f"{emitter.session_id}.jsonl")
assert events[0]["data"]["agents_file"] == "agents.yaml"
assert events[0]["data"]["yaml_path"] == str((tmp_path / "agents.yaml").resolve())
assert events[-1]["event_type"] == "session_end"
def test_agents_yaml_save_replay_key_records(traces_dir, tmp_path, monkeypatch):
result, emitter = _drive_legacy_main(
monkeypatch, tmp_path, ["praisonai", "agents.yaml"], "save_replay: true\n" + _AGENTS_YAML
)
assert result == "RESULT" and emitter.enabled
assert (traces_dir / f"{emitter.session_id}.jsonl").exists()
def test_agents_yaml_without_save_is_untraced(traces_dir, tmp_path, monkeypatch):
result, emitter = _drive_legacy_main(monkeypatch, tmp_path, ["praisonai", "agents.yaml"], _AGENTS_YAML)
assert result == "RESULT" and not emitter.enabled
def test_save_replay_is_a_known_yaml_field():
from praisonai.config.schema import YAMLConfig
assert "save_replay" in YAMLConfig.model_fields
def test_record_stamps_run_id_on_every_event(tmp_path, monkeypatch):
"""Every event in a recording shares the recording's run_id (schema 1.1 SDKs)."""
import json
import pytest
ce = pytest.importorskip("praisonaiagents.trace.context_events")
if not hasattr(ce, "_run_scope"):
pytest.skip("SDK predates schema 1.1 run scopes")
monkeypatch.setenv("PRAISONAI_TRACES_DIR", str(tmp_path))
from praisonai.replay import record
with record(run_id="run-scope-test") as rec:
ce.get_context_emitter().agent_start("Solo", {})
ce.get_context_emitter().agent_end("Solo", {})
events = [json.loads(line) for line in open(rec.path)]
assert events and {e.get("run_id") for e in events} == {"run-scope-test"}
assert ce._run_id_var.get() is None # scope closed after the block
@@ -0,0 +1,37 @@
"""Sync and async moderation must select the same model by default."""
import asyncio
import sys
from types import SimpleNamespace
from unittest.mock import AsyncMock, Mock
import pytest
from praisonai.capabilities.moderations import amoderate, moderate
@pytest.mark.parametrize("model", [None, "text-moderation-stable"])
def test_moderation_model_parity(monkeypatch, model):
response = SimpleNamespace(results=[SimpleNamespace(flagged=False)])
sync_call = Mock(return_value=response)
async_call = AsyncMock(return_value=response)
monkeypatch.setitem(
sys.modules,
"litellm",
SimpleNamespace(moderation=sync_call, amoderation=async_call),
)
kwargs = {} if model is None else {"model": model}
sync_results = moderate("A friendly greeting", **kwargs)
async_results = asyncio.run(amoderate("A friendly greeting", **kwargs))
expected_model = model or "omni-moderation-latest"
expected = {
"input": "A friendly greeting",
"model": expected_model,
"timeout": 600.0,
}
sync_call.assert_called_once_with(**expected)
async_call.assert_awaited_once_with(**expected)
assert sync_results == async_results
assert async_results[0].model == expected_model
@@ -52,6 +52,13 @@ class _MockOpenAI(BaseHTTPRequestHandler):
self.wfile.write(data)
@pytest.fixture(autouse=True)
def _project_root(tmp_path, monkeypatch):
# Task output_file writes are contained to the project root (#5770); these
# runs write under tmp_path, so make it the root.
monkeypatch.setenv("PRAISONAI_PROJECT_ROOT", str(tmp_path))
@pytest.fixture
def mock_endpoint(monkeypatch):
server = HTTPServer(("127.0.0.1", 0), _MockOpenAI)