mirror of
https://github.com/MervinPraison/PraisonAI.git
synced 2026-11-10 03:31:28 -10:00
Compare commits
145
Commits
No files matched your search
@@ -313,7 +313,11 @@ runs:
|
||||
# Add allowed tools (include GitHub tools by default)
|
||||
TOOLS="mcp__github__get_issue,mcp__github__get_issue_comments,mcp__github__update_issue,mcp__github__search_issues,mcp__github__list_issues,mcp__github__add_issue_comment,mcp__github__get_pull_request,mcp__github__get_pull_request_files,mcp__github__get_pull_request_reviews,mcp__github__get_pull_request_status,Read,Write,Edit,Bash"
|
||||
if [[ -n "$INPUT_ALLOWED_TOOLS" ]]; then
|
||||
TOOLS="$TOOLS,$INPUT_ALLOWED_TOOLS"
|
||||
# Callers pass a YAML block (one tool per line). `read` below stops at
|
||||
# the first newline, so join lines with commas or every tool after
|
||||
# the first is silently dropped.
|
||||
EXTRA_TOOLS=$(printf '%s' "$INPUT_ALLOWED_TOOLS" | tr '\n' ',' | sed 's/,,*/,/g; s/^,//; s/,$//')
|
||||
TOOLS="$TOOLS,$EXTRA_TOOLS"
|
||||
fi
|
||||
|
||||
# Remove any explicitly disallowed tools from the effective allow-list so
|
||||
|
||||
@@ -21,6 +21,24 @@ const finals = [
|
||||
];
|
||||
assert('stale when head after final', mg.isStaleFinalAfterPush(finals, '2026-06-12T09:00:00Z'));
|
||||
|
||||
// PR #6056 shape: FINAL requested after HEAD; Claude answered with a security-only
|
||||
// comment (no Phase 1 table). That FINAL never completes, so recovery must re-ask.
|
||||
{
|
||||
const head = '2026-10-27T20:40:21Z';
|
||||
const trigger = { user: { login: 'MervinPraison' }, body: mg.FINAL_CLAUDE_REVIEW_BODY, created_at: '2026-10-27T20:42:06Z' };
|
||||
const securityOnly = { user: { login: 'praisonai-triage-agent[bot]' }, body: '## Security and architecture review at `fec9380dd`. Origin: outsider\n\n**Verdict:** PASS', created_at: '2026-10-27T20:50:17Z' };
|
||||
const realFinal = { user: { login: 'praisonai-triage-agent[bot]' }, body: '#### Phase 1 — AGENTS.md review\n| Check | Result |', created_at: '2026-10-27T20:55:00Z' };
|
||||
const gateSecurity = { user: { login: 'praisonai-triage-agent[bot]' }, body: 'SECURITY_REVIEW: PASS\n| file | category |', created_at: '2026-10-27T20:50:17Z' };
|
||||
assert('FINAL answered without the review is stale (re-ask)', mg.isStaleFinalAfterPush([trigger, securityOnly], head));
|
||||
assert('FINAL answered without the review is not complete', !mg.finalClaudeCompletedOnSha([trigger, securityOnly], head));
|
||||
assert('a real FINAL reply after a security-only one completes it', !mg.isStaleFinalAfterPush([trigger, securityOnly, realFinal], head)
|
||||
&& mg.finalClaudeCompletedOnSha([trigger, securityOnly, realFinal], head));
|
||||
assert('the gate security review comment is not a FINAL answer', !mg.isStaleFinalAfterPush([trigger, gateSecurity], head));
|
||||
assert('no reply yet: FINAL in progress, not stale', !mg.isStaleFinalAfterPush([trigger], head));
|
||||
const retries = [0, 1, 2].map((i) => ({ ...trigger, created_at: `2026-10-27T2${i + 1}:00:00Z` }));
|
||||
assert('re-asking stops after the per-HEAD cap', !mg.isStaleFinalAfterPush([...retries, { ...securityOnly, created_at: '2026-10-27T23:30:00Z' }], head));
|
||||
}
|
||||
|
||||
const withRecovery = [
|
||||
...finals,
|
||||
{ user: { login: 'github-actions[bot]' }, body: '@claude FINAL architecture reviewer', created_at: '2026-06-12T09:30:00Z' },
|
||||
|
||||
@@ -322,6 +322,7 @@ function isStaleFinalAfterPush(comments, headPushedAt) {
|
||||
new Date(a.created_at) > new Date(b.created_at) ? a : b
|
||||
);
|
||||
const finalTime = new Date(latestFinal.created_at).getTime();
|
||||
if (finalAnsweredWithoutReview(comments, finalTime, headTime)) return true;
|
||||
if (headTime <= finalTime + 60000) return false;
|
||||
|
||||
// The finished reply must follow BOTH the latest FINAL trigger and HEAD, so
|
||||
@@ -342,6 +343,31 @@ function isStaleFinalAfterPush(comments, headPushedAt) {
|
||||
return !claudeSinceHead;
|
||||
}
|
||||
|
||||
// FINAL re-requests allowed for one HEAD before a maintainer has to look.
|
||||
const FINAL_RETRIES_PER_HEAD = 3;
|
||||
|
||||
/**
|
||||
* Claude answered the latest FINAL, but not with the review the gate needs (no
|
||||
* Phase 1 table / "Final architecture review" heading), e.g. a security-only
|
||||
* comment on an outsider-origin PR. That FINAL never completes, and because it
|
||||
* was requested after HEAD it is not "stale" either, so without this the PR
|
||||
* waits forever. Treat it as stale so recovery asks again (capped per HEAD).
|
||||
*/
|
||||
function finalAnsweredWithoutReview(comments, finalTime, headTime) {
|
||||
const after = comments.filter((c) => new Date(c.created_at).getTime() > finalTime);
|
||||
if (after.some(isClaudeFinalReplyComment)) return false;
|
||||
const retries = comments.filter(
|
||||
(c) => isFinalClaudeTriggerComment(c) && new Date(c.created_at).getTime() >= headTime - 60000
|
||||
).length;
|
||||
if (retries >= FINAL_RETRIES_PER_HEAD) return false;
|
||||
return after.some((c) => {
|
||||
if (!CLAUDE_AUTOMATION_LOGINS.has(c.user?.login || '') || AUTO_ACTORS.includes(c.user?.login)) return false;
|
||||
const body = c.body || '';
|
||||
if (isClaudeTriggerNoise(c) || /^\s*@claude\b/i.test(body) || /^\s*SECURITY_REVIEW:/m.test(body)) return false;
|
||||
return true;
|
||||
});
|
||||
}
|
||||
|
||||
function needsStaleFinalRecovery(comments, headPushedAt) {
|
||||
return (
|
||||
hasFinalClaudeReviewTrigger(comments) &&
|
||||
|
||||
@@ -380,7 +380,7 @@ jobs:
|
||||
if (process.env.IS_FORK === 'true') {
|
||||
preamble = `${untrusted}\n\nThis is a fork PR. Do a security review as well as your normal review:\n${checklist}\nStart your single comment with a line "Security review: PASS" or "Security review: FAIL" plus the findings (advisory: a maintainer merges fork PRs), then your normal review.`;
|
||||
} else if (isPr) {
|
||||
preamble = `${untrusted}\n\nThis PR was built from or commented on by outsiders. ${rules} Also check the diff against this list and report concerns:\n${checklist}\nThe merge gate runs an independent security review before merging.`;
|
||||
preamble = `${untrusted}\n\nThis PR was built from or commented on by outsiders. ${rules} Also check the diff against this list and report concerns:\n${checklist}\nThe merge gate runs an independent security review before merging. If the request asks for a review in a set format (the FINAL architecture review's Phase 1 table and Phase 3 verdict), keep that format exactly and add these security findings as an extra section; the merge gate only recognises the requested format.`;
|
||||
} else {
|
||||
preamble = `${untrusted}\n\nYou may implement this outsider's request, under these rules. ${rules} An independent security review runs before any merge.`;
|
||||
}
|
||||
|
||||
@@ -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>
|
||||
|
||||
@@ -16,7 +16,7 @@ RUN mkdir -p /root/.praison
|
||||
# Install Python packages (using latest versions)
|
||||
RUN pip install --no-cache-dir \
|
||||
praisonai_tools \
|
||||
"praisonai>=4.7.18" \
|
||||
"praisonai>=4.7.21" \
|
||||
"praisonai[chat]" \
|
||||
"embedchain[github,youtube]"
|
||||
|
||||
|
||||
@@ -20,7 +20,7 @@ RUN mkdir -p /root/.praison
|
||||
# Install Python packages (using latest versions)
|
||||
RUN pip install --no-cache-dir \
|
||||
praisonai_tools \
|
||||
"praisonai>=4.7.18" \
|
||||
"praisonai>=4.7.21" \
|
||||
"praisonai[ui]" \
|
||||
"praisonai[chat]" \
|
||||
"praisonai[realtime]" \
|
||||
|
||||
@@ -16,7 +16,7 @@ RUN mkdir -p /root/.praison
|
||||
# Install Python packages (using latest versions)
|
||||
RUN pip install --no-cache-dir \
|
||||
praisonai_tools \
|
||||
"praisonai>=4.7.18" \
|
||||
"praisonai>=4.7.21" \
|
||||
"praisonai[ui]" \
|
||||
"praisonai[crewai]"
|
||||
|
||||
|
||||
@@ -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)
|
||||
@@ -12,7 +12,7 @@ from collections import OrderedDict
|
||||
import inspect
|
||||
|
||||
# Decomposed agent functionality - real implementations in mixin files
|
||||
from .chat_mixin import ChatMixin
|
||||
from .chat_mixin import ChatMixin, _deadline_scope
|
||||
from .execution_mixin import ExecutionMixin
|
||||
from .memory_mixin import MemoryMixin
|
||||
from .async_memory_mixin import AsyncMemoryMixin
|
||||
@@ -463,7 +463,11 @@ class Agent(GoalLoopMixin, SteeringMixin, SandboxMixin, SkillReviewMixin, Unifie
|
||||
concurrent turns on a shared agent do not share per-turn counters."""
|
||||
if self._loop_guard is None:
|
||||
from ..escalation.loop_guard import LoopGuard, LoopGuardConfig
|
||||
self._loop_guard = LoopGuard(LoopGuardConfig(enabled=True))
|
||||
config = LoopGuardConfig(enabled=True)
|
||||
limit = getattr(self, "max_execution_time", None)
|
||||
if limit is not None:
|
||||
config.max_time_per_turn = float(limit) if limit > 0 else float("inf")
|
||||
self._loop_guard = LoopGuard(config)
|
||||
from .chat_mixin import _current_turn_loop_guard
|
||||
turn_guard = _current_turn_loop_guard(self)
|
||||
return turn_guard if turn_guard is not None else self._loop_guard
|
||||
@@ -1913,6 +1917,17 @@ class Agent(GoalLoopMixin, SteeringMixin, SandboxMixin, SkillReviewMixin, Unifie
|
||||
"Use Agent(policy=PolicyEngine(...)) for tool policy enforcement, "
|
||||
"or pass a validator via GuardrailConfig(validator=...)."
|
||||
)
|
||||
# GuardrailConfig.on_fail has a single implemented policy (retry, then
|
||||
# return the blocked message); refuse "skip"/"raise" rather than
|
||||
# silently behaving as "retry".
|
||||
if isinstance(_guardrails_config, GuardrailConfig) and (
|
||||
getattr(_guardrails_config.on_fail, "value", _guardrails_config.on_fail) != "retry"
|
||||
):
|
||||
raise ValueError(
|
||||
f"GuardrailConfig.on_fail={_guardrails_config.on_fail!r} is not supported; "
|
||||
"only 'retry' is implemented. Remove on_fail, or handle a blocked "
|
||||
"response via Agent.last_guardrail_error."
|
||||
)
|
||||
|
||||
# Fail loud when the value could not be turned into an enforceable
|
||||
# guardrail, instead of storing something that never runs. Silently
|
||||
@@ -3283,6 +3298,12 @@ Your Goal: {self.goal}
|
||||
if k in ("_db_run_context", "_db_session_init_lock"):
|
||||
# A clone starts outside the source agent's active DB turn.
|
||||
continue
|
||||
elif k == "_tool_executor":
|
||||
# Thread pools can't be copied; the clone lazily builds its own.
|
||||
continue
|
||||
elif k == "_Agent__hook_runner":
|
||||
# Holds locks; rebuilt lazily on first use.
|
||||
object.__setattr__(result, k, None)
|
||||
elif k in ("_Agent__cache_lock",):
|
||||
object.__setattr__(result, k, threading.RLock())
|
||||
elif k in ("_cost_lock", "_auto_save_lock"):
|
||||
@@ -4109,7 +4130,9 @@ Summary:"""
|
||||
step budget from ``ExecutionConfig.max_steps`` was reached and the run was
|
||||
truncated), a provider block/refusal/truncation
|
||||
(``"content_filtered" | "refused" | "length_truncated"``) derived from the
|
||||
LLM ``finish_reason``/refusal signal, or ``"error"``. Lets CLI/CI callers
|
||||
LLM ``finish_reason``/refusal signal, ``"timeout"`` (the turn passed
|
||||
``ExecutionConfig.max_execution_time`` or the ``run_autonomous``
|
||||
deadline), or ``"error"``. Lets CLI/CI callers
|
||||
branch on the terminal reason instead of parsing a magic string. Reads
|
||||
from whichever backend (OpenAI-native or LiteLLM) executed the last turn.
|
||||
"""
|
||||
@@ -4727,7 +4750,7 @@ Summary:"""
|
||||
if result.success:
|
||||
print(result.output)
|
||||
"""
|
||||
from .autonomy import AutonomyResult
|
||||
from .autonomy import AutonomyResult, _timeout_result
|
||||
import time as time_module
|
||||
from datetime import datetime, timezone
|
||||
|
||||
@@ -4766,6 +4789,11 @@ Summary:"""
|
||||
effective_promise = completion_promise
|
||||
if effective_promise is None:
|
||||
effective_promise = self.autonomy_config.get("completion_promise")
|
||||
# The promise is the only success signal, so tell the model about it
|
||||
# unless the caller already did (e.g. the `loop` CLI instructions).
|
||||
_promise_tag = f"<promise>{effective_promise}</promise>"
|
||||
if effective_promise and _promise_tag not in prompt and _promise_tag not in (self.instructions or ""):
|
||||
prompt = f"{prompt}\n\nWhen the task is fully complete, include {_promise_tag} in your response."
|
||||
|
||||
# Get clear_context from config if not explicitly set
|
||||
effective_clear_context = clear_context
|
||||
@@ -4796,17 +4824,10 @@ Summary:"""
|
||||
)
|
||||
|
||||
# Check timeout
|
||||
if timeout_seconds and (time_module.time() - start_time) > timeout_seconds:
|
||||
return AutonomyResult(
|
||||
success=False,
|
||||
output="Task timed out",
|
||||
completion_reason="timeout",
|
||||
iterations=iterations,
|
||||
stage=stage,
|
||||
actions=actions_taken,
|
||||
duration_seconds=time_module.time() - start_time,
|
||||
started_at=started_at,
|
||||
)
|
||||
_timed_out = _timeout_result(
|
||||
timeout_seconds, iterations, stage, actions_taken, start_time, started_at)
|
||||
if _timed_out is not None:
|
||||
return _timed_out
|
||||
|
||||
# G2: Check for interrupt request (cooperative cancellation) - sync version
|
||||
if self.interrupt_controller and self.interrupt_controller.is_set():
|
||||
@@ -4829,9 +4850,17 @@ Summary:"""
|
||||
# Reset the per-turn warning latch and inject any queued nudge.
|
||||
turn_prompt = self._consume_self_correction(prompt)
|
||||
try:
|
||||
response = self.chat(turn_prompt)
|
||||
with _deadline_scope(timeout_seconds and timeout_seconds - (time_module.time() - start_time)):
|
||||
response = self.chat(turn_prompt)
|
||||
except Exception as e:
|
||||
return AutonomyResult(
|
||||
# The run deadline cancels an overrunning turn; report it as a timeout.
|
||||
# A budget stop keeps its typed budget_exhausted result when a cap applies.
|
||||
return _timeout_result(
|
||||
timeout_seconds, iterations, stage, actions_taken, start_time, started_at
|
||||
) or (isinstance(e, BudgetExceededError) and self._autonomy_budget_result(
|
||||
iterations, stage, actions_taken, start_time, started_at,
|
||||
"", spend_baseline
|
||||
)) or AutonomyResult(
|
||||
success=False,
|
||||
output=str(e),
|
||||
completion_reason="error",
|
||||
@@ -4843,6 +4872,14 @@ Summary:"""
|
||||
started_at=started_at,
|
||||
)
|
||||
|
||||
# A turn that overran the run deadline or max_execution_time
|
||||
# must not count as progress.
|
||||
_timed_out = _timeout_result(
|
||||
timeout_seconds, iterations, stage, actions_taken, start_time, started_at,
|
||||
turn_timed_out=getattr(self, "_last_stop_reason", None) == "timeout")
|
||||
if _timed_out is not None:
|
||||
return _timed_out
|
||||
|
||||
response_str = str(response)
|
||||
|
||||
# Spend kill-switch: halt if cumulative cost/tokens exceed cap.
|
||||
@@ -5008,7 +5045,7 @@ Summary:"""
|
||||
# the task naturally (model stopped calling tools = done).
|
||||
# This is the same signal the CLI/wrapper path trusts.
|
||||
# ─────────────────────────────────────────────────────────────
|
||||
if self._autonomy_turn_tool_count > 0 and len(response_str) > 100:
|
||||
if not effective_promise and self._autonomy_turn_tool_count > 0 and len(response_str) > 100:
|
||||
# P3/G2: Emit completion callback
|
||||
execute_sync_callback('autonomy_complete',
|
||||
completion_reason="tool_completion",
|
||||
@@ -5045,7 +5082,8 @@ Summary:"""
|
||||
if obs is not None:
|
||||
obs.emit(_EvtType.STAGE_ESCALATE, {"from": prev_stage, "to": stage})
|
||||
|
||||
# Check for completion promise FIRST (structured signal)
|
||||
# Completion promise: when configured, it is the only success signal
|
||||
# (heuristic exits above/below are skipped in promise mode)
|
||||
if effective_promise:
|
||||
promise_tag = f"<promise>{effective_promise}</promise>"
|
||||
if promise_tag in response_str:
|
||||
@@ -5067,7 +5105,7 @@ Summary:"""
|
||||
)
|
||||
|
||||
# Check for keyword-based completion signals (word-boundary, G-COMPLETION-1 fix)
|
||||
if self._is_completion_signal(response_str):
|
||||
if not effective_promise and self._is_completion_signal(response_str):
|
||||
# P3/G2: Emit completion callback
|
||||
execute_sync_callback('autonomy_complete',
|
||||
completion_reason="goal",
|
||||
@@ -5086,7 +5124,7 @@ Summary:"""
|
||||
)
|
||||
|
||||
# For DIRECT stage, complete after first response
|
||||
if stage == "direct":
|
||||
if not effective_promise and stage == "direct":
|
||||
# P3/G2: Emit completion callback
|
||||
execute_sync_callback('autonomy_complete',
|
||||
completion_reason="goal",
|
||||
@@ -5114,7 +5152,7 @@ Summary:"""
|
||||
# Suppressed when doom recovery is active — once a doom loop
|
||||
# is detected, no_tool_calls (a success exit) should not fire;
|
||||
# the doom loop system manages termination instead.
|
||||
if (self._autonomy_turn_tool_count == 0 and iterations > 1
|
||||
if (not effective_promise and self._autonomy_turn_tool_count == 0 and iterations > 1
|
||||
and not self._doom_recovery_active):
|
||||
self._consecutive_no_tool_turns += 1
|
||||
if self._consecutive_no_tool_turns >= 2:
|
||||
@@ -5226,7 +5264,7 @@ Summary:"""
|
||||
|
||||
asyncio.run(main())
|
||||
"""
|
||||
from .autonomy import AutonomyResult
|
||||
from .autonomy import AutonomyResult, _timeout_result
|
||||
import time as time_module
|
||||
from datetime import datetime, timezone
|
||||
import asyncio
|
||||
@@ -5253,6 +5291,11 @@ Summary:"""
|
||||
effective_promise = completion_promise
|
||||
if effective_promise is None:
|
||||
effective_promise = self.autonomy_config.get("completion_promise")
|
||||
# The promise is the only success signal, so tell the model about it
|
||||
# unless the caller already did (e.g. the `loop` CLI instructions).
|
||||
_promise_tag = f"<promise>{effective_promise}</promise>"
|
||||
if effective_promise and _promise_tag not in prompt and _promise_tag not in (self.instructions or ""):
|
||||
prompt = f"{prompt}\n\nWhen the task is fully complete, include {_promise_tag} in your response."
|
||||
|
||||
# Get clear_context from config if not explicitly set
|
||||
effective_clear_context = clear_context
|
||||
@@ -5273,17 +5316,10 @@ Summary:"""
|
||||
iterations += 1
|
||||
|
||||
# Check timeout
|
||||
if timeout_seconds and (time_module.time() - start_time) > timeout_seconds:
|
||||
return AutonomyResult(
|
||||
success=False,
|
||||
output="Task timed out",
|
||||
completion_reason="timeout",
|
||||
iterations=iterations,
|
||||
stage=stage,
|
||||
actions=actions_taken,
|
||||
duration_seconds=time_module.time() - start_time,
|
||||
started_at=started_at,
|
||||
)
|
||||
_timed_out = _timeout_result(
|
||||
timeout_seconds, iterations, stage, actions_taken, start_time, started_at)
|
||||
if _timed_out is not None:
|
||||
return _timed_out
|
||||
|
||||
# G2: Check for interrupt request (cooperative cancellation)
|
||||
if self.interrupt_controller and self.interrupt_controller.is_set():
|
||||
@@ -5307,9 +5343,17 @@ Summary:"""
|
||||
# Reset the per-turn warning latch and inject any queued nudge.
|
||||
turn_prompt = self._consume_self_correction(prompt)
|
||||
try:
|
||||
response = await self.achat(turn_prompt)
|
||||
with _deadline_scope(timeout_seconds and timeout_seconds - (time_module.time() - start_time)):
|
||||
response = await self.achat(turn_prompt)
|
||||
except Exception as e:
|
||||
return AutonomyResult(
|
||||
# The run deadline cancels an overrunning turn; report it as a timeout.
|
||||
# A budget stop keeps its typed budget_exhausted result when a cap applies.
|
||||
return _timeout_result(
|
||||
timeout_seconds, iterations, stage, actions_taken, start_time, started_at
|
||||
) or (isinstance(e, BudgetExceededError) and self._autonomy_budget_result(
|
||||
iterations, stage, actions_taken, start_time, started_at,
|
||||
"", spend_baseline
|
||||
)) or AutonomyResult(
|
||||
success=False,
|
||||
output=str(e),
|
||||
completion_reason="error",
|
||||
@@ -5321,6 +5365,14 @@ Summary:"""
|
||||
started_at=started_at,
|
||||
)
|
||||
|
||||
# A turn that overran the run deadline or max_execution_time
|
||||
# must not count as progress.
|
||||
_timed_out = _timeout_result(
|
||||
timeout_seconds, iterations, stage, actions_taken, start_time, started_at,
|
||||
turn_timed_out=getattr(self, "_last_stop_reason", None) == "timeout")
|
||||
if _timed_out is not None:
|
||||
return _timed_out
|
||||
|
||||
response_str = str(response)
|
||||
|
||||
# Spend kill-switch: halt if cumulative cost/tokens exceed cap.
|
||||
@@ -5475,7 +5527,7 @@ Summary:"""
|
||||
# produced a substantive response, the inner loop completed
|
||||
# the task naturally (model stopped calling tools = done).
|
||||
# ─────────────────────────────────────────────────────────────
|
||||
if self._autonomy_turn_tool_count > 0 and len(response_str) > 100:
|
||||
if not effective_promise and self._autonomy_turn_tool_count > 0 and len(response_str) > 100:
|
||||
return AutonomyResult(
|
||||
success=True,
|
||||
output=response_str,
|
||||
@@ -5498,7 +5550,8 @@ Summary:"""
|
||||
if obs is not None:
|
||||
obs.emit(_EvtType.STAGE_ESCALATE, {"from": stage_order[idx], "to": stage})
|
||||
|
||||
# Check for completion promise FIRST (structured signal)
|
||||
# Completion promise: when configured, it is the only success signal
|
||||
# (heuristic exits above/below are skipped in promise mode)
|
||||
if effective_promise:
|
||||
promise_tag = f"<promise>{effective_promise}</promise>"
|
||||
if promise_tag in response_str:
|
||||
@@ -5514,7 +5567,7 @@ Summary:"""
|
||||
)
|
||||
|
||||
# Check for keyword-based completion signals (word-boundary, G-COMPLETION-1 fix)
|
||||
if self._is_completion_signal(response_str):
|
||||
if not effective_promise and self._is_completion_signal(response_str):
|
||||
return AutonomyResult(
|
||||
success=True,
|
||||
output=response_str,
|
||||
@@ -5527,7 +5580,7 @@ Summary:"""
|
||||
)
|
||||
|
||||
# For DIRECT stage, complete after first response
|
||||
if stage == "direct":
|
||||
if not effective_promise and stage == "direct":
|
||||
return AutonomyResult(
|
||||
success=True,
|
||||
output=response_str,
|
||||
@@ -5548,7 +5601,7 @@ Summary:"""
|
||||
# Suppressed when doom recovery is active — once a doom loop
|
||||
# is detected, no_tool_calls (a success exit) should not fire;
|
||||
# the doom loop system manages termination instead.
|
||||
if (self._autonomy_turn_tool_count == 0 and iterations > 1
|
||||
if (not effective_promise and self._autonomy_turn_tool_count == 0 and iterations > 1
|
||||
and not self._doom_recovery_active):
|
||||
self._consecutive_no_tool_turns += 1
|
||||
if self._consecutive_no_tool_turns >= 2:
|
||||
@@ -6662,10 +6715,19 @@ Summary:"""
|
||||
if hasattr(self, 'llm_instance') and self.llm_instance:
|
||||
llm = self.llm_instance
|
||||
|
||||
# rerank=True without rerank_model has no backend reranker,
|
||||
# so use the zero-dependency keyword reranker; with
|
||||
# rerank_model the vector store already reranks.
|
||||
reranker = None
|
||||
if rag_config_obj.rerank and not getattr(self._retrieval_config, "rerank_model", None):
|
||||
from praisonaiagents.rag.retriever import SimpleReranker
|
||||
reranker = SimpleReranker()
|
||||
|
||||
self._rag_instance = RAG(
|
||||
knowledge=self.knowledge,
|
||||
config=rag_config_obj,
|
||||
llm=llm,
|
||||
reranker=reranker,
|
||||
)
|
||||
except ImportError:
|
||||
logging.warning("RAG module not available. Install with: pip install 'praisonaiagents[rag]'")
|
||||
@@ -7091,6 +7153,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
|
||||
|
||||
|
||||
@@ -400,6 +400,26 @@ class AutonomyResult:
|
||||
"""
|
||||
return self.output or ""
|
||||
|
||||
|
||||
def _timeout_result(timeout_seconds, iterations, stage, actions, start_time, started_at,
|
||||
turn_timed_out=False):
|
||||
"""Return a ``timeout`` AutonomyResult once a run is past ``timeout_seconds``
|
||||
or the last turn was cut by ``max_execution_time`` (``turn_timed_out``), else None."""
|
||||
import time
|
||||
if not turn_timed_out and (not timeout_seconds or (time.time() - start_time) <= timeout_seconds):
|
||||
return None
|
||||
return AutonomyResult(
|
||||
success=False,
|
||||
output="Task timed out",
|
||||
completion_reason="timeout",
|
||||
iterations=iterations,
|
||||
stage=stage,
|
||||
actions=actions,
|
||||
duration_seconds=time.time() - start_time,
|
||||
started_at=started_at,
|
||||
)
|
||||
|
||||
|
||||
class DoomLoopTracker:
|
||||
"""Tracks actions to detect doom loops with graduated recovery.
|
||||
|
||||
|
||||
@@ -6,6 +6,7 @@ and response processing. Extracted from agent.py for maintainability.
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import contextlib
|
||||
import contextvars
|
||||
import os
|
||||
import re
|
||||
@@ -29,6 +30,14 @@ _turns_lock = threading.Lock()
|
||||
# into tool worker threads with the context, so each turn has its own counters and
|
||||
# overlapping turns on a shared agent never exhaust each other's limits.
|
||||
_turn_loop_guard = contextvars.ContextVar("praisonai_turn_loop_guard", default=None)
|
||||
# Agent whose streamed turn already records agent_start/agent_end, set while its
|
||||
# non-streaming chat() fallback runs so that turn is not recorded twice.
|
||||
_stream_traced_agent = contextvars.ContextVar("praisonai_stream_traced_agent", default=None)
|
||||
|
||||
|
||||
def _agent_end_status(db_status: str) -> str:
|
||||
"""``agent_end`` status for replay: completed, failed or cancelled."""
|
||||
return "failed" if db_status == "error" else db_status
|
||||
|
||||
|
||||
def _enter_turn(agent, cancel):
|
||||
@@ -97,20 +106,50 @@ def _sleep_unless_cancelled(seconds: float, cancel_token: Any = None, step: floa
|
||||
time.sleep(min(step, remaining))
|
||||
|
||||
|
||||
# ``time.monotonic()`` deadline for turns started in this context, set by
|
||||
# ``run_autonomous(timeout_seconds=...)`` so a single long turn cannot overrun it.
|
||||
_run_deadline = contextvars.ContextVar("praisonai_run_deadline", default=None)
|
||||
|
||||
|
||||
@contextlib.contextmanager
|
||||
def _deadline_scope(seconds):
|
||||
"""Bound every turn started in this block to ``seconds`` from now (no-op when falsy)."""
|
||||
if not seconds:
|
||||
yield
|
||||
return
|
||||
deadline = time.monotonic() + seconds
|
||||
outer = _run_deadline.get()
|
||||
token = _run_deadline.set(deadline if outer is None else min(deadline, outer))
|
||||
try:
|
||||
yield
|
||||
finally:
|
||||
_run_deadline.reset(token)
|
||||
|
||||
|
||||
class _TurnCancelToken:
|
||||
"""Track cancellation observed by one call without shared backend state."""
|
||||
"""Track cancellation observed by one call without shared backend state.
|
||||
|
||||
An optional ``deadline`` (``time.monotonic()``) makes the token report set
|
||||
once it passes, so the existing cancel checks also stop an overrunning turn.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
source: Any,
|
||||
*,
|
||||
turn_id: Optional[int] = None,
|
||||
deadline: Optional[float] = None,
|
||||
):
|
||||
self._source = source
|
||||
self._observed = False
|
||||
self._turn_id = turn_id
|
||||
self._deadline = deadline
|
||||
self.timed_out = False
|
||||
|
||||
def is_set(self) -> bool:
|
||||
if self._deadline is not None and time.monotonic() >= self._deadline:
|
||||
self.timed_out = self._observed = True
|
||||
return True
|
||||
turn_checker = getattr(self._source, "_turn_is_cancelled", None)
|
||||
if self._turn_id is not None and callable(turn_checker):
|
||||
is_set = bool(turn_checker(self._turn_id))
|
||||
@@ -121,6 +160,8 @@ class _TurnCancelToken:
|
||||
|
||||
@property
|
||||
def reason(self):
|
||||
if self.timed_out:
|
||||
return "timeout"
|
||||
return getattr(self._source, "reason", None)
|
||||
|
||||
def check(self) -> None:
|
||||
@@ -157,13 +198,22 @@ class ChatMixin:
|
||||
return get_durable_run()
|
||||
|
||||
def _turn_cancel_token(self, source: Any, *, explicit: bool):
|
||||
"""Create per-turn cancellation state while consuming default requests once."""
|
||||
"""Create per-turn cancellation state while consuming default requests once.
|
||||
|
||||
``ExecutionConfig.max_execution_time`` and any enclosing
|
||||
``run_autonomous`` timeout become the token's deadline.
|
||||
"""
|
||||
limit = getattr(self, "max_execution_time", None)
|
||||
deadline = time.monotonic() + limit if limit else None
|
||||
outer = _run_deadline.get()
|
||||
if outer is not None:
|
||||
deadline = outer if deadline is None else min(deadline, outer)
|
||||
if source is None:
|
||||
return None
|
||||
return None if deadline is None else _TurnCancelToken(None, deadline=deadline)
|
||||
begin_turn = getattr(source, "_begin_turn", None)
|
||||
if explicit or not callable(begin_turn):
|
||||
return _TurnCancelToken(source)
|
||||
return _TurnCancelToken(source, turn_id=begin_turn())
|
||||
return _TurnCancelToken(source, deadline=deadline)
|
||||
return _TurnCancelToken(source, turn_id=begin_turn(), deadline=deadline)
|
||||
|
||||
@staticmethod
|
||||
def _notify_tool_call(tool_name, tool_input, tool_result, elapsed_time=None, success=True):
|
||||
@@ -1912,6 +1962,107 @@ Your Goal: {self.goal}"""
|
||||
self._openai_client = get_openai_client(api_key=getattr(self, 'api_key', None))
|
||||
return create_llm_dispatcher(openai_client=self._openai_client, model=model)
|
||||
|
||||
def _fallback_llm_for(self, model):
|
||||
"""Build a throwaway LLM bound to fallback ``model`` for the custom-LLM path.
|
||||
|
||||
Never assigned to ``self.llm_instance``, so concurrent turns keep the
|
||||
primary model. A fallback entry that is already a model object is used
|
||||
as-is. Option params (metrics, max_iter, retries, ...) are inherited;
|
||||
connection settings (base_url, api_key, auth) only when the fallback
|
||||
targets the same provider prefix as the primary.
|
||||
"""
|
||||
if not isinstance(model, str):
|
||||
return model
|
||||
from ..llm.llm import LLM
|
||||
params = dict(getattr(self, '_llm_init_params', None) or {})
|
||||
primary = str(params.get('model') or self.llm or '')
|
||||
# Bare names (no "provider/" prefix) share the configured server.
|
||||
prefix = lambda name: name.split('/', 1)[0] if '/' in name else ''
|
||||
if prefix(primary) != prefix(model):
|
||||
for key in ('base_url', 'api_key', 'auth'):
|
||||
params.pop(key, None)
|
||||
params['model'] = model
|
||||
return LLM(**params)
|
||||
|
||||
def _classify_custom_llm_fallback(self, error, failed_model, fallback_index):
|
||||
"""Return ``(next_model, category)`` if ``error`` warrants a model fallback."""
|
||||
next_model = self._get_next_fallback_model(fallback_index)
|
||||
if next_model is None:
|
||||
return None, None
|
||||
from ..llm.error_classifier import classify_llm_error
|
||||
model_name = str(failed_model)
|
||||
provider = model_name.split('/', 1)[0] if '/' in model_name else "openai"
|
||||
classification = classify_llm_error(error, provider=provider, model=model_name)
|
||||
if not classification.should_fallback_model:
|
||||
return None, None
|
||||
logging.info(f"[{self.name}] {model_name} unavailable — falling back to {next_model}")
|
||||
return next_model, classification.error_category
|
||||
|
||||
@staticmethod
|
||||
def _track_tool_runs(llm_kwargs):
|
||||
"""Return ``(llm_kwargs, ran)``; ``ran`` becomes non-empty once a tool executes.
|
||||
|
||||
A fallback replays the whole turn from the original messages, so it is
|
||||
only safe while no tool has run (otherwise side effects would repeat).
|
||||
"""
|
||||
import functools
|
||||
import inspect
|
||||
ran = []
|
||||
fn = llm_kwargs.get('execute_tool_fn')
|
||||
if fn is None:
|
||||
return llm_kwargs, ran
|
||||
if inspect.iscoroutinefunction(fn):
|
||||
@functools.wraps(fn)
|
||||
async def tracked(*args, **kwargs):
|
||||
ran.append(True)
|
||||
return await fn(*args, **kwargs)
|
||||
else:
|
||||
@functools.wraps(fn)
|
||||
def tracked(*args, **kwargs):
|
||||
ran.append(True)
|
||||
return fn(*args, **kwargs)
|
||||
return {**llm_kwargs, 'execute_tool_fn': tracked}, ran
|
||||
|
||||
def _custom_llm_get_response(self, llm_kwargs):
|
||||
"""Call ``llm_instance.get_response`` walking ``fallback_models`` on outage."""
|
||||
llm, index, current = self.llm_instance, 0, self.llm
|
||||
stream_cb = self.stream_emitter.emit if hasattr(self, 'stream_emitter') else None
|
||||
llm_kwargs, tools_ran = self._track_tool_runs(llm_kwargs)
|
||||
while True:
|
||||
try:
|
||||
with self._llm_usage_sink(llm):
|
||||
return llm.get_response(**llm_kwargs)
|
||||
except (ToolExecutionError, BudgetExceededError, InterruptedError):
|
||||
raise
|
||||
except Exception as e:
|
||||
if tools_ran:
|
||||
raise
|
||||
next_model, category = self._classify_custom_llm_fallback(e, current, index)
|
||||
if next_model is None:
|
||||
raise
|
||||
self._emit_model_fallback(current, next_model, category, index, stream_cb)
|
||||
llm, index, current = self._fallback_llm_for(next_model), index + 1, next_model
|
||||
|
||||
async def _custom_llm_get_response_async(self, llm_kwargs):
|
||||
"""Async twin of :meth:`_custom_llm_get_response`."""
|
||||
llm, index, current = self.llm_instance, 0, self.llm
|
||||
stream_cb = self.stream_emitter.emit if hasattr(self, 'stream_emitter') else None
|
||||
llm_kwargs, tools_ran = self._track_tool_runs(llm_kwargs)
|
||||
while True:
|
||||
try:
|
||||
with self._llm_usage_sink(llm):
|
||||
return await llm.get_response_async(**llm_kwargs)
|
||||
except (ToolExecutionError, BudgetExceededError, InterruptedError):
|
||||
raise
|
||||
except Exception as e:
|
||||
if tools_ran:
|
||||
raise
|
||||
next_model, category = self._classify_custom_llm_fallback(e, current, index)
|
||||
if next_model is None:
|
||||
raise
|
||||
await self._aemit_model_fallback(current, next_model, category, index, stream_cb)
|
||||
llm, index, current = self._fallback_llm_for(next_model), index + 1, next_model
|
||||
|
||||
def _get_model_fallback_lock(self):
|
||||
"""Lazily create the per-instance lock guarding model-fallback swaps.
|
||||
|
||||
@@ -2075,6 +2226,82 @@ Your Goal: {self.goal}"""
|
||||
return max_retries
|
||||
return 2
|
||||
|
||||
def _extract_llm_usage(self, response, model=None):
|
||||
"""Return ``(prompt_tokens, completion_tokens, cost_usd)`` for one LLM response priced as ``model``."""
|
||||
usage = getattr(response, 'usage', None) if response else None
|
||||
if not usage:
|
||||
return 0, 0, 0.0
|
||||
# Chat Completions names, then the Responses API names.
|
||||
prompt_tokens = getattr(usage, 'prompt_tokens', 0) or getattr(usage, 'input_tokens', 0) or 0
|
||||
completion_tokens = getattr(usage, 'completion_tokens', 0) or getattr(usage, 'output_tokens', 0) or 0
|
||||
# Calculate cost using litellm (if available) or fallback pricing
|
||||
cost_usd = self._calculate_llm_cost(prompt_tokens, completion_tokens, response=response, model=model)
|
||||
return prompt_tokens, completion_tokens, cost_usd
|
||||
|
||||
def _record_llm_usage(self, prompt_tokens, completion_tokens, cost_usd):
|
||||
"""Add one LLM call to the agent's cost/token totals and enforce ``max_budget``.
|
||||
|
||||
The single accounting point for every path: the OpenAI-native
|
||||
completions call it directly and the custom-LLM (``provider/model``)
|
||||
tool loop reaches it through the LLM usage sink, once per call.
|
||||
"""
|
||||
# Thread-safe cost tracking (Gap 1a fix)
|
||||
with self._cost_lock:
|
||||
self._total_cost += cost_usd
|
||||
self._total_tokens_in += prompt_tokens
|
||||
self._total_tokens_out += completion_tokens
|
||||
self._llm_call_count += 1
|
||||
budget_exceeded = self._max_budget and self._total_cost >= self._max_budget
|
||||
current_cost = self._total_cost
|
||||
|
||||
if budget_exceeded:
|
||||
if self._on_budget_exceeded == "stop":
|
||||
raise BudgetExceededError(
|
||||
f"Agent '{self.name}' exceeded budget: ${current_cost:.4f} >= ${self._max_budget:.4f}",
|
||||
budget_type="cost",
|
||||
limit=self._max_budget,
|
||||
used=current_cost,
|
||||
agent_id=self.name
|
||||
)
|
||||
elif self._on_budget_exceeded == "warn":
|
||||
logging.warning(
|
||||
f"[budget] {self.name}: ${current_cost:.4f} exceeded "
|
||||
f"${self._max_budget:.4f} budget"
|
||||
)
|
||||
elif callable(self._on_budget_exceeded):
|
||||
self._on_budget_exceeded(current_cost, self._max_budget)
|
||||
|
||||
@contextlib.contextmanager
|
||||
def _llm_usage_sink(self, llm=None):
|
||||
"""Route every call ``llm`` (default: the custom LLM) makes in this block into ``_record_llm_usage``.
|
||||
|
||||
Refuses to start when ``max_budget`` is already spent (``"stop"`` mode),
|
||||
so retrying after ``BudgetExceededError`` sends no further paid request.
|
||||
"""
|
||||
if self._max_budget and self._on_budget_exceeded == "stop":
|
||||
with self._cost_lock:
|
||||
current_cost = self._total_cost
|
||||
if current_cost >= self._max_budget:
|
||||
raise BudgetExceededError(
|
||||
f"Agent '{self.name}' exceeded budget: ${current_cost:.4f} >= ${self._max_budget:.4f}",
|
||||
budget_type="cost",
|
||||
limit=self._max_budget,
|
||||
used=current_cost,
|
||||
agent_id=self.name
|
||||
)
|
||||
from ..llm.llm import _usage_sink
|
||||
llm = llm if llm is not None else self.llm_instance
|
||||
# A fallback model is priced as itself, not as the primary ``self.llm``.
|
||||
model = None if llm is self.llm_instance else getattr(llm, 'model', None)
|
||||
token = _usage_sink.set((
|
||||
llm,
|
||||
lambda response: self._record_llm_usage(*self._extract_llm_usage(response, model)),
|
||||
))
|
||||
try:
|
||||
yield
|
||||
finally:
|
||||
_usage_sink.reset(token)
|
||||
|
||||
def _chat_completion(self, messages, temperature=None, tools=None, stream=None, reasoning_steps=False, task_name=None, task_description=None, task_id=None, response_format=None, _retry_depth=0, _fallback_index=0, cancel_token=None):
|
||||
start_time = time.time()
|
||||
|
||||
@@ -2274,19 +2501,8 @@ Your Goal: {self.goal}"""
|
||||
|
||||
# Emit LLM response trace event with token usage
|
||||
_duration_ms = (time.time() - start_time) * 1000
|
||||
_prompt_tokens = 0
|
||||
_completion_tokens = 0
|
||||
_cost_usd = 0.0
|
||||
|
||||
# Extract token usage from response if available
|
||||
if final_response:
|
||||
_usage = getattr(final_response, 'usage', None)
|
||||
if _usage:
|
||||
_prompt_tokens = getattr(_usage, 'prompt_tokens', 0) or 0
|
||||
_completion_tokens = getattr(_usage, 'completion_tokens', 0) or 0
|
||||
# Calculate cost using litellm (if available) or fallback pricing
|
||||
_cost_usd = self._calculate_llm_cost(_prompt_tokens, _completion_tokens, response=final_response)
|
||||
|
||||
_prompt_tokens, _completion_tokens, _cost_usd = self._extract_llm_usage(final_response)
|
||||
|
||||
_trace_emitter.llm_response(
|
||||
self.name,
|
||||
duration_ms=_duration_ms,
|
||||
@@ -2296,33 +2512,8 @@ Your Goal: {self.goal}"""
|
||||
cost_usd=_cost_usd,
|
||||
)
|
||||
|
||||
# Budget tracking & enforcement (zero overhead when _max_budget is None)
|
||||
# Thread-safe cost tracking (Gap 1a fix)
|
||||
with self._cost_lock:
|
||||
self._total_cost += _cost_usd
|
||||
self._total_tokens_in += _prompt_tokens
|
||||
self._total_tokens_out += _completion_tokens
|
||||
self._llm_call_count += 1
|
||||
budget_exceeded = self._max_budget and self._total_cost >= self._max_budget
|
||||
current_cost = self._total_cost
|
||||
|
||||
if budget_exceeded:
|
||||
if self._on_budget_exceeded == "stop":
|
||||
raise BudgetExceededError(
|
||||
f"Agent '{self.name}' exceeded budget: ${current_cost:.4f} >= ${self._max_budget:.4f}",
|
||||
budget_type="cost",
|
||||
limit=self._max_budget,
|
||||
used=current_cost,
|
||||
agent_id=self.name
|
||||
)
|
||||
elif self._on_budget_exceeded == "warn":
|
||||
logging.warning(
|
||||
f"[budget] {self.name}: ${current_cost:.4f} exceeded "
|
||||
f"${self._max_budget:.4f} budget"
|
||||
)
|
||||
elif callable(self._on_budget_exceeded):
|
||||
self._on_budget_exceeded(current_cost, self._max_budget)
|
||||
|
||||
self._record_llm_usage(_prompt_tokens, _completion_tokens, _cost_usd)
|
||||
|
||||
# Trigger AFTER_LLM hook
|
||||
# Only build the input if an AFTER_LLM hook is actually registered
|
||||
from ..hooks import HookEvent, AfterLLMInput
|
||||
@@ -2336,6 +2527,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)
|
||||
@@ -2990,40 +3182,8 @@ Your Goal: {self.goal}"""
|
||||
)
|
||||
|
||||
# Post-call budget accounting (parity with sync _chat_completion).
|
||||
_prompt_tokens = 0
|
||||
_completion_tokens = 0
|
||||
_cost_usd = 0.0
|
||||
if final_response:
|
||||
_usage = getattr(final_response, 'usage', None)
|
||||
if _usage:
|
||||
_prompt_tokens = getattr(_usage, 'prompt_tokens', 0) or 0
|
||||
_completion_tokens = getattr(_usage, 'completion_tokens', 0) or 0
|
||||
_cost_usd = self._calculate_llm_cost(
|
||||
_prompt_tokens, _completion_tokens, response=final_response
|
||||
)
|
||||
with self._cost_lock:
|
||||
self._total_cost += _cost_usd
|
||||
self._total_tokens_in += _prompt_tokens
|
||||
self._total_tokens_out += _completion_tokens
|
||||
self._llm_call_count += 1
|
||||
budget_exceeded = self._max_budget and self._total_cost >= self._max_budget
|
||||
current_cost = self._total_cost
|
||||
if budget_exceeded:
|
||||
if self._on_budget_exceeded == "stop":
|
||||
raise BudgetExceededError(
|
||||
f"Agent '{self.name}' exceeded budget: ${current_cost:.4f} >= ${self._max_budget:.4f}",
|
||||
budget_type="cost",
|
||||
limit=self._max_budget,
|
||||
used=current_cost,
|
||||
agent_id=self.name
|
||||
)
|
||||
elif self._on_budget_exceeded == "warn":
|
||||
logging.warning(
|
||||
f"[budget] {self.name}: ${current_cost:.4f} exceeded "
|
||||
f"${self._max_budget:.4f} budget"
|
||||
)
|
||||
elif callable(self._on_budget_exceeded):
|
||||
self._on_budget_exceeded(current_cost, self._max_budget)
|
||||
_prompt_tokens, _completion_tokens, _cost_usd = self._extract_llm_usage(final_response)
|
||||
self._record_llm_usage(_prompt_tokens, _completion_tokens, _cost_usd)
|
||||
|
||||
# Trigger AFTER_LLM hook (parity with sync _chat_completion).
|
||||
from ..hooks import AfterLLMInput
|
||||
@@ -3037,6 +3197,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)
|
||||
@@ -3448,7 +3609,9 @@ Your Goal: {self.goal}"""
|
||||
# Emit context trace event (zero overhead when not set)
|
||||
from ..trace.context_events import get_context_emitter
|
||||
_trace_emitter = get_context_emitter()
|
||||
_trace_emitter.agent_start(self.name, {"role": self.role, "goal": self.goal})
|
||||
_trace_turn = _stream_traced_agent.get() is not self
|
||||
if _trace_turn:
|
||||
_trace_emitter.agent_start(self.name, {"role": self.role, "goal": self.goal})
|
||||
durable_context = None
|
||||
durable_token = None
|
||||
previous_run_state = self._get_db_run_state()
|
||||
@@ -3517,12 +3680,15 @@ Your Goal: {self.goal}"""
|
||||
end_durable_run(durable_token)
|
||||
durable_context.close()
|
||||
try:
|
||||
_trace_emitter.agent_end(self.name)
|
||||
if _trace_turn:
|
||||
_trace_emitter.agent_end(self.name, {"status": _agent_end_status(db_status)})
|
||||
finally:
|
||||
if getattr(self, "_active_turn_token", None) is _cancel:
|
||||
self._active_turn_token = None
|
||||
_exit_turn(self, _cancel_ctx)
|
||||
if _cancel is not None:
|
||||
if getattr(_cancel, "timed_out", False):
|
||||
self._last_stop_reason = "timeout"
|
||||
_cancel.close()
|
||||
|
||||
def _chat_impl(self, prompt, temperature, tools, output_json, output_pydantic, reasoning_steps, stream, task_name, task_description, task_id, config, force_retrieval, skip_retrieval, attachments, _trace_emitter, tool_choice=None, seed=None, cancel_token=None):
|
||||
@@ -3849,7 +4015,7 @@ Your Goal: {self.goal}"""
|
||||
if hasattr(self, '_drain_steering_messages') and getattr(self, '_message_steering', None) is not None:
|
||||
llm_kwargs['steering_drain'] = self._drain_steering_messages
|
||||
|
||||
response_text = self.llm_instance.get_response(**llm_kwargs)
|
||||
response_text = self._custom_llm_get_response(llm_kwargs)
|
||||
|
||||
self._add_to_chat_history("assistant", response_text)
|
||||
# Persist assistant message to DB
|
||||
@@ -3880,7 +4046,7 @@ Your Goal: {self.goal}"""
|
||||
# Rollback chat history on guardrail failure
|
||||
self._rollback_chat_history_to(chat_history_length)
|
||||
return self._guardrail_blocked_message(e)
|
||||
except ToolExecutionError:
|
||||
except (ToolExecutionError, BudgetExceededError):
|
||||
raise
|
||||
except Exception as e:
|
||||
self._last_stop_reason = "error"
|
||||
@@ -3888,7 +4054,7 @@ Your Goal: {self.goal}"""
|
||||
self._rollback_chat_history_to(chat_history_length)
|
||||
_get_display_functions()['display_error'](f"Error in LLM chat: {e}")
|
||||
return None
|
||||
except ToolExecutionError:
|
||||
except (ToolExecutionError, BudgetExceededError):
|
||||
raise
|
||||
except Exception as e:
|
||||
self._last_stop_reason = "error"
|
||||
@@ -4240,11 +4406,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
|
||||
@@ -4360,12 +4532,14 @@ Output MUST be JSON with 'reflection' and 'satisfactory'.
|
||||
await durable_context.aclose()
|
||||
finally:
|
||||
try:
|
||||
_trace_emitter.agent_end(self.name)
|
||||
_trace_emitter.agent_end(self.name, {"status": _agent_end_status(db_status)})
|
||||
finally:
|
||||
if getattr(self, "_active_turn_token", None) is _cancel:
|
||||
self._active_turn_token = None
|
||||
_exit_turn(self, _cancel_ctx)
|
||||
if _cancel is not None:
|
||||
if getattr(_cancel, "timed_out", False):
|
||||
self._last_stop_reason = "timeout"
|
||||
_cancel.close()
|
||||
|
||||
async def _achat_impl(self, prompt, temperature, tools, output_json, output_pydantic, reasoning_steps, stream, task_name, task_description, task_id, config, force_retrieval, skip_retrieval, attachments, _trace_emitter, tool_choice=None, seed=None, cancel_token=None):
|
||||
@@ -4576,7 +4750,7 @@ Output MUST be JSON with 'reflection' and 'satisfactory'.
|
||||
if hasattr(self, '_drain_steering_messages') and getattr(self, '_message_steering', None) is not None:
|
||||
llm_kwargs['steering_drain'] = self._drain_steering_messages
|
||||
|
||||
response_text = await self.llm_instance.get_response_async(**llm_kwargs)
|
||||
response_text = await self._custom_llm_get_response_async(llm_kwargs)
|
||||
|
||||
# LLM call succeeded - now it's safe to commit any compacted history
|
||||
if effective_history is not self.chat_history:
|
||||
@@ -4604,7 +4778,7 @@ Output MUST be JSON with 'reflection' and 'satisfactory'.
|
||||
# Rollback chat history on guardrail failure
|
||||
self._rollback_chat_history_to(chat_history_length)
|
||||
return self._guardrail_blocked_message(e)
|
||||
except ToolExecutionError:
|
||||
except (ToolExecutionError, BudgetExceededError):
|
||||
raise
|
||||
except Exception as e:
|
||||
self._last_stop_reason = "error"
|
||||
@@ -5118,6 +5292,9 @@ Output MUST be JSON with 'reflection' and 'satisfactory'.
|
||||
except ToolExecutionError:
|
||||
raise
|
||||
except Exception as e:
|
||||
# The custom-LLM path raises on an exhausted budget (sync parity).
|
||||
if isinstance(e, BudgetExceededError) and self._using_custom_llm:
|
||||
raise
|
||||
self._last_stop_reason = "error"
|
||||
_get_display_functions()['display_error'](f"Error in achat: {e}")
|
||||
if get_logger().getEffectiveLevel() == logging.DEBUG:
|
||||
@@ -5319,9 +5496,26 @@ Output MUST be JSON with 'reflection' and 'satisfactory'.
|
||||
# via _enter_turn) so overlapping streams on a shared agent neither
|
||||
# exhaust nor inherit each other's per-turn limits.
|
||||
guard_token = _turn_loop_guard.set([self, None])
|
||||
# Record the streamed turn like chat()/achat() do, so a live replay can
|
||||
# show the agent running while its words stream (no-op when tracing is off).
|
||||
from ..trace.context_events import get_context_emitter
|
||||
_trace_emitter = get_context_emitter()
|
||||
status = "completed"
|
||||
if _trace_emitter._enabled:
|
||||
_trace_emitter.agent_start(self.name, {"role": self.role, "goal": self.goal})
|
||||
try:
|
||||
yield from self._start_stream_turn(prompt, **kwargs)
|
||||
# _start_stream_turn returns "failed" when a guardrail or hook blocks
|
||||
# the turn, matching chat(), which records a blocked turn as failed.
|
||||
status = (yield from self._start_stream_turn(prompt, **kwargs)) or status
|
||||
except (GeneratorExit, InterruptedError):
|
||||
status = "cancelled"
|
||||
raise
|
||||
except BaseException:
|
||||
status = "failed"
|
||||
raise
|
||||
finally:
|
||||
if _trace_emitter._enabled:
|
||||
_trace_emitter.agent_end(self.name, {"status": status})
|
||||
try:
|
||||
_turn_loop_guard.reset(guard_token)
|
||||
except (ValueError, LookupError):
|
||||
@@ -5346,7 +5540,7 @@ Output MUST be JSON with 'reflection' and 'satisfactory'.
|
||||
if not _in_ok:
|
||||
logging.warning(f"Agent {getattr(self, 'name', '')}: input blocked by guardrail: {_in_err}")
|
||||
yield f"[Input blocked by guardrail: {_in_err}]"
|
||||
return
|
||||
return "failed"
|
||||
prompt = _in_prompt
|
||||
|
||||
# Same pre-turn checks as chat()/achat(): a streamed turn must not
|
||||
@@ -5367,7 +5561,7 @@ Output MUST be JSON with 'reflection' and 'satisfactory'.
|
||||
))
|
||||
if self._hook_runner.is_blocked(hook_results):
|
||||
logging.warning(f"Agent {self.name} execution blocked by BEFORE_AGENT hook")
|
||||
return
|
||||
return "failed"
|
||||
for res in hook_results:
|
||||
if res.output and res.output.modified_input and "prompt" in res.output.modified_input:
|
||||
prompt = res.output.modified_input["prompt"]
|
||||
@@ -6187,6 +6381,9 @@ Output MUST be JSON with 'reflection' and 'satisfactory'.
|
||||
else:
|
||||
fallback_kwargs.pop('stream', None)
|
||||
|
||||
# The streamed turn already records agent_start/agent_end; keep chat()
|
||||
# from recording the same turn a second time.
|
||||
_traced = _stream_traced_agent.set(self)
|
||||
try:
|
||||
return self.chat(prompt, **fallback_kwargs)
|
||||
except Exception as fallback_error:
|
||||
@@ -6199,6 +6396,8 @@ Output MUST be JSON with 'reflection' and 'satisfactory'.
|
||||
except Exception:
|
||||
pass
|
||||
raise fallback_error from streaming_error
|
||||
finally:
|
||||
_stream_traced_agent.reset(_traced)
|
||||
|
||||
def _create_llm_summarize_function(self):
|
||||
"""
|
||||
@@ -6668,6 +6867,57 @@ Output MUST be JSON with 'reflection' and 'satisfactory'.
|
||||
raise RuntimeError("Retry loop completed without returning or raising an exception")
|
||||
|
||||
async def _achat_completion_with_retry(self, messages, temperature=None, tools=None, stream=None, reasoning_steps=False, task_name=None, task_description=None, task_id=None, response_format=None, stream_callback=None, emit_events=True):
|
||||
"""Async model call with retry, recorded as ``llm_request``/``llm_response`` like the sync path.
|
||||
|
||||
``Agent.astart``/``achat`` recorded no model calls before this, so a replay
|
||||
of an async run showed the agent but never what it asked the model.
|
||||
"""
|
||||
from ..trace.context_events import get_context_emitter
|
||||
|
||||
_trace_emitter = get_context_emitter()
|
||||
call = self._achat_completion_with_retry_dispatch(
|
||||
messages, temperature, tools, stream, reasoning_steps, task_name,
|
||||
task_description, task_id, response_format,
|
||||
stream_callback=stream_callback, emit_events=emit_events,
|
||||
)
|
||||
if not _trace_emitter._enabled: # zero overhead when tracing is off
|
||||
return await call
|
||||
_trace_emitter.llm_request(
|
||||
self.name,
|
||||
messages_count=len(messages),
|
||||
tokens_used=0,
|
||||
model=self.llm if isinstance(self.llm, str) else None,
|
||||
messages=messages,
|
||||
)
|
||||
started = time.time()
|
||||
try:
|
||||
response = await call
|
||||
except Exception as e:
|
||||
_trace_emitter.llm_response(
|
||||
self.name,
|
||||
duration_ms=(time.time() - started) * 1000,
|
||||
finish_reason="error",
|
||||
response_content=str(e),
|
||||
)
|
||||
raise
|
||||
prompt_tokens = completion_tokens = 0
|
||||
cost_usd = 0.0
|
||||
usage = getattr(response, "usage", None)
|
||||
if usage:
|
||||
prompt_tokens = getattr(usage, "prompt_tokens", 0) or 0
|
||||
completion_tokens = getattr(usage, "completion_tokens", 0) or 0
|
||||
cost_usd = self._calculate_llm_cost(prompt_tokens, completion_tokens, response=response)
|
||||
_trace_emitter.llm_response(
|
||||
self.name,
|
||||
duration_ms=(time.time() - started) * 1000,
|
||||
response_content=self._extract_llm_response_content(response),
|
||||
prompt_tokens=prompt_tokens,
|
||||
completion_tokens=completion_tokens,
|
||||
cost_usd=cost_usd,
|
||||
)
|
||||
return response
|
||||
|
||||
async def _achat_completion_with_retry_dispatch(self, messages, temperature=None, tools=None, stream=None, reasoning_steps=False, task_name=None, task_description=None, task_id=None, response_format=None, stream_callback=None, emit_events=True):
|
||||
"""
|
||||
Async wrapper for _execute_unified_achat_completion that adds jittered exponential backoff retry logic.
|
||||
|
||||
|
||||
@@ -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}")
|
||||
|
||||
@@ -497,157 +497,111 @@ class ExecutionMixin:
|
||||
# The backend's stream method is async, so we need to handle that
|
||||
import asyncio
|
||||
|
||||
async def stream_wrapper():
|
||||
async for chunk in self.backend.stream(prompt, **kwargs):
|
||||
yield chunk
|
||||
|
||||
# Convert async generator to sync generator
|
||||
import threading
|
||||
import queue
|
||||
|
||||
def _to_text(item):
|
||||
# Managed backends may yield event dicts; convert them to the
|
||||
# string chunks expected by Agent and skip other event types
|
||||
# (session.status_idle, etc.).
|
||||
if isinstance(item, str):
|
||||
return item
|
||||
if isinstance(item, dict) and item.get('type') == 'agent.message':
|
||||
content = item.get('content', [])
|
||||
if isinstance(content, str):
|
||||
return content
|
||||
if isinstance(content, list):
|
||||
text_parts = []
|
||||
for block in content:
|
||||
if isinstance(block, dict) and block.get('type') == 'text':
|
||||
text_parts.append(block.get('text', ''))
|
||||
elif isinstance(block, str):
|
||||
text_parts.append(block)
|
||||
if text_parts:
|
||||
return ''.join(text_parts)
|
||||
return None
|
||||
|
||||
# Convert async generator to sync generator. The backend's async
|
||||
# generator is always driven on a private loop in a background
|
||||
# thread, whether or not the caller already has a running loop, so
|
||||
# the caller receives chunks incrementally and a backend error
|
||||
# (including RuntimeError) can never re-run the stream.
|
||||
def sync_stream():
|
||||
try:
|
||||
loop = asyncio.get_running_loop()
|
||||
# Already in async context - need to handle differently
|
||||
import concurrent.futures
|
||||
import threading
|
||||
import queue
|
||||
|
||||
result_queue = queue.Queue()
|
||||
exception_holder = [None]
|
||||
|
||||
def run_in_thread():
|
||||
new_loop = asyncio.new_event_loop()
|
||||
asyncio.set_event_loop(new_loop)
|
||||
# Bounded queue so an abandoned consumer applies back-pressure
|
||||
# instead of letting the producer fill memory unbounded.
|
||||
# ``cancel_event`` lets the producer stop early when the caller
|
||||
# closes the generator.
|
||||
bridge_queue = queue.Queue(maxsize=64)
|
||||
cancel_event = threading.Event()
|
||||
|
||||
def run_in_thread():
|
||||
new_loop = asyncio.new_event_loop()
|
||||
asyncio.set_event_loop(new_loop)
|
||||
try:
|
||||
async def pump():
|
||||
agen = self.backend.stream(prompt, **kwargs)
|
||||
try:
|
||||
async for item in agen:
|
||||
if cancel_event.is_set():
|
||||
break
|
||||
# Block until space is available or the
|
||||
# consumer signals cancellation, so we
|
||||
# never spin nor grow the queue forever.
|
||||
while not cancel_event.is_set():
|
||||
try:
|
||||
bridge_queue.put(('item', item), timeout=0.1)
|
||||
break
|
||||
except queue.Full:
|
||||
continue
|
||||
if not cancel_event.is_set():
|
||||
bridge_queue.put(('done', None))
|
||||
except Exception as e:
|
||||
if not cancel_event.is_set():
|
||||
bridge_queue.put(('error', e))
|
||||
finally:
|
||||
# Run the backend's own cleanup (e.g. stopping
|
||||
# its producer) when the consumer closes early.
|
||||
aclose = getattr(agen, 'aclose', None)
|
||||
if aclose is not None:
|
||||
try:
|
||||
await aclose()
|
||||
except Exception as e:
|
||||
logger.debug(f"Backend stream cleanup failed: {e}")
|
||||
|
||||
new_loop.run_until_complete(pump())
|
||||
finally:
|
||||
try:
|
||||
async def collect():
|
||||
try:
|
||||
async for item in self.backend.stream(prompt, **kwargs):
|
||||
result_queue.put(('item', item))
|
||||
result_queue.put(('done', None))
|
||||
except Exception as e:
|
||||
exception_holder[0] = e
|
||||
result_queue.put(('error', e))
|
||||
|
||||
new_loop.run_until_complete(collect())
|
||||
new_loop.run_until_complete(new_loop.shutdown_asyncgens())
|
||||
finally:
|
||||
new_loop.close()
|
||||
|
||||
thread = threading.Thread(target=run_in_thread)
|
||||
thread.start()
|
||||
|
||||
|
||||
# Daemon thread so a stalled/abandoned producer can never
|
||||
# block interpreter shutdown.
|
||||
thread = threading.Thread(target=run_in_thread, daemon=True)
|
||||
thread.start()
|
||||
|
||||
try:
|
||||
while True:
|
||||
msg_type, data = result_queue.get()
|
||||
msg_type, data = bridge_queue.get()
|
||||
if msg_type == 'item':
|
||||
# For managed backends, we might get event objects
|
||||
# Convert to string format expected by Agent
|
||||
if isinstance(data, dict):
|
||||
if data.get('type') == 'agent.message':
|
||||
content = data.get('content', [])
|
||||
if isinstance(content, list):
|
||||
text_parts = []
|
||||
for block in content:
|
||||
if isinstance(block, dict) and block.get('type') == 'text':
|
||||
text_parts.append(block.get('text', ''))
|
||||
elif isinstance(block, str):
|
||||
text_parts.append(block)
|
||||
if text_parts:
|
||||
yield ''.join(text_parts)
|
||||
elif isinstance(content, str):
|
||||
yield content
|
||||
# Skip other event types (session.status_idle, etc.)
|
||||
elif isinstance(data, str):
|
||||
yield data
|
||||
text = _to_text(data)
|
||||
if text is not None:
|
||||
yield text
|
||||
elif msg_type == 'done':
|
||||
break
|
||||
elif msg_type == 'error':
|
||||
raise data
|
||||
|
||||
thread.join()
|
||||
|
||||
except RuntimeError:
|
||||
# No event loop (the common script/CLI case). Drive the
|
||||
# backend's async generator from a background thread and
|
||||
# push each chunk through a queue as it arrives, so the
|
||||
# caller receives chunks incrementally instead of only
|
||||
# after the whole response has been buffered. Mirrors the
|
||||
# running-loop branch above.
|
||||
import threading
|
||||
import queue
|
||||
# Bounded queue so an abandoned consumer applies
|
||||
# back-pressure instead of letting the producer fill memory
|
||||
# unbounded. ``cancel_event`` lets the producer stop early
|
||||
# when the caller closes the generator.
|
||||
bridge_queue = queue.Queue(maxsize=64)
|
||||
cancel_event = threading.Event()
|
||||
|
||||
def run_in_thread_noloop():
|
||||
new_loop = asyncio.new_event_loop()
|
||||
asyncio.set_event_loop(new_loop)
|
||||
try:
|
||||
async def pump():
|
||||
try:
|
||||
async for item in self.backend.stream(prompt, **kwargs):
|
||||
if cancel_event.is_set():
|
||||
break
|
||||
# Block until space is available or the
|
||||
# consumer signals cancellation, so we
|
||||
# never spin nor grow the queue forever.
|
||||
while not cancel_event.is_set():
|
||||
try:
|
||||
bridge_queue.put(('item', item), timeout=0.1)
|
||||
break
|
||||
except queue.Full:
|
||||
continue
|
||||
if not cancel_event.is_set():
|
||||
bridge_queue.put(('done', None))
|
||||
except Exception as e:
|
||||
if not cancel_event.is_set():
|
||||
bridge_queue.put(('error', e))
|
||||
|
||||
new_loop.run_until_complete(pump())
|
||||
finally:
|
||||
new_loop.close()
|
||||
|
||||
# Daemon thread so a stalled/abandoned producer can never
|
||||
# block interpreter shutdown.
|
||||
thread = threading.Thread(target=run_in_thread_noloop, daemon=True)
|
||||
thread.start()
|
||||
|
||||
finally:
|
||||
# Runs on normal completion AND on early generator
|
||||
# close/GC: signal the producer to stop, drain one slot
|
||||
# so a blocked put() can observe the cancel, and join.
|
||||
cancel_event.set()
|
||||
try:
|
||||
while True:
|
||||
msg_type, data = bridge_queue.get()
|
||||
if msg_type == 'item':
|
||||
item = data
|
||||
# Same conversion logic as the running-loop branch.
|
||||
if isinstance(item, dict):
|
||||
if item.get('type') == 'agent.message':
|
||||
content = item.get('content', [])
|
||||
if isinstance(content, list):
|
||||
text_parts = []
|
||||
for block in content:
|
||||
if isinstance(block, dict) and block.get('type') == 'text':
|
||||
text_parts.append(block.get('text', ''))
|
||||
elif isinstance(block, str):
|
||||
text_parts.append(block)
|
||||
if text_parts:
|
||||
yield ''.join(text_parts)
|
||||
elif isinstance(content, str):
|
||||
yield content
|
||||
elif isinstance(item, str):
|
||||
yield item
|
||||
elif msg_type == 'done':
|
||||
break
|
||||
elif msg_type == 'error':
|
||||
raise data
|
||||
finally:
|
||||
# Runs on normal completion AND on early generator
|
||||
# close/GC: signal the producer to stop, drain one slot
|
||||
# so a blocked put() can observe the cancel, and join.
|
||||
cancel_event.set()
|
||||
try:
|
||||
bridge_queue.get_nowait()
|
||||
except queue.Empty:
|
||||
pass
|
||||
thread.join(timeout=5)
|
||||
|
||||
bridge_queue.get_nowait()
|
||||
except queue.Empty:
|
||||
pass
|
||||
thread.join(timeout=5)
|
||||
|
||||
return sync_stream()
|
||||
|
||||
except Exception as e:
|
||||
@@ -1427,6 +1381,10 @@ Write the complete compiled report:"""
|
||||
_loop_detection.record_tool_outcome(
|
||||
_ld_history, function_name, arguments, _result, _ld_config
|
||||
)
|
||||
# Per-turn tool count for autonomy completion heuristics (sync
|
||||
# parity with tool_execution.py). Counted only once the tool has
|
||||
# actually run, so a BEFORE_TOOL-blocked call cannot count.
|
||||
self._autonomy_turn_tool_count = getattr(self, '_autonomy_turn_tool_count', 0) + 1
|
||||
return _result
|
||||
|
||||
# Enforce BEFORE_TOOL/AFTER_TOOL security hooks for every async caller,
|
||||
@@ -1709,6 +1667,7 @@ Write the complete compiled report:"""
|
||||
async def _execute_tool_async_impl(self, function_name: str, arguments: Dict[str, Any], tool_call_id: Optional[str] = None, tools_override: Optional[List] = None) -> Any:
|
||||
"""Internal async tool execution implementation"""
|
||||
from ..errors import ToolExecutionError
|
||||
from .tool_execution import _CALLER_ERROR_TYPES
|
||||
try:
|
||||
logging.info(f"Executing async tool: {function_name} with arguments: {arguments}")
|
||||
|
||||
@@ -1940,20 +1899,50 @@ Write the complete compiled report:"""
|
||||
with tool_progress_channel(_progress_sink):
|
||||
return await call_target(**call_arguments)
|
||||
logging.debug(f"Executing sync function in executor: {function_name}")
|
||||
loop = asyncio.get_running_loop()
|
||||
from ..trace.context_events import copy_context_to_callable
|
||||
# copy_context_to_callable snapshots the CURRENT context (incl.
|
||||
# the injection contextvar) into the executor thread, so setting
|
||||
# it here propagates cancel_event to the sync tool body too.
|
||||
if _inject_state is not None:
|
||||
with with_injection_context(_inject_state), tool_progress_channel(_progress_sink):
|
||||
return await loop.run_in_executor(
|
||||
None, copy_context_to_callable(lambda: call_target(**call_arguments))
|
||||
)
|
||||
return await _run_sync(copy_context_to_callable(lambda: call_target(**call_arguments)))
|
||||
with tool_progress_channel(_progress_sink):
|
||||
return await loop.run_in_executor(
|
||||
None, copy_context_to_callable(lambda: call_target(**call_arguments))
|
||||
return await _run_sync(copy_context_to_callable(lambda: call_target(**call_arguments)))
|
||||
|
||||
async def _run_sync(fn):
|
||||
# With a tool timeout, a timed-out sync tool keeps its thread
|
||||
# (wait_for cannot stop it). Run it on the bounded per-agent
|
||||
# pool (sync-path parity) instead of the shared default
|
||||
# executor, so hung tools cannot starve asyncio.to_thread
|
||||
# callers elsewhere in the process.
|
||||
if not (tool_timeout and tool_timeout > 0) or not hasattr(self, "_submit_to_tool_pool"):
|
||||
return await asyncio.get_running_loop().run_in_executor(None, fn)
|
||||
# spill=True: when every worker is busy the call runs at once
|
||||
# on a counted overflow thread (executor is None) instead of
|
||||
# queueing into its timeout; refused once that cap is hit.
|
||||
future, executor, orphaned = self._submit_to_tool_pool(fn, spill=True)
|
||||
if future is None:
|
||||
logging.warning(
|
||||
f"Tool worker pool exhausted; refusing to queue '{function_name}' "
|
||||
f"behind {orphaned} hung tool call(s)."
|
||||
)
|
||||
return {
|
||||
"error": (
|
||||
f"Tool worker pool exhausted: {orphaned} earlier tool call(s) "
|
||||
f"are still running past their timeout. Refusing to queue "
|
||||
f"'{function_name}'."
|
||||
),
|
||||
"tool_pool_exhausted": True,
|
||||
"_praison_retryable": False,
|
||||
}
|
||||
try:
|
||||
return await asyncio.wrap_future(future)
|
||||
except asyncio.CancelledError:
|
||||
# Timed out (or caller cancelled) while still running:
|
||||
# retire the stuck pool so the next call gets a fresh worker.
|
||||
if executor is not None and not future.done():
|
||||
self._retire_tool_executor(executor)
|
||||
raise
|
||||
|
||||
# Circuit-breaker parity with the sync path (tool_execution.py).
|
||||
# The async tool path previously had no breaker, so a repeatedly
|
||||
@@ -1973,23 +1962,12 @@ Write the complete compiled report:"""
|
||||
_NeutralOutcome = _CircuitBreakerException
|
||||
try:
|
||||
from ..tools.circuit_breaker import (
|
||||
get_circuit_breaker,
|
||||
CircuitBreakerConfig,
|
||||
CircuitBreakerException as _CircuitBreakerException,
|
||||
_NeutralOutcome,
|
||||
)
|
||||
breaker_name = f"tool_{id(self)}_{function_name}"
|
||||
breaker = get_circuit_breaker(
|
||||
breaker_name,
|
||||
CircuitBreakerConfig(
|
||||
failure_threshold=5,
|
||||
recovery_timeout=60.0,
|
||||
timeout=30.0,
|
||||
graceful_degradation=True,
|
||||
),
|
||||
)
|
||||
if hasattr(self, "_register_breaker_finalizer"):
|
||||
self._register_breaker_finalizer(breaker_name)
|
||||
# Same per-agent breaker and policy as the sync path.
|
||||
if hasattr(self, "_get_tool_circuit_breaker"):
|
||||
breaker = self._get_tool_circuit_breaker(function_name)
|
||||
except ImportError:
|
||||
# Circuit breaker not available - fall back to direct invocation.
|
||||
logging.debug("Circuit breaker not available on async path, executing directly")
|
||||
@@ -2010,7 +1988,7 @@ Write the complete compiled report:"""
|
||||
# half-open one. Unwrapped in _invoke_guarded.
|
||||
try:
|
||||
r = await _invoke()
|
||||
except (ValueError, TypeError, AttributeError) as e:
|
||||
except _CALLER_ERROR_TYPES as e:
|
||||
raise _NeutralOutcome(e)
|
||||
if isinstance(r, dict) and r.get("error") and \
|
||||
r.get("_praison_retryable") is False:
|
||||
@@ -2051,13 +2029,7 @@ Write the complete compiled report:"""
|
||||
# _execute_tool_async_with_retry treats it as terminal.
|
||||
if loop_guard is not None:
|
||||
loop_guard.record(function_name, arguments, False, result=None)
|
||||
return {
|
||||
"error": f"Tool '{function_name}' circuit breaker open - too many recent failures",
|
||||
"circuit_open": True,
|
||||
"agent_name": getattr(self, "name", None),
|
||||
"session_id": getattr(self, "_session_id", None),
|
||||
"remediation": "Wait for recovery_timeout (60s) or investigate recent tool failures.",
|
||||
}
|
||||
return self._circuit_breaker_open_result(function_name)
|
||||
|
||||
# Apply the per-agent tool timeout (ToolConfig.timeout) so the async
|
||||
# path matches the sync path in tool_execution.py. asyncio.wait_for
|
||||
@@ -2217,7 +2189,7 @@ Write the complete compiled report:"""
|
||||
if isinstance(e, ToolExecutionError):
|
||||
retryable = e.is_retryable
|
||||
else:
|
||||
retryable = not isinstance(e, (ValueError, TypeError, AttributeError))
|
||||
retryable = not isinstance(e, _CALLER_ERROR_TYPES)
|
||||
# Private control-plane key: a tool's own result may legitimately
|
||||
# contain "retryable" (common for HTTP-API wrappers) and must not
|
||||
# collide with the framework's verdict nor leak to the model.
|
||||
|
||||
@@ -190,6 +190,39 @@ def _get_handoff_depth() -> int:
|
||||
"""Get current handoff depth."""
|
||||
return len(_get_handoff_chain())
|
||||
|
||||
def _tool_pair_safe_tail(messages: List[Any], n: int) -> List[Any]:
|
||||
"""Return the last ``n`` messages without splitting tool_calls/tool pairs.
|
||||
|
||||
A count-based slice can cut between an assistant ``tool_calls`` message and
|
||||
its ``tool`` responses; providers reject the resulting orphans. Drop tool
|
||||
results whose call was cut off, and assistant tool_calls whose responses
|
||||
are incomplete (together with their partial responses).
|
||||
"""
|
||||
if n <= 0:
|
||||
return []
|
||||
tail = list(messages[-n:])
|
||||
|
||||
def _call_ids(msg: Any) -> set:
|
||||
calls = msg.get('tool_calls') if isinstance(msg, dict) else None
|
||||
return {c.get('id') if isinstance(c, dict) else getattr(c, 'id', None) for c in calls or ()}
|
||||
|
||||
def _is_tool(msg: Any) -> bool:
|
||||
return isinstance(msg, dict) and msg.get('role') == 'tool'
|
||||
|
||||
response_ids = {m.get('tool_call_id') for m in tail if _is_tool(m)}
|
||||
dropped: set = set()
|
||||
kept = []
|
||||
for msg in tail:
|
||||
ids = _call_ids(msg)
|
||||
if ids and not ids <= response_ids:
|
||||
dropped |= ids
|
||||
continue
|
||||
kept.append(msg)
|
||||
call_ids = set().union(*(_call_ids(m) for m in kept))
|
||||
return [m for m in kept
|
||||
if not (_is_tool(m) and (m.get('tool_call_id') in dropped
|
||||
or m.get('tool_call_id') not in call_ids))]
|
||||
|
||||
def _push_handoff(agent_name: str) -> None:
|
||||
"""Push agent to handoff chain.
|
||||
|
||||
@@ -630,6 +663,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.
|
||||
@@ -681,22 +735,22 @@ class Handoff:
|
||||
# Apply context policy
|
||||
if self.config.context_policy == ContextPolicy.NONE:
|
||||
filtered_messages = []
|
||||
elif self.config.context_policy == ContextPolicy.LAST_N:
|
||||
n = self.config.max_context_messages
|
||||
elif self.config.context_policy in (ContextPolicy.LAST_N, ContextPolicy.SUMMARY):
|
||||
if self.config.preserve_system:
|
||||
system_msgs = [m for m in messages if isinstance(m, dict) and m.get('role') == 'system']
|
||||
other_msgs = [m for m in messages if not (isinstance(m, dict) and m.get('role') == 'system')]
|
||||
filtered_messages = system_msgs + other_msgs[-n:]
|
||||
base = [m for m in messages if not (isinstance(m, dict) and m.get('role') == 'system')]
|
||||
else:
|
||||
filtered_messages = messages[-n:]
|
||||
elif self.config.context_policy == ContextPolicy.SUMMARY:
|
||||
# For summary, keep system + last few messages
|
||||
if self.config.preserve_system:
|
||||
system_msgs = [m for m in messages if isinstance(m, dict) and m.get('role') == 'system']
|
||||
other_msgs = [m for m in messages if not (isinstance(m, dict) and m.get('role') == 'system')]
|
||||
filtered_messages = system_msgs + other_msgs[-3:] # Last 3 non-system
|
||||
system_msgs, base = [], list(messages)
|
||||
if self.config.context_policy == ContextPolicy.LAST_N:
|
||||
tail = _tool_pair_safe_tail(base, self.config.max_context_messages)
|
||||
else:
|
||||
filtered_messages = messages[-3:]
|
||||
# SUMMARY: the first user message (the task goal) + last few
|
||||
# messages. An LLM summary stays opt-in via ``input_filter``.
|
||||
tail = _tool_pair_safe_tail(base, 3)
|
||||
goal = next((m for m in base if isinstance(m, dict) and m.get('role') == 'user'), None)
|
||||
if goal is not None and not any(m is goal for m in tail):
|
||||
tail = [goal] + tail
|
||||
filtered_messages = system_msgs + tail
|
||||
else: # FULL
|
||||
filtered_messages = messages[:]
|
||||
|
||||
@@ -778,6 +832,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 +943,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 +1072,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 +1689,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 +1805,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)
|
||||
|
||||
@@ -14,7 +14,7 @@ import warnings
|
||||
warnings.filterwarnings("ignore", "Valid config keys have changed in V2", UserWarning)
|
||||
class ImageGenerationConfig(BaseModel):
|
||||
"""Configuration for image generation settings."""
|
||||
style: str = Field(default="natural", description="Style of the generated image")
|
||||
style: Optional[str] = Field(default=None, description="Style of the generated image (only sent when set)")
|
||||
response_format: str = Field(default="url", description="Format of the response (url or b64_json)")
|
||||
timeout: int = Field(default=600, description="Timeout in seconds for the API call")
|
||||
api_base: Optional[str] = Field(default=None, description="Optional API base URL")
|
||||
@@ -36,7 +36,7 @@ class ImageAgent(Agent):
|
||||
backstory: Optional[str] = None,
|
||||
instructions: Optional[str] = None,
|
||||
llm: Optional[Union[str, Any]] = None,
|
||||
style: str = "natural",
|
||||
style: Optional[str] = None,
|
||||
response_format: str = "url",
|
||||
timeout: int = 600,
|
||||
api_base: Optional[str] = None,
|
||||
@@ -163,10 +163,14 @@ class ImageAgent(Agent):
|
||||
# Gemini provider doesn't support response_format parameter
|
||||
# Apply this filter if provider is explicitly 'gemini' or as fallback for gemini models
|
||||
config.pop('response_format', None)
|
||||
elif custom_llm_provider == "openai" and config.get('response_format') == "url":
|
||||
# "url" is OpenAI's default for DALL-E and is rejected by models that
|
||||
# don't accept response_format (e.g. gpt-image), so don't send it
|
||||
config.pop('response_format', None)
|
||||
elif custom_llm_provider == "openai":
|
||||
if config.get('response_format') == "url":
|
||||
# "url" is OpenAI's default for DALL-E and is rejected by models that
|
||||
# don't accept response_format (e.g. gpt-image), so don't send it
|
||||
config.pop('response_format', None)
|
||||
if 'dall-e-3' not in model_name.lower():
|
||||
# Only dall-e-3 accepts 'style'; other OpenAI image models reject it
|
||||
config.pop('style', None)
|
||||
from rich.progress import Progress, SpinnerColumn, TextColumn
|
||||
with Progress(
|
||||
SpinnerColumn(),
|
||||
|
||||
@@ -38,10 +38,26 @@ _active_turn_cancel: "contextvars.ContextVar[Optional[tuple]]" = contextvars.Con
|
||||
"praisonai_active_turn_cancel", default=None
|
||||
)
|
||||
|
||||
# Exceptions a tool body raises about its *input* (bad value, missing key or
|
||||
# path). They are terminal for the same arguments and say nothing about the
|
||||
# tool's health, so they are never retried and never count toward the per-tool
|
||||
# circuit breaker. Timeouts, connection errors and other OSErrors still do.
|
||||
_CALLER_ERROR_TYPES = (
|
||||
ValueError, TypeError, AttributeError, LookupError,
|
||||
FileNotFoundError, FileExistsError, IsADirectoryError, NotADirectoryError, PermissionError,
|
||||
)
|
||||
|
||||
# Worker count for the per-agent tool ThreadPoolExecutor. Used both when the
|
||||
# pool is created and to decide (with the per-executor in-flight counter)
|
||||
# whether a fresh call would have to queue behind hung workers.
|
||||
_TOOL_POOL_MAX_WORKERS = 2
|
||||
# whether a fresh call would have to queue behind hung workers. Matches the
|
||||
# default ParallelToolCallExecutor width so a parallel batch never queues.
|
||||
_TOOL_POOL_MAX_WORKERS = 5
|
||||
|
||||
# Upper bound on overflow ("spill") threads a single agent may run at once for
|
||||
# timed async tool calls that find every pool worker busy. A hung spill thread
|
||||
# cannot be reclaimed, so without this cap leaked threads would grow with
|
||||
# concurrency instead of staying bounded (#5356).
|
||||
_MAX_SPILL_TOOL_THREADS = 8
|
||||
|
||||
# Sentinel returned as the "arguments" value when a tool call's argument string
|
||||
# cannot be parsed (e.g. truncated by max_tokens or a dropped connection).
|
||||
@@ -1077,6 +1093,110 @@ class ToolExecutionMixin:
|
||||
from .async_safety import DualLock
|
||||
return self.__dict__.setdefault("_tool_executor_lock", DualLock())
|
||||
|
||||
def _submit_to_tool_pool(self, fn, *args, spill=False):
|
||||
"""Submit ``fn(*args)`` to the bounded per-agent tool pool.
|
||||
|
||||
Shared by the sync and async tool paths so a timed-out tool never
|
||||
strands a thread in the process-wide default executor. Returns
|
||||
``(future, executor, orphaned)``; ``future`` is ``None`` when the pool
|
||||
is exhausted by hung calls and the caller must refuse fast. With
|
||||
``spill=True`` and every worker busy (but not orphan-capped), the call
|
||||
runs immediately on a counted daemon thread instead of waiting for a
|
||||
worker, and ``(future, None, orphaned)`` is returned; once
|
||||
``_MAX_SPILL_TOOL_THREADS`` spill threads are still running the call is
|
||||
refused like an exhausted pool.
|
||||
"""
|
||||
with self._get_tool_executor_lock().sync():
|
||||
if getattr(self, '_tool_executor', None) is None:
|
||||
self._tool_executor = concurrent.futures.ThreadPoolExecutor(
|
||||
max_workers=_TOOL_POOL_MAX_WORKERS,
|
||||
thread_name_prefix=f"tool-{self.name}"
|
||||
)
|
||||
# Per-executor in-flight counter so a retired pool's
|
||||
# late-finishing future can never corrupt the *new* pool's count.
|
||||
self._tool_executor._inflight = 0
|
||||
executor = self._tool_executor
|
||||
orphaned = getattr(self, '_tool_executor_orphaned', 0)
|
||||
inflight = getattr(executor, '_inflight', 0)
|
||||
# Once the orphan cap is hit we can no longer retire a stuck pool,
|
||||
# so we would otherwise queue behind hung workers and time out a
|
||||
# healthy tool that never got to start. Refuse fast and visibly
|
||||
# instead when no worker is free, and let the pool self-heal (the
|
||||
# orphan count drops as hung threads finish).
|
||||
if inflight >= _TOOL_POOL_MAX_WORKERS:
|
||||
if orphaned >= _MAX_ORPHANED_TOOL_EXECUTORS:
|
||||
return None, executor, orphaned
|
||||
if spill:
|
||||
spilled = getattr(self, '_tool_spill_inflight', 0)
|
||||
if spilled >= _MAX_SPILL_TOOL_THREADS:
|
||||
return None, executor, orphaned + spilled
|
||||
self._tool_spill_inflight = spilled + 1
|
||||
return self._start_spill_thread(fn, *args), None, orphaned
|
||||
executor._inflight = inflight + 1
|
||||
future = executor.submit(fn, *args)
|
||||
|
||||
def _on_tool_done(_f, _agent=self, _executor=executor): # noqa: ANN001
|
||||
with _agent._get_tool_executor_lock().sync():
|
||||
_executor._inflight = max(0, getattr(_executor, '_inflight', 1) - 1)
|
||||
# Orphan accounting is per-POOL, not per-future: retiring a pool
|
||||
# adds exactly one to the orphan count, so a pool may return
|
||||
# exactly one slot. Recover it only when this retired pool's
|
||||
# LAST in-flight worker exits (``_inflight`` == 0), and guard
|
||||
# with a one-shot flag so a multi-worker pool can never
|
||||
# decrement more than the single unit it contributed and steal
|
||||
# another orphaned pool's slot.
|
||||
if (
|
||||
getattr(_agent, '_tool_executor', None) is not _executor
|
||||
and _executor._inflight == 0
|
||||
and not getattr(_executor, '_orphan_recovered', False)
|
||||
):
|
||||
_executor._orphan_recovered = True
|
||||
_agent._tool_executor_orphaned = max(
|
||||
0, getattr(_agent, '_tool_executor_orphaned', 0) - 1
|
||||
)
|
||||
|
||||
future.add_done_callback(_on_tool_done)
|
||||
return future, executor, orphaned
|
||||
|
||||
def _start_spill_thread(self, fn, *args):
|
||||
"""Run ``fn(*args)`` on a daemon thread; caller already counted it."""
|
||||
future = concurrent.futures.Future()
|
||||
|
||||
def _target():
|
||||
try:
|
||||
if future.set_running_or_notify_cancel():
|
||||
try:
|
||||
future.set_result(fn(*args))
|
||||
except BaseException as e: # noqa: BLE001 - re-raised by awaiter
|
||||
future.set_exception(e)
|
||||
finally:
|
||||
with self._get_tool_executor_lock().sync():
|
||||
self._tool_spill_inflight = max(0, getattr(self, '_tool_spill_inflight', 1) - 1)
|
||||
|
||||
try:
|
||||
threading.Thread(target=_target, daemon=True, name=f"tool-{self.name}-spill").start()
|
||||
except BaseException:
|
||||
self._tool_spill_inflight = max(0, getattr(self, '_tool_spill_inflight', 1) - 1)
|
||||
raise
|
||||
return future
|
||||
|
||||
def _retire_tool_executor(self, executor):
|
||||
"""Retire a pool whose worker is stuck past its timeout.
|
||||
|
||||
The timed-out thread cannot be reclaimed. Retire this executor so the
|
||||
next call gets a fresh worker instead of queuing behind a permanently
|
||||
stuck one — but bound how many orphaned (still-blocked) pools we allow
|
||||
to accumulate: once the cap is reached we stop recycling and keep
|
||||
reusing the existing pool so repeated timeouts can't exhaust process
|
||||
resources with an unbounded number of leaked threads.
|
||||
"""
|
||||
with self._get_tool_executor_lock().sync():
|
||||
orphaned = getattr(self, '_tool_executor_orphaned', 0)
|
||||
if getattr(self, '_tool_executor', None) is executor and orphaned < _MAX_ORPHANED_TOOL_EXECUTORS:
|
||||
executor.shutdown(wait=False)
|
||||
self._tool_executor = None
|
||||
self._tool_executor_orphaned = orphaned + 1
|
||||
|
||||
def _get_turn_tools_lock(self):
|
||||
"""Return the DualLock guarding the per-turn tool buffer.
|
||||
|
||||
@@ -1342,65 +1462,16 @@ class ToolExecutionMixin:
|
||||
with tool_progress_channel(_attempt_sink):
|
||||
# Use copy_context to preserve injection context in executor thread
|
||||
ctx = contextvars.copy_context()
|
||||
started = threading.Event()
|
||||
|
||||
def execute_with_context():
|
||||
started.set()
|
||||
with with_injection_context(state):
|
||||
return self._execute_tool_with_circuit_breaker(function_name, arguments)
|
||||
|
||||
# Use reusable executor to prevent resource leaks
|
||||
pool_exhausted = False
|
||||
with self._get_tool_executor_lock().sync():
|
||||
if not hasattr(self, '_tool_executor') or self._tool_executor is None:
|
||||
self._tool_executor = concurrent.futures.ThreadPoolExecutor(
|
||||
max_workers=_TOOL_POOL_MAX_WORKERS,
|
||||
thread_name_prefix=f"tool-{self.name}"
|
||||
)
|
||||
# Per-executor in-flight counter so a
|
||||
# retired pool's late-finishing future can
|
||||
# never corrupt the *new* pool's count.
|
||||
self._tool_executor._inflight = 0
|
||||
executor = self._tool_executor
|
||||
orphaned = getattr(self, '_tool_executor_orphaned', 0)
|
||||
inflight = getattr(executor, '_inflight', 0)
|
||||
# Once the orphan cap is hit we can no longer
|
||||
# retire a stuck pool, so we would otherwise
|
||||
# queue behind hung workers and time out a
|
||||
# healthy tool that never got to start. Refuse
|
||||
# fast and visibly instead when no worker is
|
||||
# free, and let the pool self-heal (the orphan
|
||||
# count drops as hung threads finish).
|
||||
if orphaned >= _MAX_ORPHANED_TOOL_EXECUTORS and inflight >= _TOOL_POOL_MAX_WORKERS:
|
||||
pool_exhausted = True
|
||||
else:
|
||||
executor._inflight = inflight + 1
|
||||
future = executor.submit(ctx.run, execute_with_context)
|
||||
|
||||
def _on_tool_done(_f, _agent=self, _executor=executor): # noqa: ANN001
|
||||
with _agent._get_tool_executor_lock().sync():
|
||||
_executor._inflight = max(0, getattr(_executor, '_inflight', 1) - 1)
|
||||
# Orphan accounting is per-POOL, not
|
||||
# per-future: retiring a pool adds
|
||||
# exactly one to the orphan count, so
|
||||
# a pool may return exactly one slot.
|
||||
# Recover it only when this retired
|
||||
# pool's LAST in-flight worker exits
|
||||
# (``_inflight`` == 0), and guard with
|
||||
# a one-shot flag so a multi-worker
|
||||
# pool can never decrement more than
|
||||
# the single unit it contributed and
|
||||
# steal another orphaned pool's slot.
|
||||
if (
|
||||
getattr(_agent, '_tool_executor', None) is not _executor
|
||||
and _executor._inflight == 0
|
||||
and not getattr(_executor, '_orphan_recovered', False)
|
||||
):
|
||||
_executor._orphan_recovered = True
|
||||
_agent._tool_executor_orphaned = max(
|
||||
0, getattr(_agent, '_tool_executor_orphaned', 0) - 1
|
||||
)
|
||||
|
||||
future.add_done_callback(_on_tool_done)
|
||||
if pool_exhausted:
|
||||
future, executor, orphaned = self._submit_to_tool_pool(ctx.run, execute_with_context)
|
||||
if future is None:
|
||||
logging.warning(
|
||||
f"Tool worker pool exhausted; refusing to queue '{function_name}' "
|
||||
f"behind {orphaned} hung tool call(s)."
|
||||
@@ -1415,6 +1486,13 @@ class ToolExecutionMixin:
|
||||
}
|
||||
break
|
||||
try:
|
||||
# Charge tool_timeout from when the body starts,
|
||||
# not from submit, so time spent queued for a
|
||||
# worker never times out a healthy tool. The
|
||||
# queue wait is itself bounded by tool_timeout
|
||||
# so a pool stuck behind hung workers still fails.
|
||||
if not started.wait(timeout=tool_timeout):
|
||||
raise concurrent.futures.TimeoutError()
|
||||
result = future.result(timeout=tool_timeout)
|
||||
except concurrent.futures.TimeoutError:
|
||||
if _progress_active is not None:
|
||||
@@ -1427,19 +1505,7 @@ class ToolExecutionMixin:
|
||||
# inner budget never saw it, so this one still
|
||||
# earns an outer retry pass under the single budget.
|
||||
result = {"error": f"Tool timed out after {tool_timeout}s", "timeout": True, "_outer_timeout": True}
|
||||
# The timed-out thread cannot be reclaimed. Retire this
|
||||
# executor so the next call gets a fresh worker instead of
|
||||
# queuing behind a permanently stuck one — but bound how many
|
||||
# orphaned (still-blocked) threads we allow to accumulate:
|
||||
# once the cap is reached we stop recycling and keep reusing
|
||||
# the existing pool so repeated timeouts can't exhaust process
|
||||
# resources with an unbounded number of leaked threads.
|
||||
with self._get_tool_executor_lock().sync():
|
||||
orphaned = getattr(self, '_tool_executor_orphaned', 0)
|
||||
if getattr(self, '_tool_executor', None) is executor and orphaned < _MAX_ORPHANED_TOOL_EXECUTORS:
|
||||
executor.shutdown(wait=False)
|
||||
self._tool_executor = None
|
||||
self._tool_executor_orphaned = orphaned + 1
|
||||
self._retire_tool_executor(executor)
|
||||
else:
|
||||
with tool_progress_channel(_progress_sink), with_injection_context(state):
|
||||
result = self._execute_tool_with_circuit_breaker(function_name, arguments)
|
||||
@@ -1678,8 +1744,28 @@ class ToolExecutionMixin:
|
||||
# `if artifact_ref:` guards remain valid on short-result paths.
|
||||
artifact_ref = None
|
||||
|
||||
# ToolConfig.output_max_lines caps lines first, then the
|
||||
# char limit below applies; output_direction picks the end(s) kept.
|
||||
_tool_cfg = getattr(self, '_tool_config', None)
|
||||
direction = getattr(_tool_cfg, 'output_direction', 'both')
|
||||
max_lines = getattr(_tool_cfg, 'output_max_lines', None)
|
||||
capped_str = result_str
|
||||
if max_lines:
|
||||
lines = result_str.split("\n")
|
||||
if len(lines) > max_lines:
|
||||
if direction == "head":
|
||||
kept_head, kept_tail = lines[:max_lines], []
|
||||
elif direction == "tail":
|
||||
kept_head, kept_tail = [], lines[-max_lines:]
|
||||
else:
|
||||
n_tail = max_lines // 2
|
||||
kept_head = lines[:max_lines - n_tail]
|
||||
kept_tail = lines[-n_tail:] if n_tail else []
|
||||
marker = f"...[{len(lines):,} lines total, showing {max_lines:,}]..."
|
||||
capped_str = "\n".join(kept_head + [marker] + kept_tail)
|
||||
|
||||
# Check if we need to spill to artifact store
|
||||
if len(result_str) > limit:
|
||||
if len(capped_str) > limit or capped_str is not result_str:
|
||||
# Try to use artifact store if available
|
||||
if hasattr(self, '_artifact_store') and self._artifact_store is not None:
|
||||
try:
|
||||
@@ -1704,22 +1790,25 @@ class ToolExecutionMixin:
|
||||
except Exception as e:
|
||||
logging.debug(f"Failed to store artifact: {e}")
|
||||
|
||||
# Generate truncated preview
|
||||
tail_size = min(limit // 5, 2000)
|
||||
head = result_str[:limit - tail_size]
|
||||
tail = result_str[-tail_size:] if tail_size > 0 else ""
|
||||
|
||||
if len(capped_str) <= limit:
|
||||
# Only the line cap applied
|
||||
truncated = capped_str
|
||||
else:
|
||||
# Generate truncated preview
|
||||
if direction == "head":
|
||||
head, tail = capped_str[:limit], ""
|
||||
elif direction == "tail":
|
||||
head, tail = "", capped_str[-limit:]
|
||||
else:
|
||||
tail_size = min(limit // 5, 2000)
|
||||
head = capped_str[:limit - tail_size]
|
||||
tail = capped_str[-tail_size:] if tail_size > 0 else ""
|
||||
total = "chars total" if artifact_ref else "chars"
|
||||
truncated = f"{head}\n...[{len(result_str):,} {total}, showing first/last portions]...\n{tail}"
|
||||
|
||||
# If we stored an artifact, include reference in the output
|
||||
if artifact_ref:
|
||||
truncated = (
|
||||
f"{head}\n"
|
||||
f"...[{len(result_str):,} chars total, showing first/last portions]...\n"
|
||||
f"{tail}\n\n"
|
||||
f"{artifact_ref.to_inline()}"
|
||||
)
|
||||
else:
|
||||
# Fallback to simple truncation
|
||||
truncated = f"{head}\n...[{len(result_str):,} chars, showing first/last portions]...\n{tail}"
|
||||
truncated = f"{truncated}\n\n{artifact_ref.to_inline()}"
|
||||
else:
|
||||
truncated = result_str
|
||||
|
||||
@@ -1737,10 +1826,15 @@ class ToolExecutionMixin:
|
||||
else:
|
||||
truncated = self._truncate_tool_output(function_name, truncated, tool_call_id)
|
||||
|
||||
if len(truncated) < len(result_str):
|
||||
# A triggered line cap applies even when its marker makes the
|
||||
# preview no shorter than the original (many short lines).
|
||||
line_capped = capped_str is not result_str
|
||||
if line_capped or len(truncated) < len(result_str):
|
||||
logging.debug(f"Truncated {function_name} output from {len(result_str)} to {len(truncated)} chars")
|
||||
# For dicts, truncate large string fields (e.g., raw_content from search)
|
||||
if isinstance(result, dict):
|
||||
# For dicts, truncate large string fields (e.g., raw_content from search).
|
||||
# Field truncation keeps both ends, so a line cap or a head/tail
|
||||
# direction uses the string preview instead.
|
||||
if isinstance(result, dict) and direction == "both" and not line_capped:
|
||||
max_field_chars = getattr(self, 'tool_output_limit', DEFAULT_TOOL_OUTPUT_LIMIT) if not self.context_manager else None
|
||||
result = self._truncate_dict_fields(result, function_name, max_field_chars, tool_call_id)
|
||||
# Per-field truncation only shortens individual fields
|
||||
@@ -2126,7 +2220,7 @@ class ToolExecutionMixin:
|
||||
|
||||
return response
|
||||
|
||||
def _calculate_llm_cost(self, prompt_tokens: int, completion_tokens: int, response: any = None) -> float:
|
||||
def _calculate_llm_cost(self, prompt_tokens: int, completion_tokens: int, response: any = None, model: Optional[str] = None) -> float:
|
||||
"""Calculate estimated cost for LLM call.
|
||||
|
||||
Uses litellm for accurate pricing (1000+ models) when available,
|
||||
@@ -2136,6 +2230,7 @@ class ToolExecutionMixin:
|
||||
prompt_tokens: Number of tokens in the prompt
|
||||
completion_tokens: Number of tokens in the completion
|
||||
response: Optional LLM response object for more accurate cost calculation
|
||||
model: Model that answered (e.g. a fallback model); defaults to ``self.llm``
|
||||
|
||||
Returns:
|
||||
Estimated cost in USD
|
||||
@@ -2144,7 +2239,7 @@ class ToolExecutionMixin:
|
||||
return calculate_llm_cost(
|
||||
prompt_tokens=prompt_tokens,
|
||||
completion_tokens=completion_tokens,
|
||||
model=self.llm,
|
||||
model=model or self.llm,
|
||||
response=response,
|
||||
)
|
||||
|
||||
@@ -3227,23 +3322,28 @@ class ToolExecutionMixin:
|
||||
return None
|
||||
breaker_name = f"tool_{id(self)}_{function_name}"
|
||||
config = CircuitBreakerConfig(
|
||||
failure_threshold=5,
|
||||
recovery_timeout=60.0,
|
||||
timeout=30.0,
|
||||
graceful_degradation=True,
|
||||
failure_threshold=5, # Open after 5 failures
|
||||
recovery_timeout=60.0, # Wait 60s before trying half-open
|
||||
graceful_degradation=True, # Return error instead of raising exception
|
||||
)
|
||||
breaker = get_circuit_breaker(breaker_name, config)
|
||||
# Ensure the registry entry is removed the moment this Agent is
|
||||
# actually garbage-collected, regardless of whether close()/aclose()
|
||||
# was ever called. This closes the CPython id-reuse window where a
|
||||
# new Agent at the same address could inherit a stale OPEN breaker.
|
||||
self._register_breaker_finalizer(breaker_name)
|
||||
return breaker
|
||||
|
||||
def _circuit_breaker_open_result(self, function_name):
|
||||
"""Return the standard ``circuit_open`` error dict for *function_name*."""
|
||||
breaker = self._get_tool_circuit_breaker(function_name)
|
||||
recovery = breaker.config.recovery_timeout if breaker is not None else 60.0
|
||||
return {
|
||||
"error": f"Tool '{function_name}' circuit breaker open - too many recent failures",
|
||||
"circuit_open": True,
|
||||
"agent_name": getattr(self, "name", None),
|
||||
"session_id": getattr(self, "_session_id", None),
|
||||
"remediation": "Wait for recovery_timeout (60s) or investigate recent tool failures.",
|
||||
"remediation": f"Wait for recovery_timeout ({recovery:g}s) or investigate recent tool failures.",
|
||||
}
|
||||
|
||||
def _circuit_breaker_precheck(self, function_name):
|
||||
@@ -3301,7 +3401,7 @@ class ToolExecutionMixin:
|
||||
"""
|
||||
# Import circuit breaker components first (lazy import for performance)
|
||||
try:
|
||||
from ..tools.circuit_breaker import get_circuit_breaker, CircuitBreakerConfig, CircuitBreakerException, _NeutralOutcome
|
||||
from ..tools.circuit_breaker import CircuitBreakerException, _NeutralOutcome
|
||||
except ImportError:
|
||||
# Circuit breaker not available - fallback to direct execution
|
||||
logging.debug("Circuit breaker not available, falling back to direct tool execution")
|
||||
@@ -3309,29 +3409,15 @@ class ToolExecutionMixin:
|
||||
|
||||
try:
|
||||
|
||||
# Get or create circuit breaker for this tool.
|
||||
# Scope the key to this Agent instance so one agent's failing tool
|
||||
# cannot trip the breaker for another agent's same-named tool.
|
||||
breaker_name = f"tool_{id(self)}_{function_name}"
|
||||
config = CircuitBreakerConfig(
|
||||
failure_threshold=5, # Open after 5 failures
|
||||
recovery_timeout=60.0, # Wait 60s before trying half-open
|
||||
timeout=30.0, # Tool call timeout
|
||||
graceful_degradation=True # Return error instead of raising exception
|
||||
)
|
||||
breaker = get_circuit_breaker(breaker_name, config)
|
||||
|
||||
# Ensure the registry entry is removed the moment this Agent is
|
||||
# actually garbage-collected, regardless of whether close()/aclose()
|
||||
# was ever called. This closes the CPython id-reuse window where a
|
||||
# new Agent at the same address could inherit a stale OPEN breaker.
|
||||
self._register_breaker_finalizer(breaker_name)
|
||||
# Get or create this agent's circuit breaker for this tool (shared
|
||||
# with the async path so both use the same key and policy).
|
||||
breaker = self._get_tool_circuit_breaker(function_name)
|
||||
|
||||
# Execute tool through circuit breaker with failure detection wrapper
|
||||
def _tool_wrapper():
|
||||
result = self._execute_tool_impl(function_name, arguments)
|
||||
# Caller errors (bad input: ValueError/TypeError/AttributeError,
|
||||
# tagged non-retryable) say nothing about the tool's health: record
|
||||
# Caller errors (bad input: _CALLER_ERROR_TYPES, tagged
|
||||
# non-retryable) say nothing about the tool's health: record
|
||||
# them as neither failure nor success, so they can neither open the
|
||||
# breaker for other callers nor close a half-open one.
|
||||
if isinstance(result, dict) and result.get("error") and \
|
||||
@@ -3367,13 +3453,7 @@ class ToolExecutionMixin:
|
||||
except CircuitBreakerException as e:
|
||||
# Circuit breaker is open - return error dict instead of raising
|
||||
logging.warning(f"Tool '{function_name}' circuit breaker open: {e}")
|
||||
return {
|
||||
"error": f"Tool '{function_name}' circuit breaker open - too many recent failures",
|
||||
"circuit_open": True,
|
||||
"agent_name": getattr(self, "name", None),
|
||||
"session_id": getattr(self, "_session_id", None),
|
||||
"remediation": "Wait for recovery_timeout (60s) or investigate recent tool failures.",
|
||||
}
|
||||
return self._circuit_breaker_open_result(function_name)
|
||||
|
||||
def _find_declared_tool(self, function_name, tools=None):
|
||||
"""Return the tool OBJECT named ``function_name``, or ``None``.
|
||||
@@ -3823,7 +3903,7 @@ class ToolExecutionMixin:
|
||||
# model and cannot collide with a tool's own "retryable" field.
|
||||
return {
|
||||
"error": error_msg,
|
||||
"_praison_retryable": not isinstance(e, (ValueError, TypeError, AttributeError)),
|
||||
"_praison_retryable": not isinstance(e, _CALLER_ERROR_TYPES),
|
||||
}
|
||||
|
||||
# Unresolved: return a corrective, model-readable message so the model can
|
||||
|
||||
@@ -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:
|
||||
|
||||
@@ -44,6 +44,9 @@ DEFAULT_DANGEROUS_TOOLS: Dict[str, str] = {
|
||||
"acp_edit_file": "high",
|
||||
"acp_delete_file": "high",
|
||||
"execute_query": "high",
|
||||
"github_create_branch": "high",
|
||||
"github_commit_and_push": "high",
|
||||
"github_create_pull_request": "high",
|
||||
"evaluate": "medium",
|
||||
"crawl": "medium",
|
||||
"scrape_page": "medium",
|
||||
@@ -76,10 +79,11 @@ DOOM_LOOP_TARGET = "__doom_loop__"
|
||||
# or set ``PRAISONAI_TOOL_SAFETY=off``; users who want stricter
|
||||
# controls can opt into ``approval="safe"`` or ``"read_only"``.
|
||||
PERMISSION_PRESETS = {
|
||||
# "default" — blocks delete + shell/code exec. Allows read/create/edit.
|
||||
# "default" — blocks delete + shell/code exec + git push/PR. Allows read/create/edit.
|
||||
"default": frozenset({
|
||||
"execute_command", "kill_process", "execute_code", "acp_execute_command",
|
||||
"delete_file", "move_file", "copy_file", "acp_delete_file",
|
||||
"github_create_branch", "github_commit_and_push", "github_create_pull_request",
|
||||
}),
|
||||
# "safe" — blocks all dangerous tools (file writes, shell exec, etc.)
|
||||
"safe": frozenset(DEFAULT_DANGEROUS_TOOLS.keys()),
|
||||
@@ -689,29 +693,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 +780,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]:
|
||||
|
||||
@@ -246,7 +246,11 @@ class ChannelDescriptor(Protocol):
|
||||
secrets and plain config keys). A descriptor MAY additionally expose an
|
||||
optional ``setup(io)`` hook for bespoke, multi-step flows the field list
|
||||
cannot express; the wizard calls it when present and merges the returned
|
||||
values. Because it is optional, a minimal descriptor (like the example
|
||||
values. It MAY likewise expose an optional ``standalone_sender`` — a (sync or async)
|
||||
``(target: DeliveryTarget, text: str) -> None`` that raises on failure — so
|
||||
scheduled ``deliver:`` targets reach the channel even when no gateway is
|
||||
resident (``praisonai schedule tick`` from OS cron / CI / scale-to-zero).
|
||||
Because these are optional, a minimal descriptor (like the example
|
||||
below) omits it — so this Protocol is intentionally *not* used as an
|
||||
``isinstance`` type guard; consumers read attributes with ``getattr``
|
||||
fallbacks instead.
|
||||
|
||||
@@ -590,8 +590,11 @@ class GuardrailConfig:
|
||||
# With validator function
|
||||
Agent(guardrails=my_validator_fn)
|
||||
|
||||
# With string preset
|
||||
Agent(guardrails="strict") # Uses strict preset
|
||||
# With a string: an LLM-validation prompt (not a preset name)
|
||||
Agent(guardrails="Ensure the response is polite")
|
||||
|
||||
# With a preset (sets max_retries only) plus a validator
|
||||
Agent(guardrails=["strict", {"validator": my_validator_fn}])
|
||||
|
||||
# For tool policy enforcement, use the dedicated `policy` param instead:
|
||||
Agent(policy=PolicyEngine(...)) # policy strings on guardrails are NOT enforced
|
||||
@@ -600,7 +603,7 @@ class GuardrailConfig:
|
||||
Agent(guardrails=GuardrailConfig(
|
||||
validator=my_validator_fn,
|
||||
max_retries=3,
|
||||
on_fail="retry",
|
||||
on_fail="retry", # only "retry" is implemented
|
||||
))
|
||||
|
||||
# With LLM-based validation
|
||||
@@ -909,7 +912,8 @@ class ExecutionConfig:
|
||||
# ``rate_limiter`` explicitly for burst/TPM control (it takes precedence).
|
||||
max_rpm: Optional[int] = None
|
||||
|
||||
# Time limits
|
||||
# Time limits: per-turn wall-clock seconds for the loop guard.
|
||||
# None = loop-guard default (600s); 0 = no time limit.
|
||||
max_execution_time: Optional[int] = None
|
||||
|
||||
# Retry settings
|
||||
|
||||
@@ -275,9 +275,11 @@ REFLECTION_PRESETS: Dict[str, Dict[str, Any]] = {
|
||||
# =============================================================================
|
||||
|
||||
GUARDRAIL_PRESETS: Dict[str, Dict[str, Any]] = {
|
||||
"strict": {"max_retries": 5, "on_fail": "raise"},
|
||||
"permissive": {"max_retries": 1, "on_fail": "skip"},
|
||||
"safety": {"max_retries": 3, "on_fail": "retry"},
|
||||
# Presets only set retry counts; supply a validator alongside them,
|
||||
# e.g. guardrails=["strict", {"validator": fn}].
|
||||
"strict": {"max_retries": 5},
|
||||
"permissive": {"max_retries": 1},
|
||||
"safety": {"max_retries": 3},
|
||||
}
|
||||
|
||||
|
||||
|
||||
@@ -107,7 +107,8 @@ class LoopGuardConfig:
|
||||
mutating_halt_threshold: int = 7
|
||||
|
||||
# Time-based limits
|
||||
max_time_per_turn: float = 120.0 # Max seconds per chat turn
|
||||
# Max seconds per chat turn; ExecutionConfig.max_execution_time overrides (0 = no limit there).
|
||||
max_time_per_turn: float = 600.0
|
||||
|
||||
# No-progress detection
|
||||
no_progress_warn: int = 4
|
||||
|
||||
@@ -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
|
||||
@@ -11,6 +11,7 @@ import time
|
||||
import asyncio
|
||||
from praisonaiagents._logging import get_logger
|
||||
import subprocess
|
||||
import threading
|
||||
from concurrent.futures import ThreadPoolExecutor
|
||||
from typing import List, Optional, Dict, Any
|
||||
|
||||
@@ -78,25 +79,14 @@ class HookRunner:
|
||||
# misbehaving hook can only exhaust hook capacity, never the loop's
|
||||
# shared default executor that guardrails/tools/locks also rely on.
|
||||
self._hook_pool: Optional[ThreadPoolExecutor] = None
|
||||
# hook.id -> in-flight Future for a sync hook. The slot is reserved
|
||||
# before submission and only cleared when the worker finishes, so a
|
||||
# still-running (timed-out or cancelled) hook caps new launches at one
|
||||
# thread per hook. Guarded by ``_stuck_lock`` so concurrent invocations
|
||||
# of the same hook cannot both reserve the slot and race past the cap.
|
||||
# hook.id -> worker Future of a sync hook that ran past its timeout.
|
||||
# It is cleared only when that worker finishes, so a stuck hook caps
|
||||
# new launches at one leaked thread per hook. Healthy concurrent calls
|
||||
# (e.g. parallel tool calls) are not tracked and never denied. A
|
||||
# threading lock, because execute_sync() drives each call on its own
|
||||
# event loop in the caller's thread.
|
||||
self._stuck: Dict[str, Any] = {}
|
||||
self._stuck_lock_obj: Optional[asyncio.Lock] = None
|
||||
|
||||
@property
|
||||
def _stuck_lock(self) -> asyncio.Lock:
|
||||
"""Lazily create the in-flight-slot lock on first use.
|
||||
|
||||
Created lazily (not in ``__init__``) because a ``HookRunner`` may be
|
||||
constructed before any event loop exists; the lock is only needed once
|
||||
we are already inside ``async def execute`` / ``_execute_function_hook``.
|
||||
"""
|
||||
if self._stuck_lock_obj is None:
|
||||
self._stuck_lock_obj = asyncio.Lock()
|
||||
return self._stuck_lock_obj
|
||||
self._stuck_lock = threading.Lock()
|
||||
|
||||
def _release_stuck(self, hook_id: str, fut: Any) -> None:
|
||||
"""Clear the in-flight slot once the worker Future completes.
|
||||
@@ -104,8 +94,9 @@ class HookRunner:
|
||||
Only removes the entry if it still points at *this* Future, so a newer
|
||||
reserved invocation is never clobbered by an older worker draining.
|
||||
"""
|
||||
if self._stuck.get(hook_id) is fut:
|
||||
self._stuck.pop(hook_id, None)
|
||||
with self._stuck_lock:
|
||||
if self._stuck.get(hook_id) is fut:
|
||||
self._stuck.pop(hook_id, None)
|
||||
|
||||
def _get_hook_pool(self) -> ThreadPoolExecutor:
|
||||
"""Lazily create the dedicated sync-hook executor."""
|
||||
@@ -439,14 +430,9 @@ class HookRunner:
|
||||
timeout=timeout
|
||||
)
|
||||
else:
|
||||
# Reserve the single in-flight slot for this hook *under a lock*
|
||||
# and *before* submitting, so two concurrent invocations can't
|
||||
# both see an empty ``_stuck`` and both launch a worker (which
|
||||
# would exceed the intended one-worker-per-hook cap and double
|
||||
# the side effects). The slot is only cleared once the worker
|
||||
# finishes; a still-running previous worker fails this one
|
||||
# closed immediately instead of leaking another thread.
|
||||
async with self._stuck_lock:
|
||||
# A previous invocation that timed out and is still running
|
||||
# fails this one closed instead of leaking another thread.
|
||||
with self._stuck_lock:
|
||||
prev = self._stuck.get(hook.id)
|
||||
if prev is not None and not prev.done():
|
||||
duration = (time.time() - start_time) * 1000
|
||||
@@ -463,29 +449,31 @@ class HookRunner:
|
||||
),
|
||||
)
|
||||
|
||||
# Execute sync function in the dedicated hook pool with
|
||||
# timeout. ``copy_context_to_callable`` preserves
|
||||
# trace/session contextvars across the thread, matching
|
||||
# other executor sites.
|
||||
from ..trace.context_events import copy_context_to_callable
|
||||
loop = asyncio.get_event_loop()
|
||||
func = hook.func
|
||||
data = input_data
|
||||
fut = loop.run_in_executor(
|
||||
self._get_hook_pool(),
|
||||
copy_context_to_callable(lambda: func(data)),
|
||||
)
|
||||
# Reserve the slot now; a completion callback clears it only
|
||||
# when the worker actually finishes, so a timeout *or* a
|
||||
# cancellation leaves the slot held until the thread drains.
|
||||
self._stuck[hook.id] = fut
|
||||
fut.add_done_callback(
|
||||
lambda f, hid=hook.id: self._release_stuck(hid, f)
|
||||
)
|
||||
|
||||
result = await asyncio.wait_for(
|
||||
asyncio.shield(fut), timeout=timeout
|
||||
# Execute sync function in the dedicated hook pool with
|
||||
# timeout. ``copy_context_to_callable`` preserves
|
||||
# trace/session contextvars across the thread, matching
|
||||
# other executor sites. A concurrent Future (not a loop-bound
|
||||
# one) so the stuck-slot release works after the loop closes.
|
||||
from ..trace.context_events import copy_context_to_callable
|
||||
func = hook.func
|
||||
data = input_data
|
||||
fut = self._get_hook_pool().submit(
|
||||
copy_context_to_callable(lambda: func(data))
|
||||
)
|
||||
try:
|
||||
result = await asyncio.wait_for(
|
||||
asyncio.shield(asyncio.wrap_future(fut)), timeout=timeout
|
||||
)
|
||||
except (asyncio.TimeoutError, asyncio.CancelledError):
|
||||
# The worker cannot be cancelled; hold the slot until it
|
||||
# actually finishes.
|
||||
if not fut.done():
|
||||
with self._stuck_lock:
|
||||
self._stuck[hook.id] = fut
|
||||
fut.add_done_callback(
|
||||
lambda f, hid=hook.id: self._release_stuck(hid, f)
|
||||
)
|
||||
raise
|
||||
|
||||
duration = (time.time() - start_time) * 1000
|
||||
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -26,7 +26,7 @@ from pydantic import BaseModel
|
||||
import time
|
||||
import json
|
||||
import xml.etree.ElementTree as ET
|
||||
from ..errors import AgentErrorKind, FailoverDecision, IdleTimeoutBreaker, ToolExecutionError
|
||||
from ..errors import AgentErrorKind, BudgetExceededError, FailoverDecision, IdleTimeoutBreaker, ToolExecutionError
|
||||
from ..model_harness.guard import check_model_request
|
||||
# Gap 2: Tool call execution imports
|
||||
from ..tools.call_executor import ToolCall, create_tool_call_executor
|
||||
@@ -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.
|
||||
@@ -185,6 +344,11 @@ except ImportError:
|
||||
# API calls per request. This is a DRY/maintenance concern, not a billing issue.
|
||||
from ..tools.hosted import is_hosted_tool # noqa: E402
|
||||
|
||||
# ``(llm, callback)`` bound by the owning Agent for one custom-LLM turn. Every
|
||||
# tracked call made by that ``llm`` reports its response to ``callback`` so the
|
||||
# Agent's cost/token counters and ``max_budget`` cover the whole tool loop.
|
||||
_usage_sink = contextvars.ContextVar("praisonai_llm_usage_sink", default=None)
|
||||
|
||||
|
||||
class LLMResponseError(Exception):
|
||||
"""Raised when the LLM tool-calling loop fails and cannot produce a response.
|
||||
@@ -1736,15 +1900,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 +1939,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 +2021,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
|
||||
@@ -3809,7 +3993,9 @@ Respond with ONLY a valid JSON tool call in this format:
|
||||
# Mark that fallback completed successfully
|
||||
fallback_completed = True
|
||||
streaming_success = False
|
||||
|
||||
|
||||
except BudgetExceededError:
|
||||
raise
|
||||
except Exception as fallback_error:
|
||||
# If non-streaming also fails, create a graceful fallback with partial streaming data
|
||||
logging.warning(f"Non-streaming fallback also failed: {fallback_error}")
|
||||
@@ -4335,7 +4521,7 @@ Respond with ONLY a valid JSON tool call in this format:
|
||||
final_response_text = response_text.strip() if response_text else ""
|
||||
break
|
||||
|
||||
except (LLMResponseError, ToolExecutionError):
|
||||
except (LLMResponseError, ToolExecutionError, BudgetExceededError):
|
||||
raise
|
||||
except Exception as e:
|
||||
logging.error(f"Error in LLM iteration {iteration_count}: {e}")
|
||||
@@ -4682,6 +4868,8 @@ Output MUST be JSON with 'reflection' and 'satisfactory'.
|
||||
interaction_displayed = True
|
||||
return _prepare_return_value(response_text)
|
||||
continue
|
||||
except BudgetExceededError:
|
||||
raise
|
||||
except Exception as e:
|
||||
_get_display_functions()['display_error'](f"Error in LLM response: {str(e)}")
|
||||
return None
|
||||
@@ -6677,6 +6865,16 @@ Output MUST be JSON with 'reflection' and 'satisfactory'.
|
||||
|
||||
def _track_token_usage(self, response: Any, model: str) -> Optional[TokenMetrics]:
|
||||
"""Extract and track token usage from LLM response."""
|
||||
try:
|
||||
return self._collect_token_metrics(response, model)
|
||||
finally:
|
||||
# Outside the collector's error swallowing so the Agent's budget
|
||||
# check can raise BudgetExceededError and stop the tool loop.
|
||||
sink = _usage_sink.get()
|
||||
if sink is not None and sink[0] is self:
|
||||
sink[1](response)
|
||||
|
||||
def _collect_token_metrics(self, response: Any, model: str) -> Optional[TokenMetrics]:
|
||||
if not TokenMetrics or not get_token_collector:
|
||||
return None
|
||||
|
||||
@@ -7636,6 +7834,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 +7873,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 +8016,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 +8182,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,
|
||||
|
||||
@@ -127,43 +127,42 @@ class MCPToolRunner(threading.Thread):
|
||||
# Signal that initialization is complete
|
||||
self.initialized.set()
|
||||
|
||||
# Process requests
|
||||
while True:
|
||||
# Process requests. Each request runs as its own task so
|
||||
# one slow call never serializes the others behind it (the
|
||||
# session multiplexes requests by id), and the queue is read
|
||||
# with a short blocking get instead of a busy 10ms poll. The
|
||||
# bounded get keeps the reader thread from outliving the loop.
|
||||
loop = asyncio.get_running_loop()
|
||||
tasks = set()
|
||||
|
||||
def _next_item():
|
||||
try:
|
||||
# Check for new requests
|
||||
try:
|
||||
item = self.queue.get(block=False)
|
||||
if item is None: # Shutdown signal
|
||||
break
|
||||
|
||||
response_queue, kind, name, arguments, cancelled = item
|
||||
if cancelled.is_set():
|
||||
# Caller already gave up (timed out); never
|
||||
# execute a call it was told did not happen.
|
||||
continue
|
||||
try:
|
||||
if kind == "resource":
|
||||
result = await asyncio.wait_for(session.read_resource(name), timeout=self.timeout)
|
||||
elif kind == "prompt":
|
||||
result = await asyncio.wait_for(session.get_prompt(name, arguments or None), timeout=self.timeout)
|
||||
else:
|
||||
result = await asyncio.wait_for(session.call_tool(name, arguments), timeout=self.timeout)
|
||||
response_queue.put((True, result))
|
||||
except asyncio.TimeoutError:
|
||||
response_queue.put((False, f"MCP {kind} call timed out after {self.timeout} seconds (server side)"))
|
||||
except Exception as e:
|
||||
if _is_transport_closed(e):
|
||||
self._dead = f"MCP server connection lost: {str(e) or type(e).__name__}"
|
||||
response_queue.put((False, self._dead))
|
||||
break # leave stdio_client() so the child is reaped
|
||||
response_queue.put((False, str(e) or type(e).__name__))
|
||||
except queue.Empty:
|
||||
pass
|
||||
|
||||
# Give other tasks a chance to run
|
||||
await asyncio.sleep(0.01)
|
||||
except asyncio.CancelledError:
|
||||
break
|
||||
return self.queue.get(timeout=0.25)
|
||||
except queue.Empty:
|
||||
return queue.Empty
|
||||
|
||||
try:
|
||||
while True:
|
||||
item = await loop.run_in_executor(None, _next_item)
|
||||
if item is queue.Empty:
|
||||
if self._dead:
|
||||
break # leave stdio_client() so the child is reaped
|
||||
continue
|
||||
if item is None: # Shutdown signal
|
||||
break
|
||||
if self._dead:
|
||||
self.queue.put(item) # answered by _fail_pending()
|
||||
break
|
||||
task = asyncio.create_task(self._handle_request(session, item))
|
||||
tasks.add(task)
|
||||
task.add_done_callback(tasks.discard)
|
||||
except asyncio.CancelledError:
|
||||
pass
|
||||
finally:
|
||||
for task in tasks:
|
||||
task.cancel()
|
||||
if tasks:
|
||||
await asyncio.gather(*tasks, return_exceptions=True)
|
||||
except asyncio.CancelledError:
|
||||
# stop() cancelled a startup that never completed; leaving the
|
||||
# contexts above has already terminated the stdio child.
|
||||
@@ -197,6 +196,35 @@ class MCPToolRunner(threading.Thread):
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
async def _handle_request(self, session, item):
|
||||
"""Run one queued request on the session and post its result."""
|
||||
response_queue, kind, name, arguments, cancelled = item
|
||||
if cancelled.is_set():
|
||||
# Caller already gave up (timed out); never execute a call it was
|
||||
# told did not happen.
|
||||
return
|
||||
try:
|
||||
if kind == "resource":
|
||||
coro = session.read_resource(name)
|
||||
elif kind == "prompt":
|
||||
coro = session.get_prompt(name, arguments or None)
|
||||
else:
|
||||
coro = session.call_tool(name, arguments)
|
||||
result = await asyncio.wait_for(coro, timeout=self.timeout)
|
||||
response_queue.put((True, result))
|
||||
except asyncio.TimeoutError:
|
||||
response_queue.put((False, f"MCP {kind} call timed out after {self.timeout} seconds (server side)"))
|
||||
except asyncio.CancelledError:
|
||||
# Report the real cause when a sibling call saw the transport die.
|
||||
response_queue.put((False, self._dead or f"MCP {kind} call cancelled: runner shut down"))
|
||||
raise
|
||||
except Exception as e:
|
||||
if _is_transport_closed(e):
|
||||
self._dead = f"MCP server connection lost: {str(e) or type(e).__name__}"
|
||||
response_queue.put((False, self._dead))
|
||||
return
|
||||
response_queue.put((False, str(e) or type(e).__name__))
|
||||
|
||||
async def _discover_resources_and_prompts(self, session):
|
||||
"""Enumerate server resources, resource templates and prompts.
|
||||
|
||||
@@ -391,8 +419,8 @@ class MCPToolRunner(threading.Thread):
|
||||
Queues the shutdown sentinel and joins the thread so the request loop
|
||||
exits and its stdio subprocess is terminated instead of leaking a
|
||||
background thread (and its child) across the process lifetime. The
|
||||
request loop polls the queue between async operations, so once any
|
||||
in-flight call unwinds the sentinel is consumed promptly.
|
||||
request loop reads the sentinel promptly and cancels any in-flight
|
||||
calls, which are answered with an error instead of left waiting.
|
||||
|
||||
Args:
|
||||
timeout: Seconds to wait for the thread to exit. Defaults to the
|
||||
|
||||
@@ -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,15 +194,27 @@ 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
|
||||
|
||||
def _reranked(self, rerank: bool) -> bool:
|
||||
"""Whether a reranker actually ran: the local one, or a backend
|
||||
reranker configured on the knowledge (e.g. mem0 with rerank_model)."""
|
||||
if not rerank:
|
||||
return False
|
||||
if self.reranker:
|
||||
return True
|
||||
cfg = getattr(self.knowledge, "config", None)
|
||||
backend = cfg.get("reranker") if isinstance(cfg, dict) else None
|
||||
return bool(isinstance(backend, dict) and backend.get("default_rerank") and backend.get("provider"))
|
||||
|
||||
def _rerank(
|
||||
self,
|
||||
query: str,
|
||||
results: List[Dict[str, Any]],
|
||||
rerank: Optional[bool] = None,
|
||||
) -> List[Dict[str, Any]]:
|
||||
"""
|
||||
Rerank results if reranker is configured.
|
||||
@@ -213,7 +226,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 +369,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": self._reranked(eff_rerank),
|
||||
"top_k": eff_top_k,
|
||||
},
|
||||
)
|
||||
|
||||
async def aretrieve(
|
||||
self,
|
||||
@@ -498,7 +506,7 @@ class RAG:
|
||||
"num_results": len(results),
|
||||
"num_citations": len(citations),
|
||||
"retrieval_strategy": self.config.retrieval_strategy.value,
|
||||
"reranked": self.config.rerank,
|
||||
"reranked": self._reranked(self.config.rerank),
|
||||
},
|
||||
)
|
||||
|
||||
|
||||
@@ -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:
|
||||
@@ -319,7 +332,10 @@ class SimpleReranker:
|
||||
|
||||
scored = []
|
||||
for result in results:
|
||||
text = result.get('text', '').lower()
|
||||
# Knowledge adapters may return SearchResultItem objects
|
||||
if not isinstance(result, dict) and hasattr(result, 'to_dict'):
|
||||
result = result.to_dict()
|
||||
text = (result.get('text') or result.get('memory') or '').lower()
|
||||
text_words = set(text.split())
|
||||
|
||||
# Calculate overlap score
|
||||
|
||||
@@ -15,7 +15,7 @@ import time
|
||||
from typing import Dict, Iterator, List, Optional
|
||||
|
||||
from .models import ScheduleJob, RunRecord
|
||||
from .due import is_due as _is_due, is_misfire as _is_misfire, resolve_schedule_timezone
|
||||
from .due import is_due as _is_due, is_misfire as _is_misfire, lease_state as _lease_state, past_until as _past_until, resolve_schedule_timezone
|
||||
from .hook_emit import emit_schedule_add, emit_schedule_remove
|
||||
|
||||
logger = get_logger(__name__)
|
||||
@@ -187,25 +187,37 @@ class ConfigYamlScheduleStore:
|
||||
job that is due and not already leased by someone else we:
|
||||
|
||||
* advance ``last_run_at`` to ``now`` (pre-advancing the schedule so a
|
||||
later poll no longer sees it as due — at-least-once → at-most-once),
|
||||
later poll no longer sees it as due),
|
||||
* take a lease (``lease_until = now + lease_seconds``, ``lease_owner``),
|
||||
persisting both in the same atomic write,
|
||||
* remove one-shot jobs (``delete_after_run``) immediately so no other
|
||||
ticker can pick them up.
|
||||
* mark spent jobs (one-shot, ``max_runs`` met) as retiring instead of
|
||||
deleting them; the lease keeps competitors out and ``mark_run``
|
||||
removes them once the run completes.
|
||||
|
||||
A crashed run leaves a lease that expires after ``lease_seconds``; the
|
||||
job then becomes due again and is retried. Losers of the race simply do
|
||||
not see the job in their returned list.
|
||||
A lease that expires without ``complete`` is an abandoned (crashed)
|
||||
run and is recorded as ``abandoned`` history. A retiring job is then
|
||||
re-offered once (at-least-once across crashes) and retired if that
|
||||
attempt is abandoned too; a recurring job is at-most-once per slot and
|
||||
next runs at its following slot. Losers of the race simply do not see
|
||||
the job in their returned list.
|
||||
"""
|
||||
claimed: List[ScheduleJob] = []
|
||||
# One-shots auto-deleted by this claim. Emitted after the lock below;
|
||||
# a later ``remove(job.id)`` by the runner is then a no-op, so a spent
|
||||
# one-shot yields exactly one SCHEDULE_REMOVE.
|
||||
# Spent jobs retired by this claim. Emitted after the lock below.
|
||||
auto_removed: List[ScheduleJob] = []
|
||||
# (job_id, job_name) for occurrences suppressed by the misfire policy.
|
||||
# Logged as ``missed`` history after the lock so a dropped occurrence is
|
||||
# observable instead of vanishing silently.
|
||||
# (job_id, job_name, status) for occurrences suppressed by the misfire
|
||||
# policy (``missed``) or runs whose lease expired (``abandoned``).
|
||||
# Logged after the lock so a lost run is observable, not silent.
|
||||
missed: List[tuple] = []
|
||||
|
||||
def take(job: ScheduleJob) -> None:
|
||||
job._lease_until = now + lease_seconds
|
||||
job._lease_owner = owner_id
|
||||
# ``run_count`` is counted at claim; flag it so ``mark_run`` does
|
||||
# not double-count the same fire.
|
||||
job._run_counted = True
|
||||
self._held_leases[job.id] = owner_id
|
||||
claimed.append(job)
|
||||
|
||||
with self._lock, self._file_lock():
|
||||
# Re-read from disk so we observe cross-process claims/leases.
|
||||
self._reload_locked()
|
||||
@@ -213,10 +225,31 @@ class ConfigYamlScheduleStore:
|
||||
for job in list(self._jobs.values()):
|
||||
if not job.enabled:
|
||||
continue
|
||||
lease_until = getattr(job, "_lease_until", 0.0) or 0.0
|
||||
lease_owner = getattr(job, "_lease_owner", None)
|
||||
# An unexpired lease held by another owner blocks the claim.
|
||||
if lease_until > now and lease_owner != owner_id:
|
||||
state = _lease_state(job, now, owner_id)
|
||||
if state == "busy" or (
|
||||
# Our own spent-job run outlived its lease but is still in
|
||||
# flight: re-offering it would fire it twice.
|
||||
state == "abandoned"
|
||||
and getattr(job, "_retiring", 0)
|
||||
and self._held_leases.get(job.id) == owner_id
|
||||
):
|
||||
continue
|
||||
if state == "abandoned":
|
||||
missed.append((job.id, job.name, "abandoned"))
|
||||
job._lease_until = 0.0
|
||||
job._lease_owner = None
|
||||
changed = True
|
||||
if getattr(job, "_retiring", 0) == 1 and not _past_until(job, now, self._default_timezone):
|
||||
# First attempt of a spent job crashed: retry once
|
||||
# (unless its ``until`` deadline has passed meanwhile).
|
||||
job._retiring = 2
|
||||
take(job)
|
||||
continue
|
||||
if getattr(job, "_retiring", 0):
|
||||
# Run finished (or its retry was abandoned too): retire.
|
||||
auto_removed.append(self._jobs.pop(job.id))
|
||||
self._job_state.pop(job.id, None)
|
||||
changed = True
|
||||
continue
|
||||
if not _is_due(job, now, self._default_timezone):
|
||||
# A job past its ``until`` end instant is not due, but must
|
||||
@@ -235,26 +268,17 @@ class ConfigYamlScheduleStore:
|
||||
# again on the next poll.
|
||||
if _is_misfire(job, now, self._default_timezone):
|
||||
job.last_run_at = now
|
||||
missed.append((job.id, job.name))
|
||||
missed.append((job.id, job.name, "missed"))
|
||||
changed = True
|
||||
continue
|
||||
# Win the claim: pre-advance + lease atomically.
|
||||
job.last_run_at = now
|
||||
# Count the fire here (needed so ``should_retire()`` can drop a
|
||||
# spent job before a competitor reclaims it) and flag it so the
|
||||
# later ``mark_run`` does not double-count the same fire.
|
||||
job.run_count += 1
|
||||
job._run_counted = True
|
||||
job._lease_until = now + lease_seconds
|
||||
job._lease_owner = owner_id
|
||||
self._held_leases[job.id] = owner_id
|
||||
claimed.append(job)
|
||||
changed = True
|
||||
if job.should_retire(now):
|
||||
# Spent (one-shot, run-count met, or past ``until``): remove
|
||||
# now so no competitor re-claims it.
|
||||
auto_removed.append(self._jobs.pop(job.id))
|
||||
self._job_state.pop(job.id, None)
|
||||
# Spent: keep it (leased) so a crash cannot lose it.
|
||||
job._retiring = 1
|
||||
take(job)
|
||||
changed = True
|
||||
if changed:
|
||||
if not self._save():
|
||||
# The lease/last_run_at advance was NOT persisted (full
|
||||
@@ -270,10 +294,10 @@ class ConfigYamlScheduleStore:
|
||||
emit_schedule_remove(job)
|
||||
# Record suppressed occurrences so a missed run is visible in history
|
||||
# rather than silently dropped (log_run takes only the thread lock).
|
||||
for job_id, job_name in missed:
|
||||
for job_id, job_name, status in missed:
|
||||
self.log_run(
|
||||
job_id=job_id,
|
||||
status="missed",
|
||||
status=status,
|
||||
job_name=job_name,
|
||||
)
|
||||
return claimed
|
||||
|
||||
@@ -228,6 +228,26 @@ def quota_hold_from_failure(
|
||||
return now + float(retry_after) + max(float(slack_seconds), 0.0)
|
||||
|
||||
|
||||
def past_until(job: Any, now: float, default_timezone: str | None = None) -> bool:
|
||||
"""Return ``True`` when ``now`` is at or past the job's ``until`` end instant."""
|
||||
until = getattr(job, "until", None)
|
||||
if until is None:
|
||||
return False
|
||||
try:
|
||||
until_ts = localize_wall_clock(
|
||||
datetime.fromisoformat(until),
|
||||
job.schedule.tz or default_timezone,
|
||||
).timestamp()
|
||||
except (ValueError, TypeError):
|
||||
# Fail safe: a stop bound that cannot be parsed must NOT let the
|
||||
# job run past it forever. Treat it as elapsed; the claim path then
|
||||
# retires it via ``ScheduleJob.should_retire`` (which returns True
|
||||
# for an unparseable ``until``) rather than firing it.
|
||||
logger.warning("Invalid 'until' timestamp for job %s: %s", job.id, until)
|
||||
return True
|
||||
return now >= until_ts
|
||||
|
||||
|
||||
def is_due(
|
||||
job: Any,
|
||||
now: float,
|
||||
@@ -264,22 +284,8 @@ def is_due(
|
||||
if hold_until is not None and now < hold_until:
|
||||
return False
|
||||
|
||||
until = getattr(job, "until", None)
|
||||
if until is not None:
|
||||
try:
|
||||
until_ts = localize_wall_clock(
|
||||
datetime.fromisoformat(until),
|
||||
sched.tz or default_timezone,
|
||||
).timestamp()
|
||||
except (ValueError, TypeError):
|
||||
# Fail safe: a stop bound that cannot be parsed must NOT let the
|
||||
# job run past it forever. Treat the job as not-due; the claim path
|
||||
# then retires it via ``ScheduleJob.should_retire`` (which returns
|
||||
# True for an unparseable ``until``) rather than firing it.
|
||||
logger.warning("Invalid 'until' timestamp for job %s: %s", job.id, until)
|
||||
return False
|
||||
if now >= until_ts:
|
||||
return False
|
||||
if past_until(job, now, default_timezone):
|
||||
return False
|
||||
|
||||
if sched.kind == "every":
|
||||
if sched.every_seconds is None:
|
||||
@@ -341,3 +347,26 @@ def is_due(
|
||||
return now >= next_run
|
||||
|
||||
return False
|
||||
|
||||
|
||||
def lease_state(job: Any, now: float, owner_id: str) -> str | None:
|
||||
"""Classify a job's claim lease for an atomic ``claim_due``.
|
||||
|
||||
Returns:
|
||||
``"busy"`` when another owner holds an unexpired lease, or this owner
|
||||
still holds one on a spent (retiring) job whose run is in flight;
|
||||
``"abandoned"`` when a lease expired without ``complete`` — the run
|
||||
crashed or the process died mid-run; ``"done"`` when a spent job's
|
||||
run finished (lease released) but it was not removed yet; ``None``
|
||||
when the job is free to be evaluated for due-ness.
|
||||
"""
|
||||
lease_until = getattr(job, "_lease_until", 0.0) or 0.0
|
||||
lease_owner = getattr(job, "_lease_owner", None)
|
||||
retiring = getattr(job, "_retiring", 0)
|
||||
if lease_owner is not None and lease_until > now:
|
||||
if lease_owner != owner_id or retiring:
|
||||
return "busy"
|
||||
return None
|
||||
if lease_owner is not None:
|
||||
return "abandoned"
|
||||
return "done" if retiring else None
|
||||
@@ -215,7 +215,7 @@ class RunRecord:
|
||||
job_id: ID of the job that was executed.
|
||||
job_name: Human-readable name of the job.
|
||||
status: Execution status. One of ``"succeeded"``, ``"failed"``,
|
||||
``"skipped"``, ``"no_change"``, or ``"missed"``. ``"no_change"``
|
||||
``"skipped"``, ``"no_change"``, ``"missed"`` or ``"abandoned"``. ``"no_change"``
|
||||
is the stateful "monitor mode" outcome — a watched source was
|
||||
unchanged since the last tick, so the model turn was suppressed
|
||||
silently (no tokens, no delivery). It is distinct from
|
||||
@@ -225,6 +225,8 @@ class RunRecord:
|
||||
the process was down and fell outside the job's
|
||||
``misfire_grace_seconds`` — it was not run, but is recorded
|
||||
(not silently dropped) so an operator can see the gap.
|
||||
``"abandoned"`` records a claimed run whose lease expired
|
||||
without completing (crash or restart mid-run).
|
||||
result: Agent response text (truncated if very long).
|
||||
error: Error message if status is ``"failed"``.
|
||||
duration: Wall-clock seconds for execution.
|
||||
@@ -234,7 +236,7 @@ class RunRecord:
|
||||
|
||||
job_id: str
|
||||
job_name: str = ""
|
||||
status: Literal["succeeded", "failed", "skipped", "no_change", "missed"] = "succeeded"
|
||||
status: Literal["succeeded", "failed", "skipped", "no_change", "missed", "abandoned"] = "succeeded"
|
||||
result: Optional[str] = None
|
||||
error: Optional[str] = None
|
||||
duration: float = 0.0
|
||||
@@ -460,7 +462,7 @@ class ScheduleJob:
|
||||
every claim path so the retire decision is defined once. A job retires
|
||||
when it is:
|
||||
|
||||
* a one-shot (``delete_after_run``),
|
||||
* a one-shot (``delete_after_run``) that has fired,
|
||||
* at or past its ``max_runs`` budget, or
|
||||
* past its ``until`` end instant (when ``now`` is supplied).
|
||||
|
||||
@@ -478,7 +480,8 @@ class ScheduleJob:
|
||||
``until`` check is skipped (callers that only ran a job and
|
||||
want the run-count/one-shot decision).
|
||||
"""
|
||||
if self.delete_after_run:
|
||||
# A one-shot that has not fired yet is waiting, not spent.
|
||||
if self.delete_after_run and self.run_count > 0:
|
||||
return True
|
||||
if self.max_runs is not None and self.run_count >= self.max_runs:
|
||||
return True
|
||||
@@ -587,6 +590,11 @@ class ScheduleJob:
|
||||
if lease_until:
|
||||
d["lease_until"] = lease_until
|
||||
d["lease_owner"] = getattr(self, "_lease_owner", None)
|
||||
# Spent job claimed for its final run: 1 = first attempt, 2 = the one
|
||||
# retry after a crashed (abandoned) attempt. Omitted otherwise.
|
||||
retiring = getattr(self, "_retiring", 0)
|
||||
if retiring:
|
||||
d["retiring"] = retiring
|
||||
return d
|
||||
|
||||
@classmethod
|
||||
@@ -645,4 +653,5 @@ class ScheduleJob:
|
||||
# Restore atomic-claim lease metadata if present (see ``to_dict``).
|
||||
job._lease_until = d.get("lease_until", 0.0) or 0.0
|
||||
job._lease_owner = d.get("lease_owner")
|
||||
job._retiring = d.get("retiring", 0) or 0
|
||||
return job
|
||||
@@ -62,8 +62,10 @@ class ScheduleRunner:
|
||||
|
||||
Args:
|
||||
owner_id: Stable identifier for this ticker (process/host).
|
||||
lease_seconds: Lease window; a claim not completed expires after
|
||||
this so a crashed run is eventually retried.
|
||||
lease_seconds: Lease window. A claim not completed within it is
|
||||
recorded as an ``abandoned`` run: a one-shot / spent job is
|
||||
retried once (at-least-once), a recurring job runs at its next
|
||||
slot (at-most-once per slot).
|
||||
"""
|
||||
claim = getattr(self._store, "claim_due", None)
|
||||
if callable(claim):
|
||||
@@ -85,7 +87,11 @@ class ScheduleRunner:
|
||||
duration: float = 0.0,
|
||||
delivered: bool = False,
|
||||
) -> None:
|
||||
"""Update ``last_run_at``, log history, and optionally delete one-shot jobs.
|
||||
"""Update ``last_run_at``, log history, and retire spent jobs.
|
||||
|
||||
A spent job (one-shot, ``max_runs`` met) is removed after a run that
|
||||
did not fail; a failed one is kept, disabled, so it stays visible in
|
||||
``schedule list`` and can be re-enabled or removed by the user.
|
||||
|
||||
Args:
|
||||
job: The job that was executed.
|
||||
@@ -97,10 +103,9 @@ class ScheduleRunner:
|
||||
"""
|
||||
job.last_run_at = time.time()
|
||||
# Count this fire exactly once. An atomic ``claim_due`` already
|
||||
# incremented ``run_count`` when it reserved the job (and needs to, so
|
||||
# ``should_retire()`` can drop a spent job before a competitor reclaims
|
||||
# it); it flags the job with ``_run_counted`` so we do not double-count
|
||||
# here. The non-atomic ``get_due_jobs`` path never sets that flag, so
|
||||
# incremented ``run_count`` when it reserved the job (so
|
||||
# ``should_retire()`` can mark a spent job as retiring); it flags the
|
||||
# job with ``_run_counted`` so we do not double-count here. The non-atomic ``get_due_jobs`` path never sets that flag, so
|
||||
# ``mark_run`` remains the single incrementer there.
|
||||
if getattr(job, "_run_counted", False):
|
||||
job._run_counted = False
|
||||
@@ -122,9 +127,19 @@ class ScheduleRunner:
|
||||
# The run already happened (its claim was persisted), so a failed write
|
||||
# here is reported, not raised: one bad write must not stop the loop.
|
||||
try:
|
||||
if job.should_retire(job.last_run_at):
|
||||
if not job.should_retire(job.last_run_at):
|
||||
self._store.update(job)
|
||||
elif status != "failed":
|
||||
self._store.remove(job.id)
|
||||
else:
|
||||
# Keep a failed spent job visible (disabled) instead of
|
||||
# deleting it, so the user can see and re-run it. Undo the
|
||||
# spent fire (history keeps the failed row) so resuming it
|
||||
# runs it again instead of retiring it unrun.
|
||||
job.enabled = False
|
||||
job._retiring = 0
|
||||
job.run_count = max(0, job.run_count - 1)
|
||||
job.last_run_at = None
|
||||
self._store.update(job)
|
||||
except OSError as e:
|
||||
logger.error("Failed to persist run state for job %r: %s", job.id, e)
|
||||
|
||||
@@ -14,7 +14,7 @@ import time
|
||||
from typing import Dict, Iterator, List, Optional
|
||||
|
||||
from .models import ScheduleJob, RunRecord
|
||||
from .due import is_due as _is_due, is_misfire as _is_misfire
|
||||
from .due import is_due as _is_due, is_misfire as _is_misfire, lease_state as _lease_state, past_until as _past_until
|
||||
from .hook_emit import emit_schedule_add, emit_schedule_remove
|
||||
from ..utils.atomic_io import atomic_write_json
|
||||
|
||||
@@ -175,25 +175,37 @@ class FileScheduleStore:
|
||||
job that is due and not already leased by someone else we:
|
||||
|
||||
* advance ``last_run_at`` to ``now`` (pre-advancing the schedule so a
|
||||
later poll no longer sees it as due — at-least-once → at-most-once),
|
||||
later poll no longer sees it as due),
|
||||
* take a lease (``lease_until = now + lease_seconds``, ``lease_owner``),
|
||||
persisting both in the same atomic write,
|
||||
* remove one-shot jobs (``delete_after_run``) immediately so no other
|
||||
ticker can pick them up.
|
||||
* mark spent jobs (one-shot, ``max_runs`` met) as retiring instead of
|
||||
deleting them; the lease keeps competitors out and ``mark_run``
|
||||
removes them once the run completes.
|
||||
|
||||
A crashed run leaves a lease that expires after ``lease_seconds``; the
|
||||
job then becomes due again and is retried. Losers of the race simply do
|
||||
not see the job in their returned list.
|
||||
A lease that expires without ``complete`` is an abandoned (crashed)
|
||||
run and is recorded as ``abandoned`` history. A retiring job is then
|
||||
re-offered once (at-least-once across crashes) and retired if that
|
||||
attempt is abandoned too; a recurring job is at-most-once per slot and
|
||||
next runs at its following slot. Losers of the race simply do not see
|
||||
the job in their returned list.
|
||||
"""
|
||||
claimed: List[ScheduleJob] = []
|
||||
# One-shots auto-deleted by this claim. Emitted after the lock below;
|
||||
# a later ``remove(job.id)`` by the runner is then a no-op, so a spent
|
||||
# one-shot yields exactly one SCHEDULE_REMOVE.
|
||||
# Spent jobs retired by this claim. Emitted after the lock below.
|
||||
auto_removed: List[ScheduleJob] = []
|
||||
# (job_id, job_name) for occurrences suppressed by the misfire policy.
|
||||
# Logged as ``missed`` history after the lock so a dropped occurrence is
|
||||
# observable instead of vanishing silently.
|
||||
# (job_id, job_name, status) for occurrences suppressed by the misfire
|
||||
# policy (``missed``) or runs whose lease expired (``abandoned``).
|
||||
# Logged after the lock so a lost run is observable, not silent.
|
||||
missed: List[tuple] = []
|
||||
|
||||
def take(job: ScheduleJob) -> None:
|
||||
job._lease_until = now + lease_seconds
|
||||
job._lease_owner = owner_id
|
||||
# ``run_count`` is counted at claim; flag it so ``mark_run`` does
|
||||
# not double-count the same fire.
|
||||
job._run_counted = True
|
||||
self._held_leases[job.id] = owner_id
|
||||
claimed.append(job)
|
||||
|
||||
with self._lock, self._file_lock():
|
||||
# Re-read from disk so we observe cross-process claims/leases.
|
||||
self._reload_locked()
|
||||
@@ -201,10 +213,30 @@ class FileScheduleStore:
|
||||
for job in list(self._jobs.values()):
|
||||
if not job.enabled:
|
||||
continue
|
||||
lease_until = getattr(job, "_lease_until", 0.0) or 0.0
|
||||
lease_owner = getattr(job, "_lease_owner", None)
|
||||
# An unexpired lease held by another owner blocks the claim.
|
||||
if lease_until > now and lease_owner != owner_id:
|
||||
state = _lease_state(job, now, owner_id)
|
||||
if state == "busy" or (
|
||||
# Our own spent-job run outlived its lease but is still in
|
||||
# flight: re-offering it would fire it twice.
|
||||
state == "abandoned"
|
||||
and getattr(job, "_retiring", 0)
|
||||
and self._held_leases.get(job.id) == owner_id
|
||||
):
|
||||
continue
|
||||
if state == "abandoned":
|
||||
missed.append((job.id, job.name, "abandoned"))
|
||||
job._lease_until = 0.0
|
||||
job._lease_owner = None
|
||||
changed = True
|
||||
if getattr(job, "_retiring", 0) == 1 and not _past_until(job, now):
|
||||
# First attempt of a spent job crashed: retry once
|
||||
# (unless its ``until`` deadline has passed meanwhile).
|
||||
job._retiring = 2
|
||||
take(job)
|
||||
continue
|
||||
if getattr(job, "_retiring", 0):
|
||||
# Run finished (or its retry was abandoned too): retire.
|
||||
auto_removed.append(self._jobs.pop(job.id))
|
||||
changed = True
|
||||
continue
|
||||
if not _is_due(job, now):
|
||||
# A job past its ``until`` end instant is not due, but must
|
||||
@@ -224,25 +256,17 @@ class FileScheduleStore:
|
||||
# poll does not see it as due again) and record the miss,
|
||||
# without claiming or firing.
|
||||
job.last_run_at = now
|
||||
missed.append((job.id, job.name))
|
||||
missed.append((job.id, job.name, "missed"))
|
||||
changed = True
|
||||
continue
|
||||
# Win the claim: pre-advance + lease atomically.
|
||||
job.last_run_at = now
|
||||
# Count the fire here (needed so ``should_retire()`` can drop a
|
||||
# spent job before a competitor reclaims it) and flag it so the
|
||||
# later ``mark_run`` does not double-count the same fire.
|
||||
job.run_count += 1
|
||||
job._run_counted = True
|
||||
job._lease_until = now + lease_seconds
|
||||
job._lease_owner = owner_id
|
||||
self._held_leases[job.id] = owner_id
|
||||
claimed.append(job)
|
||||
changed = True
|
||||
if job.should_retire(now):
|
||||
# Spent (one-shot, run-count met, or past ``until``): remove
|
||||
# now so no competitor re-claims it.
|
||||
auto_removed.append(self._jobs.pop(job.id))
|
||||
# Spent: keep it (leased) so a crash cannot lose it.
|
||||
job._retiring = 1
|
||||
take(job)
|
||||
changed = True
|
||||
if changed:
|
||||
if not self._save():
|
||||
# The lease/last_run_at advance was NOT persisted (full
|
||||
@@ -259,10 +283,10 @@ class FileScheduleStore:
|
||||
# Record suppressed occurrences so a missed run is visible in history
|
||||
# rather than silently dropped. Logged outside the file lock (log_run
|
||||
# takes it itself; flock is not re-entrant across file handles).
|
||||
for job_id, job_name in missed:
|
||||
for job_id, job_name, status in missed:
|
||||
self.log_run(
|
||||
job_id=job_id,
|
||||
status="missed",
|
||||
status=status,
|
||||
job_name=job_name,
|
||||
)
|
||||
return claimed
|
||||
|
||||
@@ -306,6 +306,8 @@ class HierarchicalSessionStore(DefaultSessionStore):
|
||||
raw = f.read()
|
||||
data = json.loads(raw.decode("utf-8"))
|
||||
session = ExtendedSessionData.from_dict(data)
|
||||
if not self._owns_file(session_id, filepath, session):
|
||||
return ExtendedSessionData(session_id=session_id)
|
||||
# Bind the fingerprint to this parsed object and these exact bytes.
|
||||
# A later writer must not certify an older payload via its stats.
|
||||
with self._lock:
|
||||
|
||||
@@ -35,7 +35,13 @@ from typing import Any, Dict, List, Optional
|
||||
|
||||
from praisonaiagents._logging import get_logger
|
||||
|
||||
from .store import DefaultSessionStore, FileLock, SessionData
|
||||
from .store import (
|
||||
DefaultSessionStore,
|
||||
FileLock,
|
||||
SessionData,
|
||||
_legacy_session_filename,
|
||||
_safe_session_filename,
|
||||
)
|
||||
|
||||
logger = get_logger(__name__)
|
||||
|
||||
@@ -390,6 +396,22 @@ class SqliteSessionStore(DefaultSessionStore):
|
||||
except Exception:
|
||||
return set()
|
||||
|
||||
def _stored_session_id(self, filename: str) -> Optional[str]:
|
||||
"""Return the id stored in ``filename`` if it maps back to that file."""
|
||||
import json
|
||||
|
||||
try:
|
||||
with open(os.path.join(self.session_dir, filename), "r", encoding="utf-8") as f:
|
||||
sid = json.load(f).get("session_id")
|
||||
except (OSError, ValueError, AttributeError):
|
||||
return None
|
||||
if isinstance(sid, str) and filename in (
|
||||
f"{_safe_session_filename(sid)}.json",
|
||||
f"{_legacy_session_filename(sid)}.json",
|
||||
):
|
||||
return sid
|
||||
return None
|
||||
|
||||
def _reindex_all(self) -> bool:
|
||||
"""Backfill the index from existing JSON transcripts.
|
||||
|
||||
@@ -433,6 +455,10 @@ class SqliteSessionStore(DefaultSessionStore):
|
||||
if not filename.endswith(".json"):
|
||||
continue
|
||||
sid = filename[:-5]
|
||||
if "~" in sid or sid not in already:
|
||||
# Hashed or pre-#6024 legacy name (Issue #6024): the real id
|
||||
# lives in the file; loading it migrates a legacy file.
|
||||
sid = self._stored_session_id(filename) or sid
|
||||
if sid in pending:
|
||||
# Failed restore retries must not be bypassed by a synthetic
|
||||
# empty transcript from the initial/upgrade backfill path.
|
||||
|
||||
@@ -6,6 +6,7 @@ Zero dependencies beyond stdlib.
|
||||
"""
|
||||
|
||||
import copy
|
||||
import hashlib
|
||||
import json
|
||||
from praisonaiagents._logging import get_logger
|
||||
import os
|
||||
@@ -30,6 +31,27 @@ from ..paths import get_sessions_dir
|
||||
|
||||
logger = get_logger(__name__)
|
||||
|
||||
|
||||
|
||||
def _legacy_session_filename(session_id: str) -> str:
|
||||
"""Pre-#6024 lossy filename (many ids -> one name); kept only for migration."""
|
||||
return "".join(c if c.isalnum() or c in "-_" else "_" for c in session_id)
|
||||
|
||||
|
||||
def _safe_session_filename(session_id: str) -> str:
|
||||
"""Injective, filesystem-safe filename stem for a session id (Issue #6024).
|
||||
|
||||
Ids already made of ``[A-Za-z0-9_-]`` keep their exact legacy name; any
|
||||
other id gets a readable prefix plus a digest of the full id, so ids that
|
||||
differ only in ``.``/``@``/``:``/… never share a file.
|
||||
"""
|
||||
readable = _legacy_session_filename(session_id)
|
||||
if readable == session_id:
|
||||
return readable
|
||||
digest = hashlib.sha256(session_id.encode("utf-8")).hexdigest()[:16]
|
||||
return f"{readable[:64]}~{digest}"
|
||||
|
||||
|
||||
# Module-level sentinel to track if we've warned about degraded locking
|
||||
_WARNED_NO_FCNTL = False
|
||||
|
||||
@@ -884,10 +906,56 @@ class DefaultSessionStore:
|
||||
)
|
||||
|
||||
def _get_session_path(self, session_id: str) -> str:
|
||||
"""Get the file path for a session."""
|
||||
# Sanitize session_id for filesystem
|
||||
safe_id = "".join(c if c.isalnum() or c in "-_" else "_" for c in session_id)
|
||||
return os.path.join(self.session_dir, f"{safe_id}.json")
|
||||
"""Get the file path for a session (injective mapping, Issue #6024)."""
|
||||
path = self._session_file_path(session_id)
|
||||
legacy = os.path.join(self.session_dir, f"{_legacy_session_filename(session_id)}.json")
|
||||
# Migrate a pre-#6024 file only when it is provably this id's own.
|
||||
# Recheck under the destination lock so a racing worker that already
|
||||
# migrated and wrote here is never clobbered by a stale move.
|
||||
if legacy != path and not os.path.exists(path) and os.path.exists(legacy):
|
||||
try:
|
||||
with FileLock(path, self.lock_timeout):
|
||||
if not os.path.exists(path):
|
||||
with open(legacy, "r", encoding="utf-8") as f:
|
||||
owner = json.load(f).get("session_id")
|
||||
if owner == session_id:
|
||||
os.replace(legacy, path)
|
||||
except (OSError, ValueError, AttributeError):
|
||||
pass
|
||||
return path
|
||||
|
||||
def _session_file_path(self, session_id: str) -> str:
|
||||
"""Session file path without legacy migration (takes no lock)."""
|
||||
return os.path.join(self.session_dir, f"{_safe_session_filename(session_id)}.json")
|
||||
|
||||
def _owns_file(self, session_id: str, filepath: str, session: SessionData) -> bool:
|
||||
"""Return False if the loaded file belongs to another session id.
|
||||
|
||||
A foreign file (e.g. a pre-#6024 legacy collision) is never served to
|
||||
the caller; it is moved to its owner's own path (or set aside when that
|
||||
is taken) so a later write here cannot destroy the owner's history.
|
||||
Caller must hold the ``filepath`` lock. Raises ``OSError`` if the file
|
||||
cannot be moved, so writers abort instead of overwriting it.
|
||||
"""
|
||||
if not session.session_id:
|
||||
session.session_id = session_id
|
||||
return True
|
||||
if session.session_id == session_id:
|
||||
return True
|
||||
logger.warning(
|
||||
"Session file %s belongs to session %r, not %r; starting fresh",
|
||||
filepath, session.session_id, session_id,
|
||||
)
|
||||
owner_path = self._session_file_path(session.session_id)
|
||||
moved = False
|
||||
if owner_path != filepath:
|
||||
with FileLock(owner_path, self.lock_timeout):
|
||||
if not os.path.exists(owner_path):
|
||||
os.replace(filepath, owner_path)
|
||||
moved = True
|
||||
if not moved and self._quarantine_corrupt(filepath) is None:
|
||||
raise OSError(f"cannot move foreign session file {filepath}")
|
||||
return False
|
||||
|
||||
def _load_session(self, session_id: str) -> SessionData:
|
||||
"""Load session from disk with file locking."""
|
||||
@@ -952,7 +1020,10 @@ class DefaultSessionStore:
|
||||
try:
|
||||
with open(filepath, "r", encoding="utf-8") as f:
|
||||
data = json.load(f)
|
||||
return SessionData.from_dict(data)
|
||||
session = SessionData.from_dict(data)
|
||||
if not self._owns_file(session_id, filepath, session):
|
||||
return SessionData(session_id=session_id)
|
||||
return session
|
||||
except (json.JSONDecodeError, UnicodeDecodeError) as e:
|
||||
# Malformed content — the parsed data is unusable, but the raw file
|
||||
# may still hold recoverable bytes. Quarantine it aside so the next
|
||||
@@ -1122,9 +1193,7 @@ class DefaultSessionStore:
|
||||
try:
|
||||
spill_dir = self._spill_dir()
|
||||
os.makedirs(spill_dir, exist_ok=True)
|
||||
safe_id = "".join(
|
||||
c if c.isalnum() or c in "-_" else "_" for c in session_id
|
||||
)
|
||||
safe_id = _safe_session_filename(session_id)
|
||||
# A monotonic ms timestamp keeps files sortable/attributable, but a
|
||||
# short random token guarantees uniqueness so consecutive failures
|
||||
# within the same PID+millisecond never overwrite an earlier spill
|
||||
@@ -1235,13 +1304,15 @@ class DefaultSessionStore:
|
||||
spill_dir = self._spill_dir()
|
||||
if not os.path.isdir(spill_dir):
|
||||
return
|
||||
safe_id = "".join(
|
||||
c if c.isalnum() or c in "-_" else "_" for c in session_id
|
||||
# Also match pre-#6024 spill names; the session_id check below
|
||||
# rejects any spill that belongs to a colliding id.
|
||||
prefixes = (
|
||||
f"{_safe_session_filename(session_id)}.",
|
||||
f"{_legacy_session_filename(session_id)}.",
|
||||
)
|
||||
prefix = f"{safe_id}."
|
||||
candidates = sorted(
|
||||
f for f in os.listdir(spill_dir)
|
||||
if f.startswith(prefix) and f.endswith(".json")
|
||||
if f.startswith(prefixes) and f.endswith(".json")
|
||||
)
|
||||
except (IOError, OSError):
|
||||
return
|
||||
@@ -1331,7 +1402,9 @@ class DefaultSessionStore:
|
||||
new_messages = [m for _, msgs in recovered for m in msgs]
|
||||
if new_messages:
|
||||
session.messages.extend(new_messages)
|
||||
filepath = self._get_session_path(session_id)
|
||||
# Callers already hold this path's lock; the migrating lookup
|
||||
# would take it again and stall for ``lock_timeout``.
|
||||
filepath = self._session_file_path(session_id)
|
||||
session.updated_at = datetime.now(timezone.utc).isoformat()
|
||||
# Recovered turns extend the active window like any other append, so
|
||||
# they must go through the same retention policy (compact/truncate)
|
||||
|
||||
@@ -540,21 +540,21 @@ class MinimalTelemetry:
|
||||
|
||||
# Use a timeout-based flush to prevent hanging
|
||||
import threading
|
||||
import time
|
||||
import concurrent.futures
|
||||
|
||||
# Use ThreadPoolExecutor for better control
|
||||
with concurrent.futures.ThreadPoolExecutor(max_workers=1) as executor:
|
||||
flush_future = executor.submit(self._safe_flush_posthog, posthog_client)
|
||||
|
||||
try:
|
||||
flush_future.result(timeout=5.0) # 5 second timeout
|
||||
self.logger.debug("PostHog flush completed successfully")
|
||||
except concurrent.futures.TimeoutError:
|
||||
self.logger.warning("PostHog flush timed out")
|
||||
flush_future.cancel()
|
||||
except Exception as e:
|
||||
self.logger.error(f"PostHog flush failed: {e}")
|
||||
|
||||
# Daemon thread + bounded join: a ThreadPoolExecutor would join
|
||||
# a wedged flush on exit/shutdown, so the timeout never applied.
|
||||
flush_thread = threading.Thread(
|
||||
target=self._safe_flush_posthog, args=(posthog_client,),
|
||||
name="praisonai-posthog-flush", daemon=True,
|
||||
)
|
||||
flush_thread.start()
|
||||
flush_thread.join(timeout=5.0) # 5 second timeout
|
||||
if flush_thread.is_alive():
|
||||
# Consumer flush/shutdown would block on the same wedged
|
||||
# transport, so skip the remaining cleanup.
|
||||
self.logger.warning("PostHog flush timed out")
|
||||
return
|
||||
self.logger.debug("PostHog flush completed")
|
||||
|
||||
# Cleanup PostHog threads safely
|
||||
self._shutdown_posthog_threads(posthog_client)
|
||||
|
||||
@@ -62,21 +62,22 @@ class DelegationTools:
|
||||
# blocking past the documented deadline. A non-positive timeout
|
||||
# means "no bound".
|
||||
if timeout and timeout > 0:
|
||||
import concurrent.futures
|
||||
# Daemon worker: a ThreadPoolExecutor worker is joined at
|
||||
# interpreter exit, so a hung delegate would keep the process alive.
|
||||
from .call_executor import ToolTimeoutError, _run_with_timeout
|
||||
|
||||
with concurrent.futures.ThreadPoolExecutor(max_workers=1) as pool:
|
||||
future = pool.submit(
|
||||
spawn, task=task_description, agent_name=agent_type
|
||||
try:
|
||||
result = _run_with_timeout(
|
||||
lambda: spawn(task=task_description, agent_name=agent_type),
|
||||
timeout,
|
||||
)
|
||||
try:
|
||||
result = future.result(timeout=timeout)
|
||||
except concurrent.futures.TimeoutError:
|
||||
return json.dumps({
|
||||
"success": False,
|
||||
"task_description": task_description,
|
||||
"agent_type": agent_type,
|
||||
"error": f"Delegated task timed out after {timeout}s",
|
||||
}, indent=2)
|
||||
except ToolTimeoutError:
|
||||
return json.dumps({
|
||||
"success": False,
|
||||
"task_description": task_description,
|
||||
"agent_type": agent_type,
|
||||
"error": f"Delegated task timed out after {timeout}s",
|
||||
}, indent=2)
|
||||
else:
|
||||
result = spawn(task=task_description, agent_name=agent_type)
|
||||
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -3,7 +3,7 @@ import re
|
||||
import subprocess
|
||||
import logging
|
||||
import secrets
|
||||
from typing import Optional
|
||||
from typing import List, Optional
|
||||
from .decorator import tool
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
@@ -11,22 +11,22 @@ logger = logging.getLogger(__name__)
|
||||
AGENT_BRANCH_PREFIX = "praisonai/"
|
||||
|
||||
|
||||
def _git(*args: str) -> subprocess.CompletedProcess:
|
||||
"""Run a git command capturing output as text."""
|
||||
return subprocess.run(["git", *args], check=True, capture_output=True, text=True)
|
||||
def _git(*args: str, cwd: Optional[str] = None) -> subprocess.CompletedProcess:
|
||||
"""Run a git command (in ``cwd`` when given) capturing output as text."""
|
||||
return subprocess.run(["git", *args], check=True, capture_output=True, text=True, cwd=cwd)
|
||||
|
||||
|
||||
def _default_branch() -> str:
|
||||
def _default_branch(cwd: Optional[str] = None) -> str:
|
||||
"""Detect the remote default branch, falling back to main/master."""
|
||||
try:
|
||||
ref = _git("symbolic-ref", "refs/remotes/origin/HEAD").stdout.strip()
|
||||
ref = _git("symbolic-ref", "refs/remotes/origin/HEAD", cwd=cwd).stdout.strip()
|
||||
# e.g. refs/remotes/origin/main -> main
|
||||
return ref.rsplit("/", 1)[-1]
|
||||
except subprocess.CalledProcessError:
|
||||
pass
|
||||
for candidate in ("main", "master"):
|
||||
try:
|
||||
_git("show-ref", "--verify", "--quiet", f"refs/remotes/origin/{candidate}")
|
||||
_git("show-ref", "--verify", "--quiet", f"refs/remotes/origin/{candidate}", cwd=cwd)
|
||||
return candidate
|
||||
except subprocess.CalledProcessError:
|
||||
continue
|
||||
@@ -45,7 +45,7 @@ def _auto_branch_name(commit_message: str) -> str:
|
||||
return f"{AGENT_BRANCH_PREFIX}{_slugify(commit_message)}-{secrets.token_hex(3)}"
|
||||
|
||||
|
||||
def _divergence_refusal(branch: str) -> Optional[str]:
|
||||
def _divergence_refusal(branch: str, cwd: Optional[str] = None) -> Optional[str]:
|
||||
"""Return a refusal reason if pushing ``branch`` would require a force.
|
||||
|
||||
A force push is required when the remote branch exists but is not an
|
||||
@@ -54,13 +54,13 @@ def _divergence_refusal(branch: str) -> Optional[str]:
|
||||
agent-prefixed ones.
|
||||
"""
|
||||
try:
|
||||
_git("show-ref", "--verify", "--quiet", f"refs/remotes/origin/{branch}")
|
||||
_git("show-ref", "--verify", "--quiet", f"refs/remotes/origin/{branch}", cwd=cwd)
|
||||
remote_exists = True
|
||||
except subprocess.CalledProcessError:
|
||||
remote_exists = False
|
||||
if remote_exists:
|
||||
try:
|
||||
_git("merge-base", "--is-ancestor", f"origin/{branch}", "HEAD")
|
||||
_git("merge-base", "--is-ancestor", f"origin/{branch}", "HEAD", cwd=cwd)
|
||||
except subprocess.CalledProcessError:
|
||||
return (
|
||||
f"refusing to push to '{branch}': the remote branch has "
|
||||
@@ -69,7 +69,7 @@ def _divergence_refusal(branch: str) -> Optional[str]:
|
||||
return None
|
||||
|
||||
|
||||
def _branch_push_safety(branch: str) -> Optional[str]:
|
||||
def _branch_push_safety(branch: str, cwd: Optional[str] = None) -> Optional[str]:
|
||||
"""Return a refusal reason if pushing ``branch`` is unsafe, else None.
|
||||
|
||||
Rules (see github_commit_and_push docstring):
|
||||
@@ -82,14 +82,14 @@ def _branch_push_safety(branch: str) -> Optional[str]:
|
||||
(i.e. would require a force push) are refused.
|
||||
"""
|
||||
# Divergence (force-push) is refused for every branch, agent-prefixed or not.
|
||||
diverged = _divergence_refusal(branch)
|
||||
diverged = _divergence_refusal(branch, cwd=cwd)
|
||||
if diverged:
|
||||
return diverged
|
||||
|
||||
if branch.startswith(AGENT_BRANCH_PREFIX):
|
||||
return None
|
||||
|
||||
default = _default_branch()
|
||||
default = _default_branch(cwd=cwd)
|
||||
if branch == default:
|
||||
return (
|
||||
f"refusing to push to '{branch}': it is the repository "
|
||||
@@ -99,16 +99,16 @@ def _branch_push_safety(branch: str) -> Optional[str]:
|
||||
# Foreign-author check: commits on this branch since the merge-base with
|
||||
# the default branch must all be authored by the configured committer.
|
||||
try:
|
||||
committer = _git("config", "user.email").stdout.strip()
|
||||
committer = _git("config", "user.email", cwd=cwd).stdout.strip()
|
||||
except subprocess.CalledProcessError:
|
||||
committer = ""
|
||||
if committer:
|
||||
try:
|
||||
merge_base = _git(
|
||||
"merge-base", f"origin/{default}", "HEAD"
|
||||
"merge-base", f"origin/{default}", "HEAD", cwd=cwd
|
||||
).stdout.strip()
|
||||
authors = _git(
|
||||
"log", "--format=%ae", f"{merge_base}..HEAD"
|
||||
"log", "--format=%ae", f"{merge_base}..HEAD", cwd=cwd
|
||||
).stdout.split()
|
||||
foreign = {a for a in authors if a and a != committer}
|
||||
if foreign:
|
||||
@@ -123,25 +123,27 @@ def _branch_push_safety(branch: str) -> Optional[str]:
|
||||
return None
|
||||
|
||||
@tool
|
||||
def github_create_branch(branch_name: str) -> str:
|
||||
def github_create_branch(branch_name: str, repo_path: Optional[str] = None) -> str:
|
||||
"""Create and checkout a new git branch.
|
||||
|
||||
|
||||
Args:
|
||||
branch_name: The name of the branch to create and checkout.
|
||||
repo_path: Repository directory to run git in. Defaults to the
|
||||
current working directory.
|
||||
"""
|
||||
try:
|
||||
# Check if we are in a git repository
|
||||
subprocess.run(["git", "rev-parse", "--is-inside-work-tree"], check=True, capture_output=True)
|
||||
subprocess.run(["git", "rev-parse", "--is-inside-work-tree"], check=True, capture_output=True, cwd=repo_path)
|
||||
# Validate branch name to prevent misinterpretation as git options or invalid refs
|
||||
try:
|
||||
subprocess.run(
|
||||
["git", "check-ref-format", f"refs/heads/{branch_name}"],
|
||||
check=True, capture_output=True, text=True
|
||||
check=True, capture_output=True, text=True, cwd=repo_path
|
||||
)
|
||||
except subprocess.CalledProcessError as e:
|
||||
return f"Error: invalid branch name '{branch_name}': {e.stderr.strip() if e.stderr else 'branch name is not a valid git ref'}"
|
||||
subprocess.run(["git", "checkout", "-B", branch_name], check=True, capture_output=True, text=True)
|
||||
subprocess.run(["git", "checkout", "-B", branch_name], check=True, capture_output=True, text=True, cwd=repo_path)
|
||||
logger.debug(f"Branch '{branch_name}' checked out successfully.")
|
||||
logger.debug(f"Successfully checked out branch '{branch_name}'")
|
||||
return f"Successfully created and checked out branch '{branch_name}'"
|
||||
@@ -154,8 +156,15 @@ def github_commit_and_push(
|
||||
commit_message: str,
|
||||
branch: Optional[str] = None,
|
||||
allow_unsafe_branch: bool = False,
|
||||
repo_path: Optional[str] = None,
|
||||
paths: Optional[List[str]] = None,
|
||||
) -> str:
|
||||
"""Stage all changes, commit them, and push to the remote repository safely.
|
||||
"""Stage tracked changes, commit them, and push to the remote repository safely.
|
||||
|
||||
Only changes to already-tracked files under the working directory are
|
||||
staged (``git add -u -- .``); new files are committed only when listed in
|
||||
``paths``. Staging happens after the safety checks below, so a refusal
|
||||
leaves the index untouched.
|
||||
|
||||
Pushes are safe by construction. Before pushing, these rules are enforced:
|
||||
|
||||
@@ -176,25 +185,27 @@ def github_commit_and_push(
|
||||
branch: Target branch to push to. Defaults to the current branch.
|
||||
allow_unsafe_branch: Opt-in override for the default/foreign-author
|
||||
refusals. Divergence (force-push) is never allowed.
|
||||
repo_path: Repository directory to run git in. Defaults to the
|
||||
current working directory.
|
||||
paths: Extra (e.g. untracked) files to stage alongside tracked changes.
|
||||
"""
|
||||
try:
|
||||
allow_unsafe = allow_unsafe_branch or os.environ.get(
|
||||
"PRAISONAI_GIT_ALLOW_UNSAFE_BRANCH", ""
|
||||
).lower() in ("1", "true", "yes")
|
||||
|
||||
# Add all changes
|
||||
subprocess.run(["git", "add", "."], check=True, capture_output=True)
|
||||
|
||||
# Check if there are changes to commit
|
||||
status = subprocess.run(["git", "status", "--porcelain"], check=True, capture_output=True, text=True)
|
||||
if not status.stdout.strip():
|
||||
# Check if there are changes to commit (tracked changes or explicit paths)
|
||||
status = _git(
|
||||
"status", "--porcelain", "--untracked-files=no", "--", ".", cwd=repo_path
|
||||
)
|
||||
if not status.stdout.strip() and not paths:
|
||||
return "No changes to commit."
|
||||
|
||||
# Resolve target branch (defaults to current)
|
||||
current_branch = _git("rev-parse", "--abbrev-ref", "HEAD").stdout.strip()
|
||||
current_branch = _git("rev-parse", "--abbrev-ref", "HEAD", cwd=repo_path).stdout.strip()
|
||||
target_branch = branch or current_branch
|
||||
|
||||
default = _default_branch()
|
||||
default = _default_branch(cwd=repo_path)
|
||||
auto_created = False
|
||||
|
||||
# Rule 2: default branch is never pushed directly (unless overridden).
|
||||
@@ -207,7 +218,7 @@ def github_commit_and_push(
|
||||
# retry can then succeed). Adding the new commit does not change the
|
||||
# divergence/foreign-author verdicts, which are about existing history.
|
||||
if not allow_unsafe:
|
||||
reason = _branch_push_safety(target_branch)
|
||||
reason = _branch_push_safety(target_branch, cwd=repo_path)
|
||||
if reason:
|
||||
return f"Error: {reason}"
|
||||
else:
|
||||
@@ -217,14 +228,14 @@ def github_commit_and_push(
|
||||
"checks for '%s'", target_branch
|
||||
)
|
||||
# Even when overriding, divergence (force) is never permitted.
|
||||
diverged = _divergence_refusal(target_branch)
|
||||
diverged = _divergence_refusal(target_branch, cwd=repo_path)
|
||||
if diverged:
|
||||
return f"Error: {diverged}"
|
||||
|
||||
# Move to the resolved branch only after passing safety checks.
|
||||
if auto_created:
|
||||
# Fresh agent branch created from current HEAD.
|
||||
_git("checkout", "-B", target_branch)
|
||||
_git("checkout", "-B", target_branch, cwd=repo_path)
|
||||
elif target_branch != current_branch:
|
||||
# Switch to the requested branch without discarding its commits:
|
||||
# ``-B`` would reset an existing branch to HEAD, so use a plain
|
||||
@@ -233,17 +244,24 @@ def github_commit_and_push(
|
||||
try:
|
||||
_git(
|
||||
"show-ref", "--verify", "--quiet",
|
||||
f"refs/heads/{target_branch}",
|
||||
f"refs/heads/{target_branch}", cwd=repo_path,
|
||||
)
|
||||
_git("checkout", target_branch)
|
||||
_git("checkout", target_branch, cwd=repo_path)
|
||||
except subprocess.CalledProcessError:
|
||||
_git("checkout", "-b", target_branch)
|
||||
_git("checkout", "-b", target_branch, cwd=repo_path)
|
||||
|
||||
# Stage tracked changes under cwd plus explicitly listed paths (never
|
||||
# ``git add .``). The ``.`` pathspec keeps ``-u`` from reaching sibling
|
||||
# directories when run from a repository subdirectory.
|
||||
_git("add", "-u", "--", ".", cwd=repo_path)
|
||||
if paths:
|
||||
_git("add", "--", *paths, cwd=repo_path)
|
||||
|
||||
# Commit changes on the resolved branch
|
||||
subprocess.run(["git", "commit", "-m", commit_message], check=True, capture_output=True, text=True)
|
||||
_git("commit", "-m", commit_message, cwd=repo_path)
|
||||
|
||||
# Push to remote (never with --force)
|
||||
subprocess.run(["git", "push", "-u", "origin", target_branch], check=True, capture_output=True, text=True)
|
||||
_git("push", "-u", "origin", target_branch, cwd=repo_path)
|
||||
if auto_created:
|
||||
return (
|
||||
f"Successfully committed and pushed changes to new agent branch "
|
||||
@@ -255,7 +273,13 @@ def github_commit_and_push(
|
||||
return f"Error committing and pushing: {e.stderr}"
|
||||
|
||||
@tool
|
||||
def github_create_pull_request(title: str, body: str, head_branch: str, base_branch: str = "main") -> str:
|
||||
def github_create_pull_request(
|
||||
title: str,
|
||||
body: str,
|
||||
head_branch: str,
|
||||
base_branch: str = "main",
|
||||
repo_path: Optional[str] = None,
|
||||
) -> str:
|
||||
"""Create a Pull Request on GitHub using the gh CLI.
|
||||
|
||||
Args:
|
||||
@@ -267,6 +291,8 @@ def github_create_pull_request(title: str, body: str, head_branch: str, base_bra
|
||||
This is typically the feature branch you've been working on.
|
||||
base_branch: The target branch to merge your changes into. Usually the main
|
||||
development branch like "main", "master", or "develop". Defaults to "main".
|
||||
repo_path: Repository directory to run gh in. Defaults to the current
|
||||
working directory.
|
||||
|
||||
Returns:
|
||||
str: Success message with PR URL if created successfully, or error message if failed.
|
||||
@@ -282,7 +308,7 @@ def github_create_pull_request(title: str, body: str, head_branch: str, base_bra
|
||||
"""
|
||||
try:
|
||||
# Verify gh CLI is installed and authenticated
|
||||
subprocess.run(["gh", "auth", "status"], check=True, capture_output=True)
|
||||
subprocess.run(["gh", "auth", "status"], check=True, capture_output=True, cwd=repo_path)
|
||||
|
||||
# Create PR
|
||||
cmd = [
|
||||
@@ -292,7 +318,7 @@ def github_create_pull_request(title: str, body: str, head_branch: str, base_bra
|
||||
"--head", head_branch,
|
||||
"--base", base_branch
|
||||
]
|
||||
result = subprocess.run(cmd, check=True, capture_output=True, text=True)
|
||||
result = subprocess.run(cmd, check=True, capture_output=True, text=True, cwd=repo_path)
|
||||
return f"Successfully created Pull Request:\n{result.stdout.strip()}"
|
||||
except subprocess.CalledProcessError as e:
|
||||
logger.error(f"Failed to create PR: {e.stderr}")
|
||||
|
||||
@@ -148,8 +148,9 @@ def schedule_add(
|
||||
default ``PRAISONAI_SCHEDULE_TIMEZONE`` is used, then UTC.
|
||||
once: When True the job is a one-shot — it is removed automatically
|
||||
after its single successful fire (maps to ``delete_after_run``).
|
||||
Useful for ``at:`` / ``in ...`` reminders so a spent job does not
|
||||
linger in listings.
|
||||
A failed fire leaves it listed but disabled; a run lost to a crash
|
||||
is recorded as ``abandoned`` and retried once. Useful for ``at:`` /
|
||||
``in ...`` reminders so a spent job does not linger in listings.
|
||||
max_runs: Optional run-count bound for a recurring schedule — fire at
|
||||
most this many times, then auto-remove (e.g. "remind me 3 times").
|
||||
``0`` (default) means unbounded. Generalises ``once`` (``max_runs=1``).
|
||||
|
||||
@@ -66,6 +66,35 @@ def _emit_subagent_stop(
|
||||
logger.debug("SUBAGENT_STOP hook failed", exc_info=True)
|
||||
|
||||
|
||||
def _messenger_delivery(messenger: Any) -> Callable[[Any], None]:
|
||||
"""Build an ``on_complete`` callback that sends a finished job's result
|
||||
back to the turn's own chat (``"origin"``) through its outbound messenger.
|
||||
|
||||
Raises when the send is not accepted so the job manager leaves the job
|
||||
un-``delivered`` instead of recording a delivery that never happened.
|
||||
"""
|
||||
|
||||
def _deliver(info: Any) -> None:
|
||||
from .messaging_tools import _run_async
|
||||
|
||||
result = info.result
|
||||
if info.error:
|
||||
text = f"Background task failed: {info.error}"
|
||||
elif isinstance(result, dict):
|
||||
text = result.get("output") if result.get("success") else (
|
||||
f"Background task failed: {result.get('error')}"
|
||||
)
|
||||
else:
|
||||
text = result
|
||||
if not text:
|
||||
text = "Background task finished with no output."
|
||||
outcome = _run_async(messenger.send("origin", str(text)))
|
||||
if not getattr(outcome, "ok", False):
|
||||
raise RuntimeError(getattr(outcome, "summary", "") or "delivery failed")
|
||||
|
||||
return _deliver
|
||||
|
||||
|
||||
def create_subagent_tool(
|
||||
agent_factory: Optional[Callable[..., Any]] = None,
|
||||
allowed_agents: Optional[List[str]] = None,
|
||||
@@ -114,9 +143,11 @@ def create_subagent_tool(
|
||||
``deliver`` target finishes. This is the observable completion
|
||||
signal a gateway (e.g. ``BotOS``) subscribes to in order to
|
||||
route the result back to the originating chat without an active
|
||||
turn. When ``None`` (the default) background jobs remain strictly
|
||||
pull-only via ``subagent_result`` — byte-for-byte the prior
|
||||
behaviour.
|
||||
turn. When ``None`` (the default), ``deliver="origin"`` targets
|
||||
the turn's own chat, the spawning turn has an outbound messenger
|
||||
registered (a running gateway) and the send policy allows it, the
|
||||
result is sent back through that messenger; otherwise background
|
||||
jobs remain pull-only via ``subagent_result``.
|
||||
|
||||
Returns:
|
||||
Tool definition dictionary
|
||||
@@ -356,16 +387,45 @@ def create_subagent_tool(
|
||||
origin: Dict[str, Any] = {}
|
||||
job_on_complete = None
|
||||
if deliver:
|
||||
# Fill the origin from the current turn (still in scope here)
|
||||
# rather than from the model; explicit arguments still win.
|
||||
from praisonaiagents.session.context import (
|
||||
get_outbound_messenger,
|
||||
get_session_context,
|
||||
)
|
||||
|
||||
ctx = get_session_context()
|
||||
use_ctx = not chat_id or chat_id == ctx.chat_id
|
||||
origin = {
|
||||
"deliver": deliver,
|
||||
"platform": platform,
|
||||
"chat_id": chat_id,
|
||||
"thread_id": thread_id,
|
||||
"platform": platform or (ctx.platform if use_ctx else ""),
|
||||
"chat_id": chat_id or ctx.chat_id,
|
||||
"thread_id": thread_id or (ctx.thread_id if use_ctx else ""),
|
||||
"session_id": session_id,
|
||||
"user_id": ctx.user_id,
|
||||
"task": task,
|
||||
}
|
||||
if on_job_complete is not None:
|
||||
job_on_complete = on_job_complete
|
||||
else:
|
||||
# Auto-send only back to the turn's own chat and only if
|
||||
# the send policy allows it (as ``send_message`` does);
|
||||
# any other target needs an explicit ``on_job_complete``.
|
||||
from .messaging_tools import _check_send_policy
|
||||
|
||||
messenger = get_outbound_messenger()
|
||||
same_chat = (
|
||||
(not chat_id or chat_id == ctx.chat_id)
|
||||
and (not platform or platform == ctx.platform)
|
||||
and (not thread_id or thread_id == ctx.thread_id)
|
||||
)
|
||||
if (
|
||||
messenger is not None
|
||||
and deliver == "origin"
|
||||
and same_chat
|
||||
and _check_send_policy(deliver) is None
|
||||
):
|
||||
job_on_complete = _messenger_delivery(messenger)
|
||||
|
||||
job_id = job_manager.start_job(
|
||||
_job,
|
||||
|
||||
@@ -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":
|
||||
"""
|
||||
@@ -61,6 +85,16 @@ def get_context_emitter() -> "ContextTraceEmitter":
|
||||
return _default_disabled_emitter
|
||||
|
||||
|
||||
def _agent_key(name: Optional[str]) -> str:
|
||||
"""The id an agent has in traces and graphs: its name, or "Agent" when unnamed.
|
||||
|
||||
``Agent(instructions=...)`` has ``name=None``; without this its events had no
|
||||
node id and its tools were recorded as ``tool:None/<tool>``, so nothing
|
||||
matched the ``agent:Agent`` node in its graph.
|
||||
"""
|
||||
return name or "Agent"
|
||||
|
||||
|
||||
def set_context_emitter(emitter: "ContextTraceEmitter") -> contextvars.Token:
|
||||
"""
|
||||
Set the context trace emitter for the current async context.
|
||||
@@ -114,12 +148,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 +311,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 +383,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 +402,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 +427,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 +444,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 +466,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 +564,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 +604,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 +633,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 +796,7 @@ class ContextTraceEmitter:
|
||||
session_id=self._session_id,
|
||||
agent_name=agent_name,
|
||||
data=metadata or {},
|
||||
node_id=f"agent:{_agent_key(agent_name)}",
|
||||
))
|
||||
|
||||
def agent_end(
|
||||
@@ -491,6 +813,7 @@ class ContextTraceEmitter:
|
||||
session_id=self._session_id,
|
||||
agent_name=agent_name,
|
||||
data=metadata or {},
|
||||
node_id=f"agent:{_agent_key(agent_name)}",
|
||||
))
|
||||
|
||||
def message_added(
|
||||
@@ -540,6 +863,7 @@ class ContextTraceEmitter:
|
||||
"tool_name": tool_name,
|
||||
"tool_args": self._redact_dict(tool_args) if tool_args else {},
|
||||
},
|
||||
node_id=f"tool:{_agent_key(agent_name)}/{tool_name}",
|
||||
))
|
||||
|
||||
def tool_call_end(
|
||||
@@ -587,6 +911,7 @@ class ContextTraceEmitter:
|
||||
"duration_ms": duration_ms,
|
||||
"error": error,
|
||||
},
|
||||
node_id=f"tool:{_agent_key(agent_name)}/{tool_name}",
|
||||
))
|
||||
|
||||
def _smart_truncate_result(self, result: str, tool_name: str) -> str:
|
||||
@@ -686,11 +1011,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 +1040,8 @@ class ContextTraceEmitter:
|
||||
tokens_used=tokens_used,
|
||||
tokens_budget=tokens_budget,
|
||||
data=data,
|
||||
))
|
||||
node_id=f"agent:{_agent_key(agent_name)}",
|
||||
)
|
||||
|
||||
def llm_response(
|
||||
self,
|
||||
@@ -711,6 +1053,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 +1066,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 +1112,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 +1124,8 @@ class ContextTraceEmitter:
|
||||
completion_tokens=actual_completion_tokens,
|
||||
cost_usd=cost_usd,
|
||||
data=data,
|
||||
))
|
||||
node_id=f"agent:{_agent_key(agent_name)}",
|
||||
)
|
||||
|
||||
def context_snapshot(
|
||||
self,
|
||||
@@ -790,8 +1157,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 +1178,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,183 @@ 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:
|
||||
# An unnamed Agent (``Agent(instructions=...)``) has ``name=None``; it is still an agent.
|
||||
return hasattr(obj, "chat") and hasattr(obj, "name")
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# 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."""
|
||||
from ..trace.context_events import _agent_key
|
||||
|
||||
name = getattr(agent, "name", None) or type(agent).__name__ # label
|
||||
key = _agent_key(getattr(agent, "name", None)) # the id the trace uses
|
||||
agent_id = f"agent:{key}"
|
||||
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:{key}/{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 +256,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 +268,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 +285,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 +384,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,
|
||||
|
||||
@@ -707,7 +707,7 @@ class YAMLWorkflowParser:
|
||||
from ..agent import ImageAgent
|
||||
|
||||
# ImageAgent accepts: llm, style, and other image config
|
||||
style = config.get('style', 'natural')
|
||||
style = config.get('style')
|
||||
|
||||
agent = ImageAgent(
|
||||
llm=llm,
|
||||
@@ -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)
|
||||
|
||||
|
||||
@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
|
||||
|
||||
[project]
|
||||
name = "praisonaiagents"
|
||||
version = "1.7.16"
|
||||
version = "1.7.19"
|
||||
description = "Praison AI agents for completing complex tasks with Self Reflection Agents"
|
||||
readme = "README.md"
|
||||
requires-python = ">=3.10"
|
||||
|
||||
@@ -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,130 @@
|
||||
"""A configured completion promise must be the only success signal (#6027)."""
|
||||
|
||||
import asyncio
|
||||
|
||||
import pytest
|
||||
|
||||
from praisonaiagents import Agent
|
||||
from praisonaiagents.model_harness import ScriptedModel
|
||||
|
||||
NOT_DONE = (
|
||||
"I listed the files and started on the parser, but work remains: the tokenizer "
|
||||
"is still failing two tests and the CLI flag is not wired yet. Continuing next turn."
|
||||
)
|
||||
DONE_REPLY = "Fixed the tokenizer and wired the flag. <promise>DONE</promise>"
|
||||
PROMPT = "Refactor auth module" # stage=autonomous
|
||||
DIRECT_PROMPT = "Implement feature X" # stage=direct
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def _offline(monkeypatch, tmp_path):
|
||||
for key in ("OPENAI_BASE_URL", "OPENAI_API_BASE", "ANTHROPIC_BASE_URL"):
|
||||
monkeypatch.setenv(key, "http://127.0.0.1:9")
|
||||
monkeypatch.setenv("HOME", str(tmp_path))
|
||||
|
||||
|
||||
def list_files(path: str) -> str:
|
||||
"""List files in a directory."""
|
||||
return "parser.py tokenizer.py cli.py"
|
||||
|
||||
|
||||
def _script(follow_ups=1):
|
||||
return (
|
||||
[ScriptedModel.tool_call("list_files", {"path": "."})]
|
||||
+ [NOT_DONE] * follow_ups
|
||||
+ [DONE_REPLY] * 3
|
||||
)
|
||||
|
||||
|
||||
def _make(model, **autonomy):
|
||||
return Agent(
|
||||
name="worker", instructions="Do the task.", llm=model, tools=[list_files],
|
||||
output="silent", autonomy={"level": "full_auto", **autonomy},
|
||||
)
|
||||
|
||||
|
||||
def _assert_promise(result):
|
||||
assert result.success is True
|
||||
assert result.completion_reason == "promise"
|
||||
assert "<promise>DONE</promise>" in result.output
|
||||
assert result.iterations >= 2
|
||||
|
||||
|
||||
def test_sync_tool_completion_does_not_preempt_promise():
|
||||
r = _make(ScriptedModel(_script())).run_autonomous(
|
||||
PROMPT, completion_promise="DONE", max_iterations=5)
|
||||
_assert_promise(r)
|
||||
|
||||
|
||||
def test_async_waits_for_promise():
|
||||
r = asyncio.run(_make(ScriptedModel(_script(2))).run_autonomous_async(
|
||||
PROMPT, completion_promise="DONE", max_iterations=5))
|
||||
_assert_promise(r)
|
||||
|
||||
|
||||
def test_sync_direct_stage_waits_for_promise():
|
||||
r = _make(ScriptedModel([NOT_DONE, DONE_REPLY])).run_autonomous(
|
||||
DIRECT_PROMPT, completion_promise="DONE", max_iterations=5)
|
||||
_assert_promise(r)
|
||||
|
||||
|
||||
def test_async_direct_stage_waits_for_promise():
|
||||
r = asyncio.run(_make(ScriptedModel([NOT_DONE, DONE_REPLY])).run_autonomous_async(
|
||||
DIRECT_PROMPT, completion_promise="DONE", max_iterations=5))
|
||||
_assert_promise(r)
|
||||
|
||||
|
||||
def test_promise_never_emitted_ends_unsuccessfully():
|
||||
r = _make(ScriptedModel([NOT_DONE] * 5)).run_autonomous(
|
||||
PROMPT, completion_promise="DONE", max_iterations=3)
|
||||
assert r.success is False
|
||||
assert r.completion_reason == "max_iterations"
|
||||
|
||||
|
||||
def test_iterative_mode_start_honours_promise():
|
||||
agent = Agent(
|
||||
name="loop_agent", instructions="Complete the task.", llm=ScriptedModel(_script()),
|
||||
tools=[list_files], output="silent",
|
||||
autonomy={"mode": "iterative", "max_iterations": 10,
|
||||
"completion_promise": "DONE", "observe": False},
|
||||
)
|
||||
_assert_promise(agent.start(PROMPT))
|
||||
|
||||
|
||||
def test_without_promise_tool_completion_unchanged():
|
||||
r = _make(ScriptedModel(_script())).run_autonomous(PROMPT, max_iterations=5)
|
||||
assert r.success is True
|
||||
assert r.completion_reason == "tool_completion"
|
||||
assert r.iterations == 1
|
||||
|
||||
|
||||
def test_async_tool_path_counts_tools():
|
||||
r = asyncio.run(_make(ScriptedModel(_script(2))).run_autonomous_async(
|
||||
PROMPT, max_iterations=5))
|
||||
assert r.success is True
|
||||
assert r.completion_reason == "tool_completion"
|
||||
assert r.iterations == 1
|
||||
|
||||
|
||||
def test_promise_tag_is_sent_to_model():
|
||||
model = ScriptedModel([DONE_REPLY])
|
||||
_make(model).run_autonomous(PROMPT, completion_promise="DONE", max_iterations=2)
|
||||
assert "<promise>DONE</promise>" in model.requests[0].last_user_message
|
||||
|
||||
|
||||
def test_async_blocked_tool_does_not_count_as_tool_completion():
|
||||
from praisonaiagents.hooks import HookEvent, HookRegistry, HookResult
|
||||
|
||||
registry = HookRegistry()
|
||||
|
||||
@registry.on(HookEvent.BEFORE_TOOL)
|
||||
def deny(data):
|
||||
return HookResult.deny("blocked")
|
||||
|
||||
agent = Agent(
|
||||
name="worker", instructions="Do the task.", llm=ScriptedModel(_script(2)),
|
||||
tools=[list_files], output="silent", hooks=registry,
|
||||
autonomy={"level": "full_auto"},
|
||||
)
|
||||
r = asyncio.run(agent.run_autonomous_async(PROMPT, max_iterations=3))
|
||||
assert r.completion_reason != "tool_completion"
|
||||
@@ -154,9 +154,14 @@ class TestPromptReinjection:
|
||||
completion_promise="DONE"
|
||||
)
|
||||
|
||||
# All prompts should be the original, not "Continue with the task"
|
||||
# All prompts should be the original (plus the promise instruction),
|
||||
# not "Continue with the task"
|
||||
expected = (
|
||||
"Build a REST API\n\nWhen the task is fully complete, include "
|
||||
"<promise>DONE</promise> in your response."
|
||||
)
|
||||
for prompt in prompts_received:
|
||||
assert prompt == "Build a REST API", f"Expected original prompt, got: {prompt}"
|
||||
assert prompt == expected, f"Expected original prompt, got: {prompt}"
|
||||
|
||||
def test_prompt_not_continue_with_task(self):
|
||||
"""Prompt should NOT be 'Continue with the task' after first iteration."""
|
||||
|
||||
@@ -0,0 +1,169 @@
|
||||
"""Regression tests for #6072: streaming bridge, async tool pool, breaker neutrality."""
|
||||
import asyncio
|
||||
import threading
|
||||
import time
|
||||
|
||||
import pytest
|
||||
|
||||
from praisonaiagents import Agent
|
||||
|
||||
|
||||
def _agent(**kw):
|
||||
return Agent(name="t", instructions="x", llm="gpt-4o-mini", **kw)
|
||||
|
||||
|
||||
class _Backend:
|
||||
def __init__(self, fail_first=False):
|
||||
self.produced = 0
|
||||
self.calls = 0
|
||||
self.fail_first = fail_first
|
||||
|
||||
async def stream(self, prompt, **kw):
|
||||
self.calls += 1
|
||||
if self.fail_first and self.calls == 1:
|
||||
yield "a"
|
||||
raise RuntimeError("boom")
|
||||
for i in range(100000):
|
||||
self.produced += 1
|
||||
yield f"c{i}"
|
||||
await asyncio.sleep(0)
|
||||
|
||||
|
||||
def test_stream_close_in_running_loop_stops_producer():
|
||||
a = _agent()
|
||||
a.backend = _Backend()
|
||||
|
||||
async def main():
|
||||
g = a._delegate_streaming_to_backend("p")
|
||||
next(g)
|
||||
g.close()
|
||||
after_close = a.backend.produced
|
||||
time.sleep(0.3)
|
||||
return a.backend.produced - after_close
|
||||
|
||||
produced = asyncio.run(main())
|
||||
assert produced == 0
|
||||
|
||||
|
||||
def test_backend_runtime_error_does_not_rerun_stream():
|
||||
a = _agent()
|
||||
a.backend = _Backend(fail_first=True)
|
||||
|
||||
async def main():
|
||||
with pytest.raises(RuntimeError, match="boom"):
|
||||
for _ in a._delegate_streaming_to_backend("p"):
|
||||
pass
|
||||
|
||||
asyncio.run(main())
|
||||
assert a.backend.calls == 1
|
||||
|
||||
|
||||
def test_async_sync_tool_with_timeout_uses_agent_pool():
|
||||
seen = {}
|
||||
|
||||
def probe() -> str:
|
||||
seen["thread"] = threading.current_thread().name
|
||||
return "ok"
|
||||
|
||||
a = _agent(tools=[probe])
|
||||
a._tool_timeout = 5
|
||||
assert asyncio.run(a.execute_tool_async("probe", {})) == "ok"
|
||||
assert seen["thread"].startswith("tool-t")
|
||||
|
||||
|
||||
def test_async_timed_out_tool_does_not_hold_default_executor():
|
||||
def hang() -> str:
|
||||
time.sleep(2)
|
||||
return "ok"
|
||||
|
||||
a = _agent(tools=[hang])
|
||||
a._tool_timeout = 0.1
|
||||
|
||||
async def main():
|
||||
tasks = [asyncio.create_task(a.execute_tool_async("hang", {})) for _ in range(40)]
|
||||
await asyncio.gather(*tasks)
|
||||
start = time.monotonic()
|
||||
await asyncio.to_thread(lambda: 1)
|
||||
return time.monotonic() - start
|
||||
|
||||
assert asyncio.run(main()) < 1.0
|
||||
|
||||
|
||||
def test_caller_errors_do_not_open_breaker():
|
||||
def read_file(path: str) -> str:
|
||||
raise FileNotFoundError(path)
|
||||
|
||||
a = _agent(tools=[read_file])
|
||||
for i in range(7):
|
||||
result = a._execute_tool_with_circuit_breaker_impl("read_file", {"path": f"/nope{i}"})
|
||||
assert not result.get("circuit_open")
|
||||
|
||||
|
||||
def test_open_breaker_remediation_uses_live_config():
|
||||
a = _agent()
|
||||
breaker = a._get_tool_circuit_breaker("x")
|
||||
breaker.config.recovery_timeout = 7
|
||||
assert "(7s)" in a._circuit_breaker_open_result("x")["remediation"]
|
||||
|
||||
|
||||
def test_stream_close_runs_backend_cleanup():
|
||||
cleaned = threading.Event()
|
||||
|
||||
class _CleanupBackend:
|
||||
async def stream(self, prompt, **kw):
|
||||
try:
|
||||
while True:
|
||||
yield "c"
|
||||
await asyncio.sleep(0)
|
||||
finally:
|
||||
cleaned.set()
|
||||
|
||||
a = _agent()
|
||||
a.backend = _CleanupBackend()
|
||||
g = a._delegate_streaming_to_backend("p")
|
||||
next(g)
|
||||
g.close()
|
||||
assert cleaned.wait(2)
|
||||
|
||||
|
||||
def test_parallel_healthy_timed_tools_do_not_queue_into_timeout():
|
||||
def work(i: int) -> str:
|
||||
time.sleep(0.6)
|
||||
return "ok"
|
||||
|
||||
a = _agent(tools=[work])
|
||||
a._tool_timeout = 1
|
||||
|
||||
async def main():
|
||||
return await asyncio.gather(*(a.execute_tool_async("work", {"i": i}) for i in range(5)))
|
||||
|
||||
# Loop-guard may append advisory notes; what matters is no timeout.
|
||||
assert all(isinstance(r, str) and r.startswith("ok") for r in asyncio.run(main()))
|
||||
|
||||
|
||||
def test_hung_spill_threads_are_capped_across_bursts():
|
||||
from praisonaiagents.agent import tool_execution as te
|
||||
|
||||
release = threading.Event()
|
||||
a = _agent()
|
||||
started = 0
|
||||
try:
|
||||
for _ in range(5):
|
||||
# A burst of concurrent calls that all hang past their timeout;
|
||||
# each timed-out burst retires the stuck pool as the async path does.
|
||||
for _ in range(20):
|
||||
future, executor, _ = a._submit_to_tool_pool(lambda: release.wait(10), spill=True)
|
||||
started += future is not None
|
||||
a._retire_tool_executor(a._tool_executor)
|
||||
assert a._tool_spill_inflight <= te._MAX_SPILL_TOOL_THREADS
|
||||
cap = (
|
||||
te._MAX_SPILL_TOOL_THREADS
|
||||
+ (te._MAX_ORPHANED_TOOL_EXECUTORS + 1) * te._TOOL_POOL_MAX_WORKERS
|
||||
)
|
||||
assert started <= cap
|
||||
finally:
|
||||
release.set()
|
||||
deadline = time.monotonic() + 5
|
||||
while a._tool_spill_inflight and time.monotonic() < deadline:
|
||||
time.sleep(0.01)
|
||||
assert a._tool_spill_inflight == 0
|
||||
@@ -0,0 +1,199 @@
|
||||
"""Budget and cost accounting on the custom-LLM (``provider/model``) path (#6029).
|
||||
|
||||
That path hands the whole turn, tool loop included, to ``LLM.get_response*``.
|
||||
Every call it makes must reach the Agent's cost/token counters so
|
||||
``ExecutionConfig.max_budget`` and the autonomy spend caps can fire.
|
||||
``litellm.completion`` is stubbed; nothing leaves the machine.
|
||||
"""
|
||||
import asyncio
|
||||
|
||||
import pytest
|
||||
|
||||
pytest.importorskip("litellm")
|
||||
|
||||
from praisonaiagents import Agent # noqa: E402
|
||||
from praisonaiagents.config.feature_configs import ExecutionConfig # noqa: E402
|
||||
from praisonaiagents.errors import BudgetExceededError # noqa: E402
|
||||
from praisonaiagents.utils.cost_utils import calculate_llm_cost # noqa: E402
|
||||
|
||||
MODEL = "anthropic/claude-3-5-sonnet-20240620"
|
||||
USAGE = {"prompt_tokens": 1500, "completion_tokens": 1500, "total_tokens": 3000}
|
||||
|
||||
|
||||
def _litellm():
|
||||
"""The litellm module LLM will import now.
|
||||
|
||||
Other tests evict ``litellm`` from sys.modules, so a module-level binding
|
||||
can go stale and stubbing it would let the real client run.
|
||||
"""
|
||||
import litellm
|
||||
return litellm
|
||||
|
||||
|
||||
def _text(model):
|
||||
return _litellm().ModelResponse(model=model, usage=USAGE, choices=[{
|
||||
"index": 0, "finish_reason": "stop",
|
||||
"message": {"role": "assistant", "content": "ok"}}])
|
||||
|
||||
|
||||
def _tool_call(model, n):
|
||||
return _litellm().ModelResponse(model=model, usage=USAGE, choices=[{
|
||||
"index": 0, "finish_reason": "tool_calls",
|
||||
"message": {"role": "assistant", "content": None, "tool_calls": [{
|
||||
"id": f"call_{n}", "type": "function",
|
||||
"function": {"name": "step", "arguments": "{}"}}]}}])
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def calls(monkeypatch):
|
||||
"""Stub litellm; answer with a tool call while ``tool_rounds`` remain."""
|
||||
monkeypatch.setenv("ANTHROPIC_API_KEY", "sk-test")
|
||||
# Agent caches its env-derived default model per process; keep the fake
|
||||
# key's Anthropic default from leaking into later tests.
|
||||
monkeypatch.setattr(Agent, "_default_model", None)
|
||||
monkeypatch.setattr(Agent, "_default_model_checked", False)
|
||||
state = {"n": 0, "tool_rounds": 0}
|
||||
|
||||
def fake(*a, **k):
|
||||
state["n"] += 1
|
||||
if k.get("tools") and state["tool_rounds"] > 0:
|
||||
state["tool_rounds"] -= 1
|
||||
return _tool_call(k.get("model"), state["n"])
|
||||
return _text(k.get("model"))
|
||||
|
||||
async def afake(*a, **k):
|
||||
return fake(*a, **k)
|
||||
|
||||
litellm = _litellm()
|
||||
monkeypatch.setattr(litellm, "completion", fake)
|
||||
monkeypatch.setattr(litellm, "acompletion", afake)
|
||||
return state
|
||||
|
||||
|
||||
def _agent(**kw):
|
||||
return Agent(name="p", instructions="x", llm=MODEL, output="silent", **kw)
|
||||
|
||||
|
||||
def step() -> str:
|
||||
"""One step of work."""
|
||||
return "done"
|
||||
|
||||
|
||||
def test_sync_chat_counts_cost_and_tokens(calls):
|
||||
agent = _agent()
|
||||
assert agent.chat("hi") == "ok"
|
||||
summary = agent.cost_summary
|
||||
assert summary["llm_calls"] == 1
|
||||
assert summary["tokens_in"] == 1500 and summary["tokens_out"] == 1500
|
||||
assert agent.total_cost > 0
|
||||
|
||||
|
||||
def test_sync_budget_raises_between_turns(calls):
|
||||
agent = _agent(execution=ExecutionConfig(max_budget=0.05))
|
||||
with pytest.raises(BudgetExceededError):
|
||||
for i in range(5):
|
||||
agent.chat(f"hello {i}")
|
||||
assert calls["n"] < 5
|
||||
assert agent.total_cost >= 0.05
|
||||
|
||||
|
||||
def test_async_budget_raises_between_turns(calls):
|
||||
agent = _agent(execution=ExecutionConfig(max_budget=0.05))
|
||||
|
||||
async def run():
|
||||
for i in range(5):
|
||||
await agent.achat(f"hello {i}")
|
||||
|
||||
with pytest.raises(BudgetExceededError):
|
||||
asyncio.run(run())
|
||||
assert calls["n"] < 5
|
||||
assert agent.total_cost >= 0.05
|
||||
|
||||
|
||||
def test_budget_trips_mid_tool_loop(calls):
|
||||
calls["tool_rounds"] = 10 # the model would keep calling tools for 11 calls
|
||||
agent = _agent(execution=ExecutionConfig(max_budget=1.0))
|
||||
agent.chat("price one call")
|
||||
one_call = agent.total_cost
|
||||
calls["n"] = 0
|
||||
# A cap between one and two calls must stop the loop at the second call.
|
||||
agent = _agent(tools=[step], execution=ExecutionConfig(max_budget=one_call * 1.5))
|
||||
with pytest.raises(BudgetExceededError):
|
||||
agent.chat("go")
|
||||
assert calls["n"] == 2
|
||||
|
||||
|
||||
def test_warn_mode_does_not_raise(calls):
|
||||
agent = _agent(execution=ExecutionConfig(max_budget=0.001, on_budget_exceeded="warn"))
|
||||
assert agent.chat("hi") == "ok"
|
||||
assert agent.total_cost > 0.001
|
||||
|
||||
|
||||
def test_autonomy_budget_exhausted_on_provider_path(calls):
|
||||
agent = _agent(autonomy={"max_budget_usd": 0.001})
|
||||
result = agent.run_autonomous("go", max_iterations=3)
|
||||
assert result.completion_reason == "budget_exhausted"
|
||||
assert calls["n"] == 1
|
||||
|
||||
|
||||
def test_sink_ignores_other_llm_instances(calls):
|
||||
"""A call made by a different LLM object is not billed to this agent."""
|
||||
from praisonaiagents.llm.llm import LLM
|
||||
|
||||
agent = _agent()
|
||||
other = LLM(model=MODEL)
|
||||
with agent._llm_usage_sink():
|
||||
other._track_token_usage(_text(MODEL), MODEL)
|
||||
assert agent.cost_summary["llm_calls"] == 0
|
||||
|
||||
|
||||
class _Unavailable(Exception):
|
||||
status_code = 503
|
||||
|
||||
|
||||
def test_fallback_model_spend_is_counted(calls, monkeypatch):
|
||||
"""A turn answered by a ``fallback_models`` entry is still billed to the agent."""
|
||||
backup = "anthropic/claude-3-haiku-20240307"
|
||||
|
||||
def fake(*a, **k):
|
||||
calls["n"] += 1
|
||||
if k.get("model") != backup:
|
||||
raise _Unavailable("503 overloaded")
|
||||
return _text(backup)
|
||||
|
||||
async def afake(*a, **k):
|
||||
return fake(*a, **k)
|
||||
|
||||
monkeypatch.setattr(_litellm(), "completion", fake)
|
||||
monkeypatch.setattr(_litellm(), "acompletion", afake)
|
||||
for run in (lambda a: a.chat("hi"), lambda a: asyncio.run(a.achat("hi"))):
|
||||
agent = _agent(retry=False) # fail over at once, no backoff sleeps
|
||||
agent.fallback_models = [backup]
|
||||
assert run(agent) == "ok"
|
||||
assert agent.cost_summary["llm_calls"] == 1
|
||||
# Priced as the model that answered, not the (pricier) primary.
|
||||
assert agent.total_cost == pytest.approx(calculate_llm_cost(1500, 1500, backup))
|
||||
assert agent.total_cost != pytest.approx(calculate_llm_cost(1500, 1500, MODEL))
|
||||
|
||||
|
||||
def test_retry_after_exhaustion_sends_no_request(calls):
|
||||
agent = _agent(execution=ExecutionConfig(max_budget=0.001))
|
||||
with pytest.raises(BudgetExceededError):
|
||||
agent.chat("first")
|
||||
sent = calls["n"]
|
||||
with pytest.raises(BudgetExceededError):
|
||||
agent.chat("retry")
|
||||
with pytest.raises(BudgetExceededError):
|
||||
asyncio.run(agent.achat("retry async"))
|
||||
assert calls["n"] == sent
|
||||
|
||||
|
||||
def test_autonomy_budget_error_keeps_budget_exhausted(calls):
|
||||
"""A max_budget stop inside the turn still reports the autonomy cap result."""
|
||||
agent = _agent(execution=ExecutionConfig(max_budget=0.001),
|
||||
autonomy={"max_budget_usd": 0.001})
|
||||
result = agent.run_autonomous("go", max_iterations=3)
|
||||
assert result.completion_reason == "budget_exhausted"
|
||||
result = asyncio.run(_agent(execution=ExecutionConfig(max_budget=0.001),
|
||||
autonomy={"max_budget_usd": 0.001}).run_autonomous_async("go", max_iterations=3))
|
||||
assert result.completion_reason == "budget_exhausted"
|
||||
@@ -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
|
||||
@@ -388,7 +388,8 @@ def test_cancellation_during_final_save_finishes_record_and_cleanup(recorded_age
|
||||
assert db.starts[0]["run_id"] == db.ends[0]["run_id"]
|
||||
assert db.ends[0]["status"] == "completed"
|
||||
assert agent._active_turn_token is None
|
||||
emitter.agent_end.assert_called_once_with(agent.name)
|
||||
# The answer came back; only the final save was cancelled.
|
||||
emitter.agent_end.assert_called_once_with(agent.name, {"status": "completed"})
|
||||
token.close.assert_called_once()
|
||||
|
||||
|
||||
|
||||
@@ -0,0 +1,162 @@
|
||||
"""Regression tests for #5795: tool/MCP calls must not be timed out by queue wait."""
|
||||
import asyncio
|
||||
import concurrent.futures as cf
|
||||
import queue
|
||||
import threading
|
||||
import time
|
||||
from contextlib import asynccontextmanager
|
||||
from types import SimpleNamespace
|
||||
|
||||
from praisonaiagents import Agent
|
||||
from praisonaiagents.hooks import HookEvent, HookRegistry
|
||||
|
||||
|
||||
def _allow_all_hooks():
|
||||
"""Own registry (no leaked global hooks) with a sync BEFORE_TOOL hook,
|
||||
so parallel tool calls also run the hook concurrently."""
|
||||
registry = HookRegistry()
|
||||
registry.register_function(HookEvent.BEFORE_TOOL, lambda data: None)
|
||||
return registry
|
||||
|
||||
|
||||
def test_parallel_sync_tools_not_timed_out_by_queue_wait():
|
||||
def slow(n: int) -> str:
|
||||
"""slow tool"""
|
||||
time.sleep(0.5)
|
||||
return f"ok{n}"
|
||||
|
||||
agent = Agent(name="t", instructions="x", tools=[slow], llm="gpt-4o-mini", hooks=_allow_all_hooks())
|
||||
agent._tool_timeout = 0.8
|
||||
with cf.ThreadPoolExecutor(5) as ex:
|
||||
results = list(ex.map(lambda i: agent.execute_tool("slow", {"n": i}), range(5)))
|
||||
for i, result in enumerate(results):
|
||||
assert isinstance(result, str) and result.startswith(f"ok{i}"), result
|
||||
|
||||
|
||||
def test_parallel_async_sync_tools_not_timed_out_by_queue_wait():
|
||||
def slow(n: int) -> str:
|
||||
"""slow tool"""
|
||||
time.sleep(0.5)
|
||||
return f"ok{n}"
|
||||
|
||||
agent = Agent(name="a", instructions="x", tools=[slow], llm="gpt-4o-mini", hooks=_allow_all_hooks())
|
||||
agent._tool_timeout = 0.8
|
||||
|
||||
async def main():
|
||||
# Twice the pool size: the second half queues for a worker first.
|
||||
return await asyncio.gather(
|
||||
*(agent.execute_tool_async("slow", {"n": i}) for i in range(10))
|
||||
)
|
||||
|
||||
for i, result in enumerate(asyncio.run(main())):
|
||||
assert isinstance(result, str) and result.startswith(f"ok{i}"), result
|
||||
|
||||
|
||||
def test_hung_async_sync_tools_do_not_starve_default_executor():
|
||||
release = threading.Event()
|
||||
tools = []
|
||||
for i in range(10):
|
||||
def hang(x: int = 0) -> str:
|
||||
"""hang"""
|
||||
release.wait(10)
|
||||
return "x"
|
||||
hang.__name__ = f"hang{i}"
|
||||
tools.append(hang)
|
||||
|
||||
agent = Agent(name="u", instructions="x", tools=tools, llm="gpt-4o-mini", hooks=HookRegistry())
|
||||
agent._tool_timeout = 0.2
|
||||
|
||||
async def main():
|
||||
for i in range(10):
|
||||
result = await asyncio.wait_for(agent.execute_tool_async(f"hang{i}", {"x": 1}), 5)
|
||||
# Times out, or is refused fast once hung calls hit the pool cap.
|
||||
assert isinstance(result, dict) and (result.get("timeout") or result.get("tool_pool_exhausted")), result
|
||||
# The loop's default executor is still free for other to_thread users.
|
||||
assert await asyncio.wait_for(asyncio.to_thread(lambda: 42), 2) == 42
|
||||
|
||||
try:
|
||||
asyncio.run(main())
|
||||
finally:
|
||||
release.set()
|
||||
|
||||
|
||||
def test_transport_wrapper_cancels_call_on_timeout():
|
||||
from praisonaiagents.mcp.mcp_sse import SSEMCPTool
|
||||
|
||||
loop = asyncio.new_event_loop()
|
||||
threading.Thread(target=loop.run_forever, daemon=True).start()
|
||||
finished = threading.Event()
|
||||
|
||||
class _Session:
|
||||
async def call_tool(self, name, kwargs):
|
||||
await asyncio.sleep(1.0)
|
||||
finished.set()
|
||||
|
||||
try:
|
||||
tool = SSEMCPTool("t", "d", _Session(), {}, timeout=0.1)
|
||||
tool._loop = loop
|
||||
result = tool()
|
||||
assert isinstance(result, dict) and result.get("timeout"), result
|
||||
time.sleep(1.3)
|
||||
assert not finished.is_set()
|
||||
finally:
|
||||
loop.call_soon_threadsafe(loop.stop)
|
||||
|
||||
|
||||
def test_stdio_runner_runs_requests_concurrently_and_cancels_on_shutdown(monkeypatch):
|
||||
import praisonaiagents.mcp.mcp as mcp_mod
|
||||
|
||||
class _Session:
|
||||
async def initialize(self):
|
||||
pass
|
||||
|
||||
async def list_tools(self):
|
||||
return SimpleNamespace(tools=[])
|
||||
|
||||
async def call_tool(self, name, arguments):
|
||||
await asyncio.sleep(10 if name == "slow" else 0)
|
||||
return name
|
||||
|
||||
@asynccontextmanager
|
||||
async def _stdio_client(params):
|
||||
yield None, None
|
||||
|
||||
@asynccontextmanager
|
||||
async def _client_session(read, write):
|
||||
yield _Session()
|
||||
|
||||
async def _no_discovery(self, session):
|
||||
pass
|
||||
|
||||
monkeypatch.setattr(mcp_mod, "stdio_client", _stdio_client)
|
||||
monkeypatch.setattr(mcp_mod, "ClientSession", _client_session)
|
||||
monkeypatch.setattr(mcp_mod.MCPToolRunner, "_discover_resources_and_prompts", _no_discovery)
|
||||
|
||||
runner = mcp_mod.MCPToolRunner(None, timeout=30)
|
||||
assert runner.initialized.wait(5) and runner._init_error is None
|
||||
|
||||
slow_q, fast_q = queue.Queue(), queue.Queue()
|
||||
runner.queue.put((slow_q, "tool", "slow", {}, threading.Event()))
|
||||
runner.queue.put((fast_q, "tool", "fast", {}, threading.Event()))
|
||||
# The fast call is not serialized behind the in-flight slow one.
|
||||
assert fast_q.get(timeout=2) == (True, "fast")
|
||||
|
||||
runner.queue.put(None)
|
||||
ok, message = slow_q.get(timeout=2)
|
||||
assert not ok and "shut down" in message
|
||||
runner.join(2)
|
||||
assert not runner.is_alive()
|
||||
|
||||
|
||||
def test_agent_deepcopy_after_async_tool_use():
|
||||
import copy
|
||||
|
||||
def echo(x: int) -> int:
|
||||
"""echo"""
|
||||
return x
|
||||
|
||||
agent = Agent(name="d", instructions="x", tools=[echo], llm="gpt-4o-mini")
|
||||
assert asyncio.run(agent.execute_tool_async("echo", {"x": 1})) == 1
|
||||
clone = copy.deepcopy(agent)
|
||||
assert "_tool_executor" not in clone.__dict__
|
||||
assert asyncio.run(clone.execute_tool_async("echo", {"x": 2})) == 2
|
||||
@@ -0,0 +1,93 @@
|
||||
"""Loop guard per-turn time limit is configurable via ExecutionConfig (#6026)."""
|
||||
import asyncio
|
||||
|
||||
import pytest
|
||||
|
||||
import praisonaiagents.escalation.loop_guard as lg
|
||||
from praisonaiagents import Agent
|
||||
from praisonaiagents.config.feature_configs import ExecutionConfig
|
||||
from praisonaiagents.escalation.loop_guard import GuardAction, LoopGuardConfig
|
||||
from praisonaiagents.model_harness import ScriptedModel
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def clock(monkeypatch):
|
||||
state = {"t": 1_000_000.0}
|
||||
|
||||
class FakeTime:
|
||||
@staticmethod
|
||||
def time():
|
||||
return state["t"]
|
||||
|
||||
monkeypatch.setattr(lg, "time", FakeTime)
|
||||
return state
|
||||
|
||||
|
||||
def _agent(clock, tool_seconds, execution=None, calls=None):
|
||||
calls = [] if calls is None else calls
|
||||
|
||||
def run_tests(suite: str) -> str:
|
||||
"""Run a (slow) test suite."""
|
||||
calls.append(suite)
|
||||
clock["t"] += tool_seconds
|
||||
return f"{suite}: 412 passed"
|
||||
|
||||
model = ScriptedModel([
|
||||
ScriptedModel.tool_call("run_tests", {"suite": "unit"}),
|
||||
ScriptedModel.tool_call("run_tests", {"suite": "integration"}),
|
||||
"All suites passed.",
|
||||
])
|
||||
kwargs = {"execution": execution} if execution is not None else {}
|
||||
return Agent(name="coder", instructions="Run the tests.", llm=model,
|
||||
tools=[run_tests], output="silent", **kwargs)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("use_async", [False, True])
|
||||
def test_default_turn_with_130s_tool_completes(clock, use_async):
|
||||
agent = _agent(clock, 130.0)
|
||||
if use_async:
|
||||
out = asyncio.run(agent.achat("Run both suites"))
|
||||
else:
|
||||
out = agent.chat("Run both suites")
|
||||
assert "All suites passed" in str(out)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("use_async", [False, True])
|
||||
def test_configured_limit_halts(clock, use_async):
|
||||
calls = []
|
||||
agent = _agent(clock, 130.0, ExecutionConfig(max_execution_time=60), calls)
|
||||
if use_async:
|
||||
asyncio.run(agent.achat("Run both suites"))
|
||||
else:
|
||||
with pytest.raises(Exception, match="Turn time limit exceeded"):
|
||||
agent.chat("Run both suites")
|
||||
# The second tool call is refused by the loop guard on both paths.
|
||||
assert calls == ["unit"]
|
||||
|
||||
|
||||
@pytest.mark.parametrize("use_async", [False, True])
|
||||
def test_zero_disables_time_halt(clock, use_async):
|
||||
agent = _agent(clock, 10_000.0, ExecutionConfig(max_execution_time=0))
|
||||
if use_async:
|
||||
out = asyncio.run(agent.achat("Run both suites"))
|
||||
else:
|
||||
out = agent.chat("Run both suites")
|
||||
assert "All suites passed" in str(out)
|
||||
|
||||
|
||||
def test_ensure_loop_guard_reads_max_execution_time():
|
||||
agent = Agent(name="a", instructions="x", llm=ScriptedModel(["ok"]),
|
||||
output="silent", execution=ExecutionConfig(max_execution_time=900))
|
||||
assert agent._ensure_loop_guard().config.max_time_per_turn == 900.0
|
||||
|
||||
|
||||
def test_real_loop_still_halts_without_time_limit():
|
||||
guard = lg.LoopGuard(LoopGuardConfig(max_time_per_turn=float("inf")))
|
||||
actions = []
|
||||
for _ in range(20):
|
||||
decision = guard.check("read_file", {"path": "a.txt"})
|
||||
actions.append(decision.action)
|
||||
if decision.action == GuardAction.HALT:
|
||||
break
|
||||
guard.record("read_file", {"path": "a.txt"}, success=True, result="same")
|
||||
assert actions[-1] == GuardAction.HALT
|
||||
@@ -0,0 +1,126 @@
|
||||
"""``ExecutionConfig.max_execution_time`` and ``run_autonomous(timeout_seconds=)`` (#6029).
|
||||
|
||||
Both caps become a deadline on the turn's cancel token, so the LLM tool loop
|
||||
stops at its next boundary instead of running to completion.
|
||||
``litellm.completion`` is stubbed; nothing leaves the machine.
|
||||
"""
|
||||
import asyncio
|
||||
import time
|
||||
|
||||
import pytest
|
||||
|
||||
pytest.importorskip("litellm")
|
||||
|
||||
from praisonaiagents import Agent # noqa: E402
|
||||
from praisonaiagents.config.feature_configs import ExecutionConfig # noqa: E402
|
||||
from praisonaiagents.escalation.loop_guard import LoopGuardConfig # noqa: E402
|
||||
|
||||
MODEL = "anthropic/claude-3-5-sonnet-20240620"
|
||||
USAGE = {"prompt_tokens": 1, "completion_tokens": 1, "total_tokens": 2}
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def calls(monkeypatch):
|
||||
"""Stub litellm so the model asks for ``slow_step`` on every call."""
|
||||
# Import here, not at module level: other tests evict ``litellm`` from
|
||||
# sys.modules, and LLM then imports a fresh module object at call time.
|
||||
import litellm
|
||||
from litellm import ModelResponse
|
||||
|
||||
monkeypatch.setenv("ANTHROPIC_API_KEY", "sk-test")
|
||||
# Agent caches its env-derived default model per process; keep the fake
|
||||
# key's Anthropic default from leaking into later tests.
|
||||
monkeypatch.setattr(Agent, "_default_model", None)
|
||||
monkeypatch.setattr(Agent, "_default_model_checked", False)
|
||||
state = {"n": 0}
|
||||
|
||||
def fake(*a, **k):
|
||||
state["n"] += 1
|
||||
if not k.get("tools"):
|
||||
return ModelResponse(model=k.get("model"), usage=USAGE, choices=[{
|
||||
"index": 0, "finish_reason": "stop",
|
||||
"message": {"role": "assistant", "content": "done"}}])
|
||||
return ModelResponse(model=k.get("model"), usage=USAGE, choices=[{
|
||||
"index": 0, "finish_reason": "tool_calls",
|
||||
"message": {"role": "assistant", "content": None, "tool_calls": [{
|
||||
"id": f"call_{state['n']}", "type": "function",
|
||||
"function": {"name": "slow_step", "arguments": "{}"}}]}}])
|
||||
|
||||
async def afake(*a, **k):
|
||||
return fake(*a, **k)
|
||||
|
||||
monkeypatch.setattr(litellm, "completion", fake)
|
||||
monkeypatch.setattr(litellm, "acompletion", afake)
|
||||
return state
|
||||
|
||||
|
||||
def slow_step() -> str:
|
||||
"""A slow step."""
|
||||
time.sleep(0.4)
|
||||
return "ok"
|
||||
|
||||
|
||||
def _agent(**kw):
|
||||
return Agent(name="t", instructions="x", llm=MODEL, tools=[slow_step],
|
||||
output="silent", **kw)
|
||||
|
||||
|
||||
def test_max_execution_time_stops_tool_loop(calls):
|
||||
agent = _agent(execution=ExecutionConfig(max_execution_time=1, max_iter=50))
|
||||
t0 = time.time()
|
||||
out = agent.chat("go")
|
||||
elapsed = time.time() - t0
|
||||
assert elapsed < 3
|
||||
assert "timeout" in (out or "")
|
||||
assert agent.last_stop_reason == "timeout"
|
||||
assert calls["n"] < 10
|
||||
|
||||
|
||||
def test_max_execution_time_stops_tool_loop_async(calls):
|
||||
agent = _agent(execution=ExecutionConfig(max_execution_time=1, max_iter=50))
|
||||
t0 = time.time()
|
||||
out = asyncio.run(agent.achat("go"))
|
||||
assert time.time() - t0 < 3
|
||||
assert "timeout" in (out or "")
|
||||
assert agent.last_stop_reason == "timeout"
|
||||
|
||||
|
||||
def test_unset_max_execution_time_creates_no_deadline():
|
||||
agent = Agent(name="t", instructions="x", output="silent")
|
||||
assert agent._turn_cancel_token(None, explicit=False) is None
|
||||
|
||||
|
||||
def test_loop_guard_turn_limit_follows_max_execution_time():
|
||||
agent = Agent(name="t", instructions="x", output="silent",
|
||||
execution=ExecutionConfig(max_execution_time=600))
|
||||
assert agent._ensure_loop_guard().config.max_time_per_turn == 600
|
||||
default = Agent(name="d", instructions="x", output="silent")
|
||||
assert default._ensure_loop_guard().config.max_time_per_turn == LoopGuardConfig.max_time_per_turn
|
||||
|
||||
|
||||
def test_run_autonomous_timeout_bounds_the_running_turn(calls):
|
||||
agent = _agent(autonomy=True, execution=ExecutionConfig(max_iter=50))
|
||||
t0 = time.time()
|
||||
result = agent.run_autonomous("go", max_iterations=1, timeout_seconds=1)
|
||||
assert time.time() - t0 < 3
|
||||
assert result.completion_reason == "timeout"
|
||||
assert not result.success
|
||||
|
||||
|
||||
def test_run_autonomous_async_timeout_bounds_the_running_turn(calls):
|
||||
agent = _agent(autonomy=True, execution=ExecutionConfig(max_iter=50))
|
||||
t0 = time.time()
|
||||
result = asyncio.run(agent.run_autonomous_async("go", max_iterations=1, timeout_seconds=1))
|
||||
assert time.time() - t0 < 3
|
||||
assert result.completion_reason == "timeout"
|
||||
|
||||
|
||||
def test_run_autonomous_reports_turn_max_execution_time(calls):
|
||||
"""A turn cut by max_execution_time is a failed timeout, not a goal."""
|
||||
agent = _agent(autonomy=True, execution=ExecutionConfig(max_execution_time=1, max_iter=50))
|
||||
result = agent.run_autonomous("go", max_iterations=3)
|
||||
assert result.completion_reason == "timeout"
|
||||
assert not result.success
|
||||
result = asyncio.run(agent.run_autonomous_async("go", max_iterations=3))
|
||||
assert result.completion_reason == "timeout"
|
||||
assert not result.success
|
||||
@@ -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)
|
||||
@@ -0,0 +1,44 @@
|
||||
"""GitHub write tools must go through the approval system (issue #6023)."""
|
||||
import pytest
|
||||
|
||||
from praisonaiagents.approval import get_approval_registry
|
||||
from praisonaiagents.approval.registry import PERMISSION_PRESETS
|
||||
|
||||
GITHUB_WRITE_TOOLS = (
|
||||
"github_create_branch",
|
||||
"github_commit_and_push",
|
||||
"github_create_pull_request",
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("name", GITHUB_WRITE_TOOLS)
|
||||
def test_registered_as_high_risk(name):
|
||||
reg = get_approval_registry()
|
||||
assert reg.is_required(name)
|
||||
# high (not critical) so YAML/declarative approve lists still apply
|
||||
assert reg.get_risk_level(name) == "high"
|
||||
|
||||
|
||||
@pytest.mark.parametrize("name", GITHUB_WRITE_TOOLS)
|
||||
def test_in_deny_presets(name):
|
||||
for preset in ("default", "safe", "read_only"):
|
||||
assert name in PERMISSION_PRESETS[preset]
|
||||
assert name not in PERMISSION_PRESETS["full"]
|
||||
|
||||
|
||||
def test_unattended_agent_denied_by_default(monkeypatch):
|
||||
monkeypatch.delenv("PRAISONAI_TOOL_SAFETY", raising=False)
|
||||
from praisonaiagents import Agent
|
||||
|
||||
agent = Agent(instructions="t", llm="gpt-4o-mini", approval=False)
|
||||
for name in ("github_commit_and_push", "github_create_pull_request"):
|
||||
result = agent._check_tool_approval_sync(name, {})
|
||||
assert isinstance(result, dict) and result.get("permission_denied")
|
||||
|
||||
|
||||
def test_explicit_full_preset_allows(monkeypatch):
|
||||
monkeypatch.setenv("PRAISONAI_TOOL_SAFETY", "off")
|
||||
from praisonaiagents import Agent
|
||||
|
||||
agent = Agent(instructions="t", llm="gpt-4o-mini", approval=False)
|
||||
assert "github_commit_and_push" not in (agent._perm_deny or ())
|
||||
@@ -0,0 +1,176 @@
|
||||
"""Config options that were accepted but never read must either work or fail loud.
|
||||
|
||||
Covers GuardrailConfig.on_fail, ToolConfig.output_max_lines/output_direction and
|
||||
rerank=True on the Agent's RAG pipeline (issue #6042). No network calls.
|
||||
"""
|
||||
import pytest
|
||||
|
||||
from praisonaiagents import Agent
|
||||
from praisonaiagents.config import GuardrailConfig, ToolConfig
|
||||
from praisonaiagents.config.presets import GUARDRAIL_PRESETS
|
||||
|
||||
|
||||
def _ok(task_output):
|
||||
return True, task_output
|
||||
|
||||
|
||||
# --- GuardrailConfig.on_fail ------------------------------------------------
|
||||
|
||||
@pytest.mark.parametrize("on_fail", ["skip", "raise"])
|
||||
def test_unsupported_on_fail_fails_loud(on_fail):
|
||||
with pytest.raises(ValueError, match="on_fail"):
|
||||
Agent(name="g", instructions="x", llm="gpt-4o-mini",
|
||||
guardrails=GuardrailConfig(validator=_ok, on_fail=on_fail))
|
||||
|
||||
|
||||
def test_retry_on_fail_still_accepted():
|
||||
a = Agent(name="g", instructions="x", llm="gpt-4o-mini",
|
||||
guardrails=GuardrailConfig(validator=_ok, max_retries=2, on_fail="retry"))
|
||||
assert a.max_guardrail_retries == 2
|
||||
|
||||
|
||||
def test_presets_carry_no_unimplemented_on_fail():
|
||||
assert all("on_fail" not in p for p in GUARDRAIL_PRESETS.values())
|
||||
a = Agent(name="g", instructions="x", llm="gpt-4o-mini",
|
||||
guardrails=["strict", {"validator": _ok}])
|
||||
assert a.max_guardrail_retries == 5
|
||||
assert a.guardrail is _ok
|
||||
|
||||
|
||||
# --- ToolConfig.output_max_lines / output_direction --------------------------
|
||||
|
||||
def many_lines() -> str:
|
||||
"""Return 200 short lines."""
|
||||
return "\n".join(f"line {i}" for i in range(200))
|
||||
|
||||
|
||||
def _run_tool(cfg):
|
||||
a = Agent(name="t", instructions="x", llm="gpt-4o-mini", tools=[many_lines],
|
||||
tool_config=cfg, output="silent")
|
||||
return str(a.execute_tool("many_lines", {}))
|
||||
|
||||
|
||||
def test_default_tool_config_does_not_cap_lines():
|
||||
out = _run_tool(ToolConfig())
|
||||
assert out.splitlines()[-1] == "line 199"
|
||||
assert out.count("\n") == 199
|
||||
|
||||
|
||||
def test_output_max_lines_head():
|
||||
lines = _run_tool(ToolConfig(output_max_lines=5, output_direction="head")).splitlines()
|
||||
assert lines[:5] == [f"line {i}" for i in range(5)]
|
||||
assert "line 5" not in lines and "line 199" not in lines
|
||||
assert "200 lines total" in lines[5]
|
||||
|
||||
|
||||
def test_output_max_lines_tail():
|
||||
lines = _run_tool(ToolConfig(output_max_lines=5, output_direction="tail")).splitlines()
|
||||
assert lines[-5:] == [f"line {i}" for i in range(195, 200)]
|
||||
assert "line 0" not in lines
|
||||
|
||||
|
||||
def test_output_max_lines_both():
|
||||
lines = _run_tool(ToolConfig(output_max_lines=4)).splitlines()
|
||||
assert lines[:2] == ["line 0", "line 1"]
|
||||
assert lines[-2:] == ["line 198", "line 199"]
|
||||
assert len(lines) == 5 # 4 kept lines + marker
|
||||
|
||||
|
||||
def test_output_direction_applies_to_char_cap():
|
||||
out = _run_tool(ToolConfig(output_limit=50, output_direction="head"))
|
||||
assert out.startswith("line 0\nline 1")
|
||||
assert "line 199" not in out
|
||||
|
||||
|
||||
def short_lines() -> str:
|
||||
"""Return ten one-character lines."""
|
||||
return "\n".join("abcdefghij")
|
||||
|
||||
|
||||
def big_dict() -> dict:
|
||||
"""Return a dict with one large field."""
|
||||
return {"body": "START" + "x" * 45000 + "END"}
|
||||
|
||||
|
||||
def test_line_cap_applies_to_short_lines():
|
||||
a = Agent(name="t", instructions="x", llm="gpt-4o-mini", tools=[short_lines],
|
||||
tool_config=ToolConfig(output_max_lines=5, output_direction="head"), output="silent")
|
||||
lines = str(a.execute_tool("short_lines", {})).splitlines()
|
||||
assert lines[:5] == list("abcde")
|
||||
assert "j" not in lines
|
||||
|
||||
|
||||
def test_output_direction_applies_to_dict_results():
|
||||
a = Agent(name="t", instructions="x", llm="gpt-4o-mini", tools=[big_dict],
|
||||
tool_config=ToolConfig(output_limit=40000, output_direction="head"), output="silent")
|
||||
out = str(a.execute_tool("big_dict", {}))
|
||||
assert "START" in out and "END" not in out
|
||||
|
||||
|
||||
# --- rerank=True on Agent.rag -----------------------------------------------
|
||||
|
||||
class StubKnowledge:
|
||||
"""Minimal Knowledge stand-in: search() returns hits worst-first."""
|
||||
|
||||
def add(self, *a, **k):
|
||||
pass
|
||||
|
||||
def search(self, query, **kw):
|
||||
return {"results": [
|
||||
{"memory": "bananas are yellow", "score": 0.9, "metadata": {"source": "a"}},
|
||||
{"memory": "paris is the capital of france", "score": 0.8, "metadata": {"source": "b"}},
|
||||
]}
|
||||
|
||||
|
||||
def _rag(retrieval_config=None):
|
||||
a = Agent(name="r", instructions="x", llm="gpt-4o-mini", knowledge=StubKnowledge(),
|
||||
retrieval_config=retrieval_config)
|
||||
return a.rag
|
||||
|
||||
|
||||
def test_rerank_true_reranks_and_reports_it():
|
||||
pack = _rag({"rerank": True}).retrieve("capital of france")
|
||||
assert pack.metadata["reranked"] is True
|
||||
assert pack.context.index("paris") < pack.context.index("bananas")
|
||||
|
||||
|
||||
def test_rerank_false_reports_not_reranked():
|
||||
rag = _rag()
|
||||
assert rag.reranker is None
|
||||
assert rag.retrieve("capital of france").metadata["reranked"] is False
|
||||
|
||||
|
||||
def test_rerank_override_without_reranker_reports_false():
|
||||
assert _rag().retrieve("q", rerank=True).metadata["reranked"] is False
|
||||
|
||||
|
||||
def test_rerank_model_leaves_reranking_to_backend():
|
||||
rag = _rag({"rerank": True, "rerank_model": "cohere/rerank-v3"})
|
||||
assert rag.reranker is None
|
||||
|
||||
|
||||
def test_rerank_model_backend_reranking_reported():
|
||||
class BackendRerankKnowledge(StubKnowledge):
|
||||
config = {"reranker": {"default_rerank": True, "provider": "cohere"}}
|
||||
|
||||
a = Agent(name="r", instructions="x", llm="gpt-4o-mini", knowledge=BackendRerankKnowledge(),
|
||||
retrieval_config={"rerank": True, "rerank_model": "cohere/rerank-v3"})
|
||||
assert a.rag.reranker is None
|
||||
assert a.rag.retrieve("q").metadata["reranked"] is True
|
||||
assert a.rag.retrieve("q", rerank=False).metadata["reranked"] is False
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_rerank_true_async():
|
||||
pack = await _rag({"rerank": True}).aretrieve("capital of france")
|
||||
assert pack.metadata["reranked"] is True
|
||||
assert pack.context.index("paris") < pack.context.index("bananas")
|
||||
|
||||
|
||||
def test_simple_reranker_accepts_search_result_items():
|
||||
from praisonaiagents.knowledge.models import SearchResultItem
|
||||
from praisonaiagents.rag.retriever import SimpleReranker
|
||||
hits = [SearchResultItem(text="bananas are yellow", score=0.9),
|
||||
SearchResultItem(text="paris is the capital of france", score=0.8)]
|
||||
ranked = SimpleReranker().rerank("capital of france", hits)
|
||||
assert ranked[0]["text"].startswith("paris")
|
||||
@@ -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]
|
||||
@@ -275,3 +275,113 @@ def test_async_fallback_passes_override_and_leaves_shared_state(monkeypatch):
|
||||
assert used[0] is not original_dispatcher
|
||||
assert agent.llm == "gpt-4o-mini" # shared model untouched
|
||||
assert agent._unified_dispatcher is original_dispatcher # shared cache untouched
|
||||
|
||||
|
||||
# --- Custom-LLM path (Issue #6030) ---------------------------------------------
|
||||
|
||||
class _Unavailable(Exception):
|
||||
status_code = 503
|
||||
|
||||
|
||||
class _StubLLM:
|
||||
def __init__(self, model, fail, calls):
|
||||
self.model, self._fail, self._calls = model, fail, calls
|
||||
|
||||
def _respond(self):
|
||||
self._calls.append(self.model)
|
||||
if self._fail:
|
||||
raise self._fail
|
||||
return f"answer from {self.model}"
|
||||
|
||||
def get_response(self, **kwargs):
|
||||
return self._respond()
|
||||
|
||||
async def get_response_async(self, **kwargs):
|
||||
return self._respond()
|
||||
|
||||
|
||||
class _ToolThenFailLLM(_StubLLM):
|
||||
"""Runs a tool, then fails, like a 503 on the follow-up model request."""
|
||||
|
||||
def get_response(self, execute_tool_fn=None, **kwargs):
|
||||
execute_tool_fn("send_email", {})
|
||||
return self._respond()
|
||||
|
||||
|
||||
def _custom_agent(primary_error, fallbacks):
|
||||
from praisonaiagents import Agent
|
||||
|
||||
calls = []
|
||||
primary = _StubLLM("primary-model", primary_error, calls)
|
||||
agent = Agent(name="c", instructions="x", llm=primary, output="silent")
|
||||
agent.fallback_models = [_StubLLM(m, None, calls) for m in fallbacks]
|
||||
seen = []
|
||||
agent._hook_runner.registry.register_function(
|
||||
HookEvent.MODEL_FALLBACK, lambda inp: seen.append(inp) or HookResult.allow()
|
||||
)
|
||||
return agent, primary, calls, seen
|
||||
|
||||
|
||||
def test_custom_llm_sync_falls_back_on_outage():
|
||||
agent, primary, calls, seen = _custom_agent(_Unavailable("503 overloaded"), ["backup"])
|
||||
assert agent.chat("hello") == "answer from backup"
|
||||
assert calls == ["primary-model", "backup"]
|
||||
assert len(seen) == 1 and seen[0].fallback_index == 0
|
||||
assert agent.llm_instance is primary # shared state untouched
|
||||
|
||||
|
||||
def test_custom_llm_async_falls_back_on_outage():
|
||||
agent, primary, calls, seen = _custom_agent(_Unavailable("503 overloaded"), ["backup"])
|
||||
assert asyncio.run(agent.achat("hello")) == "answer from backup"
|
||||
assert calls == ["primary-model", "backup"]
|
||||
assert len(seen) == 1
|
||||
assert agent.llm_instance is primary
|
||||
|
||||
|
||||
def test_custom_llm_non_fallback_error_does_not_walk_chain():
|
||||
err = Exception("401 Unauthorized: invalid api key")
|
||||
err.status_code = 401
|
||||
agent, _, calls, seen = _custom_agent(err, ["backup"])
|
||||
assert agent.chat("hello") is None
|
||||
assert calls == ["primary-model"]
|
||||
assert seen == []
|
||||
|
||||
|
||||
def test_fallback_llm_for_drops_connection_settings_across_providers(monkeypatch):
|
||||
import praisonaiagents.llm.llm as llm_mod
|
||||
|
||||
built = []
|
||||
monkeypatch.setattr(llm_mod, "LLM", lambda **kw: built.append(kw) or kw)
|
||||
agent = _DummyAgent()
|
||||
agent.llm = "ollama/llama3"
|
||||
agent._llm_init_params = {"model": "ollama/llama3", "base_url": "http://h", "api_key": "k", "max_iter": 3}
|
||||
agent._fallback_llm_for("ollama/qwen")
|
||||
agent._fallback_llm_for("openai/gpt-4o-mini")
|
||||
assert built[0] == {"model": "ollama/qwen", "base_url": "http://h", "api_key": "k", "max_iter": 3}
|
||||
assert built[1] == {"model": "openai/gpt-4o-mini", "max_iter": 3}
|
||||
assert agent._llm_init_params["model"] == "ollama/llama3"
|
||||
|
||||
|
||||
def test_fallback_llm_for_keeps_connection_settings_for_bare_names(monkeypatch):
|
||||
import praisonaiagents.llm.llm as llm_mod
|
||||
|
||||
built = []
|
||||
monkeypatch.setattr(llm_mod, "LLM", lambda **kw: built.append(kw) or kw)
|
||||
agent = _DummyAgent()
|
||||
agent.llm = "model-a"
|
||||
agent._llm_init_params = {"model": "model-a", "base_url": "http://custom/v1", "api_key": "k"}
|
||||
agent._fallback_llm_for("model-b")
|
||||
assert built[0] == {"model": "model-b", "base_url": "http://custom/v1", "api_key": "k"}
|
||||
|
||||
|
||||
def test_custom_llm_does_not_fall_back_after_a_tool_ran():
|
||||
from praisonaiagents import Agent
|
||||
|
||||
calls, sent = [], []
|
||||
primary = _ToolThenFailLLM("primary-model", _Unavailable("503 overloaded"), calls)
|
||||
agent = Agent(name="c", instructions="x", llm=primary, output="silent")
|
||||
agent.fallback_models = [_StubLLM("backup", None, calls)]
|
||||
agent.execute_tool = lambda name, args, **kw: sent.append(name)
|
||||
assert agent.chat("hello") is None # existing error handling, no replay
|
||||
assert calls == ["primary-model"]
|
||||
assert sent == ["send_email"]
|
||||
@@ -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,119 @@
|
||||
"""Distinct session ids must never share a session file (Issue #6024)."""
|
||||
|
||||
import json
|
||||
import os
|
||||
|
||||
import pytest
|
||||
|
||||
from praisonaiagents.session.store import DefaultSessionStore
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def store(tmp_path):
|
||||
return DefaultSessionStore(session_dir=str(tmp_path))
|
||||
|
||||
|
||||
@pytest.mark.parametrize("a,b", [
|
||||
("alice.smith", "alice_smith"),
|
||||
("a:b", "a_b"),
|
||||
("[email protected]", "[email protected]"),
|
||||
("[email protected]", "u_tag_x_com"),
|
||||
("a/b", "a b"),
|
||||
("telegram:acct:default:chat:-1001:", "telegram_acct_default_chat_-1001_"),
|
||||
])
|
||||
def test_colliding_ids_get_distinct_paths(store, a, b):
|
||||
assert store._get_session_path(a) != store._get_session_path(b)
|
||||
|
||||
|
||||
def test_safe_ids_keep_legacy_filename(store):
|
||||
sid = "user_123-abc" + "x" * 100
|
||||
assert os.path.basename(store._get_session_path(sid)) == f"{sid}.json"
|
||||
|
||||
|
||||
def test_colliding_users_have_separate_histories(store):
|
||||
store.add_user_message("[email protected]", "My PIN is 4471.")
|
||||
assert store.get_chat_history("[email protected]") == []
|
||||
store.add_user_message("[email protected]", "hello")
|
||||
assert [m["content"] for m in store.get_chat_history("[email protected]")] == ["My PIN is 4471."]
|
||||
assert [m["content"] for m in store.get_chat_history("[email protected]")] == ["hello"]
|
||||
|
||||
|
||||
def _write_legacy(store, stem, owner, content):
|
||||
path = os.path.join(store.session_dir, f"{stem}.json")
|
||||
with open(path, "w", encoding="utf-8") as f:
|
||||
json.dump({"session_id": owner, "messages": [
|
||||
{"role": "user", "content": content, "timestamp": 1.0}]}, f)
|
||||
return path
|
||||
|
||||
|
||||
def test_legacy_file_migrated_only_for_owner(store):
|
||||
legacy = _write_legacy(store, "alice_smith", "alice.smith", "secret")
|
||||
# A colliding unsafe id never adopts it.
|
||||
assert store.get_chat_history("alice:smith") == []
|
||||
assert os.path.exists(legacy)
|
||||
# The real owner migrates it on first read.
|
||||
assert [m["content"] for m in store.get_chat_history("alice.smith")] == ["secret"]
|
||||
assert not os.path.exists(legacy)
|
||||
|
||||
|
||||
def test_foreign_file_never_served_and_not_destroyed(store):
|
||||
# Safe id whose filename holds another id's unmigrated legacy transcript.
|
||||
_write_legacy(store, "alice_smith", "alice.smith", "secret")
|
||||
assert store.get_chat_history("alice_smith") == []
|
||||
store.add_user_message("alice_smith", "mine")
|
||||
assert [m["content"] for m in store.get_chat_history("alice_smith")] == ["mine"]
|
||||
# The owner's history was relocated rather than overwritten.
|
||||
assert [m["content"] for m in store.get_chat_history("alice.smith")] == ["secret"]
|
||||
|
||||
|
||||
def test_failed_foreign_move_aborts_write(store, monkeypatch, tmp_path):
|
||||
monkeypatch.setattr(store, "_spill_dir", lambda: str(tmp_path / "spill"))
|
||||
store.add_user_message("alice.smith", "new") # owner path taken
|
||||
legacy = _write_legacy(store, "alice_smith", "alice.smith", "secret")
|
||||
monkeypatch.setattr(store, "_quarantine_corrupt", lambda p: None)
|
||||
assert store.add_user_message("alice_smith", "mine") is False
|
||||
with open(legacy, encoding="utf-8") as f:
|
||||
assert json.load(f)["session_id"] == "alice.smith"
|
||||
|
||||
|
||||
def test_foreign_file_set_aside_when_owner_path_taken(store):
|
||||
store.add_user_message("alice.smith", "new")
|
||||
legacy = _write_legacy(store, "alice_smith", "alice.smith", "secret")
|
||||
store.add_user_message("alice_smith", "mine")
|
||||
assert [m["content"] for m in store.get_chat_history("alice_smith")] == ["mine"]
|
||||
assert [m["content"] for m in store.get_chat_history("alice.smith")] == ["new"]
|
||||
aside = [p for p in os.listdir(store.session_dir) if p.startswith("alice_smith.json.corrupt-")]
|
||||
assert len(aside) == 1 and os.path.exists(legacy)
|
||||
|
||||
|
||||
def test_sqlite_backfill_indexes_hashed_filenames(tmp_path):
|
||||
from praisonaiagents.session.sqlite_store import SqliteSessionStore
|
||||
|
||||
DefaultSessionStore(session_dir=str(tmp_path)).add_user_message(
|
||||
"[email protected]", "narwhal record")
|
||||
store = SqliteSessionStore(session_dir=str(tmp_path), db_path=":memory:")
|
||||
assert [h.session_id for h in store.search("narwhal")] == ["[email protected]"]
|
||||
|
||||
|
||||
def test_sqlite_backfill_indexes_legacy_collision_files(tmp_path):
|
||||
from praisonaiagents.session.sqlite_store import SqliteSessionStore
|
||||
|
||||
legacy = DefaultSessionStore(session_dir=str(tmp_path))
|
||||
_write_legacy(legacy, "alice_smith", "alice.smith", "walrus record")
|
||||
store = SqliteSessionStore(session_dir=str(tmp_path), db_path=":memory:")
|
||||
assert [h.session_id for h in store.search("walrus")] == ["alice.smith"]
|
||||
|
||||
|
||||
def test_spill_recovery_does_not_wait_on_own_lock(store, monkeypatch, tmp_path):
|
||||
import time
|
||||
|
||||
spill = tmp_path / "spill"
|
||||
spill.mkdir()
|
||||
monkeypatch.setattr(store, "_spill_dir", lambda: str(spill))
|
||||
# Legacy name of "alice.smith" belongs to another id, so it is never migrated.
|
||||
_write_legacy(store, "alice_smith", "alice:smith", "theirs")
|
||||
(spill / "alice_smith.1.json").write_text(json.dumps({"session_id": "alice.smith", "messages": [
|
||||
{"role": "user", "content": "spilled", "timestamp": 2.0}]}), encoding="utf-8")
|
||||
start = time.monotonic()
|
||||
assert [m["content"] for m in store.get_chat_history("alice.smith")] == ["spilled"]
|
||||
assert time.monotonic() - start < store.lock_timeout
|
||||
@@ -0,0 +1,36 @@
|
||||
"""Telemetry shutdown must return promptly when the PostHog flush is wedged."""
|
||||
|
||||
import threading
|
||||
import time
|
||||
from unittest.mock import MagicMock
|
||||
|
||||
from praisonaiagents.telemetry.telemetry import MinimalTelemetry
|
||||
|
||||
|
||||
class _FastJoinThread(threading.Thread):
|
||||
def join(self, timeout=None):
|
||||
return super().join(timeout=0.2 if timeout else timeout)
|
||||
|
||||
|
||||
def test_shutdown_returns_when_flush_hangs(monkeypatch):
|
||||
release = threading.Event()
|
||||
client = MagicMock()
|
||||
client.flush.side_effect = lambda: release.wait(10)
|
||||
|
||||
telemetry = MinimalTelemetry(enabled=False)
|
||||
telemetry.enabled = True
|
||||
telemetry._shutdown_lock = threading.Lock()
|
||||
telemetry._shutdown_complete = False
|
||||
telemetry._posthog = client
|
||||
monkeypatch.setattr(telemetry, "flush", lambda: None)
|
||||
monkeypatch.setattr(threading, "Thread", _FastJoinThread)
|
||||
|
||||
start = time.monotonic()
|
||||
try:
|
||||
telemetry.shutdown()
|
||||
finally:
|
||||
release.set()
|
||||
assert time.monotonic() - start < 2
|
||||
# Cleanup that would block on the same wedged transport is skipped.
|
||||
client.shutdown.assert_not_called()
|
||||
assert telemetry._posthog is None
|
||||
@@ -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,66 @@
|
||||
"""Handoff context policies must keep tool-call pairs intact and keep the goal."""
|
||||
|
||||
from types import SimpleNamespace
|
||||
|
||||
from praisonaiagents.agent.handoff import (
|
||||
ContextPolicy,
|
||||
Handoff,
|
||||
HandoffConfig,
|
||||
_tool_pair_safe_tail,
|
||||
)
|
||||
|
||||
CALL = {"id": "c1", "type": "function", "function": {"name": "search", "arguments": "{}"}}
|
||||
HISTORY = [
|
||||
{"role": "system", "content": "sys"},
|
||||
{"role": "user", "content": "GOAL: book flight"},
|
||||
{"role": "assistant", "content": "ok, searching"},
|
||||
{"role": "assistant", "content": None, "tool_calls": [CALL]},
|
||||
{"role": "tool", "tool_call_id": "c1", "content": "results"},
|
||||
{"role": "assistant", "content": "found 3 options"},
|
||||
]
|
||||
|
||||
|
||||
def _prepare(policy, n=10):
|
||||
target = SimpleNamespace(name="tgt")
|
||||
handoff = Handoff(agent=target, config=HandoffConfig(context_policy=policy, max_context_messages=n))
|
||||
source = SimpleNamespace(name="src", chat_history=list(HISTORY))
|
||||
return handoff._prepare_context(source, {}).messages
|
||||
|
||||
|
||||
def _assert_pairs_valid(messages):
|
||||
call_ids = {c["id"] for m in messages for c in (m.get("tool_calls") or [])}
|
||||
response_ids = {m["tool_call_id"] for m in messages if m.get("role") == "tool"}
|
||||
assert call_ids == response_ids
|
||||
|
||||
|
||||
def test_last_n_drops_orphan_tool_result():
|
||||
msgs = _prepare(ContextPolicy.LAST_N, n=2)
|
||||
_assert_pairs_valid(msgs)
|
||||
assert [m["role"] for m in msgs] == ["system", "assistant"]
|
||||
|
||||
|
||||
def test_last_n_keeps_complete_pair():
|
||||
msgs = _prepare(ContextPolicy.LAST_N, n=3)
|
||||
_assert_pairs_valid(msgs)
|
||||
assert any(m.get("tool_calls") for m in msgs)
|
||||
|
||||
|
||||
def test_summary_keeps_goal_and_valid_pairs():
|
||||
msgs = _prepare(ContextPolicy.SUMMARY)
|
||||
_assert_pairs_valid(msgs)
|
||||
assert msgs[0]["role"] == "system"
|
||||
assert msgs[1]["content"] == "GOAL: book flight"
|
||||
assert msgs[-1]["content"] == "found 3 options"
|
||||
|
||||
|
||||
def test_summary_does_not_duplicate_goal_in_tail():
|
||||
target = SimpleNamespace(name="tgt")
|
||||
handoff = Handoff(agent=target, config=HandoffConfig(context_policy=ContextPolicy.SUMMARY))
|
||||
source = SimpleNamespace(name="src", chat_history=[{"role": "user", "content": "hi"}])
|
||||
assert handoff._prepare_context(source, {}).messages == [{"role": "user", "content": "hi"}]
|
||||
|
||||
|
||||
def test_tail_drops_tool_calls_without_responses():
|
||||
msgs = HISTORY[1:4] # ends on an unanswered tool_calls message
|
||||
assert _tool_pair_safe_tail(msgs, 3) == HISTORY[1:3]
|
||||
assert _tool_pair_safe_tail(msgs, 0) == []
|
||||
@@ -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,47 @@
|
||||
"""OpenAI image calls must only forward style to models that accept it (dall-e-3)."""
|
||||
|
||||
import pytest
|
||||
|
||||
from praisonaiagents import ImageAgent
|
||||
|
||||
|
||||
@pytest.mark.parametrize("llm,kwargs,expected", [
|
||||
("openai/dall-e-2", {}, None),
|
||||
("openai/gpt-image-1", {}, None),
|
||||
("openai/gpt-image-1", {"style": "vivid"}, None),
|
||||
(None, {}, None),
|
||||
("openai/dall-e-3", {}, None),
|
||||
("openai/dall-e-3", {"style": "natural"}, "natural"),
|
||||
("openai/dall-e-3", {"style": "vivid"}, "vivid"),
|
||||
])
|
||||
def test_style_forwarding(llm, kwargs, expected):
|
||||
agent = ImageAgent(llm=llm, verbose=False)
|
||||
received = {}
|
||||
|
||||
def fake_image_generation(**call_kwargs):
|
||||
received.update(call_kwargs)
|
||||
return {"data": [{"url": "https://example.com/image.png"}]}
|
||||
|
||||
agent._litellm = fake_image_generation
|
||||
try:
|
||||
agent.generate("A small green square icon", **kwargs)
|
||||
if expected is None:
|
||||
assert "style" not in received
|
||||
else:
|
||||
assert received["style"] == expected
|
||||
finally:
|
||||
agent.close()
|
||||
|
||||
|
||||
@pytest.mark.parametrize("config,expected", [
|
||||
({}, None),
|
||||
({"style": "vivid"}, "vivid"),
|
||||
])
|
||||
def test_yaml_image_agent_style(config, expected):
|
||||
from praisonaiagents.workflows.yaml_parser import YAMLWorkflowParser
|
||||
|
||||
agent = YAMLWorkflowParser()._create_image_agent("img", "openai/dall-e-3", config)
|
||||
try:
|
||||
assert agent.image_config.style == expected
|
||||
finally:
|
||||
agent.close()
|
||||
@@ -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"
|
||||
@@ -7,7 +7,7 @@ Tests:
|
||||
- Concurrent claims (threads) each get disjoint jobs — no double-fire
|
||||
- complete() releases the lease
|
||||
- Expired leases become claimable again (crash recovery)
|
||||
- One-shot jobs (delete_after_run) are removed on claim
|
||||
- One-shot jobs (delete_after_run) stay leased on claim, retired on completion
|
||||
- Lease metadata round-trips through to_dict/from_dict
|
||||
- ScheduleRunner.claim_due_jobs falls back to get_due_jobs without claim
|
||||
"""
|
||||
@@ -78,7 +78,7 @@ class TestClaimDue:
|
||||
reclaimed = store.claim_due(t1, owner_id="B", lease_seconds=5)
|
||||
assert len(reclaimed) == 1
|
||||
|
||||
def test_one_shot_removed_on_claim(self):
|
||||
def test_one_shot_retired_on_completion(self):
|
||||
with tempfile.TemporaryDirectory() as d:
|
||||
store = FileScheduleStore(store_dir=d)
|
||||
job = ScheduleJob(
|
||||
@@ -90,6 +90,10 @@ class TestClaimDue:
|
||||
store.add(job)
|
||||
claimed = store.claim_due(time.time(), owner_id="A")
|
||||
assert len(claimed) == 1
|
||||
# Kept (leased, retiring) so a crash mid-run cannot lose it.
|
||||
kept = store.get(job.id)
|
||||
assert kept is not None and kept._retiring == 1
|
||||
ScheduleRunner(store).mark_run(claimed[0])
|
||||
assert store.get(job.id) is None
|
||||
|
||||
def test_concurrent_claims_no_double_fire(self):
|
||||
@@ -335,7 +339,7 @@ class TestConfigYamlClaimDue:
|
||||
reclaimed = store.claim_due(t1, owner_id="B", lease_seconds=5)
|
||||
assert len(reclaimed) == 1
|
||||
|
||||
def test_one_shot_removed_on_claim(self):
|
||||
def test_one_shot_retired_on_completion(self):
|
||||
with tempfile.TemporaryDirectory() as d:
|
||||
store = self._make_store(d)
|
||||
job = ScheduleJob(
|
||||
@@ -347,6 +351,10 @@ class TestConfigYamlClaimDue:
|
||||
store.add(job)
|
||||
claimed = store.claim_due(time.time(), owner_id="A")
|
||||
assert len(claimed) == 1
|
||||
# Kept (leased, retiring) so a crash mid-run cannot lose it.
|
||||
kept = store.get(job.id)
|
||||
assert kept is not None and kept._retiring == 1
|
||||
ScheduleRunner(store).mark_run(claimed[0])
|
||||
assert store.get(job.id) is None
|
||||
|
||||
def test_concurrent_claims_no_double_fire(self):
|
||||
|
||||
@@ -0,0 +1,183 @@
|
||||
"""
|
||||
Crash recovery for claimed schedule jobs (issue #6036).
|
||||
|
||||
A spent job (one-shot / ``max_runs`` met) is no longer deleted at claim time:
|
||||
it stays leased and is retired on completion. A lease that expires without
|
||||
``complete`` is recorded as an ``abandoned`` run; a one-shot is retried once,
|
||||
a recurring job waits for its next slot. Covers both stores.
|
||||
"""
|
||||
|
||||
import os
|
||||
import time
|
||||
from datetime import datetime, timedelta, timezone
|
||||
|
||||
import pytest
|
||||
|
||||
from praisonaiagents.scheduler.config_store import ConfigYamlScheduleStore
|
||||
from praisonaiagents.scheduler.models import Schedule, ScheduleJob
|
||||
from praisonaiagents.scheduler.runner import ScheduleRunner
|
||||
from praisonaiagents.scheduler.store import FileScheduleStore
|
||||
|
||||
|
||||
@pytest.fixture(params=["file", "yaml"])
|
||||
def make_store(request, tmp_path):
|
||||
if request.param == "file":
|
||||
return lambda: FileScheduleStore(store_dir=str(tmp_path))
|
||||
path = os.path.join(str(tmp_path), "config.yaml")
|
||||
return lambda: ConfigYamlScheduleStore(config_path=path)
|
||||
|
||||
|
||||
def _one_shot(name="remind-once"):
|
||||
past = (datetime.now(timezone.utc) - timedelta(seconds=5)).isoformat()
|
||||
return ScheduleJob(
|
||||
name=name,
|
||||
message="ping",
|
||||
schedule=Schedule(kind="at", at=past),
|
||||
delete_after_run=True,
|
||||
)
|
||||
|
||||
|
||||
def _statuses(store, job_id):
|
||||
return [r.status for r in store.get_history(job_id=job_id)]
|
||||
|
||||
|
||||
def test_claim_keeps_one_shot_until_completion(make_store):
|
||||
store = make_store()
|
||||
job = _one_shot()
|
||||
store.add(job)
|
||||
claimed = store.claim_due(time.time(), owner_id="A", lease_seconds=300)
|
||||
assert [j.id for j in claimed] == [job.id]
|
||||
kept = make_store().get(job.id)
|
||||
assert kept is not None and kept._retiring == 1
|
||||
assert kept._lease_owner == "A"
|
||||
# A competitor cannot claim it while leased.
|
||||
assert make_store().claim_due(time.time(), owner_id="B") == []
|
||||
ScheduleRunner(store).mark_run(claimed[0])
|
||||
store.complete(job.id, "A")
|
||||
assert make_store().get(job.id) is None
|
||||
|
||||
|
||||
def test_crashed_one_shot_is_retried_once(make_store):
|
||||
store = make_store()
|
||||
job = _one_shot()
|
||||
store.add(job)
|
||||
t0 = time.time()
|
||||
assert store.claim_due(t0, owner_id="A", lease_seconds=1)
|
||||
del store # crash: no mark_run / complete
|
||||
|
||||
restarted = make_store()
|
||||
retry = restarted.claim_due(t0 + 5, owner_id="B", lease_seconds=1)
|
||||
assert [j.id for j in retry] == [job.id]
|
||||
assert _statuses(restarted, job.id) == ["abandoned"]
|
||||
# Retry also crashes: logged and retired, never offered a third time.
|
||||
third = make_store()
|
||||
assert third.claim_due(t0 + 10, owner_id="C", lease_seconds=1) == []
|
||||
assert third.get(job.id) is None
|
||||
assert _statuses(make_store(), job.id) == ["abandoned", "abandoned"]
|
||||
|
||||
|
||||
def test_retried_one_shot_completes_once(make_store):
|
||||
store = make_store()
|
||||
job = _one_shot()
|
||||
store.add(job)
|
||||
t0 = time.time()
|
||||
store.claim_due(t0, owner_id="A", lease_seconds=1)
|
||||
restarted = make_store()
|
||||
retry = restarted.claim_due(t0 + 5, owner_id="B", lease_seconds=300)
|
||||
ScheduleRunner(restarted).mark_run(retry[0])
|
||||
restarted.complete(job.id, "B")
|
||||
final = make_store()
|
||||
assert final.get(job.id) is None
|
||||
assert sorted(_statuses(final, job.id)) == ["abandoned", "succeeded"]
|
||||
|
||||
|
||||
def test_crashed_recurring_job_logs_abandoned_and_waits_for_next_slot(make_store):
|
||||
store = make_store()
|
||||
job = ScheduleJob(
|
||||
name="hourly", message="digest",
|
||||
schedule=Schedule(kind="every", every_seconds=3600),
|
||||
)
|
||||
store.add(job)
|
||||
t0 = time.time()
|
||||
assert store.claim_due(t0, owner_id="A", lease_seconds=1)
|
||||
|
||||
restarted = make_store()
|
||||
assert restarted.claim_due(t0 + 5, owner_id="B") == []
|
||||
assert _statuses(restarted, job.id) == ["abandoned"]
|
||||
assert restarted.get(job.id) is not None
|
||||
# Next slot fires normally.
|
||||
assert len(restarted.claim_due(t0 + 3601, owner_id="B")) == 1
|
||||
|
||||
|
||||
def test_failed_one_shot_stays_listed_disabled(make_store):
|
||||
store = make_store()
|
||||
job = _one_shot()
|
||||
store.add(job)
|
||||
claimed = store.claim_due(time.time(), owner_id="A")
|
||||
ScheduleRunner(store).mark_run(claimed[0], status="failed", error="provider 503")
|
||||
store.complete(job.id, "A")
|
||||
after = make_store()
|
||||
kept = after.get(job.id)
|
||||
assert kept is not None and kept.enabled is False
|
||||
assert _statuses(after, job.id) == ["failed"]
|
||||
# Disabled: never re-claimed or silently retired.
|
||||
assert after.claim_due(time.time() + 10, owner_id="B") == []
|
||||
assert after.get(job.id) is not None
|
||||
|
||||
|
||||
def test_one_shot_in_flight_past_lease_not_refired_by_same_process(make_store):
|
||||
store = make_store()
|
||||
job = _one_shot()
|
||||
store.add(job)
|
||||
t0 = time.time()
|
||||
assert store.claim_due(t0, owner_id="A", lease_seconds=1)
|
||||
# Same ticker polls again while its run is still going.
|
||||
assert store.claim_due(t0 + 5, owner_id="A", lease_seconds=1) == []
|
||||
|
||||
|
||||
def test_future_one_shot_survives_poll_before_its_time(make_store):
|
||||
store = make_store()
|
||||
future = (datetime.now(timezone.utc) + timedelta(hours=1)).isoformat()
|
||||
job = ScheduleJob(
|
||||
name="later", message="ping",
|
||||
schedule=Schedule(kind="at", at=future), delete_after_run=True,
|
||||
)
|
||||
store.add(job)
|
||||
assert store.claim_due(time.time(), owner_id="A") == []
|
||||
assert make_store().get(job.id) is not None
|
||||
|
||||
|
||||
def test_resumed_failed_one_shot_runs_again(make_store):
|
||||
store = make_store()
|
||||
job = _one_shot()
|
||||
store.add(job)
|
||||
claimed = store.claim_due(time.time(), owner_id="A")
|
||||
ScheduleRunner(store).mark_run(claimed[0], status="failed", error="boom")
|
||||
store.complete(job.id, "A")
|
||||
resumed = make_store()
|
||||
kept = resumed.get(job.id)
|
||||
kept.enabled = True
|
||||
resumed.update(kept)
|
||||
retry = resumed.claim_due(time.time() + 10, owner_id="B")
|
||||
assert [j.id for j in retry] == [job.id]
|
||||
ScheduleRunner(resumed).mark_run(retry[0])
|
||||
resumed.complete(job.id, "B")
|
||||
assert make_store().get(job.id) is None
|
||||
|
||||
|
||||
def test_crashed_spent_job_not_retried_past_until(make_store):
|
||||
store = make_store()
|
||||
t0 = time.time()
|
||||
until = datetime.fromtimestamp(t0 + 3, timezone.utc).isoformat()
|
||||
job = ScheduleJob(
|
||||
name="bounded", message="ping",
|
||||
schedule=Schedule(kind="every", every_seconds=3600),
|
||||
max_runs=1, until=until,
|
||||
)
|
||||
store.add(job)
|
||||
assert store.claim_due(t0, owner_id="A", lease_seconds=1)
|
||||
# Crash, then restart after the deadline: logged and retired, not re-run.
|
||||
restarted = make_store()
|
||||
assert restarted.claim_due(t0 + 5, owner_id="B") == []
|
||||
assert restarted.get(job.id) is None
|
||||
assert _statuses(make_store(), job.id) == ["abandoned"]
|
||||
@@ -241,9 +241,9 @@ class TestFileScheduleStoreEmits:
|
||||
|
||||
|
||||
class TestOneShotAutoDelete:
|
||||
"""A spent ``delete_after_run`` job is deleted inside ``claim_due``,
|
||||
bypassing ``remove()`` entirely — it must still announce itself, exactly
|
||||
once (the runner's later ``remove(job.id)`` finds nothing)."""
|
||||
"""A spent ``delete_after_run`` job is kept (leased) by ``claim_due`` and
|
||||
retired by the runner's ``mark_run`` once the run completes — announcing
|
||||
itself exactly once."""
|
||||
|
||||
def test_claim_due_auto_delete_fires_once(self, tmp_path):
|
||||
from praisonaiagents.scheduler.store import FileScheduleStore
|
||||
@@ -258,12 +258,15 @@ class TestOneShotAutoDelete:
|
||||
with _Recorder(HookEvent.SCHEDULE_REMOVE) as rec:
|
||||
claimed = store.claim_due(now=time.time() + 5, owner_id="w1")
|
||||
assert [j.id for j in claimed] == [job.id]
|
||||
# The runner's post-run cleanup must not double-announce it.
|
||||
assert store.remove(job.id) is False
|
||||
# Not deleted at claim: a crash mid-run must not lose it.
|
||||
assert rec.fired == []
|
||||
ScheduleRunner(store).mark_run(claimed[0])
|
||||
store.complete(job.id, "w1")
|
||||
# A later poll finds nothing left to retire.
|
||||
assert store.claim_due(now=time.time() + 10, owner_id="w1") == []
|
||||
|
||||
assert rec.fired == ["schedule_remove"]
|
||||
assert rec.payloads[0]["job_name"] == "one-shot"
|
||||
assert ScheduleRunner is not None # cleanup path lives here
|
||||
|
||||
def test_recurring_claim_does_not_fire(self, tmp_path):
|
||||
from praisonaiagents.scheduler.store import FileScheduleStore
|
||||
|
||||
@@ -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
|
||||
@@ -101,15 +101,19 @@ def test_delegate_task_timeout_exceeded(monkeypatch):
|
||||
|
||||
def factory(*, agent_factory=None, **kwargs):
|
||||
def spawn(task, agent_name=None, **kw):
|
||||
time.sleep(2)
|
||||
time.sleep(4)
|
||||
return {"success": True, "output": "too late", "agent_name": agent_name}
|
||||
return {"function": spawn}
|
||||
|
||||
monkeypatch.setattr(delegation_tools, "create_subagent_tool", factory)
|
||||
start = time.monotonic()
|
||||
raw = DelegationTools().delegate_task("Slow task", agent_type="research", timeout=1)
|
||||
elapsed = time.monotonic() - start
|
||||
data = json.loads(raw)
|
||||
assert data["success"] is False
|
||||
assert "timed out" in data["error"]
|
||||
# The caller must not wait for the abandoned worker to finish.
|
||||
assert elapsed < 3
|
||||
|
||||
|
||||
def test_delegate_task_requires_approval_registered():
|
||||
|
||||
@@ -5,11 +5,13 @@ import pytest
|
||||
from praisonaiagents.tools import github_tools
|
||||
|
||||
|
||||
def _fake_git_factory(responses, calls):
|
||||
def _fake_git_factory(responses, calls, cwds=None):
|
||||
"""Build a fake subprocess.run that records calls and serves canned output."""
|
||||
cwds = cwds if cwds is not None else []
|
||||
|
||||
def fake_run(cmd, check=False, capture_output=False, text=False):
|
||||
def fake_run(cmd, check=False, capture_output=False, text=False, cwd=None):
|
||||
calls.append(cmd)
|
||||
cwds.append(cwd)
|
||||
assert cmd[0] == "git"
|
||||
key = tuple(cmd[1:])
|
||||
# Match on a prefix of the git args
|
||||
@@ -256,3 +258,50 @@ def test_create_branch_accepts_slash_names(monkeypatch):
|
||||
checkouts = [c for c in calls if c[:2] == ["git", "checkout"]]
|
||||
assert checkouts, "expected a checkout for a valid branch name"
|
||||
assert ["git", "checkout", "-B", "feature/new-thing"] in checkouts
|
||||
|
||||
|
||||
def test_repo_path_used_as_cwd_for_every_git_call(monkeypatch):
|
||||
responses = {
|
||||
("status", "--porcelain"): ("M file\n", 0),
|
||||
("rev-parse", "--abbrev-ref", "HEAD"): ("praisonai/work\n", 0),
|
||||
("symbolic-ref", "refs/remotes/origin/HEAD"): ("refs/remotes/origin/main\n", 0),
|
||||
}
|
||||
calls, cwds = [], []
|
||||
monkeypatch.setattr(subprocess, "run", _fake_git_factory(responses, calls, cwds))
|
||||
result = _run_tool(github_tools.github_commit_and_push, "fix", repo_path="/repo")
|
||||
assert "Successfully" in result
|
||||
assert calls and all(c == "/repo" for c in cwds)
|
||||
|
||||
|
||||
def test_never_stages_untracked_with_add_dot(monkeypatch):
|
||||
responses = {
|
||||
("status", "--porcelain"): ("M file\n", 0),
|
||||
("rev-parse", "--abbrev-ref", "HEAD"): ("praisonai/work\n", 0),
|
||||
("symbolic-ref", "refs/remotes/origin/HEAD"): ("refs/remotes/origin/main\n", 0),
|
||||
}
|
||||
calls = _call(monkeypatch, responses)
|
||||
_run_tool(github_tools.github_commit_and_push, "fix", paths=["new.py"])
|
||||
adds = [c for c in calls if c[:2] == ["git", "add"]]
|
||||
# scoped to cwd so a subdirectory call cannot stage sibling edits
|
||||
assert ["git", "add", "-u", "--", "."] in adds
|
||||
assert ["git", "add", "-u"] not in adds
|
||||
assert ["git", "add", "--", "new.py"] in adds
|
||||
assert ["git", "add", "."] not in adds
|
||||
# untracked files are not considered when checking for changes
|
||||
status = [c for c in calls if c[:2] == ["git", "status"]]
|
||||
assert all("--untracked-files=no" in c for c in status)
|
||||
|
||||
|
||||
def test_refusal_leaves_index_untouched(monkeypatch):
|
||||
responses = {
|
||||
("status", "--porcelain"): ("M file\n", 0),
|
||||
("rev-parse", "--abbrev-ref", "HEAD"): ("release/v2\n", 0),
|
||||
("symbolic-ref", "refs/remotes/origin/HEAD"): ("refs/remotes/origin/main\n", 0),
|
||||
("config", "user.email"): ("[email protected]\n", 0),
|
||||
("merge-base", "origin/main", "HEAD"): ("abc\n", 0),
|
||||
("log", "--format=%ae", "abc..HEAD"): ("[email protected]\n", 0),
|
||||
}
|
||||
calls = _call(monkeypatch, responses)
|
||||
result = _run_tool(github_tools.github_commit_and_push, "fix", branch="release/v2")
|
||||
assert "Error" in result
|
||||
assert not [c for c in calls if c[:2] == ["git", "add"]]
|
||||
@@ -0,0 +1,178 @@
|
||||
"""spawn_subagent(background=True, deliver=...) inside a gateway turn (#6037).
|
||||
|
||||
The model cannot supply the origin chat, so it must come from the turn's
|
||||
session context, and the result must reach the turn's outbound messenger.
|
||||
"""
|
||||
|
||||
import time
|
||||
|
||||
from praisonaiagents.background.job_manager import get_job_manager
|
||||
from praisonaiagents.gateway.protocols import DeliveryResult
|
||||
from praisonaiagents.session.context import (
|
||||
clear_outbound_messenger,
|
||||
clear_session_context,
|
||||
register_outbound_messenger,
|
||||
set_session_context,
|
||||
)
|
||||
from praisonaiagents.tools.subagent_tool import create_subagent_tool
|
||||
|
||||
|
||||
class _StubAgent:
|
||||
def chat(self, prompt):
|
||||
return "long job done"
|
||||
|
||||
|
||||
class _FakeMessenger:
|
||||
def __init__(self, ok=True):
|
||||
self.ok = ok
|
||||
self.sent = []
|
||||
|
||||
async def send(self, target, text, *, media=None):
|
||||
self.sent.append((target, text))
|
||||
return DeliveryResult(ok=self.ok, target=target)
|
||||
|
||||
def list_targets(self):
|
||||
return []
|
||||
|
||||
|
||||
def _spawn_in_turn(messenger, deliver="origin", **kwargs):
|
||||
tool = create_subagent_tool(agent_factory=lambda **kw: _StubAgent())
|
||||
ctx = set_session_context(
|
||||
platform="telegram", chat_id="100", thread_id="7", user_id="alice"
|
||||
)
|
||||
msg = register_outbound_messenger(messenger)
|
||||
try:
|
||||
return tool["function"](
|
||||
task="summarise", background=True, deliver=deliver, **kwargs
|
||||
)
|
||||
finally:
|
||||
clear_outbound_messenger(msg)
|
||||
clear_session_context(ctx)
|
||||
|
||||
|
||||
def _wait_info(job_id):
|
||||
jm = get_job_manager()
|
||||
jm.get_result(job_id, timeout=5)
|
||||
for _ in range(50):
|
||||
info = jm.get_job_info(job_id)
|
||||
if info.delivered:
|
||||
break
|
||||
time.sleep(0.02)
|
||||
return info
|
||||
|
||||
|
||||
def test_origin_filled_from_session_context_and_delivered_once():
|
||||
messenger = _FakeMessenger()
|
||||
handle = _spawn_in_turn(messenger)
|
||||
info = _wait_info(handle["job_id"])
|
||||
|
||||
assert info.origin["platform"] == "telegram"
|
||||
assert info.origin["chat_id"] == "100"
|
||||
assert info.origin["thread_id"] == "7"
|
||||
assert info.origin["user_id"] == "alice"
|
||||
assert info.delivered is True
|
||||
assert len(messenger.sent) == 1
|
||||
target, text = messenger.sent[0]
|
||||
assert target == "origin"
|
||||
assert "long job done" in text
|
||||
|
||||
|
||||
def test_explicit_origin_arguments_win():
|
||||
messenger = _FakeMessenger()
|
||||
handle = _spawn_in_turn(messenger, platform="slack", chat_id="C9")
|
||||
info = _wait_info(handle["job_id"])
|
||||
assert info.origin["platform"] == "slack"
|
||||
assert info.origin["chat_id"] == "C9"
|
||||
assert info.origin["thread_id"] == ""
|
||||
assert info.origin["user_id"] == "alice"
|
||||
# A different chat is not the turn's chat: never send it to chat "100".
|
||||
assert messenger.sent == []
|
||||
assert info.delivered is False
|
||||
|
||||
|
||||
def test_explicit_other_thread_is_not_auto_sent():
|
||||
messenger = _FakeMessenger()
|
||||
handle = _spawn_in_turn(messenger, thread_id="99")
|
||||
info = _wait_info(handle["job_id"])
|
||||
assert info.origin["thread_id"] == "99"
|
||||
assert messenger.sent == []
|
||||
assert info.delivered is False
|
||||
|
||||
|
||||
def test_non_origin_target_is_not_auto_sent():
|
||||
messenger = _FakeMessenger()
|
||||
handle = _spawn_in_turn(messenger, deliver="telegram:999")
|
||||
info = _wait_info(handle["job_id"])
|
||||
assert messenger.sent == []
|
||||
assert info.delivered is False
|
||||
|
||||
|
||||
def test_send_policy_denial_blocks_auto_send():
|
||||
from praisonaiagents.gateway.protocols import SendDecision
|
||||
from praisonaiagents.session.context import (
|
||||
clear_send_policy,
|
||||
register_send_policy,
|
||||
)
|
||||
|
||||
class _Deny:
|
||||
def evaluate(self, target, **kw):
|
||||
return SendDecision(allow=False, reason="no")
|
||||
|
||||
messenger = _FakeMessenger()
|
||||
token = register_send_policy(_Deny())
|
||||
try:
|
||||
handle = _spawn_in_turn(messenger)
|
||||
finally:
|
||||
clear_send_policy(token)
|
||||
info = _wait_info(handle["job_id"])
|
||||
assert messenger.sent == []
|
||||
assert info.delivered is False
|
||||
|
||||
|
||||
def test_rejected_send_is_not_marked_delivered():
|
||||
messenger = _FakeMessenger(ok=False)
|
||||
handle = _spawn_in_turn(messenger)
|
||||
jm = get_job_manager()
|
||||
jm.get_result(handle["job_id"], timeout=5)
|
||||
time.sleep(0.1)
|
||||
assert len(messenger.sent) == 1
|
||||
assert jm.get_job_info(handle["job_id"]).delivered is False
|
||||
|
||||
|
||||
def test_without_gateway_stays_pull_only():
|
||||
tool = create_subagent_tool(agent_factory=lambda **kw: _StubAgent())
|
||||
handle = tool["function"](task="t", background=True, deliver="origin")
|
||||
jm = get_job_manager()
|
||||
jm.get_result(handle["job_id"], timeout=5)
|
||||
time.sleep(0.05)
|
||||
info = jm.get_job_info(handle["job_id"])
|
||||
assert info.delivered is False
|
||||
assert info.origin["chat_id"] == ""
|
||||
|
||||
|
||||
def test_restating_own_chat_keeps_turn_thread_and_platform():
|
||||
messenger = _FakeMessenger()
|
||||
handle = _spawn_in_turn(messenger, chat_id="100")
|
||||
info = _wait_info(handle["job_id"])
|
||||
assert info.origin["platform"] == "telegram"
|
||||
assert info.origin["thread_id"] == "7"
|
||||
assert len(messenger.sent) == 1
|
||||
|
||||
|
||||
def test_empty_output_still_notifies_before_marking_delivered():
|
||||
class _Silent:
|
||||
def chat(self, prompt):
|
||||
return ""
|
||||
|
||||
tool = create_subagent_tool(agent_factory=lambda **kw: _Silent())
|
||||
messenger = _FakeMessenger()
|
||||
ctx = set_session_context(platform="telegram", chat_id="100", user_id="alice")
|
||||
msg = register_outbound_messenger(messenger)
|
||||
try:
|
||||
handle = tool["function"](task="t", background=True, deliver="origin")
|
||||
finally:
|
||||
clear_outbound_messenger(msg)
|
||||
clear_session_context(ctx)
|
||||
info = _wait_info(handle["job_id"])
|
||||
assert info.delivered is True
|
||||
assert len(messenger.sent) == 1
|
||||
@@ -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,260 @@
|
||||
"""Unnamed agents in traces and graphs, agent_end status, and async model-call events.
|
||||
|
||||
Deterministic: the model call is replaced, so nothing reaches a provider.
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
from types import SimpleNamespace
|
||||
from unittest.mock import patch
|
||||
|
||||
import pytest
|
||||
|
||||
from praisonaiagents.trace.context_events import (
|
||||
ContextListSink,
|
||||
ContextTraceEmitter,
|
||||
reset_context_emitter,
|
||||
set_context_emitter,
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def sink(monkeypatch):
|
||||
monkeypatch.setenv("OPENAI_API_KEY", "sk-test-not-used")
|
||||
s = ContextListSink()
|
||||
token = set_context_emitter(ContextTraceEmitter(sink=s, session_id="t", enabled=True))
|
||||
yield s
|
||||
reset_context_emitter(token)
|
||||
|
||||
|
||||
def _of(sink, kind):
|
||||
return [e for e in sink.get_events() if e.event_type.value == kind]
|
||||
|
||||
|
||||
def _fake_completion(text="Solar sails ride sunlight."):
|
||||
message = SimpleNamespace(content=text, tool_calls=None, role="assistant", reasoning_content=None)
|
||||
return SimpleNamespace(
|
||||
choices=[SimpleNamespace(message=message, finish_reason="stop", index=0)],
|
||||
usage=SimpleNamespace(prompt_tokens=11, completion_tokens=7, total_tokens=18),
|
||||
model="gpt-4o-mini",
|
||||
)
|
||||
|
||||
|
||||
def search(query: str) -> str:
|
||||
"""Search the web."""
|
||||
return "result"
|
||||
|
||||
|
||||
# --- to_graph and ids ---------------------------------------------------------
|
||||
|
||||
def test_to_graph_accepts_an_unnamed_agent(monkeypatch):
|
||||
monkeypatch.setenv("OPENAI_API_KEY", "sk-test-not-used")
|
||||
from praisonaiagents import Agent
|
||||
from praisonaiagents.workflows.diagram import to_graph
|
||||
|
||||
graph = to_graph(Agent(instructions="Answer briefly.", tools=[search]))
|
||||
ids = [n["id"] for n in graph["nodes"]]
|
||||
assert ids[:3] == ["start:begin", "agent:Agent", "tool:Agent/search"] or set(ids) >= {
|
||||
"start:begin", "agent:Agent", "tool:Agent/search", "end:done"}
|
||||
assert next(n for n in graph["nodes"] if n["id"] == "agent:Agent")["label"] == "Agent"
|
||||
|
||||
|
||||
def test_named_agent_graph_is_unchanged(monkeypatch):
|
||||
monkeypatch.setenv("OPENAI_API_KEY", "sk-test-not-used")
|
||||
from praisonaiagents import Agent
|
||||
from praisonaiagents.workflows.diagram import to_graph
|
||||
|
||||
ids = [n["id"] for n in to_graph(Agent(name="Writer", instructions="x", tools=[search]))["nodes"]]
|
||||
assert "agent:Writer" in ids and "tool:Writer/search" in ids
|
||||
|
||||
|
||||
def test_emitter_ids_for_an_unnamed_agent():
|
||||
s = ContextListSink()
|
||||
em = ContextTraceEmitter(sink=s, session_id="t", enabled=True)
|
||||
em.agent_start(None)
|
||||
em.tool_call_start(None, "search", {"query": "x"})
|
||||
em.tool_call_end(None, "search", "ok", 1.0)
|
||||
em.agent_end(None)
|
||||
assert [e.node_id for e in s.get_events()] == [
|
||||
"agent:Agent", "tool:Agent/search", "tool:Agent/search", "agent:Agent"]
|
||||
|
||||
|
||||
def test_unnamed_agent_run_matches_its_graph(sink):
|
||||
from praisonaiagents import Agent
|
||||
from praisonaiagents.workflows.diagram import to_graph
|
||||
|
||||
agent = Agent(instructions="Answer briefly.")
|
||||
with patch.object(agent, "_chat_completion", return_value=_fake_completion()):
|
||||
agent.chat("Why do solar sails work?")
|
||||
graph_ids = {n["id"] for n in to_graph(agent)["nodes"]}
|
||||
for kind in ("agent_start", "agent_end"):
|
||||
(event,) = _of(sink, kind)
|
||||
assert event.agent_name is None # the name is not invented on the event
|
||||
assert event.node_id == "agent:Agent" and event.node_id in graph_ids
|
||||
|
||||
|
||||
# --- agent_end status ---------------------------------------------------------
|
||||
|
||||
def test_agent_end_records_completed(sink):
|
||||
from praisonaiagents import Agent
|
||||
|
||||
agent = Agent(name="Writer", instructions="x")
|
||||
with patch.object(agent, "_chat_completion", return_value=_fake_completion()):
|
||||
assert agent.chat("hi")
|
||||
assert _of(sink, "agent_end")[0].data == {"status": "completed"}
|
||||
|
||||
|
||||
def test_agent_end_records_failed(sink):
|
||||
from praisonaiagents import Agent
|
||||
|
||||
agent = Agent(name="Writer", instructions="x")
|
||||
with patch.object(agent, "_chat_completion", side_effect=RuntimeError("provider down")):
|
||||
try:
|
||||
result = agent.chat("hi")
|
||||
except Exception:
|
||||
result = None
|
||||
assert result is None
|
||||
assert _of(sink, "agent_end")[0].data == {"status": "failed"}
|
||||
|
||||
|
||||
# --- async model calls --------------------------------------------------------
|
||||
|
||||
def test_astart_records_the_model_call(sink):
|
||||
from praisonaiagents import Agent
|
||||
|
||||
agent = Agent(name="Writer", instructions="x", llm="gpt-4o-mini")
|
||||
|
||||
async def fake_dispatch(*args, **kwargs):
|
||||
return _fake_completion()
|
||||
|
||||
with patch.object(agent, "_achat_completion_with_retry_dispatch", side_effect=fake_dispatch):
|
||||
answer = asyncio.run(agent.astart("Why do solar sails work?"))
|
||||
assert "sunlight" in str(answer)
|
||||
(req,) = _of(sink, "llm_request")
|
||||
(resp,) = _of(sink, "llm_response")
|
||||
assert req.node_id == resp.node_id == "agent:Writer"
|
||||
assert resp.prompt_tokens == 11 and resp.completion_tokens == 7 and resp.cost_usd > 0
|
||||
assert "sunlight" in (resp.data.get("response_content") or "")
|
||||
assert _of(sink, "agent_end")[0].data == {"status": "completed"}
|
||||
|
||||
|
||||
def test_async_model_error_is_recorded(sink):
|
||||
from praisonaiagents import Agent
|
||||
|
||||
agent = Agent(name="Writer", instructions="x", llm="gpt-4o-mini")
|
||||
|
||||
async def boom(*args, **kwargs):
|
||||
raise RuntimeError("rate limited")
|
||||
|
||||
with patch.object(agent, "_achat_completion_with_retry_dispatch", side_effect=boom):
|
||||
try:
|
||||
asyncio.run(agent.astart("hi"))
|
||||
except Exception:
|
||||
pass
|
||||
responses = _of(sink, "llm_response")
|
||||
assert responses and responses[0].data.get("finish_reason") == "error"
|
||||
assert "rate limited" in (responses[0].data.get("response_content") or "")
|
||||
|
||||
|
||||
def test_async_path_records_nothing_when_tracing_is_off(monkeypatch):
|
||||
monkeypatch.setenv("OPENAI_API_KEY", "sk-test-not-used")
|
||||
from praisonaiagents import Agent
|
||||
|
||||
s = ContextListSink()
|
||||
token = set_context_emitter(ContextTraceEmitter(sink=s, session_id="t", enabled=False))
|
||||
try:
|
||||
agent = Agent(name="Writer", instructions="x", llm="gpt-4o-mini")
|
||||
|
||||
async def fake_dispatch(*args, **kwargs):
|
||||
return _fake_completion()
|
||||
|
||||
with patch.object(agent, "_achat_completion_with_retry_dispatch", side_effect=fake_dispatch) as d:
|
||||
asyncio.run(agent.astart("hi"))
|
||||
assert d.called and s.get_events() == []
|
||||
finally:
|
||||
reset_context_emitter(token)
|
||||
|
||||
|
||||
# --- streamed turns (start(stream=True)) ---------------------------------------
|
||||
|
||||
def _fake_stream(self, prompt, **kwargs):
|
||||
yield "Solar "
|
||||
yield "sails."
|
||||
|
||||
|
||||
def test_streamed_turn_records_agent_start_and_end(sink):
|
||||
from praisonaiagents import Agent
|
||||
|
||||
agent = Agent(instructions="x")
|
||||
with patch.object(type(agent), "_start_stream_turn", _fake_stream):
|
||||
text = "".join(agent._start_stream("hi"))
|
||||
assert text == "Solar sails."
|
||||
kinds = [(e.event_type.value, e.node_id, e.data.get("status")) for e in sink.get_events()]
|
||||
assert kinds == [("agent_start", "agent:Agent", None), ("agent_end", "agent:Agent", "completed")]
|
||||
|
||||
|
||||
def test_streamed_turn_stopped_early_is_cancelled(sink):
|
||||
from praisonaiagents import Agent
|
||||
|
||||
agent = Agent(name="Writer", instructions="x")
|
||||
with patch.object(type(agent), "_start_stream_turn", _fake_stream):
|
||||
gen = agent._start_stream("hi")
|
||||
next(gen)
|
||||
gen.close() # the page stopped reading
|
||||
assert _of(sink, "agent_end")[0].data == {"status": "cancelled"}
|
||||
|
||||
|
||||
def test_streamed_turn_that_raises_is_failed(sink):
|
||||
from praisonaiagents import Agent
|
||||
|
||||
def boom(self, prompt, **kwargs):
|
||||
yield "x"
|
||||
raise RuntimeError("provider down")
|
||||
|
||||
agent = Agent(name="Writer", instructions="x")
|
||||
with patch.object(type(agent), "_start_stream_turn", boom):
|
||||
with pytest.raises(RuntimeError):
|
||||
list(agent._start_stream("hi"))
|
||||
assert _of(sink, "agent_end")[0].data == {"status": "failed"}
|
||||
|
||||
|
||||
def test_streamed_turn_records_nothing_when_tracing_is_off(monkeypatch):
|
||||
monkeypatch.setenv("OPENAI_API_KEY", "sk-test-not-used")
|
||||
from praisonaiagents import Agent
|
||||
|
||||
s = ContextListSink()
|
||||
token = set_context_emitter(ContextTraceEmitter(sink=s, session_id="t", enabled=False))
|
||||
try:
|
||||
agent = Agent(name="Writer", instructions="x")
|
||||
with patch.object(type(agent), "_start_stream_turn", _fake_stream):
|
||||
assert "".join(agent._start_stream("hi")) == "Solar sails."
|
||||
assert s.get_events() == []
|
||||
finally:
|
||||
reset_context_emitter(token)
|
||||
|
||||
|
||||
def test_streamed_turn_blocked_by_guardrail_is_failed(sink):
|
||||
from praisonaiagents import Agent
|
||||
|
||||
agent = Agent(name="Writer", instructions="x")
|
||||
with patch.object(agent, "_validate_input_with_guardrail", return_value=(False, None, "no")):
|
||||
assert "blocked" in "".join(agent._start_stream("hi"))
|
||||
assert _of(sink, "agent_end")[0].data == {"status": "failed"}
|
||||
|
||||
|
||||
def test_streamed_turn_fallback_records_the_turn_once(sink):
|
||||
from praisonaiagents import Agent
|
||||
|
||||
def falls_back(self, prompt, **kwargs):
|
||||
yield self._stream_fallback_chat(prompt, kwargs, RuntimeError("stream broke"))
|
||||
|
||||
agent = Agent(name="Writer", instructions="x")
|
||||
with patch.object(type(agent), "_start_stream_impl", falls_back), \
|
||||
patch.object(agent, "_chat_completion", return_value=_fake_completion()):
|
||||
assert "".join(agent._start_stream("hi")) == "Solar sails ride sunlight."
|
||||
kinds = [e.event_type.value for e in sink.get_events() if e.event_type.value.startswith("agent_")]
|
||||
assert kinds == ["agent_start", "agent_end"]
|
||||
assert _of(sink, "agent_end")[0].data == {"status": "completed"}
|
||||
# chat() outside a stream still records its own turn.
|
||||
with patch.object(agent, "_chat_completion", return_value=_fake_completion()):
|
||||
agent.chat("again")
|
||||
assert len(_of(sink, "agent_start")) == 2
|
||||
@@ -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"
|
||||
}
|
||||
Loaded 100 of 204 files, more files were not shown because too many files have changed in this diff.
Show more
Reference in new issue
Block a user