mirror of
https://github.com/MervinPraison/PraisonAI.git
synced 2026-11-12 18:47:44 -10:00
Compare commits
65
Commits
No files matched your search
@@ -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: |
|
||||
|
||||
@@ -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>
|
||||
|
||||
@@ -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.
|
||||
@@ -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>"""
|
||||
@@ -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")
|
||||
@@ -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)
|
||||
|
||||
@@ -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):
|
||||
"""
|
||||
|
||||
@@ -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",
|
||||
|
||||
Generated
+3
-3
@@ -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" },
|
||||
|
||||
@@ -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
|
||||
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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",
|
||||
|
||||
@@ -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.
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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)
|
||||
@@ -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
|
||||
@@ -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,
|
||||
)
|
||||
@@ -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")
|
||||
@@ -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)
|
||||
@@ -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
|
||||
@@ -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
|
||||
|
||||
|
||||
@@ -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.",
|
||||
))
|
||||
@@ -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)
|
||||
|
||||
Reference in new issue
Block a user