mirror of
https://github.com/MervinPraison/PraisonAI.git
synced 2026-11-07 17:33:33 -10:00
Compare commits
232
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
|
||||
|
||||
@@ -94,3 +94,17 @@ runs:
|
||||
cd "$WRAPPER_PKG"
|
||||
pip install -e "${WRAPPER_SUFFIX}"
|
||||
fi
|
||||
|
||||
# crawl4ai depends on unclecode-litellm, a fork that writes into the same
|
||||
# litellm/ directory. Installed together, the two overwrite each other's
|
||||
# files in whatever order they land, so imports of litellm fail at random.
|
||||
# Reinstall the resolved litellm last so its own files win.
|
||||
if python -m pip show unclecode-litellm >/dev/null 2>&1; then
|
||||
LITELLM_VERSION="$(python -m pip show litellm | awk '/^Version:/{print $2}')"
|
||||
if [ "${{ inputs.use-uv }}" = "true" ]; then
|
||||
uv pip install --system --reinstall --no-deps "litellm==${LITELLM_VERSION}"
|
||||
else
|
||||
pip install --force-reinstall --no-deps "litellm==${LITELLM_VERSION}"
|
||||
fi
|
||||
python -c "import litellm"
|
||||
fi
|
||||
@@ -178,13 +178,84 @@ assert('no source issue', o.sourceIssueNumbers({ head: { ref: 'feat/x' }, body:
|
||||
await o.touchedByRun(runGh([{ event: 'committed', committer: { date: '2026-10-17T10:07:00Z' } }]), 'o', 'r', item(), SINCE, APP));
|
||||
assert('App editing the body during the run is touched (GraphQL login has no [bot])',
|
||||
await o.touchedByRun(runGh([], [{ editedAt: '2026-10-17T10:08:00Z', editor: { login: 'praisonai-triage-agent' } }]), 'o', 'r', item(), SINCE, APP));
|
||||
assert('routine label sync by the App during the run is not touched (shared App identity)',
|
||||
!(await o.touchedByRun(runGh([{ event: 'unlabeled', actor: { login: 'praisonai-triage-agent[bot]' }, label: { name: 'pipeline/blocked:ci' }, created_at: '2026-10-17T10:06:00Z' }]), 'o', 'r', item(), SINCE, APP)));
|
||||
assert('the App adding the claude label during the run is touched (it starts a Claude run)',
|
||||
await o.touchedByRun(runGh([{ event: 'labeled', actor: { login: 'praisonai-triage-agent[bot]' }, label: { name: 'claude' }, created_at: '2026-10-17T10:06:00Z' }]), 'o', 'r', item(), SINCE, APP));
|
||||
const appLabel = (event, name) => runGh([{ event, actor: { login: 'praisonai-triage-agent[bot]' }, label: { name }, created_at: '2026-10-17T10:06:00Z' }]);
|
||||
assert('the App removing no-auto-merge during the run is touched (it loosens the gate)',
|
||||
await o.touchedByRun(appLabel('unlabeled', 'no-auto-merge'), 'o', 'r', item(), SINCE, APP));
|
||||
assert('the App removing a manual-only label during the run is touched',
|
||||
await o.touchedByRun(appLabel('unlabeled', 'needs-manual-review'), 'o', 'r', item(), SINCE, APP));
|
||||
assert('the App adding a maintainer-accept override during the run is touched',
|
||||
await o.touchedByRun(appLabel('labeled', 'maintainer-accept-core-tools'), 'o', 'r', item(), SINCE, APP));
|
||||
assert('the App adding no-auto-merge (tightening) is not touched',
|
||||
!(await o.touchedByRun(appLabel('labeled', 'no-auto-merge'), 'o', 'r', item(), SINCE, APP)));
|
||||
assert('the gate releasing its own lock label is not touched',
|
||||
!(await o.touchedByRun(appLabel('unlabeled', 'claude-merge-gate-active'), 'o', 'r', item(), SINCE, APP)));
|
||||
assert('the App force-pushing during the run is touched',
|
||||
await o.touchedByRun(runGh([{ event: 'head_ref_force_pushed', actor: { login: 'praisonai-triage-agent[bot]' }, created_at: '2026-10-17T10:06:00Z' }]), 'o', 'r', item(), SINCE, APP));
|
||||
assert('a maintainer comment during the run is not touched',
|
||||
!(await o.touchedByRun(runGh([{ event: 'commented', user: { login: 'MervinPraison' }, created_at: '2026-10-17T10:06:00Z' }]), 'o', 'r', item(), SINCE, APP)));
|
||||
assert('App activity after the run ended is not touched',
|
||||
!(await o.touchedByRun(runGh([{ event: 'commented', user: { login: 'praisonai-triage-agent[bot]' }, created_at: '2026-10-20T10:00:00Z' }]), 'o', 'r', item(), SINCE, APP, '2026-10-17T10:30:00Z')));
|
||||
assert('App activity before the run is not touched',
|
||||
!(await o.touchedByRun(runGh([{ event: 'labeled', actor: { login: 'praisonai-triage-agent[bot]' }, created_at: '2026-10-17T09:00:00Z' }]), 'o', 'r', item(), SINCE, APP)));
|
||||
const editFail = await (async () => { const g = runGh([]); g.graphql = async () => { throw new Error('rate'); }; return o.touchedByRun(g, 'o', 'r', item(), SINCE, APP).then(() => null, (e) => e); })();
|
||||
assert('an edit-history lookup failure fails closed (throws)', editFail instanceof Error);
|
||||
|
||||
{
|
||||
const jobs = [{ name: 'issue-triage', completed_at: '2026-10-17T10:00:01Z' }, { name: 'claude-response', completed_at: '2026-10-17T10:20:00Z' }];
|
||||
const g = { rest: { actions: { listJobsForWorkflowRun: 'jobs' } }, paginate: async () => jobs };
|
||||
assert('the Claude job end time bounds the run window', (await o.jobCompletedAt(g, 'o', 'r', 1, 'claude-response')) === '2026-10-17T10:20:00Z');
|
||||
const hub = { rest: g.rest, paginate: async () => [{ name: 'claude / claude-response', completed_at: '2026-10-17T10:21:00Z' }] };
|
||||
assert('a reusable workflow job (caller / job) is found', (await o.jobCompletedAt(hub, 'o', 'r', 1, 'claude-response')) === '2026-10-17T10:21:00Z');
|
||||
assert('an unfinished or missing job gives no bound', (await o.jobCompletedAt(g, 'o', 'r', 1, 'nope')) === null);
|
||||
const bad = { rest: { actions: { listJobsForWorkflowRun: 'jobs' } }, paginate: async () => { throw new Error('403'); } };
|
||||
assert('a jobs lookup failure gives no bound (caller counts up to now)', (await o.jobCompletedAt(bad, 'o', 'r', 1, 'claude-response')) === null);
|
||||
}
|
||||
|
||||
// A request mentioning another repo's numbers as bare #N (Docs issues cite PraisonAI PRs this way)
|
||||
// must not make the run outsider: a number that names nothing here carries no content.
|
||||
r = await o.referencedIssueReasons(fakeGithub({ issueError: 404 }), 'o', 'r', 'see upstream #5987 and #6366');
|
||||
assert('a referenced number that does not exist here (404) is not outsider influence', r.length === 0);
|
||||
{
|
||||
// The issue exists (writer-owned) but its comments call 404: that must not silently skip the comment check.
|
||||
const g = fakeGithub({ issues: { 12: maintainerIssue } });
|
||||
g.paginate = async (method) => { if (method === 'issues.listComments') throw Object.assign(new Error('nf'), { status: 404 }); return []; };
|
||||
r = await o.referencedIssueReasons(g, 'o', 'r', 'see #12');
|
||||
assert('a 404 after the issue was found still fails closed (only a missing issue is ignored)', r.some((x) => x.includes('#12 is unavailable')));
|
||||
}
|
||||
r = await o.referencedIssueReasons(fakeGithub({ issueError: 410 }), 'o', 'r', 'see #12');
|
||||
assert('a referenced issue that was deleted (410) still fails closed', r.some((x) => x.includes('unavailable (HTTP 410)')));
|
||||
r = await o.referencedIssueReasons(
|
||||
fakeGithub({ issues: { 500: maintainerIssue }, timeline: [{ event: 'labeled', label: { name: o.OUTSIDER_ORIGIN_LABEL } }] }), 'o', 'r', 'implement #500'
|
||||
);
|
||||
assert('a referenced issue flagged earlier stays outsider (outsider comment since deleted)', r.some((x) => x.includes('#500 was flagged')));
|
||||
r = await o.referencedIssueReasons(fakeGithub({ issues: { 500: { user: { login: 'stranger' } } } }), 'o', 'r', 'implement #500');
|
||||
assert('a referenced outsider issue still counts', r.some((x) => x.includes('referenced issue #500 opened by @stranger')));
|
||||
r = await o.outsiderIssueReasons(
|
||||
fakeGithub({ issues: { 600: maintainerIssue } }), 'o', 'r', 600, 'Docs for upstream #5987 (PraisonAI)'
|
||||
);
|
||||
assert('maintainer issue citing missing numbers is not outsider', r.length === 0);
|
||||
r = await o.outsiderOriginReasons(fakeGithub({ issueError: 404 }), 'o', 'r', internalPr(), []);
|
||||
assert('a missing SOURCE issue (404) still fails closed', r.some((x) => x.includes('unavailable (HTTP 404)')));
|
||||
|
||||
// A person with write access can clear a wrong flag; bots cannot.
|
||||
const lab = (event, login, t) => ({ event, label: { name: o.OUTSIDER_ORIGIN_LABEL }, actor: { login }, created_at: t });
|
||||
r = await o.outsiderOriginReasons(fakeGithub({ issue: maintainerIssue, timeline: [
|
||||
lab('labeled', 'praisonai-triage-agent[bot]', '2026-10-19T09:30:00Z'), lab('unlabeled', 'MervinPraison', '2026-10-26T10:00:00Z')] }), 'o', 'r', internalPr(), []);
|
||||
assert('a writer removing the flag clears it', !r.some((x) => x.includes('flagged')));
|
||||
r = await o.outsiderOriginReasons(fakeGithub({ issue: maintainerIssue, timeline: [
|
||||
lab('labeled', 'praisonai-triage-agent[bot]', '2026-10-19T09:30:00Z'), lab('unlabeled', 'praisonai-triage-agent[bot]', '2026-10-26T10:00:00Z')] }), 'o', 'r', internalPr(), []);
|
||||
assert('a bot removing the flag does not clear it', r.some((x) => x.includes('flagged')));
|
||||
r = await o.outsiderOriginReasons(fakeGithub({ issue: maintainerIssue, timeline: [
|
||||
lab('labeled', 'praisonai-triage-agent[bot]', '2026-10-19T09:30:00Z'), lab('unlabeled', 'stranger', '2026-10-26T10:00:00Z')] }), 'o', 'r', internalPr(), []);
|
||||
assert('a non-writer removing the flag does not clear it', r.some((x) => x.includes('flagged')));
|
||||
r = await o.outsiderOriginReasons(fakeGithub({ issue: maintainerIssue, timeline: [
|
||||
lab('labeled', 'x[bot]', '2026-10-19T09:30:00Z'), lab('unlabeled', 'MervinPraison', '2026-10-20T10:00:00Z'), lab('labeled', 'x[bot]', '2026-10-21T10:00:00Z')] }), 'o', 'r', internalPr(), []);
|
||||
assert('flagged again after a clear is flagged', r.some((x) => x.includes('flagged')));
|
||||
|
||||
r = await o.outsiderOriginReasons(fakeGithub({ issueError: 410 }), 'o', 'r', internalPr(), []);
|
||||
assert('a deleted or transferred source issue (410) counts as unknown, so outsider', r.some((x) => x.includes('unavailable (HTTP 410)')));
|
||||
|
||||
|
||||
@@ -26,6 +26,15 @@ const TRUSTED_AUTOMATION = new Set(['Copilot']);
|
||||
const SECURITY_STATUS_CONTEXT = 'security/outsider-review';
|
||||
const SECURITY_STATUS_CREATOR = 'github-actions[bot]';
|
||||
const OUTSIDER_ORIGIN_LABEL = 'origin/outsider';
|
||||
// Timeline events that change content Claude or the merge gate reads.
|
||||
const CONTENT_EVENTS = new Set(['commented', 'reviewed', 'line-commented', 'commit-commented', 'head_ref_force_pushed', 'renamed']);
|
||||
// Label changes that loosen what Claude or the merge gate does (see merge-gate.js):
|
||||
// adding a trigger or an override, or removing a stop or manual-only label.
|
||||
// Pipeline status labels and the gate's own lock are routine and do not count.
|
||||
const LOOSENING_LABELS = {
|
||||
labeled: new Set(['claude', 'merge-gate-ci-only', 'maintainer-accept-example', 'maintainer-accept-core-tools']),
|
||||
unlabeled: new Set(['no-auto-merge', 'claude-conflict-pending', 'auto-merged-by-gate', 'security', 'breaking-change', 'needs-manual-review', 'release']),
|
||||
};
|
||||
|
||||
// Build, dependency, CI and install-time surface. Code here runs on developer
|
||||
// machines, in CI with secrets, or at install time, so an outsider-origin
|
||||
@@ -164,12 +173,24 @@ async function outsiderOriginReasons(github, owner, repo, pr, comments = []) {
|
||||
return reasons;
|
||||
}
|
||||
|
||||
/**
|
||||
* Was this ever flagged? Timeline events cannot be edited, so deleting the
|
||||
* outsider content does not clear it. Only a person with write access can clear
|
||||
* a wrong flag, by removing the label after it was last added (bots cannot:
|
||||
* every Claude run acts as the App).
|
||||
*/
|
||||
async function wasLabelled(github, owner, repo, number, label) {
|
||||
try {
|
||||
const events = await github.paginate(github.rest.issues.listEventsForTimeline, {
|
||||
const events = (await github.paginate(github.rest.issues.listEventsForTimeline, {
|
||||
owner, repo, issue_number: number, per_page: 100,
|
||||
});
|
||||
return events.some((e) => e.event === 'labeled' && e.label?.name === label);
|
||||
})).filter((e) => (e.event === 'labeled' || e.event === 'unlabeled') && e.label?.name === label);
|
||||
const lastAdded = events.map((e) => e.event).lastIndexOf('labeled');
|
||||
if (lastAdded < 0) return false;
|
||||
const isWriter = makePermissionLookup(github, owner, repo);
|
||||
for (const e of events.slice(lastAdded + 1)) {
|
||||
if (e.event === 'unlabeled' && !isAutomation(e.actor?.login) && (await isWriter(e.actor.login))) return false;
|
||||
}
|
||||
return true;
|
||||
} catch (err) {
|
||||
if (err.status === 404 || err.status === 410) return false;
|
||||
throw err;
|
||||
@@ -195,10 +216,12 @@ async function maintainerApprovedHead(github, owner, repo, prNumber, headSha) {
|
||||
}
|
||||
|
||||
/** Outsider authorship of an issue and its comments ([] for a PR number). */
|
||||
async function issueReasons(github, owner, repo, issueNumber, isWriter, label = 'issue') {
|
||||
async function issueReasons(github, owner, repo, issueNumber, isWriter, label = 'issue', { missingOk = false } = {}) {
|
||||
const reasons = [];
|
||||
let found = false;
|
||||
try {
|
||||
const { data: issue } = await github.rest.issues.get({ owner, repo, issue_number: issueNumber });
|
||||
found = true;
|
||||
if (issue.pull_request) {
|
||||
// A PR as the source: its head and author matter too.
|
||||
const { data: src } = await github.rest.pulls.get({ owner, repo, pull_number: issueNumber });
|
||||
@@ -214,6 +237,11 @@ async function issueReasons(github, owner, repo, issueNumber, isWriter, label =
|
||||
if (outsiders.length) reasons.push(`${label} #${issueNumber} has comments from ${outsiders.map((l) => '@' + l).join(', ')}`);
|
||||
} catch (err) {
|
||||
if (err.status !== 404 && err.status !== 410) throw err;
|
||||
// A number merely mentioned that names nothing here (e.g. another repo's #N) has no
|
||||
// content to inject; a deleted issue (410) did, so it still fails closed.
|
||||
// Only the issue lookup itself may come back missing; a 404 from a later call
|
||||
// (comments, the PR behind it) would otherwise skip checks and fail open.
|
||||
if (missingOk && err.status === 404 && !found) return reasons;
|
||||
// Deleted, transferred or never existed: who wrote it is unknown, so fail closed.
|
||||
reasons.push(`${label} #${issueNumber} is unavailable (HTTP ${err.status}); provenance unknown`);
|
||||
}
|
||||
@@ -226,16 +254,20 @@ async function issueReasons(github, owner, repo, issueNumber, isWriter, label =
|
||||
* commits to it, or edited its text. Unrelated activity by people does not count,
|
||||
* so flagging what an outsider run touched does not sweep up other work.
|
||||
*/
|
||||
async function touchedByRun(github, owner, repo, item, since, appLogins) {
|
||||
async function touchedByRun(github, owner, repo, item, since, appLogins, until = new Date().toISOString()) {
|
||||
const t = new Date(since).getTime();
|
||||
const after = (d) => !!d && new Date(d).getTime() >= t;
|
||||
const end = new Date(until).getTime();
|
||||
const after = (d) => !!d && new Date(d).getTime() >= t && new Date(d).getTime() <= end;
|
||||
const isApp = (login) => !!login && (appLogins.has(login) || appLogins.has(`${login}[bot]`));
|
||||
if (isApp(item.user?.login) && after(item.created_at)) return true;
|
||||
const events = await github.paginate(github.rest.issues.listEventsForTimeline, {
|
||||
owner, repo, issue_number: item.number, per_page: 100,
|
||||
});
|
||||
for (const e of events) {
|
||||
if (isApp(e.actor?.login || e.user?.login) && after(e.created_at || e.submitted_at || e.updated_at)) return true;
|
||||
// Only actions that change what Claude or the gate reads. Label syncing by the
|
||||
// pipeline also runs as the App and must not count as the outsider run.
|
||||
const content = CONTENT_EVENTS.has(e.event) || !!LOOSENING_LABELS[e.event]?.has(e.label?.name);
|
||||
if (content && isApp(e.actor?.login || e.user?.login) && after(e.created_at || e.submitted_at || e.updated_at)) return true;
|
||||
if (e.event === 'committed' && after(e.committer?.date)) return true;
|
||||
}
|
||||
let data;
|
||||
@@ -258,6 +290,22 @@ async function touchedByRun(github, owner, repo, item, since, appLogins) {
|
||||
return edits.some((n) => after(n.editedAt) && isApp(n.editor?.login));
|
||||
}
|
||||
|
||||
/**
|
||||
* When a job of this workflow run finished (null if unknown). It bounds what counts
|
||||
* as that job's activity, so work by later runs is not attributed to it.
|
||||
*/
|
||||
async function jobCompletedAt(github, owner, repo, runId, jobName) {
|
||||
try {
|
||||
const jobs = await github.paginate(github.rest.actions.listJobsForWorkflowRun, {
|
||||
owner, repo, run_id: runId, per_page: 100,
|
||||
});
|
||||
// A reusable workflow's jobs are named '<caller job> / <job>'.
|
||||
return jobs.find((j) => j.name === jobName || j.name.endsWith(` / ${jobName}`))?.completed_at || null;
|
||||
} catch {
|
||||
return null;
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Why an issue Claude is asked to work on is outsider-influenced ([] when not):
|
||||
* the issue itself, anything it was flagged as before, and every issue the
|
||||
@@ -267,7 +315,7 @@ async function outsiderIssueReasons(github, owner, repo, issueNumber, requestTex
|
||||
const isWriter = makePermissionLookup(github, owner, repo);
|
||||
const reasons = await issueReasons(github, owner, repo, issueNumber, isWriter);
|
||||
for (const n of referencedIssueNumbers(requestText, owner, repo)) {
|
||||
if (n !== issueNumber) reasons.push(...(await issueReasons(github, owner, repo, n, isWriter, 'referenced issue')));
|
||||
if (n !== issueNumber) reasons.push(...(await issueReasons(github, owner, repo, n, isWriter, 'referenced issue', { missingOk: true })));
|
||||
}
|
||||
for (const n of [issueNumber, ...referencedIssueNumbers(requestText, owner, repo)]) {
|
||||
if (await wasLabelled(github, owner, repo, n, OUTSIDER_ORIGIN_LABEL)) {
|
||||
@@ -278,6 +326,20 @@ async function outsiderIssueReasons(github, owner, repo, issueNumber, requestTex
|
||||
return reasons;
|
||||
}
|
||||
|
||||
/** Outsider influence from issues a request points at ("implement #500"); missing numbers are ignored. */
|
||||
async function referencedIssueReasons(github, owner, repo, text) {
|
||||
const isWriter = makePermissionLookup(github, owner, repo);
|
||||
const reasons = [];
|
||||
for (const n of referencedIssueNumbers(text, owner, repo)) {
|
||||
reasons.push(...(await issueReasons(github, owner, repo, n, isWriter, 'referenced issue', { missingOk: true })));
|
||||
// Sticky like everywhere else: a deleted outsider comment must not clear the flag.
|
||||
if (await wasLabelled(github, owner, repo, n, OUTSIDER_ORIGIN_LABEL)) {
|
||||
reasons.push(`referenced issue #${n} was flagged ${OUTSIDER_ORIGIN_LABEL} earlier`);
|
||||
}
|
||||
}
|
||||
return reasons;
|
||||
}
|
||||
|
||||
/** Every #N, owner/repo#N (this repo) and issue/PR URL in this repo mentioned in text. */
|
||||
function referencedIssueNumbers(text, owner, repo) {
|
||||
const here = `${owner}/${repo}`.toLowerCase();
|
||||
@@ -418,7 +480,9 @@ module.exports = {
|
||||
outsiderOriginReasons,
|
||||
outsiderIssueReasons,
|
||||
referencedIssueNumbers,
|
||||
referencedIssueReasons,
|
||||
touchedByRun,
|
||||
jobCompletedAt,
|
||||
outsiderSensitivePaths,
|
||||
riskFlags,
|
||||
formatRiskFlags,
|
||||
|
||||
@@ -11,9 +11,52 @@ on:
|
||||
description: 'Fail if performance regresses'
|
||||
required: false
|
||||
default: 'true'
|
||||
pull_request:
|
||||
paths:
|
||||
- 'src/praisonai-agents/praisonaiagents/agent/**'
|
||||
- 'src/praisonai-agents/praisonaiagents/llm/**'
|
||||
- 'src/praisonai-agents/benchmarks/simple_benchmark.py'
|
||||
- 'src/praisonai-agents/benchmarks/thresholds.json'
|
||||
- '.github/workflows/benchmark.yml'
|
||||
|
||||
permissions:
|
||||
contents: read
|
||||
|
||||
jobs:
|
||||
instantiation-gate:
|
||||
# Fails the PR when Agent instantiation exceeds benchmarks/thresholds.json.
|
||||
runs-on: ubuntu-latest
|
||||
timeout-minutes: 15
|
||||
steps:
|
||||
- name: Checkout repository
|
||||
uses: actions/checkout@v4
|
||||
with:
|
||||
persist-credentials: false
|
||||
|
||||
- name: Set up Python 3.12
|
||||
uses: actions/setup-python@v5
|
||||
with:
|
||||
python-version: '3.12'
|
||||
cache: 'pip'
|
||||
|
||||
- name: Install praisonaiagents
|
||||
run: |
|
||||
python -m pip install --upgrade pip
|
||||
pip install -e src/praisonai-agents
|
||||
|
||||
- name: Check instantiation threshold
|
||||
run: |
|
||||
cd src/praisonai-agents
|
||||
python benchmarks/simple_benchmark.py --skip-crewai --json --fail-over-threshold > simple.json
|
||||
python - <<'PY' >> "$GITHUB_STEP_SUMMARY"
|
||||
import json
|
||||
data = json.load(open('simple.json'))
|
||||
budget = json.load(open('benchmarks/thresholds.json'))['praisonai_instantiation_us_max']
|
||||
print(f"PraisonAI instantiation: {data['results']['PraisonAI']:.2f} us (budget {budget} us)")
|
||||
PY
|
||||
|
||||
benchmark:
|
||||
if: github.event_name == 'workflow_dispatch'
|
||||
runs-on: ubuntu-latest
|
||||
timeout-minutes: 30
|
||||
|
||||
@@ -44,7 +87,8 @@ jobs:
|
||||
python benchmarks/import_time.py 2>&1 | tee import_time_output.txt
|
||||
|
||||
# Extract import time
|
||||
IMPORT_TIME=$(grep -oP 'Median:\s+\K[\d.]+' import_time_output.txt || echo "0")
|
||||
IMPORT_TIME=$(grep -oP 'Median:\s+\K[\d.]+' import_time_output.txt | head -1)
|
||||
[ -n "$IMPORT_TIME" ] || { echo "::error::Could not parse import time"; exit 1; }
|
||||
echo "import_time=$IMPORT_TIME" >> $GITHUB_OUTPUT
|
||||
|
||||
# Check threshold (200ms)
|
||||
@@ -66,7 +110,8 @@ jobs:
|
||||
python benchmarks/memory_usage.py 2>&1 | tee memory_output.txt
|
||||
|
||||
# Extract peak memory
|
||||
PEAK_MEMORY=$(grep -oP 'Peak:\s+\K[\d.]+' memory_output.txt || echo "0")
|
||||
PEAK_MEMORY=$(grep -oP 'Peak:\s+\K[\d.]+' memory_output.txt | head -1)
|
||||
[ -n "$PEAK_MEMORY" ] || { echo "::error::Could not parse peak memory"; exit 1; }
|
||||
echo "peak_memory=$PEAK_MEMORY" >> $GITHUB_OUTPUT
|
||||
|
||||
# Check threshold (30MB)
|
||||
@@ -83,10 +128,8 @@ jobs:
|
||||
id: simple
|
||||
run: |
|
||||
cd src/praisonai-agents
|
||||
python benchmarks/simple_benchmark.py 2>&1 | tee simple_output.txt
|
||||
|
||||
# Extract PraisonAI time
|
||||
PRAISON_TIME=$(grep -oP 'PraisonAI\s+\K[\d.]+' simple_output.txt | head -1 || echo "0")
|
||||
python benchmarks/simple_benchmark.py --json > simple.json
|
||||
PRAISON_TIME=$(python -c "import json; print(f\"{json.load(open('simple.json'))['results']['PraisonAI']:.2f}\")")
|
||||
echo "praison_time=$PRAISON_TIME" >> $GITHUB_OUTPUT
|
||||
|
||||
- name: Run Tools Benchmark
|
||||
@@ -96,7 +139,8 @@ jobs:
|
||||
python benchmarks/tools_benchmark.py 2>&1 | tee tools_output.txt
|
||||
|
||||
# Extract PraisonAI time with tools
|
||||
PRAISON_TOOLS_TIME=$(grep -oP 'PraisonAI\s+\K[\d.]+' tools_output.txt | head -1 || echo "0")
|
||||
PRAISON_TOOLS_TIME=$(grep -oP 'PraisonAI\s+\K[\d.]+' tools_output.txt | head -1)
|
||||
[ -n "$PRAISON_TOOLS_TIME" ] || { echo "::error::Could not parse tools benchmark"; exit 1; }
|
||||
echo "praison_tools_time=$PRAISON_TOOLS_TIME" >> $GITHUB_OUTPUT
|
||||
|
||||
- name: Run Parameter Impact Benchmark
|
||||
@@ -133,7 +177,7 @@ jobs:
|
||||
echo "|--------|-------|-----------|--------|" >> $GITHUB_STEP_SUMMARY
|
||||
echo "| Import Time | ${{ steps.import_time.outputs.import_time }}ms | <200ms | ${{ steps.import_time.outputs.import_passed == 'true' && '✅' || '❌' }} |" >> $GITHUB_STEP_SUMMARY
|
||||
echo "| Peak Memory | ${{ steps.memory.outputs.peak_memory }}MB | <30MB | ${{ steps.memory.outputs.memory_passed == 'true' && '✅' || '❌' }} |" >> $GITHUB_STEP_SUMMARY
|
||||
echo "| Agent Instantiation | ${{ steps.simple.outputs.praison_time }}μs | - | ℹ️ |" >> $GITHUB_STEP_SUMMARY
|
||||
echo "| Agent Instantiation | ${{ steps.simple.outputs.praison_time }}μs | see thresholds.json | ℹ️ |" >> $GITHUB_STEP_SUMMARY
|
||||
echo "| Agent with Tools | ${{ steps.tools.outputs.praison_tools_time }}μs | - | ℹ️ |" >> $GITHUB_STEP_SUMMARY
|
||||
echo "" >> $GITHUB_STEP_SUMMARY
|
||||
echo "### Lazy Import Check" >> $GITHUB_STEP_SUMMARY
|
||||
|
||||
@@ -353,9 +353,7 @@ jobs:
|
||||
const { data: pr } = await github.rest.pulls.get({ owner, repo, pull_number: number });
|
||||
const comments = await github.paginate(github.rest.issues.listComments, { owner, repo, issue_number: number, per_page: 100 });
|
||||
origin = await o.outsiderOriginReasons(github, owner, repo, pr, comments);
|
||||
for (const n of o.referencedIssueNumbers(request, owner, repo)) {
|
||||
origin.push(...(await o.outsiderIssueReasons(github, owner, repo, n)));
|
||||
}
|
||||
origin.push(...(await o.referencedIssueReasons(github, owner, repo, request)));
|
||||
} else {
|
||||
origin = await o.outsiderIssueReasons(github, owner, repo, number, request);
|
||||
}
|
||||
@@ -560,13 +558,17 @@ jobs:
|
||||
const o = require(path.join(process.env.GITHUB_WORKSPACE, '.github/scripts/outsider-review.js'));
|
||||
const { owner, repo } = context.repo;
|
||||
const since = process.env.STARTED_AT || new Date(0).toISOString();
|
||||
// Claude's job is over (its token revoked, its processes killed); later
|
||||
// activity belongs to other runs. Unknown: count everything up to now.
|
||||
const until = (await o.jobCompletedAt(github, owner, repo, context.runId, 'claude-response')) || new Date().toISOString();
|
||||
core.info(`Outsider run window: ${since} .. ${until}`);
|
||||
// Issues and PRs created, edited, commented on or pushed to since then.
|
||||
const touched = await github.paginate(github.rest.issues.listForRepo, { owner, repo, state: 'all', since, per_page: 100 });
|
||||
const app = new Set(['praisonai-triage-agent[bot]']);
|
||||
for (const item of touched) {
|
||||
let flag = true;
|
||||
try {
|
||||
flag = await o.touchedByRun(github, owner, repo, item, since, app);
|
||||
flag = await o.touchedByRun(github, owner, repo, item, since, app, until);
|
||||
} catch (err) {
|
||||
core.warning(`#${item.number}: ${err.message}; flagging it to be safe`);
|
||||
}
|
||||
|
||||
@@ -96,11 +96,11 @@ jobs:
|
||||
--head "$branch" \
|
||||
--title "$msg" \
|
||||
--body "$(cat <<EOF
|
||||
Automated sync from [\`MervinPraison/PraisonAI\`](https://github.com/MervinPraison/PraisonAI) \`${SOURCE_PATH}/\` into the npm mirror repo.
|
||||
Automated sync from [\`MervinPraison/PraisonAI\`](https://github.com/MervinPraison/PraisonAI) \`${SOURCE_PATH}/\` into the npm mirror repo.
|
||||
|
||||
Canonical development is in the monorepo; this PR updates the public \`praisonai-js\` checkout.
|
||||
Canonical development is in the monorepo; this PR updates the public \`praisonai-js\` checkout.
|
||||
|
||||
Excludes from rsync source: \`node_modules/\`, \`dist/\`, \`coverage/\`. \`.github/\` on praisonai-js is unchanged.
|
||||
EOF
|
||||
)"
|
||||
Excludes from rsync source: \`node_modules/\`, \`dist/\`, \`coverage/\`. \`.github/\` on praisonai-js is unchanged.
|
||||
EOF
|
||||
)"
|
||||
echo "Opened PR on ${TARGET_REPO} (${branch} → ${TARGET_BRANCH})."
|
||||
@@ -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: |
|
||||
|
||||
@@ -1002,11 +1002,13 @@ agent.start('Write a movie script about a robot in Mars');
|
||||
|
||||
## ⚡ Performance
|
||||
|
||||
PraisonAI is built for speed, with agent instantiation in around 14μs. This reduces overhead, improves responsiveness, and helps multi-agent systems scale efficiently in real-world production workloads.
|
||||
PraisonAI keeps framework overhead microsecond-class: creating an `Agent` (one tool, no LLM call) averaged **46 μs** in our 2026-10-15 run, so per-agent setup is negligible next to model latency in multi-agent workloads.
|
||||
|
||||
| Performance Metric | PraisonAI |
|
||||
|--------------------|-----------|
|
||||
| Avg Instantiation Time | **14 μs** |
|
||||
| Performance Metric | PraisonAI | Measured |
|
||||
|--------------------|-----------|----------|
|
||||
| Avg Agent instantiation (1 tool, 1000 iterations) | **46 μs** | 2026-10-15 · praisonaiagents 1.7.12 |
|
||||
|
||||
Methodology, comparisons and how to reproduce: [benchmarks/BENCHMARK_RESULTS.md](https://github.com/MervinPraison/PraisonAI/blob/main/src/praisonai-agents/benchmarks/BENCHMARK_RESULTS.md) · [docs](https://docs.praison.ai/docs/developers/benchmarks)
|
||||
|
||||
---
|
||||
|
||||
|
||||
@@ -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.16" \
|
||||
"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.16" \
|
||||
"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.16" \
|
||||
"praisonai>=4.7.21" \
|
||||
"praisonai[ui]" \
|
||||
"praisonai[crewai]"
|
||||
|
||||
|
||||
@@ -5,7 +5,6 @@ FastContext provides rapid code search capabilities.
|
||||
Use it directly from the context module for code search.
|
||||
|
||||
FastContext provides:
|
||||
- 10-20x faster code search than traditional methods
|
||||
- Parallel execution (8 concurrent searches)
|
||||
- Automatic caching for repeated queries
|
||||
|
||||
|
||||
@@ -2,7 +2,7 @@
|
||||
Example 1: Basic Fast Context Usage
|
||||
|
||||
This example demonstrates the basic usage of Fast Context for rapid code search.
|
||||
Fast Context is 10-20x faster than traditional sequential search because it:
|
||||
Fast Context speeds up code search because it:
|
||||
- Executes up to 8 tool calls in parallel
|
||||
- Limits search to 4 turns maximum
|
||||
- Uses a lightweight model optimized for search
|
||||
|
||||
@@ -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)
|
||||
@@ -1,280 +0,0 @@
|
||||
# Performance Benchmark CI
|
||||
# Runs import time and instantiation benchmarks on push/PR
|
||||
# Fails if import time exceeds 200ms threshold
|
||||
|
||||
name: Performance Benchmark
|
||||
|
||||
on:
|
||||
push:
|
||||
branches: [main, develop]
|
||||
paths:
|
||||
- 'praisonaiagents/**'
|
||||
- 'benchmarks/**'
|
||||
- '.github/workflows/benchmark.yml'
|
||||
pull_request:
|
||||
branches: [main, develop]
|
||||
paths:
|
||||
- 'praisonaiagents/**'
|
||||
- 'benchmarks/**'
|
||||
- '.github/workflows/benchmark.yml'
|
||||
workflow_dispatch: # Allow manual trigger
|
||||
|
||||
jobs:
|
||||
benchmark:
|
||||
runs-on: ubuntu-latest
|
||||
timeout-minutes: 10
|
||||
|
||||
steps:
|
||||
- uses: actions/checkout@v4
|
||||
|
||||
- name: Set up Python
|
||||
uses: actions/setup-python@v5
|
||||
with:
|
||||
python-version: '3.11'
|
||||
cache: 'pip'
|
||||
|
||||
- name: Install package
|
||||
run: |
|
||||
pip install -e .
|
||||
pip install pytest
|
||||
|
||||
- name: Measure Import Time
|
||||
id: import_time
|
||||
run: |
|
||||
# Measure import time (full Agent import)
|
||||
IMPORT_TIME=$(python -c "
|
||||
import time
|
||||
import sys
|
||||
|
||||
# Clear any cached modules
|
||||
for key in list(sys.modules.keys()):
|
||||
if 'praisonai' in key:
|
||||
del sys.modules[key]
|
||||
|
||||
start = time.perf_counter()
|
||||
from praisonaiagents import Agent
|
||||
elapsed = (time.perf_counter() - start) * 1000
|
||||
print(f'{elapsed:.1f}')
|
||||
")
|
||||
|
||||
echo "import_time=$IMPORT_TIME" >> $GITHUB_OUTPUT
|
||||
echo "Import time: ${IMPORT_TIME}ms"
|
||||
|
||||
- name: Check Import Time Threshold
|
||||
run: |
|
||||
IMPORT_TIME=${{ steps.import_time.outputs.import_time }}
|
||||
THRESHOLD=200
|
||||
|
||||
echo "Import time: ${IMPORT_TIME}ms"
|
||||
echo "Threshold: ${THRESHOLD}ms"
|
||||
|
||||
# Compare using bc for floating point
|
||||
if (( $(echo "$IMPORT_TIME > $THRESHOLD" | bc -l) )); then
|
||||
echo "❌ FAIL: Import time ${IMPORT_TIME}ms exceeds threshold ${THRESHOLD}ms"
|
||||
exit 1
|
||||
else
|
||||
echo "✅ PASS: Import time ${IMPORT_TIME}ms is within threshold ${THRESHOLD}ms"
|
||||
fi
|
||||
|
||||
- name: Measure Instantiation Time
|
||||
run: |
|
||||
python -c "
|
||||
import time
|
||||
from praisonaiagents import Agent
|
||||
|
||||
# Warmup
|
||||
for _ in range(5):
|
||||
Agent(name='Test', output='silent')
|
||||
|
||||
# Measure
|
||||
times = []
|
||||
for _ in range(50):
|
||||
start = time.perf_counter()
|
||||
Agent(name='Test', output='silent')
|
||||
times.append((time.perf_counter() - start) * 1e6)
|
||||
|
||||
avg = sum(times) / len(times)
|
||||
min_t = min(times)
|
||||
max_t = max(times)
|
||||
|
||||
print(f'Instantiation time:')
|
||||
print(f' Average: {avg:.2f}μs')
|
||||
print(f' Min: {min_t:.2f}μs')
|
||||
print(f' Max: {max_t:.2f}μs')
|
||||
"
|
||||
|
||||
- name: Verify Lazy Imports
|
||||
run: |
|
||||
python -c "
|
||||
import sys
|
||||
|
||||
# Clear any cached modules
|
||||
for key in list(sys.modules.keys()):
|
||||
if 'praisonai' in key or 'litellm' in key:
|
||||
del sys.modules[key]
|
||||
|
||||
# Import package
|
||||
import praisonaiagents
|
||||
|
||||
# Check heavy deps are NOT loaded
|
||||
heavy_deps = ['litellm', 'chromadb', 'mem0', 'rich']
|
||||
loaded = []
|
||||
for dep in heavy_deps:
|
||||
if dep in sys.modules:
|
||||
loaded.append(dep)
|
||||
|
||||
if loaded:
|
||||
print(f'⚠️ WARNING: Heavy deps loaded eagerly: {loaded}')
|
||||
# Don't fail, just warn - some deps may be needed
|
||||
else:
|
||||
print('✅ All heavy deps are lazy loaded')
|
||||
"
|
||||
|
||||
- name: Run Quick Benchmark (if available)
|
||||
continue-on-error: true
|
||||
run: |
|
||||
if [ -f "benchmarks/quick_benchmark.py" ]; then
|
||||
python benchmarks/quick_benchmark.py --fast || true
|
||||
else
|
||||
echo "Quick benchmark not found, skipping"
|
||||
fi
|
||||
|
||||
# Real API benchmark - only runs when API keys are available
|
||||
# Triggered manually or on main branch with secrets
|
||||
real-api-benchmark:
|
||||
runs-on: ubuntu-latest
|
||||
timeout-minutes: 15
|
||||
if: github.event_name == 'workflow_dispatch' || (github.ref == 'refs/heads/main' && github.event_name == 'push')
|
||||
|
||||
steps:
|
||||
- uses: actions/checkout@v4
|
||||
|
||||
- name: Set up Python
|
||||
uses: actions/setup-python@v5
|
||||
with:
|
||||
python-version: '3.11'
|
||||
cache: 'pip'
|
||||
|
||||
- name: Install package
|
||||
run: |
|
||||
pip install -e .
|
||||
pip install pytest
|
||||
|
||||
- name: Check API Key Availability
|
||||
id: check_keys
|
||||
run: |
|
||||
if [ -n "${{ secrets.OPENAI_API_KEY }}" ]; then
|
||||
echo "has_openai=true" >> $GITHUB_OUTPUT
|
||||
echo "✅ OpenAI API key available"
|
||||
else
|
||||
echo "has_openai=false" >> $GITHUB_OUTPUT
|
||||
echo "⚠️ OpenAI API key not available, skipping real API tests"
|
||||
fi
|
||||
|
||||
- name: Run Real API Benchmark
|
||||
if: steps.check_keys.outputs.has_openai == 'true'
|
||||
env:
|
||||
OPENAI_API_KEY: ${{ secrets.OPENAI_API_KEY }}
|
||||
run: |
|
||||
python -c "
|
||||
import time
|
||||
import os
|
||||
|
||||
# Verify API key is set
|
||||
if not os.environ.get('OPENAI_API_KEY'):
|
||||
print('❌ OPENAI_API_KEY not set')
|
||||
exit(1)
|
||||
|
||||
from praisonaiagents import Agent
|
||||
|
||||
# Create a simple agent
|
||||
agent = Agent(
|
||||
name='BenchmarkAgent',
|
||||
instructions='You are a helpful assistant. Be concise.',
|
||||
llm='gpt-4o-mini',
|
||||
output='silent'
|
||||
)
|
||||
|
||||
# Warmup run
|
||||
print('Running warmup...')
|
||||
agent.chat('Say hello in one word')
|
||||
|
||||
# Benchmark runs
|
||||
print('Running benchmark (3 iterations)...')
|
||||
times = []
|
||||
for i in range(3):
|
||||
start = time.perf_counter()
|
||||
response = agent.chat('What is 2+2? Answer with just the number.')
|
||||
elapsed = time.perf_counter() - start
|
||||
times.append(elapsed)
|
||||
print(f' Run {i+1}: {elapsed:.2f}s')
|
||||
|
||||
avg = sum(times) / len(times)
|
||||
print(f'')
|
||||
print(f'Real API Benchmark Results:')
|
||||
print(f' Average response time: {avg:.2f}s')
|
||||
print(f' Min: {min(times):.2f}s')
|
||||
print(f' Max: {max(times):.2f}s')
|
||||
|
||||
# Verify response quality
|
||||
if '4' in response:
|
||||
print('✅ Response quality check passed')
|
||||
else:
|
||||
print(f'⚠️ Unexpected response: {response}')
|
||||
"
|
||||
|
||||
- name: Run Multi-Agent Benchmark
|
||||
if: steps.check_keys.outputs.has_openai == 'true'
|
||||
env:
|
||||
OPENAI_API_KEY: ${{ secrets.OPENAI_API_KEY }}
|
||||
run: |
|
||||
python -c "
|
||||
import time
|
||||
import os
|
||||
|
||||
from praisonaiagents import Agent, Agents, Task
|
||||
|
||||
# Create agents
|
||||
researcher = Agent(
|
||||
name='Researcher',
|
||||
instructions='You research topics. Be brief.',
|
||||
llm='gpt-4o-mini',
|
||||
output='silent'
|
||||
)
|
||||
|
||||
writer = Agent(
|
||||
name='Writer',
|
||||
instructions='You write summaries. Be brief.',
|
||||
llm='gpt-4o-mini',
|
||||
output='silent'
|
||||
)
|
||||
|
||||
# Create tasks
|
||||
task1 = Task(
|
||||
description='List 2 facts about Python programming',
|
||||
agent=researcher,
|
||||
expected_output='2 brief facts'
|
||||
)
|
||||
|
||||
task2 = Task(
|
||||
description='Summarize the facts in one sentence',
|
||||
agent=writer,
|
||||
expected_output='One sentence summary'
|
||||
)
|
||||
|
||||
# Run workflow
|
||||
print('Running multi-agent workflow...')
|
||||
start = time.perf_counter()
|
||||
|
||||
agents = Agents(agents=[researcher, writer], tasks=[task1, task2])
|
||||
result = agents.start()
|
||||
|
||||
elapsed = time.perf_counter() - start
|
||||
|
||||
print(f'')
|
||||
print(f'Multi-Agent Benchmark Results:')
|
||||
print(f' Total workflow time: {elapsed:.2f}s')
|
||||
print(f' Agents: 2')
|
||||
print(f' Tasks: 2')
|
||||
print('✅ Multi-agent workflow completed successfully')
|
||||
"
|
||||
@@ -1,70 +1,47 @@
|
||||
# PraisonAI Agents - Benchmark Results
|
||||
|
||||
**Generated:** 2026-06-17
|
||||
**Methodology:** 10 warmup runs + 1000 measured iterations (`performance_benchmark.py` config)
|
||||
**Config:** `Agent(name=..., instructions=..., model="gpt-4o-mini", tools=[get_weather], output="silent")`
|
||||
**Generated:** 2026-10-15 (single run of `performance_benchmark.py`)
|
||||
**Methodology:** 10 warmup runs + 1000 measured iterations per framework
|
||||
**Setup (every framework):** one agent, model `gpt-4o-mini`, one `get_weather` tool, no LLM call. Tool wrapping (`function_tool`, `@tool`, ...) happens inside the timed call. Frameworks that were not installed are omitted.
|
||||
|
||||
## Agent Instantiation Time
|
||||
|
||||
| Framework | Avg Time (μs) | Min (μs) | Median (μs) | Relative |
|
||||
|-----------|---------------|----------|-------------|----------|
|
||||
| OpenAI Agents SDK | 5.05 | — | — | 1.00x (fastest) |
|
||||
| Agno | 5.27 | — | — | 1.04x |
|
||||
| **PraisonAI** | **13.62** | **12.58** | **12.96** | **2.69x** |
|
||||
| PraisonAI (LiteLLM) | 14.49 | 13.37 | 13.71 | 2.86x |
|
||||
|
||||
### PraisonAI detail (1000 iter)
|
||||
|
||||
| Variant | Avg (μs) | Min (μs) | Median (μs) | Std (μs) |
|
||||
|---------|----------|----------|-------------|----------|
|
||||
| With tools (`gpt-4o-mini`) | 13.62 | 12.58 | 12.96 | 2.73 |
|
||||
| With tools (`openai/gpt-4o-mini`) | 14.49 | 13.37 | 13.71 | 4.14 |
|
||||
| Minimal, no tools | 13.40 | — | — | — |
|
||||
|
||||
## Historical baseline
|
||||
|
||||
| Source | Config | Avg (μs) | Notes |
|
||||
|--------|--------|----------|-------|
|
||||
| `performance_benchmark.py` (Dec 2025) | tools + 1000 iter | **3.20** | Pre-init-bloat baseline |
|
||||
| `simple_benchmark.py` (Dec 2025) | no tools, 100 iter | **3.77** | Saved in prior `BENCHMARK_RESULTS.md` |
|
||||
| Current (Jun 2026) | tools + 1000 iter | **13.62** | After lazy-init optimisations (was ~31 μs pre-fix) |
|
||||
| Framework | Avg (μs) | Min (μs) | Median (μs) | Std (μs) | Avg memory (KiB) | Relative |
|
||||
|-----------|----------|----------|-------------|----------|------------------|----------|
|
||||
| Agno | 9.51 | 9.16 | 9.33 | 1.00 | 6.83 | 1.00x (fastest) |
|
||||
| PraisonAI | 45.55 | 42.63 | 44.26 | 4.11 | 20.55 | 4.79x |
|
||||
| PraisonAI (LiteLLM) | 46.94 | 44.65 | 45.82 | 3.26 | 20.81 | 4.94x |
|
||||
| OpenAI Agents SDK | 591.02 | 529.26 | 566.69 | 60.62 | 60.19 | 62.18x |
|
||||
| PydanticAI | 1,021.99 | 882.02 | 935.63 | 1,674.03 | 51.22 | 108x |
|
||||
| LangGraph | 3,232.91 | 2,836.27 | 3,183.02 | 233.32 | 89.66 | 340x |
|
||||
|
||||
## Package Versions
|
||||
|
||||
| Package | Version |
|
||||
|---------|---------|
|
||||
| PraisonAI | 1.6.60 |
|
||||
| Agno | 2.4.2 |
|
||||
| OpenAI Agents SDK | 0.7.0 |
|
||||
| praisonaiagents | 1.7.12 |
|
||||
| agno | 3.1.3 |
|
||||
| pydantic-ai-slim | 2.56.0 |
|
||||
| openai-agents | 0.23.1 |
|
||||
| langgraph | 1.2.15 |
|
||||
|
||||
## Machine
|
||||
|
||||
| Field | Value |
|
||||
|-------|-------|
|
||||
| Platform | Linux-6.17.0-1022-azure-x86_64-with-glibc2.39 |
|
||||
| Processor | x86_64 |
|
||||
| CPUs | 4 |
|
||||
| Python | 3.12.3 |
|
||||
| Load avg (1 min) at start | 0.85 |
|
||||
|
||||
## How to Reproduce
|
||||
|
||||
PraisonAI-only (fast; avoids CrewAI/LangGraph import hang):
|
||||
|
||||
```bash
|
||||
cd praisonai-agents
|
||||
PYTHONWARNINGS='ignore::DeprecationWarning' python -c "
|
||||
import time, statistics
|
||||
from typing import Literal
|
||||
from praisonaiagents import Agent
|
||||
|
||||
def get_weather(city: Literal['nyc', 'sf']):
|
||||
return 'cloudy' if city == 'nyc' else 'sunny'
|
||||
|
||||
def bench(fn, n=1000, w=10):
|
||||
for _ in range(w): fn()
|
||||
times = []
|
||||
for _ in range(n):
|
||||
s = time.perf_counter(); fn(); times.append(time.perf_counter()-s)
|
||||
print(f'avg={statistics.mean(times)*1e6:.2f} μs')
|
||||
|
||||
bench(lambda: Agent(name='Test', instructions='Hi', model='gpt-4o-mini', tools=[get_weather], output='silent'))
|
||||
"
|
||||
cd src/praisonai-agents
|
||||
pip install -e . agno openai-agents "pydantic-ai-slim[openai]" langgraph langchain-openai crewai
|
||||
# OpenAI-backed constructors need a key to exist; no request is made
|
||||
OPENAI_API_KEY=sk-dummy PYTHONWARNINGS='ignore::DeprecationWarning' python benchmarks/performance_benchmark.py --save --iterations 1000 --warmup 10 --only praisonai praisonai_litellm agno pydantic openai_agents langgraph
|
||||
```
|
||||
|
||||
Full multi-framework benchmark:
|
||||
|
||||
```bash
|
||||
cd praisonai-agents
|
||||
PYTHONWARNINGS='ignore::DeprecationWarning' python benchmarks/performance_benchmark.py --save
|
||||
```
|
||||
Older, superseded result files live in [`archive/`](archive/).
|
||||
@@ -6,7 +6,7 @@
|
||||
|
||||
## Summary
|
||||
|
||||
Instantiation regressed from **~3.2 μs** (Dec 2025) to **~31 μs** (~10×). After fixes, steady-state is **~16 μs** (~2× historical). PraisonAI is competitive with Agno (~5 μs) on warmed benchmarks using `model=`.
|
||||
Instantiation regressed from **~3.2 μs** (Dec 2025) to **~31 μs** (~10×). After fixes, steady-state is **~16 μs** (~2× historical). Agno measured ~5 μs on the same machine at the time, so PraisonAI was still ~3× slower; see `BENCHMARK_RESULTS.md` for current numbers.
|
||||
|
||||
## Before vs After Fixes
|
||||
|
||||
@@ -53,10 +53,45 @@ Instantiation regressed from **~3.2 μs** (Dec 2025) to **~31 μs** (~10×). Aft
|
||||
|
||||
The residual ~13 μs is primarily **`Agent.__init__` size** (~1200 lines of consolidated config resolution added since v0.1.5). Further optimisation requires a dedicated slim fast-path for `Agent(name=..., output="silent")` without full config object churn.
|
||||
|
||||
## Re-measure (2026-10-15)
|
||||
|
||||
Quiet machine (GitHub-hosted `ubuntu` x86_64 runner, 4 vCPU, load 0.77, Python 3.12.3). Same config as the
|
||||
June run (`performance_benchmark.py`: one tool, 10–20 warmup + 1000 iterations). Both versions were run on the
|
||||
same machine, so the two columns are directly comparable. The June numbers in `BENCHMARK_RESULTS.md` came from a
|
||||
faster machine.
|
||||
|
||||
| Code | Avg (μs) | Min (μs) | Median (μs) |
|
||||
|------|----------|----------|-------------|
|
||||
| praisonaiagents 1.6.60 (`709a03878`) | 20.65 | 19.85 | 20.30 |
|
||||
| `main` @ `a84e42aa9` (1.7.12 dev) | 47.92 | 45.71 | 47.11 |
|
||||
|
||||
The slowdown is real: **~2.3×** on identical hardware. Median, sampled along `main` first-parent history:
|
||||
|
||||
| Date | Commit | Median (μs) |
|
||||
|------|--------|-------------|
|
||||
| 2026-06-17 | `709a03878` (1.6.60) | 20.3 |
|
||||
| 2026-06-24 | `bb2d35ba6` | 24.2 |
|
||||
| 2026-07-01 | `f45c49b32` | 26.7 |
|
||||
| 2026-07-15 | `b1086c275` | 30.5 |
|
||||
| 2026-08-05 | `aadb62acf` | 29.7 |
|
||||
| 2026-08-19 | `d32fe29b4` | 34.0 |
|
||||
| 2026-08-26 | `7746520df` | **3,063** (transient ms-scale regression, since fixed) |
|
||||
| 2026-09-03 | `17d8726c3` | **3,012** |
|
||||
| 2026-09-10 | `8d2b4cdbc` | 48.3 |
|
||||
| 2026-10-02 | `a2b52f669` | 60.2 |
|
||||
| 2026-10-14 | `a84e42aa9` | 47.2 |
|
||||
|
||||
`git bisect` with a 33 μs cut-off has no single culprit: it lands on a wrapper-only commit, and neighbouring
|
||||
commits differ by less than the noise. The cost is gradual creep. cProfile on `main` puts ~43% of construction
|
||||
time in the body of `Agent.__init__`, and no callee costs more than ~10%. The fix is the slim construction path
|
||||
described above, tracked in #5835. A ~3 ms regression shipped for two weeks unnoticed; the PR gate
|
||||
in `.github/workflows/benchmark.yml` (`simple_benchmark.py --json --fail-over-threshold`, 150 μs budget) now
|
||||
catches regressions of that size.
|
||||
|
||||
## How to reproduce
|
||||
|
||||
```bash
|
||||
cd praisonai-agents
|
||||
cd src/praisonai-agents
|
||||
PYTHONWARNINGS='ignore::DeprecationWarning' python benchmarks/simple_benchmark.py
|
||||
PYTHONWARNINGS='ignore::DeprecationWarning' python benchmarks/performance_benchmark.py
|
||||
PYTHONWARNINGS='ignore::DeprecationWarning' python benchmarks/performance_benchmark.py --only praisonai
|
||||
```
|
||||
+2
@@ -1,3 +1,5 @@
|
||||
> **Superseded.** Historical result from 2026-01-13, kept for reference only. Package versions and methodology differ from current runs; do not cite. Current results: [../BENCHMARK_RESULTS.md](../BENCHMARK_RESULTS.md).
|
||||
|
||||
# PraisonAI Agents - Real Execution Benchmark
|
||||
|
||||
**Generated:** 2026-01-13 06:07:28
|
||||
+2
@@ -1,3 +1,5 @@
|
||||
> **Superseded.** Historical result from 2026-01-13, kept for reference only. Package versions and methodology differ from current runs; do not cite. Current results: [../BENCHMARK_RESULTS.md](../BENCHMARK_RESULTS.md).
|
||||
|
||||
# PraisonAI Agents - Real-World Execution Benchmark
|
||||
|
||||
**Generated:** 2026-01-13 14:42:59
|
||||
+2
@@ -1,3 +1,5 @@
|
||||
> **Superseded.** Historical result from 2025-12-18, kept for reference only. Package versions and methodology differ from current runs; do not cite. Current results: [../BENCHMARK_RESULTS.md](../BENCHMARK_RESULTS.md).
|
||||
|
||||
# PraisonAI Agents - Benchmark Results
|
||||
|
||||
**Generated:** 2025-12-18 13:23:43
|
||||
+2
@@ -1,3 +1,5 @@
|
||||
> **Superseded.** Historical result from 2026-01-23, kept for reference only. Package versions and methodology differ from current runs; do not cite. Current results: [../BENCHMARK_RESULTS.md](../BENCHMARK_RESULTS.md).
|
||||
|
||||
# PraisonAI Agents - Tools Benchmark Results
|
||||
|
||||
**Generated:** 2026-01-23 11:50:18
|
||||
@@ -6,6 +6,8 @@ multiple AI agent frameworks using industry-standard methodology.
|
||||
|
||||
Usage:
|
||||
python benchmarks/performance_benchmark.py
|
||||
python benchmarks/performance_benchmark.py --save # rewrite BENCHMARK_RESULTS.md
|
||||
python benchmarks/performance_benchmark.py --only praisonai
|
||||
"""
|
||||
|
||||
import gc
|
||||
@@ -134,39 +136,14 @@ def run_benchmark(func: Callable, name: str, iterations: int = 1000, warmup: int
|
||||
|
||||
|
||||
# ============================================================================
|
||||
# Framework Imports
|
||||
# Framework factories
|
||||
#
|
||||
# Each builder imports its framework lazily and returns a zero-arg callable
|
||||
# that creates one agent with one tool. A missing framework raises ImportError
|
||||
# and is skipped, so the script runs with only praisonaiagents installed.
|
||||
# ============================================================================
|
||||
from praisonaiagents import Agent as PraisonAgent
|
||||
from agno.agent import Agent as AgnoAgent
|
||||
from agno.models.openai import OpenAIChat
|
||||
|
||||
# Additional frameworks for comparison
|
||||
try:
|
||||
from pydantic_ai import Agent as PydanticAgent
|
||||
PYDANTIC_AVAILABLE = True
|
||||
except ImportError:
|
||||
PYDANTIC_AVAILABLE = False
|
||||
|
||||
try:
|
||||
from crewai.agent import Agent as CrewAgent
|
||||
from crewai.tools import tool as crewai_tool
|
||||
CREWAI_AVAILABLE = True
|
||||
except ImportError:
|
||||
CREWAI_AVAILABLE = False
|
||||
|
||||
try:
|
||||
from langchain_core.tools import tool as langchain_tool
|
||||
from langchain_openai import ChatOpenAI
|
||||
from langgraph.prebuilt import create_react_agent
|
||||
LANGGRAPH_AVAILABLE = True
|
||||
except ImportError:
|
||||
LANGGRAPH_AVAILABLE = False
|
||||
|
||||
try:
|
||||
from agents import Agent as OpenAIAgent, function_tool
|
||||
OPENAI_AGENTS_AVAILABLE = True
|
||||
except ImportError:
|
||||
OPENAI_AGENTS_AVAILABLE = False
|
||||
INSTRUCTIONS = "Be concise, reply with one sentence."
|
||||
|
||||
|
||||
def get_weather(city: Literal["nyc", "sf"]):
|
||||
@@ -180,195 +157,181 @@ def get_weather(city: Literal["nyc", "sf"]):
|
||||
tools = [get_weather]
|
||||
|
||||
|
||||
def instantiate_praisonai_agent():
|
||||
"""Instantiate a PraisonAI agent with one tool."""
|
||||
return PraisonAgent(
|
||||
name="Test Agent",
|
||||
instructions="Be concise, reply with one sentence.",
|
||||
model="gpt-4o-mini",
|
||||
tools=tools,
|
||||
output="silent"
|
||||
def build_praisonai():
|
||||
from praisonaiagents import Agent as PraisonAgent
|
||||
return lambda: PraisonAgent(
|
||||
name="Test Agent", instructions=INSTRUCTIONS, model="gpt-4o-mini",
|
||||
tools=tools, output="silent",
|
||||
)
|
||||
|
||||
|
||||
def instantiate_praisonai_litellm_agent():
|
||||
"""Instantiate a PraisonAI agent with LiteLLM backend."""
|
||||
return PraisonAgent(
|
||||
name="Test Agent",
|
||||
instructions="Be concise, reply with one sentence.",
|
||||
model="openai/gpt-4o-mini",
|
||||
tools=tools,
|
||||
output="silent"
|
||||
def build_praisonai_litellm():
|
||||
from praisonaiagents import Agent as PraisonAgent
|
||||
return lambda: PraisonAgent(
|
||||
name="Test Agent", instructions=INSTRUCTIONS, model="openai/gpt-4o-mini",
|
||||
tools=tools, output="silent",
|
||||
)
|
||||
|
||||
|
||||
def instantiate_agno_agent():
|
||||
"""Instantiate an Agno agent with one tool."""
|
||||
return AgnoAgent(model=OpenAIChat(id="gpt-4o-mini"), tools=tools)
|
||||
def build_agno():
|
||||
from agno.agent import Agent as AgnoAgent
|
||||
from agno.models.openai import OpenAIChat
|
||||
return lambda: AgnoAgent(model=OpenAIChat(id="gpt-4o-mini"), instructions=INSTRUCTIONS, tools=tools)
|
||||
|
||||
|
||||
def instantiate_pydantic_agent():
|
||||
"""Instantiate a PydanticAI agent with one tool."""
|
||||
agent = PydanticAgent("openai:gpt-4o-mini", system_prompt="Be concise, reply with one sentence.")
|
||||
@agent.tool_plain
|
||||
def get_weather_pydantic(city: Literal["nyc", "sf"]):
|
||||
"""Use this to get weather information."""
|
||||
if city == "nyc":
|
||||
return "It might be cloudy in nyc"
|
||||
elif city == "sf":
|
||||
return "It's always sunny in sf"
|
||||
return agent
|
||||
def build_pydantic():
|
||||
from pydantic_ai import Agent as PydanticAgent
|
||||
|
||||
def create():
|
||||
agent = PydanticAgent("openai:gpt-4o-mini", system_prompt=INSTRUCTIONS)
|
||||
agent.tool_plain(get_weather)
|
||||
return agent
|
||||
return create
|
||||
|
||||
|
||||
def instantiate_openai_agents():
|
||||
"""Instantiate an OpenAI Agents SDK agent with one tool."""
|
||||
return OpenAIAgent(
|
||||
name="Test agent",
|
||||
instructions="Be concise, reply with one sentence.",
|
||||
model="gpt-4o-mini",
|
||||
def build_openai_agents():
|
||||
from agents import Agent as OpenAIAgent, function_tool
|
||||
return lambda: OpenAIAgent(
|
||||
name="Test agent", instructions=INSTRUCTIONS, model="gpt-4o-mini",
|
||||
tools=[function_tool(get_weather)],
|
||||
)
|
||||
|
||||
|
||||
def instantiate_langgraph_agent():
|
||||
"""Instantiate a LangGraph agent with one tool."""
|
||||
@langchain_tool
|
||||
def get_weather_lg(city: Literal["nyc", "sf"]):
|
||||
"""Use this to get weather information."""
|
||||
if city == "nyc":
|
||||
return "It might be cloudy in nyc"
|
||||
elif city == "sf":
|
||||
return "It's always sunny in sf"
|
||||
return create_react_agent(model=ChatOpenAI(model="gpt-4o-mini"), tools=[get_weather_lg])
|
||||
|
||||
|
||||
def instantiate_crewai_agent():
|
||||
"""Instantiate a CrewAI agent with one tool."""
|
||||
@crewai_tool("Weather Tool")
|
||||
def get_weather_crew(city: Literal["nyc", "sf"]):
|
||||
"""Use this to get weather information."""
|
||||
if city == "nyc":
|
||||
return "It might be cloudy in nyc"
|
||||
elif city == "sf":
|
||||
return "It's always sunny in sf"
|
||||
return CrewAgent(
|
||||
llm="gpt-4o-mini",
|
||||
role="Test Agent",
|
||||
goal="Be concise, reply with one sentence.",
|
||||
tools=[get_weather_crew],
|
||||
backstory="Test",
|
||||
def build_langgraph():
|
||||
from langchain_core.tools import tool as langchain_tool
|
||||
from langchain_openai import ChatOpenAI
|
||||
from langgraph.prebuilt import create_react_agent
|
||||
return lambda: create_react_agent(
|
||||
model=ChatOpenAI(model="gpt-4o-mini"), tools=[langchain_tool(get_weather)],
|
||||
)
|
||||
|
||||
|
||||
def get_package_versions():
|
||||
"""Get version numbers for all benchmarked packages."""
|
||||
def build_crewai():
|
||||
from crewai.agent import Agent as CrewAgent
|
||||
from crewai.tools import tool as crewai_tool
|
||||
|
||||
def create():
|
||||
return CrewAgent(
|
||||
llm="gpt-4o-mini", role="Test Agent", goal=INSTRUCTIONS,
|
||||
tools=[crewai_tool("Weather Tool")(get_weather)], backstory="Test",
|
||||
)
|
||||
return create
|
||||
|
||||
|
||||
# key -> (display name, PyPI package, builder)
|
||||
FRAMEWORKS = {
|
||||
'praisonai': ('PraisonAI', 'praisonaiagents', build_praisonai),
|
||||
'praisonai_litellm': ('PraisonAI (LiteLLM)', 'praisonaiagents', build_praisonai_litellm),
|
||||
'agno': ('Agno', 'agno', build_agno),
|
||||
'pydantic': ('PydanticAI', 'pydantic-ai-slim', build_pydantic),
|
||||
'openai_agents': ('OpenAI Agents SDK', 'openai-agents', build_openai_agents),
|
||||
'langgraph': ('LangGraph', 'langgraph', build_langgraph),
|
||||
'crewai': ('CrewAI', 'crewai', build_crewai),
|
||||
}
|
||||
|
||||
|
||||
def get_package_versions(keys=None):
|
||||
"""Get version numbers for benchmarked packages."""
|
||||
from importlib.metadata import version as get_version
|
||||
|
||||
packages = {
|
||||
'PraisonAI': 'praisonaiagents',
|
||||
'Agno': 'agno',
|
||||
'PydanticAI': 'pydantic-ai',
|
||||
'OpenAI Agents SDK': 'openai-agents',
|
||||
'LangGraph': 'langgraph',
|
||||
'CrewAI': 'crewai'
|
||||
}
|
||||
|
||||
|
||||
versions = {}
|
||||
for display_name, pkg_name in packages.items():
|
||||
for key in keys or FRAMEWORKS:
|
||||
display_name, pkg_name, _ = FRAMEWORKS[key]
|
||||
if pkg_name in versions:
|
||||
continue
|
||||
try:
|
||||
versions[display_name] = get_version(pkg_name)
|
||||
versions[pkg_name] = get_version(pkg_name)
|
||||
except Exception:
|
||||
versions[display_name] = 'not installed'
|
||||
|
||||
versions[pkg_name] = 'not installed'
|
||||
return versions
|
||||
|
||||
|
||||
def save_benchmark_results(results: dict, baseline):
|
||||
"""Save benchmark results to BENCHMARK_RESULTS.md and update README.md"""
|
||||
def get_machine_info() -> dict:
|
||||
"""Describe the machine so results are traceable."""
|
||||
import os
|
||||
import re
|
||||
from datetime import datetime
|
||||
|
||||
benchmarks_dir = os.path.dirname(__file__)
|
||||
filepath = os.path.join(benchmarks_dir, 'BENCHMARK_RESULTS.md')
|
||||
|
||||
# Get package versions
|
||||
versions = get_package_versions()
|
||||
|
||||
display_names = {
|
||||
'agno': 'Agno',
|
||||
'praisonai': 'PraisonAI',
|
||||
'praisonai_litellm': 'PraisonAI (LiteLLM)',
|
||||
'pydantic': 'PydanticAI',
|
||||
'openai_agents': 'OpenAI Agents SDK',
|
||||
'langgraph': 'LangGraph',
|
||||
'crewai': 'CrewAI'
|
||||
import platform
|
||||
|
||||
load = os.getloadavg()[0] if hasattr(os, 'getloadavg') else None
|
||||
return {
|
||||
'Platform': platform.platform(),
|
||||
'Processor': platform.processor() or platform.machine(),
|
||||
'CPUs': os.cpu_count(),
|
||||
'Python': platform.python_version(),
|
||||
'Load avg (1 min) at start': f'{load:.2f}' if load is not None else 'n/a',
|
||||
}
|
||||
|
||||
# Build table rows
|
||||
table_rows = []
|
||||
sorted_results = sorted(results.items(), key=lambda x: x[1].avg_run_time)
|
||||
|
||||
for name, result in sorted_results:
|
||||
display_name = display_names.get(name, name)
|
||||
avg_us = result.avg_run_time * 1e6
|
||||
ratio = result.avg_run_time / baseline.avg_run_time if baseline.avg_run_time > 0 else 0
|
||||
|
||||
if name == 'praisonai':
|
||||
table_rows.append(f'| **{display_name}** | **{avg_us:.2f}** | **1.00x (fastest)** |')
|
||||
else:
|
||||
avg_str = f'{avg_us:,.2f}' if avg_us >= 1000 else f'{avg_us:.2f}'
|
||||
ratio_str = f'{ratio:,.0f}x' if ratio >= 100 else f'{ratio:.2f}x'
|
||||
table_rows.append(f'| {display_name} | {avg_str} | {ratio_str} |')
|
||||
|
||||
# Save to BENCHMARK_RESULTS.md
|
||||
|
||||
|
||||
def rank_results(results: dict) -> list:
|
||||
"""Return ``(key, result, ratio)`` sorted fastest first, ratio vs the fastest."""
|
||||
ranked = sorted(results.items(), key=lambda x: x[1].avg_run_time)
|
||||
fastest = ranked[0][1].avg_run_time if ranked else 0
|
||||
return [(k, r, r.avg_run_time / fastest if fastest > 0 else 0) for k, r in ranked]
|
||||
|
||||
|
||||
def save_benchmark_results(results: dict, machine: dict, iterations: int, warmup: int,
|
||||
filename: str = 'BENCHMARK_RESULTS.md') -> str:
|
||||
"""Write every framework's results from this single run to ``BENCHMARK_RESULTS.md``."""
|
||||
import os
|
||||
from datetime import datetime
|
||||
|
||||
filepath = os.path.join(os.path.dirname(__file__), filename)
|
||||
versions = get_package_versions(results.keys())
|
||||
|
||||
def us(seconds):
|
||||
value = seconds * 1e6
|
||||
return f'{value:,.2f}' if value >= 1000 else f'{value:.2f}'
|
||||
|
||||
lines = [
|
||||
'# PraisonAI Agents - Benchmark Results', '',
|
||||
f'**Generated:** {datetime.now().strftime("%Y-%m-%d")} (single run of `performance_benchmark.py`)',
|
||||
f'**Methodology:** {warmup} warmup runs + {iterations} measured iterations per framework',
|
||||
f'**Setup (every framework):** one agent, model `gpt-4o-mini`, one `get_weather` tool, '
|
||||
'no LLM call. Tool wrapping (`function_tool`, `@tool`, ...) happens inside the timed call. '
|
||||
'Frameworks that were not installed are omitted.',
|
||||
'', '## Agent Instantiation Time', '',
|
||||
'| Framework | Avg (μs) | Min (μs) | Median (μs) | Std (μs) | Avg memory (KiB) | Relative |',
|
||||
'|-----------|----------|----------|-------------|----------|------------------|----------|',
|
||||
]
|
||||
for rank, (key, r, ratio) in enumerate(rank_results(results)):
|
||||
relative = '1.00x (fastest)' if rank == 0 else (f'{ratio:,.0f}x' if ratio >= 100 else f'{ratio:.2f}x')
|
||||
lines.append(
|
||||
f'| {FRAMEWORKS[key][0]} | {us(r.avg_run_time)} | {us(r.min_run_time)} | '
|
||||
f'{us(r.median_run_time)} | {us(r.std_dev_run_time)} | '
|
||||
f'{r.avg_memory_usage * 1024:.2f} | {relative} |'
|
||||
)
|
||||
lines += ['', '## Package Versions', '', '| Package | Version |', '|---------|---------|']
|
||||
lines += [f'| {pkg} | {ver} |' for pkg, ver in versions.items()]
|
||||
lines += ['', '## Machine', '', '| Field | Value |', '|-------|-------|']
|
||||
lines += [f'| {k} | {v} |' for k, v in machine.items()]
|
||||
lines += [
|
||||
'', '## How to Reproduce', '', '```bash', 'cd src/praisonai-agents',
|
||||
'pip install -e . agno openai-agents "pydantic-ai-slim[openai]" langgraph langchain-openai crewai',
|
||||
"# OpenAI-backed constructors need a key to exist; no request is made",
|
||||
"OPENAI_API_KEY=sk-dummy PYTHONWARNINGS='ignore::DeprecationWarning' "
|
||||
f"python benchmarks/performance_benchmark.py --save --iterations {iterations} "
|
||||
f"--warmup {warmup} --only {' '.join(results)}",
|
||||
'```', '',
|
||||
'Older, superseded result files live in [`archive/`](archive/).', '',
|
||||
]
|
||||
with open(filepath, 'w') as f:
|
||||
f.write('# PraisonAI Agents - Benchmark Results\n\n')
|
||||
f.write(f'**Generated:** {datetime.now().strftime("%Y-%m-%d %H:%M:%S")}\n')
|
||||
f.write('**Methodology:** 10 warmup runs + 1000 measured iterations\n\n')
|
||||
f.write('## Agent Instantiation Time\n\n')
|
||||
f.write('| Framework | Avg Time (μs) | Relative |\n')
|
||||
f.write('|-----------|---------------|----------|\n')
|
||||
f.write('\n'.join(table_rows) + '\n')
|
||||
f.write('\n## Package Versions\n\n')
|
||||
f.write('| Package | Version |\n')
|
||||
f.write('|---------|--------|\n')
|
||||
for pkg, ver in versions.items():
|
||||
f.write(f'| {pkg} | {ver} |\n')
|
||||
f.write('\n## How to Reproduce\n\n')
|
||||
f.write('```bash\n')
|
||||
f.write('cd praisonai-agents\n')
|
||||
f.write('python benchmarks/performance_benchmark.py\n')
|
||||
f.write('```\n')
|
||||
|
||||
f.write('\n'.join(lines))
|
||||
|
||||
print(f'\nResults saved to: {filepath}')
|
||||
|
||||
# Update README.md
|
||||
readme_path = os.path.join(benchmarks_dir, '..', '..', '..', 'README.md')
|
||||
if os.path.exists(readme_path):
|
||||
with open(readme_path, 'r') as f:
|
||||
content = f.read()
|
||||
|
||||
new_table = '''| Framework | Avg Time (μs) | Relative |
|
||||
|-----------|---------------|----------|
|
||||
''' + '\n'.join(table_rows)
|
||||
|
||||
pattern = r'(\| Framework \| Avg Time \(μs\) \| Relative \|\n\|[-|]+\|\n)(\|[^\n]+\|\n)+'
|
||||
|
||||
if re.search(pattern, content):
|
||||
content = re.sub(pattern, new_table + '\n', content)
|
||||
|
||||
with open(readme_path, 'w') as f:
|
||||
f.write(content)
|
||||
|
||||
print(f'README.md updated: {readme_path}')
|
||||
return filepath
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
def main():
|
||||
parser = argparse.ArgumentParser(description='PraisonAI Agents - Comprehensive Performance Benchmark')
|
||||
parser.add_argument('--save', action='store_true', help='Save results to file')
|
||||
parser.add_argument('--save', action='store_true', help='Save results to BENCHMARK_RESULTS.md')
|
||||
parser.add_argument('--iterations', type=int, default=1000, help='Measured iterations (default: 1000)')
|
||||
parser.add_argument('--warmup', type=int, default=10, help='Warmup runs (default: 10)')
|
||||
parser.add_argument('--only', nargs='+', choices=list(FRAMEWORKS), help='Benchmark only these frameworks')
|
||||
args = parser.parse_args()
|
||||
|
||||
if args.iterations <= 0 or args.warmup < 0:
|
||||
parser.error('--iterations must be > 0 and --warmup must be >= 0')
|
||||
|
||||
machine = get_machine_info()
|
||||
|
||||
print("="*70)
|
||||
print("PraisonAI Agents - Comprehensive Performance Benchmark")
|
||||
print("="*70)
|
||||
@@ -376,156 +339,48 @@ if __name__ == "__main__":
|
||||
print(" 1. Agent instantiation time (how fast an agent object is created)")
|
||||
print(" 2. Memory footprint (how much memory an agent uses)")
|
||||
print("\nMethodology:")
|
||||
print(" - 10 warmup runs (not measured)")
|
||||
print(" - 1000 measured iterations")
|
||||
print(f" - {args.warmup} warmup runs (not measured)")
|
||||
print(f" - {args.iterations} measured iterations")
|
||||
print(" - tracemalloc for memory measurement")
|
||||
print(" - Baseline subtraction for accuracy")
|
||||
print("\nFrameworks available:")
|
||||
print(" - PraisonAI: ✅")
|
||||
print(f" - Agno: ✅")
|
||||
print(f" - PydanticAI: {'✅' if PYDANTIC_AVAILABLE else '❌'}")
|
||||
print(f" - OpenAI Agents SDK: {'✅' if OPENAI_AGENTS_AVAILABLE else '❌'}")
|
||||
print(f" - LangGraph: {'✅' if LANGGRAPH_AVAILABLE else '❌'}")
|
||||
print(f" - CrewAI: {'✅' if CREWAI_AVAILABLE else '❌'}")
|
||||
|
||||
print("\nMachine:")
|
||||
for k, v in machine.items():
|
||||
print(f" - {k}: {v}")
|
||||
|
||||
results = {}
|
||||
|
||||
# Test PraisonAI first
|
||||
try:
|
||||
print("\n" + "="*70)
|
||||
print("Testing PraisonAI Agents")
|
||||
print("="*70)
|
||||
results['praisonai'] = run_benchmark(
|
||||
instantiate_praisonai_agent,
|
||||
"PraisonAI",
|
||||
iterations=1000,
|
||||
warmup=10
|
||||
)
|
||||
except Exception as e:
|
||||
print(f"Error testing PraisonAI: {e}")
|
||||
|
||||
try:
|
||||
print("\n" + "="*70)
|
||||
print("Testing PraisonAI Agents (LiteLLM)")
|
||||
print("="*70)
|
||||
results['praisonai_litellm'] = run_benchmark(
|
||||
instantiate_praisonai_litellm_agent,
|
||||
"PraisonAI (LiteLLM)",
|
||||
iterations=1000,
|
||||
warmup=10
|
||||
)
|
||||
except Exception as e:
|
||||
print(f"Error testing PraisonAI (LiteLLM): {e}")
|
||||
|
||||
# Test Agno
|
||||
try:
|
||||
print("\n" + "="*70)
|
||||
print("Testing Agno")
|
||||
print("="*70)
|
||||
results['agno'] = run_benchmark(
|
||||
instantiate_agno_agent,
|
||||
"Agno",
|
||||
iterations=1000,
|
||||
warmup=10
|
||||
)
|
||||
except Exception as e:
|
||||
print(f"Error testing Agno: {e}")
|
||||
|
||||
# Test PydanticAI
|
||||
if PYDANTIC_AVAILABLE:
|
||||
for key in args.only or FRAMEWORKS:
|
||||
display_name, _, builder = FRAMEWORKS[key]
|
||||
try:
|
||||
print("\n" + "="*70)
|
||||
print("Testing PydanticAI")
|
||||
print("="*70)
|
||||
results['pydantic'] = run_benchmark(
|
||||
instantiate_pydantic_agent,
|
||||
"PydanticAI Agent Instantiation",
|
||||
iterations=1000,
|
||||
warmup=10
|
||||
)
|
||||
except Exception as e:
|
||||
print(f"Error testing PydanticAI: {e}")
|
||||
|
||||
# Test OpenAI Agents SDK
|
||||
if OPENAI_AGENTS_AVAILABLE:
|
||||
create = builder()
|
||||
except ImportError:
|
||||
print(f"\nSkipping {display_name}: not installed")
|
||||
continue
|
||||
try:
|
||||
print("\n" + "="*70)
|
||||
print("Testing OpenAI Agents SDK")
|
||||
print("="*70)
|
||||
results['openai_agents'] = run_benchmark(
|
||||
instantiate_openai_agents,
|
||||
"OpenAI Agents SDK Instantiation",
|
||||
iterations=1000,
|
||||
warmup=10
|
||||
)
|
||||
results[key] = run_benchmark(create, display_name, iterations=args.iterations, warmup=args.warmup)
|
||||
except Exception as e:
|
||||
print(f"Error testing OpenAI Agents SDK: {e}")
|
||||
|
||||
# Test LangGraph
|
||||
if LANGGRAPH_AVAILABLE:
|
||||
try:
|
||||
print("\n" + "="*70)
|
||||
print("Testing LangGraph")
|
||||
print("="*70)
|
||||
results['langgraph'] = run_benchmark(
|
||||
instantiate_langgraph_agent,
|
||||
"LangGraph Agent Instantiation",
|
||||
iterations=1000,
|
||||
warmup=10
|
||||
)
|
||||
except Exception as e:
|
||||
print(f"Error testing LangGraph: {e}")
|
||||
|
||||
# Test CrewAI
|
||||
if CREWAI_AVAILABLE:
|
||||
try:
|
||||
print("\n" + "="*70)
|
||||
print("Testing CrewAI")
|
||||
print("="*70)
|
||||
results['crewai'] = run_benchmark(
|
||||
instantiate_crewai_agent,
|
||||
"CrewAI Agent Instantiation",
|
||||
iterations=1000,
|
||||
warmup=10
|
||||
)
|
||||
except Exception as e:
|
||||
print(f"Error testing CrewAI: {e}")
|
||||
|
||||
print(f"Error testing {display_name}: {e}")
|
||||
|
||||
if not results:
|
||||
print("\nNo framework could be benchmarked.")
|
||||
return 1
|
||||
|
||||
# Final Comparison Table
|
||||
print("\n" + "="*70)
|
||||
print("RESULTS SUMMARY")
|
||||
print("="*70)
|
||||
|
||||
if 'praisonai' in results:
|
||||
baseline = results['praisonai']
|
||||
|
||||
print(f"\n{'Framework':<25} {'Avg Time (μs)':<15} {'Avg Memory (KiB)':<18} {'Relative':<15}")
|
||||
print("-" * 73)
|
||||
|
||||
# Sort by time
|
||||
sorted_results = sorted(results.items(), key=lambda x: x[1].avg_run_time)
|
||||
|
||||
for name, result in sorted_results:
|
||||
time_ratio = result.avg_run_time / baseline.avg_run_time if baseline.avg_run_time > 0 else 0
|
||||
|
||||
time_str = f"{time_ratio:.2f}x"
|
||||
|
||||
display_name = {
|
||||
'agno': 'Agno',
|
||||
'praisonai': 'PraisonAI',
|
||||
'praisonai_litellm': 'PraisonAI (LiteLLM)',
|
||||
'pydantic': 'PydanticAI',
|
||||
'openai_agents': 'OpenAI Agents SDK',
|
||||
'langgraph': 'LangGraph',
|
||||
'crewai': 'CrewAI'
|
||||
}.get(name, name)
|
||||
|
||||
print(f"{display_name:<25} {result.avg_run_time*1e6:<15.2f} {result.avg_memory_usage*1024:<18.2f} {time_str:<15}")
|
||||
|
||||
print("\n" + "="*70)
|
||||
|
||||
# Save results to files
|
||||
if args.save:
|
||||
save_benchmark_results(results, baseline)
|
||||
else:
|
||||
print('\nResults not saved (use --save flag to save results to file)')
|
||||
print(f"\n{'Framework':<25} {'Avg Time (μs)':<15} {'Median (μs)':<13} {'Avg Memory (KiB)':<18} {'Relative':<15}")
|
||||
print("-" * 86)
|
||||
for key, result, ratio in rank_results(results):
|
||||
print(f"{FRAMEWORKS[key][0]:<25} {result.avg_run_time*1e6:<15.2f} {result.median_run_time*1e6:<13.2f} "
|
||||
f"{result.avg_memory_usage*1024:<18.2f} {ratio:.2f}x")
|
||||
print("\n" + "="*70)
|
||||
|
||||
if args.save:
|
||||
save_benchmark_results(results, machine, args.iterations, args.warmup)
|
||||
else:
|
||||
print('\nResults not saved (use --save flag to save results to file)')
|
||||
return 0
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
raise SystemExit(main())
|
||||
@@ -188,15 +188,17 @@ def get_package_versions():
|
||||
return versions
|
||||
|
||||
|
||||
def save_results(results: dict, filename: str = 'BENCHMARK_RESULTS.md', *,
|
||||
def save_results(results: dict, filename: str = 'SIMPLE_BENCHMARK_RESULTS.md', *,
|
||||
iterations: int = ITERATIONS, log=print):
|
||||
"""Save benchmark results to a markdown file and update README.
|
||||
"""Save benchmark results to a markdown file.
|
||||
|
||||
Writes ``SIMPLE_BENCHMARK_RESULTS.md`` by default so the curated
|
||||
``BENCHMARK_RESULTS.md`` (from ``performance_benchmark.py``) is never
|
||||
overwritten.
|
||||
|
||||
Status messages are emitted via ``log`` (defaults to ``print``). Pass a
|
||||
stderr logger when stdout must stay machine-readable (e.g. ``--json``).
|
||||
"""
|
||||
import os
|
||||
import re
|
||||
from datetime import datetime
|
||||
|
||||
filepath = os.path.join(os.path.dirname(__file__), filename)
|
||||
@@ -217,7 +219,6 @@ def save_results(results: dict, filename: str = 'BENCHMARK_RESULTS.md', *,
|
||||
ratio_str = f'{ratio:,.0f}x' if ratio >= 100 else f'{ratio:.2f}x'
|
||||
table_rows.append(f'| {name} | {avg_str} | {ratio_str} |')
|
||||
|
||||
# Save to BENCHMARK_RESULTS.md
|
||||
with open(filepath, 'w') as f:
|
||||
f.write('# PraisonAI Agents - Benchmark Results\n\n')
|
||||
f.write(f'**Generated:** {datetime.now().strftime("%Y-%m-%d %H:%M:%S")}\n')
|
||||
@@ -239,39 +240,9 @@ def save_results(results: dict, filename: str = 'BENCHMARK_RESULTS.md', *,
|
||||
f.write('```\n')
|
||||
|
||||
log(f'\nResults saved to: {filepath}')
|
||||
|
||||
# Also update the main README.md
|
||||
readme_path = os.path.join(os.path.dirname(__file__), '..', '..', '..', 'README.md')
|
||||
if os.path.exists(readme_path):
|
||||
update_readme(readme_path, table_rows, log=log)
|
||||
|
||||
return filepath
|
||||
|
||||
|
||||
def update_readme(readme_path: str, table_rows: list, *, log=print):
|
||||
"""Update the performance section in README.md with latest results."""
|
||||
import re
|
||||
|
||||
with open(readme_path, 'r') as f:
|
||||
content = f.read()
|
||||
|
||||
# Build the new table
|
||||
new_table = '''| Framework | Avg Time (μs) | Relative |
|
||||
|-----------|---------------|----------|
|
||||
''' + '\n'.join(table_rows)
|
||||
|
||||
# Pattern to match the performance table
|
||||
pattern = r'(\| Framework \| Avg Time \(μs\) \| Relative \|\n\|[-|]+\|\n)(\|[^\n]+\|\n)+'
|
||||
|
||||
if re.search(pattern, content):
|
||||
content = re.sub(pattern, new_table + '\n', content)
|
||||
|
||||
with open(readme_path, 'w') as f:
|
||||
f.write(content)
|
||||
|
||||
log(f'README.md updated: {readme_path}')
|
||||
|
||||
|
||||
def main():
|
||||
parser = argparse.ArgumentParser(description='PraisonAI Agents - Performance Benchmark')
|
||||
parser.add_argument('--save', action='store_true', help='Save results to file')
|
||||
|
||||
@@ -1,5 +1,5 @@
|
||||
{
|
||||
"praisonai_instantiation_us_max": 450,
|
||||
"praisonai_instantiation_us_max": 150,
|
||||
"iterations": 100,
|
||||
"warmup": 10
|
||||
}
|
||||
@@ -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
|
||||
@@ -457,11 +457,20 @@ class Agent(GoalLoopMixin, SteeringMixin, SandboxMixin, SkillReviewMixin, Unifie
|
||||
self._wire_deferred_history_callback()
|
||||
|
||||
def _ensure_loop_guard(self):
|
||||
"""Lazy-create loop guard on first tool execution."""
|
||||
"""Lazy-create loop guard on first tool execution.
|
||||
|
||||
Inside a chat turn, returns that turn's own guard (same config) so
|
||||
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))
|
||||
return self._loop_guard
|
||||
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
|
||||
|
||||
def _ensure_loop_detector(self):
|
||||
"""Lazy-create the result-aware tool-loop detector on first tool call.
|
||||
@@ -1908,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
|
||||
@@ -3278,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"):
|
||||
@@ -4104,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.
|
||||
"""
|
||||
@@ -4722,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
|
||||
|
||||
@@ -4761,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
|
||||
@@ -4791,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():
|
||||
@@ -4824,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",
|
||||
@@ -4838,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.
|
||||
@@ -5003,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",
|
||||
@@ -5040,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:
|
||||
@@ -5062,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",
|
||||
@@ -5081,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",
|
||||
@@ -5109,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:
|
||||
@@ -5221,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
|
||||
@@ -5248,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
|
||||
@@ -5268,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():
|
||||
@@ -5302,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",
|
||||
@@ -5316,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.
|
||||
@@ -5470,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,
|
||||
@@ -5493,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:
|
||||
@@ -5509,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,
|
||||
@@ -5522,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,
|
||||
@@ -5543,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:
|
||||
@@ -6524,7 +6582,8 @@ Summary:"""
|
||||
# In async context - prefer async hook if available
|
||||
if hasattr(self._memory_instance, 'aon_memory_write'):
|
||||
# Schedule async hook without blocking (fire and forget)
|
||||
asyncio.create_task(self._memory_instance.aon_memory_write(action, memory_type, content, metadata))
|
||||
from ..utils.async_bridge import spawn_background
|
||||
spawn_background(self._memory_instance.aon_memory_write(action, memory_type, content, metadata), "aon_memory_write")
|
||||
elif hasattr(self._memory_instance, 'on_memory_write'):
|
||||
self._memory_instance.on_memory_write(action, memory_type, content, metadata)
|
||||
except RuntimeError:
|
||||
@@ -6656,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]'")
|
||||
@@ -7085,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,8 @@ and response processing. Extracted from agent.py for maintainability.
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import contextlib
|
||||
import contextvars
|
||||
import os
|
||||
import re
|
||||
import time
|
||||
@@ -24,12 +26,25 @@ logger = logging.getLogger(__name__)
|
||||
|
||||
_turns_lock = threading.Lock()
|
||||
|
||||
# Per-turn loop guard: ``[agent, LoopGuard | None]`` bound by ``_enter_turn``. Copied
|
||||
# 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):
|
||||
"""Mark a ``chat``/``achat`` turn in flight and bind its cancel token to this context."""
|
||||
with _turns_lock:
|
||||
agent._turns_in_flight = getattr(agent, "_turns_in_flight", 0) + 1
|
||||
return _active_turn_cancel.set((agent, cancel))
|
||||
return _active_turn_cancel.set((agent, cancel)), _turn_loop_guard.set([agent, None])
|
||||
|
||||
|
||||
def _exit_turn(agent, ctx_token):
|
||||
@@ -38,18 +53,36 @@ def _exit_turn(agent, ctx_token):
|
||||
return
|
||||
with _turns_lock:
|
||||
agent._turns_in_flight = max(getattr(agent, "_turns_in_flight", 1) - 1, 0)
|
||||
try:
|
||||
_active_turn_cancel.reset(ctx_token)
|
||||
except (ValueError, LookupError):
|
||||
_active_turn_cancel.set(None)
|
||||
cancel_token, guard_token = ctx_token
|
||||
for var, token in ((_active_turn_cancel, cancel_token), (_turn_loop_guard, guard_token)):
|
||||
try:
|
||||
var.reset(token)
|
||||
except (ValueError, LookupError):
|
||||
var.set(None)
|
||||
|
||||
|
||||
def _current_turn_loop_guard(agent):
|
||||
"""Return this turn's loop guard for ``agent`` (lazily created, sharing the agent
|
||||
guard's config), or ``None`` when not inside one of ``agent``'s chat turns."""
|
||||
slot = _turn_loop_guard.get()
|
||||
if slot is None or slot[0] is not agent:
|
||||
return None
|
||||
with _turns_lock:
|
||||
if slot[1] is None:
|
||||
from ..escalation.loop_guard import LoopGuard
|
||||
slot[1] = LoopGuard(agent._loop_guard.config)
|
||||
return slot[1]
|
||||
|
||||
|
||||
def _reset_loop_guard_for_turn(agent):
|
||||
"""Reset per-turn loop-guard limits unless another turn on this agent is still
|
||||
running (the guard is shared; resetting it would restart that turn's limits).
|
||||
Never force-creates the guard."""
|
||||
"""Start fresh per-turn loop-guard limits. Inside a chat turn this drops the turn's
|
||||
own guard; outside one it resets the agent-level guard. Never force-creates a guard."""
|
||||
slot = _turn_loop_guard.get()
|
||||
if slot is not None and slot[0] is agent:
|
||||
slot[1] = None
|
||||
return
|
||||
guard = getattr(agent, "_loop_guard", None)
|
||||
if guard is not None and getattr(agent, "_turns_in_flight", 0) <= 1:
|
||||
if guard is not None:
|
||||
guard.reset_turn()
|
||||
|
||||
|
||||
@@ -73,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))
|
||||
@@ -97,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:
|
||||
@@ -133,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):
|
||||
@@ -1888,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.
|
||||
|
||||
@@ -2051,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()
|
||||
|
||||
@@ -2250,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,
|
||||
@@ -2272,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
|
||||
@@ -2312,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)
|
||||
@@ -2966,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
|
||||
@@ -3013,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)
|
||||
@@ -3424,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()
|
||||
@@ -3493,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):
|
||||
@@ -3825,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
|
||||
@@ -3856,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"
|
||||
@@ -3864,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"
|
||||
@@ -4216,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
|
||||
@@ -4336,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):
|
||||
@@ -4552,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:
|
||||
@@ -4580,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"
|
||||
@@ -5094,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:
|
||||
@@ -5291,6 +5492,36 @@ Output MUST be JSON with 'reflection' and 'satisfactory'.
|
||||
|
||||
def _start_stream(self, prompt: str, **kwargs) -> Generator[str, None, None]:
|
||||
"""Own the durable context for the complete streaming generator."""
|
||||
# A stream is a turn too: give it its own loop guard (as chat/achat get
|
||||
# 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:
|
||||
# _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):
|
||||
_turn_loop_guard.set(None)
|
||||
|
||||
def _start_stream_turn(self, prompt: str, **kwargs) -> Generator[str, None, None]:
|
||||
# Tools go to their sandbox before the first model call, once per
|
||||
# agent. Placed on chat/achat/_start_stream rather than run/start
|
||||
# because AgentTeam and AgentFlow call these three directly -- an
|
||||
@@ -5309,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
|
||||
@@ -5330,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"]
|
||||
@@ -6150,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:
|
||||
@@ -6162,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):
|
||||
"""
|
||||
@@ -6631,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}")
|
||||
@@ -1777,7 +1777,7 @@ Focus on creating clear, structured metadata that will help with intelligent fil
|
||||
|
||||
request = urllib.request.Request(api_url, headers=headers)
|
||||
|
||||
with urllib.request.urlopen(request) as response:
|
||||
with urllib.request.urlopen(request, timeout=15) as response:
|
||||
file_data = json.loads(response.read())
|
||||
|
||||
if file_data.get('content') and file_data.get('encoding') == 'base64':
|
||||
@@ -2350,7 +2350,7 @@ This PRP provides:
|
||||
import urllib.request
|
||||
import json
|
||||
|
||||
with urllib.request.urlopen(api_url) as response:
|
||||
with urllib.request.urlopen(api_url, timeout=15) as response:
|
||||
data = json.loads(response.read())
|
||||
|
||||
file_structure = []
|
||||
|
||||
@@ -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
|
||||
@@ -1970,24 +1959,15 @@ Write the complete compiled report:"""
|
||||
# (a never-raised local class never matches a real exception).
|
||||
class _CircuitBreakerException(Exception):
|
||||
pass
|
||||
_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")
|
||||
@@ -2002,7 +1982,17 @@ Write the complete compiled report:"""
|
||||
super().__init__(error_dict.get("error", "Tool execution failed"))
|
||||
|
||||
async def _invoke_for_breaker():
|
||||
r = await _invoke()
|
||||
# Caller errors (bad input) say nothing about the tool's health:
|
||||
# recorded as neither failure nor success (sync parity), so they
|
||||
# can neither open the breaker for other callers nor close a
|
||||
# half-open one. Unwrapped in _invoke_guarded.
|
||||
try:
|
||||
r = await _invoke()
|
||||
except _CALLER_ERROR_TYPES as e:
|
||||
raise _NeutralOutcome(e)
|
||||
if isinstance(r, dict) and r.get("error") and \
|
||||
r.get("_praison_retryable") is False:
|
||||
raise _NeutralOutcome(r)
|
||||
if isinstance(r, dict) and r.get("error") and \
|
||||
not r.get("approval_denied") and \
|
||||
not r.get("permission_denied") and \
|
||||
@@ -2026,6 +2016,11 @@ Write the complete compiled report:"""
|
||||
except _ToolFailure as tf:
|
||||
# Failure was counted by the breaker; return the original dict.
|
||||
return tf.error_dict
|
||||
except _NeutralOutcome as n:
|
||||
neutral = n.value
|
||||
if isinstance(neutral, BaseException):
|
||||
raise neutral
|
||||
return neutral
|
||||
|
||||
def _circuit_open_result():
|
||||
# Record the rejection as a failure so repeated open-circuit
|
||||
@@ -2034,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
|
||||
@@ -2200,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.
|
||||
|
||||
@@ -630,6 +630,27 @@ class Handoff:
|
||||
None, lambda: self.agent.chat(prompt, tools=effective_tools)
|
||||
)
|
||||
|
||||
def _trace_handoff(self, source_agent: 'Agent', context: Optional[Dict[str, Any]], reason: Optional[str] = None) -> None:
|
||||
"""Record the hand-off on the replay trace (no-op when tracing is off)."""
|
||||
from ..trace.context_events import get_context_emitter
|
||||
emitter = get_context_emitter()
|
||||
if not emitter._enabled:
|
||||
return
|
||||
context = dict(context or {})
|
||||
if reason is None:
|
||||
reason = context.get("reason") or context.get("task") or context.get("prompt")
|
||||
policy = getattr(self.config, "context_policy", None)
|
||||
emitter.agent_handoff(
|
||||
getattr(source_agent, "name", None) or "agent",
|
||||
getattr(self.agent, "name", None),
|
||||
reason=emitter._preview(reason) if reason is not None else None,
|
||||
context_passed={
|
||||
"context_policy": getattr(policy, "value", policy),
|
||||
"args": context,
|
||||
},
|
||||
kind="handoff",
|
||||
)
|
||||
|
||||
def _check_safety(self, source_agent: 'Agent') -> None:
|
||||
"""
|
||||
Check safety constraints before handoff.
|
||||
@@ -778,6 +799,7 @@ class Handoff:
|
||||
# Track handoff chain
|
||||
_push_handoff(source_agent.name)
|
||||
pushed = True
|
||||
self._trace_handoff(source_agent, kwargs, reason=prompt)
|
||||
|
||||
# Execute on_handoff callback
|
||||
self._execute_callback(self.config.on_handoff or self.on_handoff, source_agent, kwargs)
|
||||
@@ -888,6 +910,7 @@ class Handoff:
|
||||
self._check_safety(source_agent)
|
||||
_push_handoff(source_agent.name)
|
||||
pushed = True
|
||||
self._trace_handoff(source_agent, kwargs, reason=prompt)
|
||||
|
||||
# Execute callback
|
||||
self._execute_callback(self.config.on_handoff or self.on_handoff, source_agent, kwargs)
|
||||
@@ -1016,6 +1039,7 @@ class Handoff:
|
||||
# Track handoff chain
|
||||
_push_handoff(source_agent.name)
|
||||
pushed = True
|
||||
self._trace_handoff(source_agent, kwargs, reason=None)
|
||||
|
||||
# Execute on_handoff callback
|
||||
self._execute_callback(self.config.on_handoff or self.on_handoff, source_agent, kwargs)
|
||||
@@ -1632,6 +1656,7 @@ class TypedHandoff(Handoff, Generic[T]):
|
||||
# Track handoff chain
|
||||
_push_handoff(source_agent.name)
|
||||
pushed = True
|
||||
self._trace_handoff(source_agent, kwargs, reason=None)
|
||||
|
||||
# Execute on_handoff callback
|
||||
self._execute_callback(self.config.on_handoff or self.on_handoff, source_agent, kwargs)
|
||||
@@ -1747,6 +1772,7 @@ class TypedHandoff(Handoff, Generic[T]):
|
||||
self._check_safety(source_agent)
|
||||
_push_handoff(source_agent.name)
|
||||
pushed = True
|
||||
self._trace_handoff(source_agent, kwargs, reason=None)
|
||||
|
||||
# Execute callback
|
||||
self._execute_callback(self.config.on_handoff or self.on_handoff, source_agent, kwargs)
|
||||
|
||||
@@ -70,6 +70,9 @@ class ImageAgent(Agent):
|
||||
api_version=api_version
|
||||
)
|
||||
|
||||
# Only an explicit llm=/model= names an image model; the base Agent's chat default must not leak into image calls
|
||||
self._image_llm_set = llm is not None or kwargs.get("model") is not None
|
||||
|
||||
# Lazy load litellm
|
||||
self._litellm = None
|
||||
# Configure logging based on verbose level
|
||||
@@ -124,15 +127,21 @@ class ImageAgent(Agent):
|
||||
"Please install it with: pip install litellm"
|
||||
)
|
||||
return self._litellm
|
||||
|
||||
def _image_model(self, default: str) -> str:
|
||||
"""Return the explicit ``llm=`` model name, else an image-capable ``default``."""
|
||||
if not self._image_llm_set:
|
||||
return default
|
||||
model_info = self.llm_model
|
||||
return model_info.model if hasattr(model_info, 'model') else str(model_info)
|
||||
|
||||
def generate_image(self, prompt: str, **kwargs) -> Dict[str, Any]:
|
||||
"""Generate an image based on the provided prompt."""
|
||||
# Merge default config with any provided kwargs
|
||||
config = self.image_config.dict(exclude_none=True)
|
||||
config.update(kwargs)
|
||||
|
||||
# Get the model name robustly from the parent Agent's property
|
||||
model_info = self.llm_model
|
||||
model_name = model_info.model if hasattr(model_info, 'model') else str(model_info)
|
||||
model_name = self._image_model("dall-e-3")
|
||||
|
||||
# Use the model name in config
|
||||
config['model'] = model_name
|
||||
@@ -154,6 +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":
|
||||
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(),
|
||||
@@ -260,7 +277,7 @@ class ImageAgent(Agent):
|
||||
import litellm
|
||||
litellm.telemetry = False
|
||||
|
||||
model_name = self.llm or self.model or "dall-e-2"
|
||||
model_name = self._image_model("dall-e-2")
|
||||
|
||||
config = {
|
||||
"model": model_name,
|
||||
@@ -335,7 +352,7 @@ class ImageAgent(Agent):
|
||||
import litellm
|
||||
litellm.telemetry = False
|
||||
|
||||
model_name = self.llm or self.model or "dall-e-2"
|
||||
model_name = self._image_model("dall-e-2")
|
||||
|
||||
config = {
|
||||
"model": model_name,
|
||||
|
||||
@@ -301,10 +301,9 @@ class RouterAgent(Agent):
|
||||
Returns:
|
||||
Task execution result
|
||||
"""
|
||||
# Estimate context size in tokens (rough estimate: ~4 chars per token)
|
||||
# This is a simplified heuristic; actual tokenization varies by model
|
||||
text_length = len(task_description) + (len(context) if context else 0)
|
||||
context_size = text_length // 4 # Approximate token count
|
||||
# Estimate context size in tokens (canonical heuristic, CJK-aware)
|
||||
from ..context.tokens import estimate_tokens_heuristic
|
||||
context_size = estimate_tokens_heuristic(task_description + (context or ""))
|
||||
|
||||
# Select the best model for this task
|
||||
selected_model = self._select_model_for_task(
|
||||
|
||||
@@ -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,
|
||||
)
|
||||
|
||||
@@ -2174,22 +2269,25 @@ class ToolExecutionMixin:
|
||||
Returns:
|
||||
Estimated minimum cost in USD for the call.
|
||||
"""
|
||||
prompt_chars = 0
|
||||
texts = []
|
||||
if messages:
|
||||
for _msg in messages:
|
||||
_content = _msg.get("content") if isinstance(_msg, dict) else None
|
||||
if isinstance(_content, str):
|
||||
prompt_chars += len(_content)
|
||||
texts.append(_content)
|
||||
elif isinstance(_content, list):
|
||||
for _part in _content:
|
||||
if isinstance(_part, dict):
|
||||
_text = _part.get("text")
|
||||
if isinstance(_text, str):
|
||||
prompt_chars += len(_text)
|
||||
texts.append(_text)
|
||||
elif _content is not None:
|
||||
prompt_chars += len(str(_content))
|
||||
texts.append(str(_content))
|
||||
|
||||
prompt_tokens = prompt_chars // 4
|
||||
from ..context.tokens import estimate_tokens_heuristic
|
||||
prompt_text = "".join(texts)
|
||||
prompt_chars = len(prompt_text)
|
||||
prompt_tokens = estimate_tokens_heuristic(prompt_text)
|
||||
if max_tokens:
|
||||
completion_tokens = int(max_tokens)
|
||||
elif prompt_chars > 0:
|
||||
@@ -3224,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):
|
||||
@@ -3298,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
|
||||
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")
|
||||
@@ -3306,27 +3409,20 @@ 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: _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 \
|
||||
result.get("_praison_retryable") is False:
|
||||
raise _NeutralOutcome(result)
|
||||
# Convert error dicts to exceptions so circuit breaker can detect failures
|
||||
# Don't treat approval/permission denials as circuit breaker failures
|
||||
if isinstance(result, dict) and result.get("error") and \
|
||||
@@ -3345,6 +3441,8 @@ class ToolExecutionMixin:
|
||||
|
||||
try:
|
||||
return breaker.call(_tool_wrapper)
|
||||
except _NeutralOutcome as n:
|
||||
return n.value
|
||||
except Exception as e:
|
||||
# Check if this is our sentinel exception
|
||||
if hasattr(e, 'error_dict'):
|
||||
@@ -3355,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``.
|
||||
@@ -3811,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:
|
||||
|
||||
@@ -336,17 +336,27 @@ class PermissionAllowlist:
|
||||
if not perm.session_only
|
||||
}
|
||||
}
|
||||
from ..utils.atomic_io import atomic_write_json
|
||||
os.makedirs(os.path.dirname(filepath) or ".", exist_ok=True)
|
||||
with open(filepath, "w") as f:
|
||||
json.dump(data, f, indent=2)
|
||||
atomic_write_json(filepath, data, indent=2)
|
||||
|
||||
@classmethod
|
||||
def load(cls, filepath: str) -> "PermissionAllowlist":
|
||||
allowlist = cls()
|
||||
if not os.path.exists(filepath):
|
||||
return allowlist
|
||||
with open(filepath, "r") as f:
|
||||
data = json.load(f)
|
||||
try:
|
||||
with open(filepath, "r") as f:
|
||||
data = json.load(f)
|
||||
except json.JSONDecodeError:
|
||||
# Fail closed: a corrupt allowlist grants nothing.
|
||||
from ..utils.atomic_io import quarantine_corrupt_file
|
||||
try:
|
||||
dest = quarantine_corrupt_file(filepath)
|
||||
logger.warning("Corrupt permission allowlist %s moved to %s; starting empty", filepath, dest)
|
||||
except OSError:
|
||||
logger.warning("Corrupt permission allowlist %s; ignoring it", filepath)
|
||||
return allowlist
|
||||
for tool_name, tool_data in data.get("tools", {}).items():
|
||||
allowlist.add_tool(tool_name, paths=tool_data.get("allowed_paths", []), session_only=tool_data.get("session_only", False))
|
||||
return allowlist
|
||||
|
||||
@@ -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()),
|
||||
@@ -91,6 +95,24 @@ PERMISSION_PRESETS = {
|
||||
"off": frozenset(),
|
||||
}
|
||||
|
||||
# Optional host hook (Issue #5912): a runner that bounds concurrent work (e.g.
|
||||
# the gateway's admission gate) binds a zero-arg factory here for the duration
|
||||
# of an admitted turn. The registry enters ``factory()`` — a sync *and* async
|
||||
# context manager — around the backend's human wait so a turn parked on a
|
||||
# person does not occupy a run slot. Unbound (the default) is a no-op.
|
||||
HUMAN_WAIT_SCOPE: contextvars.ContextVar[Optional[Callable[[], object]]] = contextvars.ContextVar(
|
||||
"praisonai_approval_human_wait_scope", default=None
|
||||
)
|
||||
|
||||
|
||||
def _human_wait_scope() -> object:
|
||||
factory = HUMAN_WAIT_SCOPE.get()
|
||||
if factory is None:
|
||||
import contextlib
|
||||
return contextlib.nullcontext()
|
||||
return factory()
|
||||
|
||||
|
||||
class ApprovalRegistry:
|
||||
"""Per-agent approval configuration.
|
||||
|
||||
@@ -671,28 +693,41 @@ class ApprovalRegistry:
|
||||
liveness=liveness,
|
||||
)
|
||||
|
||||
# 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,
|
||||
@@ -745,25 +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:
|
||||
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]:
|
||||
|
||||
@@ -12,6 +12,7 @@ from enum import Enum
|
||||
from typing import Any, Callable, Dict, List, Optional, Protocol, runtime_checkable
|
||||
from dataclasses import dataclass, field
|
||||
from concurrent.futures import CancelledError, ThreadPoolExecutor, Future
|
||||
from concurrent.futures import TimeoutError as FuturesTimeoutError
|
||||
|
||||
from praisonaiagents._logging import get_logger
|
||||
|
||||
@@ -148,12 +149,18 @@ class BackgroundJobManager:
|
||||
"""
|
||||
if self._store is None:
|
||||
return
|
||||
try:
|
||||
self._store.upsert(job)
|
||||
except Exception as e: # noqa: BLE001 — persistence must never crash worker
|
||||
logger.debug(
|
||||
"Persisting job %s failed (non-fatal): %s", job.job_id, e,
|
||||
)
|
||||
# Check-and-write under the lock: a finished job whose id was reused
|
||||
# must not overwrite the newer job's persisted record.
|
||||
with self._lock:
|
||||
current = self._jobs.get(job.job_id)
|
||||
if current is not None and current is not job:
|
||||
return
|
||||
try:
|
||||
self._store.upsert(job)
|
||||
except Exception as e: # noqa: BLE001 — persistence must never crash worker
|
||||
logger.debug(
|
||||
"Persisting job %s failed (non-fatal): %s", job.job_id, e,
|
||||
)
|
||||
|
||||
@property
|
||||
def active_jobs(self) -> int:
|
||||
@@ -192,14 +199,21 @@ class BackgroundJobManager:
|
||||
|
||||
Returns:
|
||||
The job ID
|
||||
|
||||
Raises:
|
||||
ValueError: If ``job_id`` belongs to a job that is still pending
|
||||
or running. Ids of terminal jobs may be reused.
|
||||
"""
|
||||
job_id = job_id or str(uuid.uuid4())[:8]
|
||||
job_id = job_id or uuid.uuid4().hex[:12]
|
||||
|
||||
# Opportunistically evict old terminal jobs so long-lived hosts
|
||||
# (gateways, bots, servers) don't retain every finished result.
|
||||
now = time.time()
|
||||
if now - self._last_gc > self._GC_INTERVAL:
|
||||
self._last_gc = now
|
||||
with self._lock:
|
||||
run_gc = now - self._last_gc > self._GC_INTERVAL
|
||||
if run_gc:
|
||||
self._last_gc = now
|
||||
if run_gc:
|
||||
self.cleanup_completed()
|
||||
|
||||
job_info = JobInfo(
|
||||
@@ -208,8 +222,19 @@ class BackgroundJobManager:
|
||||
origin=dict(origin) if origin else {},
|
||||
)
|
||||
|
||||
# Result handle owned by the manager (not the executor) so cancel_job
|
||||
# can wake get_result waiters even while ``func`` is still running.
|
||||
done: Future = Future()
|
||||
with self._lock:
|
||||
existing = self._jobs.get(job_id)
|
||||
if existing is not None and existing.status in (
|
||||
JobStatus.PENDING, JobStatus.RUNNING,
|
||||
):
|
||||
raise ValueError(f"Job '{job_id}' is already active")
|
||||
# Publish job and future together so a concurrent cancel_job never
|
||||
# sees the new job paired with the old future.
|
||||
self._jobs[job_id] = job_info
|
||||
self._futures[job_id] = done
|
||||
# Persist the PENDING record before it starts so a crash between spawn
|
||||
# and first run still leaves a recoverable trace (reconciled as LOST).
|
||||
self._persist(job_info)
|
||||
@@ -224,10 +249,7 @@ class BackgroundJobManager:
|
||||
a missing hooks package or a raising subscriber never affects the
|
||||
job's own result.
|
||||
"""
|
||||
with self._lock:
|
||||
info = self._jobs.get(job_id)
|
||||
if info is None:
|
||||
return
|
||||
info = job_info
|
||||
try:
|
||||
from praisonaiagents.hooks import fire_hook
|
||||
|
||||
@@ -256,10 +278,7 @@ class BackgroundJobManager:
|
||||
_emit_job_completed()
|
||||
if on_complete is None:
|
||||
return
|
||||
with self._lock:
|
||||
info = self._jobs.get(job_id)
|
||||
if info is None:
|
||||
return
|
||||
info = job_info
|
||||
try:
|
||||
on_complete(info)
|
||||
# Record successful deliver-back so restart reconciliation does
|
||||
@@ -273,37 +292,88 @@ class BackgroundJobManager:
|
||||
job_id, e,
|
||||
)
|
||||
|
||||
def _run_job():
|
||||
# The worker writes only to its own ``job_info`` (never re-looked-up by
|
||||
# id) so a later job reusing the id, or eviction, cannot be corrupted.
|
||||
job = job_info
|
||||
|
||||
def _abort(error: BaseException, cancelled: bool = False) -> None:
|
||||
"""Move a still-active job to a terminal state and resolve ``done``."""
|
||||
with self._lock:
|
||||
job = self._jobs[job_id]
|
||||
if job.status not in (JobStatus.PENDING, JobStatus.RUNNING):
|
||||
return
|
||||
job.status = JobStatus.CANCELLED if cancelled else JobStatus.FAILED
|
||||
job.completed_at = time.time()
|
||||
if not cancelled:
|
||||
# Bare SystemExit()/KeyboardInterrupt() stringify to "".
|
||||
job.error = str(error) or type(error).__name__
|
||||
if cancelled:
|
||||
done.cancel()
|
||||
else:
|
||||
done.set_exception(error)
|
||||
self._persist(job)
|
||||
# Failures must reach deliver-back/hook subscribers like any other
|
||||
# failure; cancellations stay silent, matching cancel_job().
|
||||
if not cancelled:
|
||||
_fire_complete()
|
||||
|
||||
def _execute():
|
||||
with self._lock:
|
||||
if job.status == JobStatus.CANCELLED: # cancelled while pending
|
||||
return
|
||||
job.status = JobStatus.RUNNING
|
||||
job.started_at = time.time()
|
||||
self._persist(job)
|
||||
|
||||
|
||||
try:
|
||||
result = func()
|
||||
with self._lock:
|
||||
job = self._jobs[job_id]
|
||||
job.status = JobStatus.COMPLETED
|
||||
job.completed_at = time.time()
|
||||
job.result = result
|
||||
self._persist(job)
|
||||
_fire_complete()
|
||||
return result
|
||||
except Exception as e:
|
||||
with self._lock:
|
||||
job = self._jobs[job_id]
|
||||
if job.status == JobStatus.CANCELLED:
|
||||
return
|
||||
job.status = JobStatus.FAILED
|
||||
job.completed_at = time.time()
|
||||
job.error = str(e)
|
||||
self._persist(job)
|
||||
try:
|
||||
_fire_complete()
|
||||
finally:
|
||||
done.set_exception(e)
|
||||
return
|
||||
with self._lock:
|
||||
# Cancelled while running: discard the late outcome.
|
||||
if job.status == JobStatus.CANCELLED:
|
||||
return
|
||||
job.status = JobStatus.COMPLETED
|
||||
job.completed_at = time.time()
|
||||
job.result = result
|
||||
self._persist(job)
|
||||
try:
|
||||
_fire_complete()
|
||||
finally:
|
||||
# Resolve even if on_complete raised a BaseException.
|
||||
done.set_result(result)
|
||||
|
||||
def _run_job():
|
||||
try:
|
||||
_execute()
|
||||
except BaseException as e:
|
||||
# ``func`` raised a non-Exception (e.g. SystemExit); never
|
||||
# strand get_result callers.
|
||||
_abort(e)
|
||||
raise
|
||||
|
||||
future = self._executor.submit(_run_job)
|
||||
with self._lock:
|
||||
self._futures[job_id] = future
|
||||
|
||||
|
||||
try:
|
||||
executor_future = self._executor.submit(_run_job)
|
||||
except BaseException as e:
|
||||
# e.g. submit after shutdown(): fail the job so its id is reusable
|
||||
# and get_result raises instead of blocking forever.
|
||||
_abort(e)
|
||||
raise
|
||||
# A queued work item dropped by the executor never runs _run_job.
|
||||
executor_future.add_done_callback(
|
||||
lambda f: _abort(CancelledError(), cancelled=True) if f.cancelled() else None
|
||||
)
|
||||
|
||||
return job_id
|
||||
|
||||
def get_status(self, job_id: str) -> JobStatus:
|
||||
@@ -362,7 +432,10 @@ class BackgroundJobManager:
|
||||
Raises:
|
||||
KeyError: If job_id is not found
|
||||
TimeoutError: If timeout is reached before completion
|
||||
RuntimeError: If job failed
|
||||
CancelledError: If the job was cancelled
|
||||
Exception: If the job failed, the original exception is re-raised
|
||||
(``RuntimeError`` with the stored error once evicted from
|
||||
memory)
|
||||
"""
|
||||
with self._lock:
|
||||
future = self._futures.get(job_id)
|
||||
@@ -379,12 +452,18 @@ class BackgroundJobManager:
|
||||
|
||||
try:
|
||||
return future.result(timeout=timeout)
|
||||
except TimeoutError:
|
||||
except FuturesTimeoutError:
|
||||
# Distinct from builtin TimeoutError on Python 3.10.
|
||||
raise TimeoutError(f"Job {job_id} did not complete within {timeout} seconds")
|
||||
|
||||
def cancel_job(self, job_id: str) -> bool:
|
||||
"""
|
||||
Attempt to cancel a job.
|
||||
|
||||
A pending job never runs. A running job is marked ``CANCELLED``
|
||||
immediately and its eventual result is discarded (no ``on_complete``
|
||||
delivery); Python threads cannot be killed, so the worker thread is
|
||||
only released once ``func`` returns.
|
||||
|
||||
Args:
|
||||
job_id: The job ID to cancel
|
||||
@@ -393,19 +472,18 @@ class BackgroundJobManager:
|
||||
True if cancellation was successful, False otherwise
|
||||
"""
|
||||
with self._lock:
|
||||
if job_id not in self._futures:
|
||||
job = self._jobs.get(job_id)
|
||||
future = self._futures.get(job_id)
|
||||
if future is None or job is None or job.status not in (
|
||||
JobStatus.PENDING, JobStatus.RUNNING,
|
||||
):
|
||||
return False
|
||||
future = self._futures[job_id]
|
||||
cancelled = future.cancel()
|
||||
if cancelled:
|
||||
self._jobs[job_id].status = JobStatus.CANCELLED
|
||||
self._jobs[job_id].completed_at = time.time()
|
||||
job = self._jobs[job_id]
|
||||
else:
|
||||
job = None
|
||||
if job is not None:
|
||||
self._persist(job)
|
||||
return cancelled
|
||||
job.status = JobStatus.CANCELLED
|
||||
job.completed_at = time.time()
|
||||
# Wakes any get_result waiter with CancelledError right away.
|
||||
future.cancel()
|
||||
self._persist(job)
|
||||
return True
|
||||
|
||||
def reconcile_on_start(
|
||||
self,
|
||||
@@ -560,6 +638,14 @@ class BackgroundJobManager:
|
||||
if job_id in self._futures:
|
||||
del self._futures[job_id]
|
||||
removed += 1
|
||||
|
||||
# Prune the durable store with the same age policy so it stays bounded.
|
||||
prune = getattr(self._store, "delete_older_than", None)
|
||||
if prune is not None:
|
||||
try:
|
||||
prune(now - max_age)
|
||||
except Exception as e: # noqa: BLE001 — pruning is best-effort
|
||||
logger.debug("Pruning background job store failed (non-fatal): %s", e)
|
||||
|
||||
return removed
|
||||
|
||||
|
||||
@@ -160,6 +160,29 @@ class SqliteBackgroundJobStore:
|
||||
).fetchall()
|
||||
return [self._row_to_job(r) for r in rows]
|
||||
|
||||
def delete_older_than(self, ts: float) -> int:
|
||||
"""Delete terminal jobs completed before ``ts``; return rows removed.
|
||||
|
||||
Completed jobs with an undelivered ``origin`` are kept so restart
|
||||
reconciliation can still re-deliver them.
|
||||
"""
|
||||
with self._lock:
|
||||
cur = self._conn.execute(
|
||||
"""
|
||||
DELETE FROM background_jobs
|
||||
WHERE completed_at IS NOT NULL AND completed_at < ?
|
||||
AND status NOT IN (?, ?)
|
||||
AND NOT (status = ? AND delivered = 0 AND origin IS NOT NULL)
|
||||
""",
|
||||
(
|
||||
ts,
|
||||
JobStatus.PENDING.value,
|
||||
JobStatus.RUNNING.value,
|
||||
JobStatus.COMPLETED.value,
|
||||
),
|
||||
)
|
||||
return cur.rowcount
|
||||
|
||||
def close(self) -> None:
|
||||
"""Close the underlying database connection."""
|
||||
with self._lock:
|
||||
|
||||
@@ -307,6 +307,8 @@ def resolve_ingress_admission(
|
||||
# forwards an unset policy does not silently start replying to all group
|
||||
# traffic (a security-relevant regression).
|
||||
policy = (group_policy or "mention_only").lower()
|
||||
if policy not in _GROUP_POLICIES:
|
||||
policy = "mention_only"
|
||||
|
||||
if policy == "command_only":
|
||||
if is_command:
|
||||
@@ -339,7 +341,7 @@ def resolve_ingress_admission(
|
||||
gate=GATE_GROUP_POLICY,
|
||||
)
|
||||
|
||||
# respond_all (or any unknown policy): admit everything.
|
||||
# respond_all: admit everything.
|
||||
return IngressDecision(
|
||||
admit=True, reason_code=REASON_ALLOWED, gate=GATE_GROUP_POLICY
|
||||
)
|
||||
@@ -50,13 +50,15 @@ Example::
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import contextlib
|
||||
import re
|
||||
import time
|
||||
from abc import ABC, abstractmethod
|
||||
from dataclasses import dataclass, field
|
||||
from enum import Enum
|
||||
from typing import Any, Dict, List, Literal, Optional, Union
|
||||
|
||||
from .protocols import PlatformCapabilities
|
||||
from .protocols import HealthResult, PlatformCapabilities, ProbeResult
|
||||
|
||||
__all__ = [
|
||||
"SendErrorKind",
|
||||
@@ -469,6 +471,12 @@ class BasePlatformAdapter(ABC):
|
||||
typing heartbeat, edit/delete fallbacks, formatting). Override any default
|
||||
only when the platform can do better.
|
||||
|
||||
The base also satisfies the runtime surface ``Bot``/``BotOS`` drive
|
||||
(``start`` / ``stop`` / ``send_message`` / ``probe`` / ``health``) with
|
||||
defaults built on those primitives, so a minimal subclass runs under the
|
||||
gateway as-is. ``Bot`` builds adapters with ``token=`` / ``agent=`` /
|
||||
``config=`` keyword arguments, so a subclass ``__init__`` should accept them.
|
||||
|
||||
Declare platform features by setting the class attribute
|
||||
:attr:`capabilities` to a :class:`PlatformCapabilities` instance. The
|
||||
default delivery path reads capabilities via ``getattr`` so it degrades
|
||||
@@ -550,20 +558,110 @@ class BasePlatformAdapter(ABC):
|
||||
``start`` should run its inbound source (poll/listen/socket) and only
|
||||
return when the channel is stopped cleanly, raising on an unexpected
|
||||
drop so the supervisor can reconnect with backoff. Concrete adapters
|
||||
override this with their platform run loop.
|
||||
with a blocking platform loop override this.
|
||||
|
||||
Default: :meth:`connect` (passing ``is_reconnect`` on supervised
|
||||
restarts), raise if it fails so the supervisor backs off, then hold
|
||||
until :meth:`stop` — suitable for push/callback transports whose
|
||||
``connect`` registers their inbound listener.
|
||||
"""
|
||||
raise NotImplementedError(
|
||||
"adapter must implement start() to run its inbound loop"
|
||||
)
|
||||
is_reconnect = getattr(self, "_base_started_at", None) is not None
|
||||
# Created before connecting so a stop() racing connect() is not lost.
|
||||
stop_event = self._base_stop_event = asyncio.Event()
|
||||
# Until connect() returns, a stop()'s disconnect may precede the
|
||||
# connection, so teardown is still owed.
|
||||
stopped_early = True
|
||||
try:
|
||||
if not await self.connect(is_reconnect=is_reconnect):
|
||||
raise ConnectionError(f"{type(self).__name__}.connect() failed")
|
||||
stopped_early = stop_event.is_set()
|
||||
self._base_started_at = time.time()
|
||||
self._base_connected = True
|
||||
await stop_event.wait()
|
||||
finally:
|
||||
self._base_connected = False
|
||||
# Cancelled by the supervisor, failed, or stopped mid-connect:
|
||||
# stop() did not (effectively) disconnect, so do it here.
|
||||
if stopped_early or not stop_event.is_set():
|
||||
stop_event.set()
|
||||
with contextlib.suppress(Exception):
|
||||
await self.disconnect()
|
||||
|
||||
async def stop(self) -> None:
|
||||
"""Signal the inbound loop to stop and tear down the connection.
|
||||
|
||||
Default delegates to :meth:`disconnect`. Adapters with an explicit
|
||||
stop event override this to unblock their run loop first.
|
||||
Default releases the default :meth:`start` wait, then delegates to
|
||||
:meth:`disconnect`. Adapters with their own run loop override this to
|
||||
unblock it first.
|
||||
"""
|
||||
event = getattr(self, "_base_stop_event", None)
|
||||
if event is not None:
|
||||
event.set()
|
||||
await self.disconnect()
|
||||
|
||||
def _base_running(self) -> bool:
|
||||
return getattr(self, "_base_connected", False)
|
||||
|
||||
async def send_message(
|
||||
self,
|
||||
channel_id: Any,
|
||||
content: Union[str, Dict[str, Any]],
|
||||
reply_to: Optional[str] = None,
|
||||
thread_id: Optional[str] = None,
|
||||
) -> SendResult:
|
||||
"""Runtime send entry point (``Bot.send_message``).
|
||||
|
||||
Default routes through :meth:`deliver` so formatting, chunking, typing
|
||||
and retry apply; ``thread_id`` is forwarded in ``metadata``.
|
||||
"""
|
||||
return await self.deliver(
|
||||
channel_id,
|
||||
content,
|
||||
reply_to=reply_to,
|
||||
metadata={"thread_id": thread_id} if thread_id else None,
|
||||
)
|
||||
|
||||
async def probe(self) -> ProbeResult:
|
||||
"""Check connectivity/credentials without processing messages.
|
||||
|
||||
Default: ok while the default :meth:`start` is running; otherwise a
|
||||
:meth:`connect` / :meth:`disconnect` round-trip.
|
||||
"""
|
||||
platform = getattr(self, "platform_name", "") or ""
|
||||
if self._base_running():
|
||||
return ProbeResult(ok=True, platform=platform)
|
||||
t0 = time.monotonic()
|
||||
error: Optional[str] = None
|
||||
try:
|
||||
ok = bool(await self.connect())
|
||||
if not ok:
|
||||
error = "connect() returned False"
|
||||
except Exception as exc:
|
||||
ok, error = False, str(exc)
|
||||
finally:
|
||||
# Also on failure/cancellation (preflight timeout), so a partially
|
||||
# opened client or listener is released.
|
||||
with contextlib.suppress(Exception):
|
||||
await self.disconnect()
|
||||
return ProbeResult(
|
||||
ok=ok,
|
||||
platform=platform,
|
||||
elapsed_ms=(time.monotonic() - t0) * 1000,
|
||||
error=error,
|
||||
)
|
||||
|
||||
async def health(self) -> HealthResult:
|
||||
"""Report health. Default reflects the default run state and uptime."""
|
||||
running = self._base_running()
|
||||
started = getattr(self, "_base_started_at", None)
|
||||
return HealthResult(
|
||||
ok=running,
|
||||
platform=getattr(self, "platform_name", "") or "",
|
||||
is_running=running,
|
||||
uptime_seconds=(time.time() - started) if running and started else None,
|
||||
error=None if running else "not running",
|
||||
)
|
||||
|
||||
# ------------------------------------------------------------------ #
|
||||
# Identity canonicalization seam. #
|
||||
# #
|
||||
|
||||
@@ -8,7 +8,7 @@ import asyncio
|
||||
from praisonaiagents._logging import get_logger
|
||||
import threading
|
||||
from typing import Any, Callable, Dict, List, Optional, Set, Union
|
||||
from ..utils.async_bridge import run_coroutine_from_any_context
|
||||
from ..utils.async_bridge import run_coroutine_from_any_context, spawn_background
|
||||
from dataclasses import dataclass, field
|
||||
from .event import Event, EventType
|
||||
|
||||
@@ -235,11 +235,14 @@ class EventBus:
|
||||
if sub.is_async:
|
||||
# Schedule async callback safely
|
||||
try:
|
||||
loop = asyncio.get_running_loop()
|
||||
loop.create_task(sub.callback(event))
|
||||
asyncio.get_running_loop()
|
||||
except RuntimeError:
|
||||
# No running loop, use safe async bridge
|
||||
run_coroutine_from_any_context(sub.callback(event))
|
||||
else:
|
||||
# Keep a strong ref and log failures under this
|
||||
# subscriber's id, like the sync path below.
|
||||
spawn_background(sub.callback(event), f"subscriber {sub.id}")
|
||||
else:
|
||||
sub.callback(event)
|
||||
except Exception as e:
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -708,12 +708,17 @@ def _config_write_target() -> Path:
|
||||
|
||||
def _dump_config_file(path: Path, data: Dict[str, Any]) -> None:
|
||||
"""Write ``data`` to ``path`` as YAML or TOML based on its suffix."""
|
||||
from ..utils.atomic_io import atomic_write_text
|
||||
|
||||
# Write through a symlinked config so the shared target gets the update,
|
||||
# and keep an existing file's mode (the atomic temp file is owner-only).
|
||||
path = path.resolve()
|
||||
mode = path.stat().st_mode & 0o7777 if path.exists() else None
|
||||
path.parent.mkdir(parents=True, exist_ok=True)
|
||||
if path.suffix.lower() in (".yaml", ".yml"):
|
||||
import yaml
|
||||
|
||||
with open(path, "w", encoding="utf-8") as f:
|
||||
yaml.safe_dump(data, f, sort_keys=False, default_flow_style=False)
|
||||
text = yaml.safe_dump(data, sort_keys=False, default_flow_style=False)
|
||||
else:
|
||||
try:
|
||||
import tomli_w
|
||||
@@ -727,8 +732,10 @@ def _dump_config_file(path: Path, data: Dict[str, Any]) -> None:
|
||||
"Install tomli-w (pip install tomli-w), or convert the config "
|
||||
"to .praisonai/config.yaml."
|
||||
) from None
|
||||
with open(path, "wb") as f:
|
||||
tomli_w.dump(data, f)
|
||||
text = tomli_w.dumps(data)
|
||||
atomic_write_text(str(path), text)
|
||||
if mode is not None:
|
||||
os.chmod(path, mode)
|
||||
|
||||
|
||||
def set_plugin_enabled(name: str, enabled: bool) -> Path:
|
||||
|
||||
@@ -16,6 +16,7 @@ from typing import Any, Callable, Dict, Optional, Set, Type, Union
|
||||
import inspect
|
||||
|
||||
from .parse_utils import (
|
||||
canonical_preset_key,
|
||||
detect_url_scheme,
|
||||
is_path_like,
|
||||
make_preset_error,
|
||||
@@ -345,12 +346,12 @@ def _resolve_string(
|
||||
if presets:
|
||||
# Case-insensitive lookup, whitespace-tolerant and treating -/_ as
|
||||
# interchangeable so the accepted spellings match the closed-set guard
|
||||
# in validate_preset_string (which uses the same normalization). Without
|
||||
# this, a value that passed construction-time validation (e.g.
|
||||
# " sliding_window " or "sliding-window") would be rejected here.
|
||||
value_norm = value.strip().lower().replace("-", "_")
|
||||
# in validate_preset_string. Both go through canonical_preset_key, the
|
||||
# single source of truth for the spelling rule, so validation and
|
||||
# resolution cannot disagree about what counts as the same preset.
|
||||
value_norm = canonical_preset_key(value)
|
||||
for preset_key in presets:
|
||||
if preset_key.strip().lower().replace("-", "_") == value_norm:
|
||||
if canonical_preset_key(preset_key) == value_norm:
|
||||
return _apply_preset(preset_key, presets, config_class)
|
||||
|
||||
# Not a valid preset - raise helpful error
|
||||
|
||||
@@ -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},
|
||||
}
|
||||
|
||||
|
||||
|
||||
@@ -157,8 +157,8 @@ class FileIndex:
|
||||
"entries": {k: v.to_dict() for k, v in self.entries.items()}
|
||||
}
|
||||
|
||||
with open(self.index_path, 'w', encoding='utf-8') as f:
|
||||
json.dump(data, f, indent=2)
|
||||
from praisonaiagents.utils.atomic_io import atomic_write_json
|
||||
atomic_write_json(self.index_path, data, indent=2)
|
||||
|
||||
return True
|
||||
except Exception as e:
|
||||
|
||||
@@ -659,50 +659,52 @@ class ContextManager:
|
||||
Returns:
|
||||
Tuple of (token_count, metrics)
|
||||
"""
|
||||
# Heuristic is O(len(text)); hashing the whole text to cache it costs
|
||||
# far more than recomputing it, so only cache the tiktoken path.
|
||||
if not (validate or self.config.estimation_mode == EstimationMode.VALIDATED):
|
||||
return estimate_tokens_heuristic(text), None
|
||||
|
||||
# Check cache
|
||||
cache_key = hashlib.sha256(text.encode()).hexdigest()
|
||||
if cache_key in self._estimation_cache:
|
||||
return self._estimation_cache[cache_key], None
|
||||
|
||||
|
||||
heuristic = estimate_tokens_heuristic(text)
|
||||
metrics = None
|
||||
|
||||
if validate or self.config.estimation_mode == EstimationMode.VALIDATED:
|
||||
try:
|
||||
from .tokens import estimate_tokens_accurate
|
||||
accurate = estimate_tokens_accurate(text, self.model)
|
||||
|
||||
error_pct = abs(heuristic - accurate) / accurate * 100 if accurate > 0 else 0
|
||||
|
||||
metrics = EstimationMetrics(
|
||||
heuristic_estimate=heuristic,
|
||||
accurate_estimate=accurate,
|
||||
error_pct=error_pct,
|
||||
estimator_used=EstimationMode.VALIDATED,
|
||||
|
||||
try:
|
||||
from .tokens import estimate_tokens_accurate
|
||||
accurate = estimate_tokens_accurate(text, self.model)
|
||||
|
||||
error_pct = abs(heuristic - accurate) / accurate * 100 if accurate > 0 else 0
|
||||
|
||||
metrics = EstimationMetrics(
|
||||
heuristic_estimate=heuristic,
|
||||
accurate_estimate=accurate,
|
||||
error_pct=error_pct,
|
||||
estimator_used=EstimationMode.VALIDATED,
|
||||
)
|
||||
|
||||
self._estimation_metrics = metrics
|
||||
|
||||
# Log if mismatch exceeds threshold
|
||||
if self.config.log_estimation_mismatch and error_pct > self.config.mismatch_threshold_pct:
|
||||
import logging
|
||||
get_logger(__name__).warning(
|
||||
f"Token estimation mismatch: heuristic={heuristic}, accurate={accurate}, error={error_pct:.1f}%"
|
||||
)
|
||||
|
||||
self._estimation_metrics = metrics
|
||||
|
||||
# Log if mismatch exceeds threshold
|
||||
if self.config.log_estimation_mismatch and error_pct > self.config.mismatch_threshold_pct:
|
||||
import logging
|
||||
get_logger(__name__).warning(
|
||||
f"Token estimation mismatch: heuristic={heuristic}, accurate={accurate}, error={error_pct:.1f}%"
|
||||
)
|
||||
|
||||
# Use accurate if validated mode
|
||||
if self.config.estimation_mode == EstimationMode.ACCURATE:
|
||||
result = accurate
|
||||
else:
|
||||
result = heuristic
|
||||
except ImportError:
|
||||
|
||||
# Use accurate if validated mode
|
||||
if self.config.estimation_mode == EstimationMode.ACCURATE:
|
||||
result = accurate
|
||||
else:
|
||||
result = heuristic
|
||||
metrics = EstimationMetrics(
|
||||
heuristic_estimate=heuristic,
|
||||
estimator_used=EstimationMode.HEURISTIC,
|
||||
)
|
||||
else:
|
||||
except ImportError:
|
||||
result = heuristic
|
||||
metrics = EstimationMetrics(
|
||||
heuristic_estimate=heuristic,
|
||||
estimator_used=EstimationMode.HEURISTIC,
|
||||
)
|
||||
|
||||
# Cache result
|
||||
if len(self._estimation_cache) < self._cache_max_size:
|
||||
|
||||
@@ -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
|
||||
@@ -228,6 +228,10 @@ _LAZY_IMPORTS = {
|
||||
"StopResult": ("praisonaiagents.gateway.protocols", "StopResult"),
|
||||
# Schema-validated inbound frame codec (Issue #2831)
|
||||
"HelloParams": ("praisonaiagents.gateway.protocols", "HelloParams"),
|
||||
# Device-identity handshake contract (Issue #5998)
|
||||
"DeviceAuth": ("praisonaiagents.gateway.protocols", "DeviceAuth"),
|
||||
"DeviceIdentityProtocol": ("praisonaiagents.gateway.protocols", "DeviceIdentityProtocol"),
|
||||
"device_id_from_public_key": ("praisonaiagents.gateway.protocols", "device_id_from_public_key"),
|
||||
"HelloResult": ("praisonaiagents.gateway.protocols", "HelloResult"),
|
||||
"HelloError": ("praisonaiagents.gateway.protocols", "HelloError"),
|
||||
"ConnectErrorCode": ("praisonaiagents.gateway.protocols", "ConnectErrorCode"),
|
||||
@@ -586,6 +590,9 @@ __all__ = [
|
||||
"StopResult",
|
||||
# Schema-validated inbound frame codec (Issue #2831)
|
||||
"HelloParams",
|
||||
"DeviceAuth",
|
||||
"DeviceIdentityProtocol",
|
||||
"device_id_from_public_key",
|
||||
"HelloResult",
|
||||
"HelloError",
|
||||
"ConnectErrorCode",
|
||||
|
||||
@@ -11,6 +11,8 @@ The ``praisonai`` wrapper provides backward-compatible shims.
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import base64
|
||||
import hashlib
|
||||
import math
|
||||
import time
|
||||
import uuid
|
||||
@@ -347,10 +349,88 @@ class SessionObserverProtocol(Protocol):
|
||||
...
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Device identity (Issue #5998)
|
||||
#
|
||||
# A per-device keypair signs a server-issued challenge nonce so the gateway can
|
||||
# derive a stable ``device_id`` and approve/revoke devices individually instead
|
||||
# of sharing one gateway-wide secret. Core only fixes the wire shape, the
|
||||
# fingerprint and the exact bytes that are signed; signature verification,
|
||||
# device storage and per-device tokens live in the wrapper gateway.
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
def device_id_from_public_key(public_key: str) -> str:
|
||||
"""Return the stable device id: hex SHA-256 of the raw public-key bytes.
|
||||
|
||||
``public_key`` must be standard base64 (RFC 4648, padded, no whitespace) of
|
||||
the raw key bytes. Hashing the decoded bytes keeps the id independent of
|
||||
how a client spells the encoding. Raises ``ValueError`` on invalid base64.
|
||||
"""
|
||||
return hashlib.sha256(base64.b64decode(public_key, validate=True)).hexdigest()
|
||||
|
||||
|
||||
@dataclass
|
||||
class DeviceAuth:
|
||||
"""Device-identity proof carried in a ``hello`` frame.
|
||||
|
||||
Attributes:
|
||||
device_id: Fingerprint of ``public_key`` (see :func:`device_id_from_public_key`).
|
||||
public_key: Standard base64 (padded) of the device's raw public key.
|
||||
signature: Standard base64 (padded) signature over :meth:`signing_payload`.
|
||||
signed_at: Integer Unix seconds at which the device signed the challenge
|
||||
(an integer so every client language renders identical signed bytes).
|
||||
nonce: The server-issued challenge nonce being answered.
|
||||
"""
|
||||
device_id: str
|
||||
public_key: str
|
||||
signature: str
|
||||
signed_at: int
|
||||
nonce: str
|
||||
|
||||
def signing_payload(self, agent_id: str) -> bytes:
|
||||
"""Canonical bytes the device signs and the gateway verifies.
|
||||
|
||||
Binds the nonce to this device, the signing time and the target agent so
|
||||
a captured proof cannot be replayed for another device or agent.
|
||||
"""
|
||||
return "\n".join((
|
||||
"praisonai-device-v1",
|
||||
self.device_id,
|
||||
self.nonce,
|
||||
str(int(self.signed_at)),
|
||||
agent_id,
|
||||
)).encode("utf-8")
|
||||
|
||||
|
||||
@runtime_checkable
|
||||
class DeviceIdentityProtocol(Protocol):
|
||||
"""Seam for per-device challenge/response auth (implemented in the wrapper).
|
||||
|
||||
Implementations must approve only on an exact device-key match — never on
|
||||
network reachability alone — and record how each device was approved.
|
||||
"""
|
||||
|
||||
def issue_challenge(self) -> str:
|
||||
"""Return a fresh single-use challenge nonce for a connecting client."""
|
||||
...
|
||||
|
||||
def verify(self, device: DeviceAuth, *, agent_id: str) -> Optional[str]:
|
||||
"""Verify ``device`` against an issued nonce; return its role, or ``None`` if rejected."""
|
||||
...
|
||||
|
||||
def list_devices(self) -> List[Dict[str, Any]]:
|
||||
"""List known devices (id, role, approval path, last seen)."""
|
||||
...
|
||||
|
||||
def revoke(self, device_id: str) -> bool:
|
||||
"""Revoke one device; return ``False`` if it was unknown."""
|
||||
...
|
||||
|
||||
|
||||
@dataclass
|
||||
class HelloParams:
|
||||
"""Parameters for initiating a versioned handshake.
|
||||
|
||||
|
||||
Attributes:
|
||||
agent_id: The agent to connect to
|
||||
protocol_min: Minimum protocol version the client supports
|
||||
@@ -358,6 +438,7 @@ class HelloParams:
|
||||
capabilities: Optional list of capability tokens the client supports
|
||||
session_id: Optional session to resume
|
||||
since: Optional cursor for event replay
|
||||
device: Optional device-identity proof (Issue #5998)
|
||||
"""
|
||||
agent_id: str
|
||||
protocol_min: int
|
||||
@@ -365,6 +446,7 @@ class HelloParams:
|
||||
capabilities: List[str] = field(default_factory=list)
|
||||
session_id: Optional[str] = None
|
||||
since: Optional[int] = None
|
||||
device: Optional[DeviceAuth] = None
|
||||
|
||||
type: str = field(default="hello", init=False)
|
||||
|
||||
@@ -897,6 +979,9 @@ def _decode_hello(data: Dict[str, Any]) -> "HelloParams":
|
||||
if since is not None:
|
||||
since = _coerce_int(since, field_name="since")
|
||||
|
||||
device_data = data.get("device")
|
||||
device = _decode_device(device_data) if device_data is not None else None
|
||||
|
||||
return HelloParams(
|
||||
agent_id=agent_id,
|
||||
protocol_min=client_min,
|
||||
@@ -904,6 +989,60 @@ def _decode_hello(data: Dict[str, Any]) -> "HelloParams":
|
||||
capabilities=capabilities,
|
||||
session_id=_as_opt_str(data.get("session_id")),
|
||||
since=since,
|
||||
device=device,
|
||||
)
|
||||
|
||||
|
||||
def _decode_device(data: Any) -> DeviceAuth:
|
||||
"""Validate a ``hello.device`` block into :class:`DeviceAuth`.
|
||||
|
||||
Shape errors are configuration errors; a ``device_id`` that is not the
|
||||
fingerprint of ``public_key`` is rejected as unauthorized so a client can
|
||||
never claim another device's id.
|
||||
"""
|
||||
if not isinstance(data, dict):
|
||||
raise FrameDecodeError(
|
||||
HelloError(
|
||||
code=ConnectErrorCode.CONFIGURATION_ERROR,
|
||||
message="Field 'device' must be an object",
|
||||
next_step=ConnectRecoveryStep.DO_NOT_RETRY,
|
||||
)
|
||||
)
|
||||
public_key = _require_str(data.get("public_key"), field_name="device.public_key")
|
||||
device_id = _require_str(data.get("device_id"), field_name="device.device_id")
|
||||
signed_at = data.get("signed_at")
|
||||
if isinstance(signed_at, bool) or not isinstance(signed_at, int):
|
||||
raise FrameDecodeError(
|
||||
HelloError(
|
||||
code=ConnectErrorCode.CONFIGURATION_ERROR,
|
||||
message="Field 'device.signed_at' must be an integer (Unix seconds)",
|
||||
next_step=ConnectRecoveryStep.DO_NOT_RETRY,
|
||||
)
|
||||
)
|
||||
try:
|
||||
fingerprint = device_id_from_public_key(public_key)
|
||||
except ValueError:
|
||||
raise FrameDecodeError(
|
||||
HelloError(
|
||||
code=ConnectErrorCode.CONFIGURATION_ERROR,
|
||||
message="Field 'device.public_key' must be padded standard base64",
|
||||
next_step=ConnectRecoveryStep.DO_NOT_RETRY,
|
||||
)
|
||||
) from None
|
||||
if device_id != fingerprint:
|
||||
raise FrameDecodeError(
|
||||
HelloError(
|
||||
code=ConnectErrorCode.AUTH_UNAUTHORIZED,
|
||||
message="device.device_id does not match the public key fingerprint",
|
||||
next_step=ConnectRecoveryStep.REAUTHENTICATE,
|
||||
)
|
||||
)
|
||||
return DeviceAuth(
|
||||
device_id=device_id,
|
||||
public_key=public_key,
|
||||
signature=_require_str(data.get("signature"), field_name="device.signature"),
|
||||
signed_at=signed_at,
|
||||
nonce=_require_str(data.get("nonce"), field_name="device.nonce"),
|
||||
)
|
||||
|
||||
|
||||
|
||||
@@ -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
|
||||
|
||||
|
||||
@@ -417,9 +417,9 @@ class FileTracker:
|
||||
return
|
||||
|
||||
try:
|
||||
os.makedirs(os.path.dirname(self._state_file), exist_ok=True)
|
||||
with open(self._state_file, "w", encoding="utf-8") as f:
|
||||
json.dump(self._tracked, f, indent=2)
|
||||
os.makedirs(os.path.dirname(self._state_file) or ".", exist_ok=True)
|
||||
from ..utils.atomic_io import atomic_write_json
|
||||
atomic_write_json(self._state_file, self._tracked, indent=2)
|
||||
except (OSError, IOError):
|
||||
pass
|
||||
|
||||
|
||||
@@ -803,6 +803,9 @@ class Knowledge:
|
||||
same path as local files. Failures are surfaced to the caller so a
|
||||
single bad URL in a batch is visible, matching local-file behaviour.
|
||||
"""
|
||||
from ..tools.url_safety import is_safe_http_url
|
||||
if not is_safe_http_url(url):
|
||||
raise ValueError(f"Refusing to fetch unsafe URL: {url}")
|
||||
result = self.markdown.convert(url)
|
||||
content = getattr(result, 'text_content', None)
|
||||
if not content:
|
||||
|
||||
@@ -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
|
||||
@@ -44,6 +44,173 @@ _TOOL_ARGUMENTS_PARSE_FAILED = object()
|
||||
_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.
|
||||
# 600s matches the OpenAI SDK default used by the bare-model-name path and is
|
||||
# a read-gap bound for streams. Pass ``timeout=None`` (or 0) to opt out.
|
||||
_DEFAULT_LLM_TIMEOUT_S = 600
|
||||
_TIMEOUT_UNSET = object()
|
||||
|
||||
|
||||
# ``litellm.success_callback``/``_async_success_callback``/``callbacks`` are
|
||||
# process-global lists shared by every LLM instance. Guard the read-modify-write
|
||||
# in ``_setup_event_tracking`` so concurrent instances don't corrupt them.
|
||||
@@ -177,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.
|
||||
@@ -234,6 +406,19 @@ class TokenUsage:
|
||||
'audio_output_tokens': self.audio_output_tokens,
|
||||
}
|
||||
|
||||
# error_classifier category -> (text-derived kinds that may refine it, default kind).
|
||||
_CATEGORY_KINDS = {
|
||||
"rate_limit": ({"rate_limit", "billing"}, "rate_limit"),
|
||||
"auth": ({"auth", "auth_permanent"}, "auth"),
|
||||
"context_limit": ({"context_overflow"}, "context_overflow"),
|
||||
"invalid_request": (
|
||||
{"context_overflow", "model_not_found", "auth", "auth_permanent", "billing", "format_error"},
|
||||
"format_error",
|
||||
),
|
||||
"transient": ({"overloaded", "idle_timeout"}, "overloaded"),
|
||||
"permanent": ({"billing"}, "billing"),
|
||||
}
|
||||
|
||||
class LLM:
|
||||
"""
|
||||
Easy to use wrapper for language models. Supports multiple providers like OpenAI,
|
||||
@@ -493,7 +678,7 @@ Respond with ONLY a valid JSON tool call in this format:
|
||||
def __init__(
|
||||
self,
|
||||
model: str,
|
||||
timeout: Optional[int] = None,
|
||||
timeout: Optional[int] = _TIMEOUT_UNSET,
|
||||
temperature: Optional[float] = None,
|
||||
top_p: Optional[float] = None,
|
||||
n: Optional[int] = None,
|
||||
@@ -535,7 +720,7 @@ Respond with ONLY a valid JSON tool call in this format:
|
||||
)
|
||||
|
||||
self.model = model
|
||||
self.timeout = timeout
|
||||
self.timeout = _DEFAULT_LLM_TIMEOUT_S if timeout is _TIMEOUT_UNSET else timeout
|
||||
self.temperature = temperature
|
||||
self.top_p = top_p
|
||||
self.n = n
|
||||
@@ -886,11 +1071,12 @@ Respond with ONLY a valid JSON tool call in this format:
|
||||
else:
|
||||
return False
|
||||
|
||||
def _parse_retry_delay(self, error_message: str) -> float:
|
||||
"""Parse retry delay from rate limit error message.
|
||||
def _parse_retry_delay(self, error_message: Union[Exception, str]) -> float:
|
||||
"""Parse retry delay from a rate limit error.
|
||||
|
||||
Args:
|
||||
error_message: The error message from the API
|
||||
error_message: The provider exception (preferred, so the structured
|
||||
``Retry-After`` response header is honoured) or its message
|
||||
|
||||
Returns:
|
||||
Retry delay in seconds, or default if not found
|
||||
@@ -902,7 +1088,8 @@ Respond with ONLY a valid JSON tool call in this format:
|
||||
|
||||
from .error_classifier import extract_retry_after
|
||||
|
||||
delay = extract_retry_after(Exception(error_message), cap_seconds=max_delay)
|
||||
error = error_message if isinstance(error_message, BaseException) else Exception(error_message)
|
||||
delay = extract_retry_after(error, cap_seconds=max_delay)
|
||||
return self._retry_delay if delay is None else delay
|
||||
|
||||
def _is_rate_limit_error(self, error: Exception) -> bool:
|
||||
@@ -939,15 +1126,35 @@ Respond with ONLY a valid JSON tool call in this format:
|
||||
def classify_error_kind(self, error: Exception) -> AgentErrorKind:
|
||||
"""
|
||||
Classify error into typed AgentErrorKind instead of freeform strings.
|
||||
|
||||
Replaces scattered regex checks with systematic classification.
|
||||
|
||||
|
||||
The HTTP status code / exception type (via ``error_classifier``) is
|
||||
authoritative; message text only refines the kind within that category,
|
||||
so e.g. a 400 whose body happens to contain "500" is never retried as
|
||||
an overload.
|
||||
|
||||
Args:
|
||||
error: Exception to classify
|
||||
|
||||
|
||||
Returns:
|
||||
AgentErrorKind: Typed error classification
|
||||
"""
|
||||
text_kind = self._classify_error_kind_from_text(error)
|
||||
from .error_classifier import _classify_by_type_and_code, _extract_status_code
|
||||
category = _classify_by_type_and_code(error)
|
||||
if category is None:
|
||||
return text_kind
|
||||
allowed, default = _CATEGORY_KINDS[category.value]
|
||||
if text_kind in allowed:
|
||||
return text_kind
|
||||
# A 404 from a completion endpoint means the model/deployment is missing.
|
||||
if _extract_status_code(error) == 404 or type(error).__name__ == "NotFoundError":
|
||||
return "model_not_found"
|
||||
if default == "overloaded" and "timeout" in type(error).__name__.lower():
|
||||
return "idle_timeout"
|
||||
return default
|
||||
|
||||
def _classify_error_kind_from_text(self, error: Exception) -> AgentErrorKind:
|
||||
"""Message-substring fallback used when no status code / type is known."""
|
||||
error_str = str(error).lower()
|
||||
error_type = type(error).__name__.lower()
|
||||
|
||||
@@ -969,7 +1176,7 @@ Respond with ONLY a valid JSON tool call in this format:
|
||||
|
||||
# Billing/quota issues (must be checked before generic 429/rate-limit)
|
||||
if any(indicator in error_str for indicator in [
|
||||
"insufficient quota", "quota exceeded", "billing", "credit",
|
||||
"insufficient quota", "quota exceeded", "billing", "credit balance", "insufficient credit",
|
||||
"payment required", "subscription required", "subscription expired",
|
||||
"plan limit"
|
||||
]):
|
||||
@@ -1101,7 +1308,7 @@ Respond with ONLY a valid JSON tool call in this format:
|
||||
|
||||
# Rate limiting - extract retry delay
|
||||
if error_kind == "rate_limit":
|
||||
backoff = self._parse_retry_delay(str(error)) # Returns seconds
|
||||
backoff = self._parse_retry_delay(error) # Returns seconds
|
||||
if backoff == 0: # No specific delay found, use exponential backoff
|
||||
backoff = min(2 ** (attempt - 1), 60) # 1s, 2s, 4s, ... cap at 60s
|
||||
return FailoverDecision(
|
||||
@@ -1693,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):
|
||||
@@ -1722,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:
|
||||
@@ -1794,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
|
||||
@@ -3766,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}")
|
||||
@@ -4292,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}")
|
||||
@@ -4639,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
|
||||
@@ -6634,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
|
||||
|
||||
@@ -7593,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.
|
||||
@@ -7606,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.
|
||||
@@ -7735,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,
|
||||
@@ -7888,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,
|
||||
|
||||
@@ -154,10 +154,19 @@ def register_approval_callback(callback_fn):
|
||||
add_display_callback = register_display_callback
|
||||
add_approval_callback = register_approval_callback
|
||||
|
||||
def _filter_callback_kwargs(callback, kwargs):
|
||||
"""Filter kwargs to what the callback accepts to maintain backward compatibility."""
|
||||
import inspect
|
||||
sig = inspect.signature(callback)
|
||||
if any(p.kind == inspect.Parameter.VAR_KEYWORD for p in sig.parameters.values()):
|
||||
return kwargs
|
||||
return {k: v for k, v in kwargs.items() if k in sig.parameters}
|
||||
|
||||
def execute_sync_callback(display_type: str, **kwargs):
|
||||
"""Execute synchronous callback for a given display type without displaying anything.
|
||||
|
||||
This function is used to trigger callbacks even when verbose=False.
|
||||
Callback errors are logged and suppressed so observers cannot abort the run.
|
||||
|
||||
Args:
|
||||
display_type (str): Type of display event
|
||||
@@ -167,61 +176,37 @@ def execute_sync_callback(display_type: str, **kwargs):
|
||||
callback = sync_display_callbacks.get(display_type)
|
||||
|
||||
if callback:
|
||||
import inspect
|
||||
sig = inspect.signature(callback)
|
||||
|
||||
# Filter kwargs to what the callback accepts to maintain backward compatibility
|
||||
if any(p.kind == inspect.Parameter.VAR_KEYWORD for p in sig.parameters.values()):
|
||||
# Callback accepts **kwargs, so pass all arguments
|
||||
supported_kwargs = kwargs
|
||||
else:
|
||||
# Only pass arguments that the callback signature supports
|
||||
supported_kwargs = {k: v for k, v in kwargs.items() if k in sig.parameters}
|
||||
|
||||
callback(**supported_kwargs)
|
||||
try:
|
||||
callback(**_filter_callback_kwargs(callback, kwargs))
|
||||
except Exception:
|
||||
logging.exception("Display callback for %r failed; ignored", display_type)
|
||||
|
||||
async def execute_callback(display_type: str, **kwargs):
|
||||
"""Execute both sync and async callbacks for a given display type.
|
||||
|
||||
Callback errors are logged and suppressed so observers cannot abort the run.
|
||||
|
||||
Args:
|
||||
display_type (str): Type of display event
|
||||
**kwargs: Arguments to pass to the callback functions
|
||||
"""
|
||||
import inspect
|
||||
|
||||
# Execute synchronous callback if registered
|
||||
with _callbacks_lock:
|
||||
sync_callback = sync_display_callbacks.get(display_type)
|
||||
async_callback = async_display_callbacks.get(display_type)
|
||||
|
||||
if sync_callback:
|
||||
sig = inspect.signature(sync_callback)
|
||||
|
||||
# Filter kwargs to what the callback accepts to maintain backward compatibility
|
||||
if any(p.kind == inspect.Parameter.VAR_KEYWORD for p in sig.parameters.values()):
|
||||
# Callback accepts **kwargs, so pass all arguments
|
||||
supported_kwargs = kwargs
|
||||
else:
|
||||
# Only pass arguments that the callback signature supports
|
||||
supported_kwargs = {k: v for k, v in kwargs.items() if k in sig.parameters}
|
||||
|
||||
loop = asyncio.get_event_loop()
|
||||
await loop.run_in_executor(None, lambda: sync_callback(**supported_kwargs))
|
||||
try:
|
||||
supported_kwargs = _filter_callback_kwargs(sync_callback, kwargs)
|
||||
loop = asyncio.get_running_loop()
|
||||
await loop.run_in_executor(None, lambda: sync_callback(**supported_kwargs))
|
||||
except Exception:
|
||||
logging.exception("Display callback for %r failed; ignored", display_type)
|
||||
|
||||
# Execute asynchronous callback if registered
|
||||
if async_callback:
|
||||
callback = async_callback
|
||||
sig = inspect.signature(callback)
|
||||
|
||||
# Filter kwargs to what the callback accepts to maintain backward compatibility
|
||||
if any(p.kind == inspect.Parameter.VAR_KEYWORD for p in sig.parameters.values()):
|
||||
# Callback accepts **kwargs, so pass all arguments
|
||||
supported_kwargs = kwargs
|
||||
else:
|
||||
# Only pass arguments that the callback signature supports
|
||||
supported_kwargs = {k: v for k, v in kwargs.items() if k in sig.parameters}
|
||||
|
||||
await callback(**supported_kwargs)
|
||||
try:
|
||||
await async_callback(**_filter_callback_kwargs(async_callback, kwargs))
|
||||
except Exception:
|
||||
logging.exception("Async display callback for %r failed; ignored", display_type)
|
||||
|
||||
def _clean_display_content(content: str, max_length: int = 20000) -> str:
|
||||
"""Helper function to clean and truncate content for display."""
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -28,15 +28,35 @@ _LOCALHOST_HOSTS = {"localhost", "127.0.0.1", "::1"}
|
||||
|
||||
|
||||
def _require_https(url: str) -> None:
|
||||
"""Reject non-HTTPS metadata URLs (loopback HTTP is allowed for testing)."""
|
||||
"""Reject OAuth URLs that are not HTTPS or that point at internal hosts.
|
||||
|
||||
Loopback hosts are allowed (over HTTP too) for local testing. Any other
|
||||
host must pass :func:`~praisonaiagents.tools.url_safety.is_safe_http_url`,
|
||||
so private, link-local (cloud metadata) and ambiguous numeric hosts are
|
||||
refused even over HTTPS.
|
||||
"""
|
||||
parsed = urlparse(url)
|
||||
if parsed.scheme == "https":
|
||||
if parsed.scheme in ("http", "https") and (parsed.hostname or "") in _LOCALHOST_HOSTS:
|
||||
return
|
||||
if parsed.scheme == "http" and (parsed.hostname or "") in _LOCALHOST_HOSTS:
|
||||
return
|
||||
raise ValueError(
|
||||
f"Refusing to fetch OAuth metadata over insecure transport: {url!r}"
|
||||
)
|
||||
if parsed.scheme != "https":
|
||||
raise ValueError(
|
||||
f"Refusing OAuth request over insecure transport: {url!r}"
|
||||
)
|
||||
from ..tools.url_safety import is_safe_http_url
|
||||
|
||||
if not is_safe_http_url(url):
|
||||
raise ValueError(
|
||||
f"Refusing OAuth request to a private/internal address: {url!r}"
|
||||
)
|
||||
|
||||
|
||||
def _require_success(resp: Any) -> None:
|
||||
"""Raise unless ``resp`` is 2xx (credential POSTs don't follow redirects)."""
|
||||
resp.raise_for_status()
|
||||
if not 200 <= resp.status_code < 300:
|
||||
raise ValueError(
|
||||
f"Unexpected HTTP {resp.status_code} from OAuth endpoint {resp.url!r}"
|
||||
)
|
||||
|
||||
|
||||
def parse_www_authenticate(header: Optional[str]) -> Dict[str, str]:
|
||||
@@ -62,30 +82,31 @@ def parse_www_authenticate(header: Optional[str]) -> Dict[str, str]:
|
||||
return params
|
||||
|
||||
|
||||
def _get_json(url: str, timeout: float = 10.0) -> Dict[str, Any]:
|
||||
def _get_json(url: str, timeout: float = 10.0, max_redirects: int = 3) -> Dict[str, Any]:
|
||||
"""Fetch and parse a JSON metadata document over HTTPS.
|
||||
|
||||
Redirects are followed but every hop — including the final destination — is
|
||||
re-validated with :func:`_require_https` so a redirect cannot escape the
|
||||
HTTPS/loopback guard to an internal or plaintext endpoint (SSRF-safe).
|
||||
Redirects are followed manually and every hop is validated with
|
||||
:func:`_require_https` *before* it is requested, so a redirect cannot
|
||||
reach an internal or plaintext endpoint (SSRF-safe).
|
||||
"""
|
||||
_require_https(url)
|
||||
import requests # lazy import — optional dependency
|
||||
|
||||
resp = requests.get(
|
||||
url,
|
||||
headers={"Accept": "application/json"},
|
||||
timeout=timeout,
|
||||
allow_redirects=True,
|
||||
)
|
||||
# Re-validate every redirect hop and the final URL: allow_redirects only
|
||||
# checks the initial URL, so a 30x to http:// or an internal host would
|
||||
# otherwise bypass the guard.
|
||||
for previous in getattr(resp, "history", None) or []:
|
||||
_require_https(previous.url)
|
||||
_require_https(resp.url)
|
||||
resp.raise_for_status()
|
||||
return resp.json()
|
||||
current = url
|
||||
for _ in range(max_redirects + 1):
|
||||
_require_https(current)
|
||||
resp = requests.get(
|
||||
current,
|
||||
headers={"Accept": "application/json"},
|
||||
timeout=timeout,
|
||||
allow_redirects=False,
|
||||
)
|
||||
location = resp.headers.get("Location")
|
||||
if resp.status_code in (301, 302, 303, 307, 308) and location:
|
||||
current = urljoin(current, location)
|
||||
continue
|
||||
resp.raise_for_status()
|
||||
return resp.json()
|
||||
raise ValueError(f"Too many redirects fetching OAuth metadata: {url!r}")
|
||||
|
||||
|
||||
def _wellknown_candidates(server_url: str) -> List[str]:
|
||||
|
||||
@@ -268,8 +268,11 @@ class MCPOAuthProvider:
|
||||
data=data,
|
||||
headers={"Accept": "application/json"},
|
||||
timeout=timeout,
|
||||
allow_redirects=False, # never re-send credentials to a redirect target
|
||||
)
|
||||
resp.raise_for_status()
|
||||
from .oauth_discovery import _require_success
|
||||
|
||||
_require_success(resp)
|
||||
return self._normalize_tokens(
|
||||
resp.json(), token_endpoint=metadata["token_endpoint"]
|
||||
)
|
||||
@@ -285,7 +288,7 @@ class MCPOAuthProvider:
|
||||
return False
|
||||
|
||||
try:
|
||||
from .oauth_discovery import _require_https
|
||||
from .oauth_discovery import _require_https, _require_success
|
||||
|
||||
# Defense in depth: never POST a refresh token / client secret to a
|
||||
# persisted endpoint that isn't HTTPS (or loopback).
|
||||
@@ -306,8 +309,9 @@ class MCPOAuthProvider:
|
||||
data=data,
|
||||
headers={"Accept": "application/json"},
|
||||
timeout=10.0,
|
||||
allow_redirects=False, # never re-send credentials to a redirect target
|
||||
)
|
||||
resp.raise_for_status()
|
||||
_require_success(resp)
|
||||
new_tokens = self._normalize_tokens(
|
||||
resp.json(), token_endpoint=token_endpoint
|
||||
)
|
||||
|
||||
@@ -10,26 +10,13 @@ returned for persistence via :meth:`MCPAuthStorage.set_client_info`.
|
||||
"""
|
||||
|
||||
from typing import Any, Dict, List, Optional
|
||||
from urllib.parse import urlparse
|
||||
|
||||
from praisonaiagents._logging import get_logger
|
||||
|
||||
from .oauth_discovery import _require_https, _require_success
|
||||
|
||||
logger = get_logger("mcp-oauth-registration")
|
||||
|
||||
_LOCALHOST_HOSTS = {"localhost", "127.0.0.1", "::1"}
|
||||
|
||||
|
||||
def _require_https(url: str) -> None:
|
||||
"""Reject non-HTTPS registration URLs (loopback HTTP allowed for testing)."""
|
||||
parsed = urlparse(url)
|
||||
if parsed.scheme == "https":
|
||||
return
|
||||
if parsed.scheme == "http" and (parsed.hostname or "") in _LOCALHOST_HOSTS:
|
||||
return
|
||||
raise ValueError(
|
||||
f"Refusing to register client over insecure transport: {url!r}"
|
||||
)
|
||||
|
||||
|
||||
def register_client(
|
||||
registration_endpoint: str,
|
||||
@@ -75,8 +62,9 @@ def register_client(
|
||||
json=payload,
|
||||
headers={"Content-Type": "application/json", "Accept": "application/json"},
|
||||
timeout=timeout,
|
||||
allow_redirects=False, # never re-send credentials to a redirect target
|
||||
)
|
||||
resp.raise_for_status()
|
||||
_require_success(resp)
|
||||
client_info = resp.json()
|
||||
|
||||
if not client_info.get("client_id"):
|
||||
|
||||
@@ -8,6 +8,7 @@ LAZY IMPORT: sqlite3 is only imported when this adapter is instantiated.
|
||||
"""
|
||||
|
||||
import os
|
||||
import re
|
||||
import sqlite3
|
||||
import json
|
||||
from praisonaiagents._logging import get_logger
|
||||
@@ -23,6 +24,39 @@ from ..protocols import MemoryProtocol
|
||||
logger = get_logger(__name__)
|
||||
|
||||
|
||||
_SEARCH_STOP_WORDS = frozenset({
|
||||
"a", "an", "and", "are", "as", "at", "be", "by", "can", "do", "does", "for",
|
||||
"from", "has", "have", "how", "i", "in", "is", "it", "me", "my", "of", "on",
|
||||
"or", "the", "to", "was", "were", "what", "when", "where", "which", "who",
|
||||
"why", "with", "you", "your",
|
||||
})
|
||||
|
||||
|
||||
def _execute_search(conn, table, query, limit, user_id=None):
|
||||
"""Match the whole query or any significant word, best matches first.
|
||||
|
||||
A plain ``LIKE '%query%'`` misses natural-language queries such as
|
||||
"Which mode does the user prefer?", so rows are also matched per word and
|
||||
ranked by exact-phrase hit, then number of matched words, then recency.
|
||||
"""
|
||||
terms = list(dict.fromkeys(
|
||||
w for w in re.findall(r"\w+", (query or "").lower())
|
||||
if w not in _SEARCH_STOP_WORDS
|
||||
))[:32] # keep well under SQLite's bound-parameter limit
|
||||
patterns = [f"%{query}%"] + [f"%{t}%" for t in terms]
|
||||
match = " OR ".join(["content LIKE ?"] * len(patterns))
|
||||
# An empty score would leave ``ORDER BY 0``, which SQLite reads as a column index
|
||||
score = " + ".join(["(content LIKE ?)"] * len(terms)) + " DESC, " if terms else ""
|
||||
user_filter = " AND memory_user_id(metadata) = ?" if user_id is not None else ""
|
||||
params = patterns + ([str(user_id)] if user_id is not None else []) + patterns + [limit]
|
||||
return conn.execute(
|
||||
f"SELECT id, content, metadata, timestamp FROM {table} "
|
||||
f"WHERE ({match}){user_filter} "
|
||||
f"ORDER BY (content LIKE ?) DESC, {score}timestamp DESC LIMIT ?",
|
||||
params,
|
||||
)
|
||||
|
||||
|
||||
def _metadata_user_id(raw):
|
||||
"""Match the existing Python JSON reader, including non-finite values."""
|
||||
try:
|
||||
@@ -206,14 +240,7 @@ class SqliteMemoryAdapter:
|
||||
"""Search short-term memory."""
|
||||
with self._write_lock if self.short_db == ":memory:" else nullcontext():
|
||||
conn = self._get_stm_conn()
|
||||
user_id = kwargs.get("user_id")
|
||||
user_filter = " AND memory_user_id(metadata) = ?" if user_id is not None else ""
|
||||
params = (f"%{query}%", str(user_id), limit) if user_id is not None else (f"%{query}%", limit)
|
||||
cursor = conn.execute(
|
||||
"SELECT id, content, metadata, timestamp FROM short_term_memory "
|
||||
"WHERE content LIKE ?" + user_filter + " ORDER BY timestamp DESC LIMIT ?",
|
||||
params
|
||||
)
|
||||
cursor = _execute_search(conn, "short_term_memory", query, limit, kwargs.get("user_id"))
|
||||
|
||||
results = []
|
||||
for row in cursor.fetchall():
|
||||
@@ -250,14 +277,7 @@ class SqliteMemoryAdapter:
|
||||
"""Search long-term memory."""
|
||||
with self._write_lock if self.long_db == ":memory:" else nullcontext():
|
||||
conn = self._get_ltm_conn()
|
||||
user_id = kwargs.get("user_id")
|
||||
user_filter = " AND memory_user_id(metadata) = ?" if user_id is not None else ""
|
||||
params = (f"%{query}%", str(user_id), limit) if user_id is not None else (f"%{query}%", limit)
|
||||
cursor = conn.execute(
|
||||
"SELECT id, content, metadata, timestamp FROM long_term_memory "
|
||||
"WHERE content LIKE ?" + user_filter + " ORDER BY timestamp DESC LIMIT ?",
|
||||
params
|
||||
)
|
||||
cursor = _execute_search(conn, "long_term_memory", query, limit, kwargs.get("user_id"))
|
||||
|
||||
results = []
|
||||
for row in cursor.fetchall():
|
||||
|
||||
@@ -35,6 +35,7 @@ import os
|
||||
import logging
|
||||
from praisonaiagents._logging import get_logger
|
||||
from pathlib import Path
|
||||
from ..utils.safe_path import safe_child
|
||||
from typing import Any, Dict, List, Optional, Literal
|
||||
from dataclasses import dataclass, field
|
||||
|
||||
@@ -366,7 +367,7 @@ class DocsManager:
|
||||
docs_dir = self.workspace_path / self.DOCS_DIR_NAME.replace("/", os.sep)
|
||||
|
||||
docs_dir.mkdir(parents=True, exist_ok=True)
|
||||
file_path = docs_dir / f"{name}.md"
|
||||
file_path = safe_child(docs_dir, name, ".md")
|
||||
|
||||
# Build frontmatter
|
||||
import yaml
|
||||
|
||||
@@ -23,6 +23,7 @@ import sys
|
||||
import hashlib
|
||||
import logging
|
||||
from praisonaiagents._logging import get_logger
|
||||
from praisonaiagents.utils.safe_path import safe_child
|
||||
import threading
|
||||
import weakref
|
||||
from collections import Counter
|
||||
@@ -142,6 +143,15 @@ class FileMemory:
|
||||
raise ValueError("user_id must not contain path separators or parent references")
|
||||
safe = user_id.strip()
|
||||
return safe or "default"
|
||||
|
||||
def _episodic_file(self, date: str) -> Path:
|
||||
"""Map a YYYY-MM-DD date to its episodic file, rejecting anything else."""
|
||||
try:
|
||||
# Normalise so "2026-3-5" lands in the same file recent reads look for
|
||||
date = datetime.strptime(date, "%Y-%m-%d").strftime("%Y-%m-%d")
|
||||
except (TypeError, ValueError):
|
||||
raise ValueError(f"date must be YYYY-MM-DD, got {date!r}") from None
|
||||
return safe_child(self.episodic_path, date, ".json")
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
@@ -753,10 +763,7 @@ class FileMemory:
|
||||
if date is None:
|
||||
date = datetime.now().strftime("%Y-%m-%d")
|
||||
|
||||
episodic_file = self.episodic_path / f"{date}.json"
|
||||
|
||||
# Load existing episodic memories for this date
|
||||
episodes = self._read_json(episodic_file, [])
|
||||
episodic_file = self._episodic_file(date)
|
||||
|
||||
# Create new item
|
||||
item = MemoryItem(
|
||||
@@ -765,8 +772,11 @@ class FileMemory:
|
||||
metadata=metadata or {}
|
||||
)
|
||||
|
||||
episodes.append(item.to_dict())
|
||||
self._write_json(episodic_file, episodes)
|
||||
# Read-modify-write under the store lock so concurrent adds are not lost
|
||||
with self._lock:
|
||||
episodes = self._read_json(episodic_file, [])
|
||||
episodes.append(item.to_dict())
|
||||
self._write_json(episodic_file, episodes)
|
||||
|
||||
self._log(f"Added episodic memory for {date}: {content[:50]}...")
|
||||
|
||||
@@ -788,7 +798,7 @@ class FileMemory:
|
||||
List of episodic memories
|
||||
"""
|
||||
if date:
|
||||
episodic_file = self.episodic_path / f"{date}.json"
|
||||
episodic_file = self._episodic_file(date)
|
||||
data = self._read_json(episodic_file, [])
|
||||
return [MemoryItem.from_dict(item) for item in data]
|
||||
|
||||
@@ -813,7 +823,8 @@ class FileMemory:
|
||||
try:
|
||||
file_date = datetime.strptime(episodic_file.stem, "%Y-%m-%d")
|
||||
if file_date < cutoff_date:
|
||||
episodic_file.unlink()
|
||||
with self._lock:
|
||||
episodic_file.unlink(missing_ok=True)
|
||||
removed += 1
|
||||
except ValueError:
|
||||
continue
|
||||
@@ -1193,16 +1204,17 @@ class FileMemory:
|
||||
for i in range(30)
|
||||
]
|
||||
|
||||
for check_date in dates_to_check:
|
||||
episodic_file = self.episodic_path / f"{check_date}.json"
|
||||
if episodic_file.exists():
|
||||
data = self._read_json(episodic_file, [])
|
||||
for i, item in enumerate(data):
|
||||
if item.get("id") == memory_id:
|
||||
del data[i]
|
||||
self._write_json(episodic_file, data)
|
||||
self._log(f"Deleted episodic memory: {memory_id}")
|
||||
return True
|
||||
with self._lock:
|
||||
for check_date in dates_to_check:
|
||||
episodic_file = self._episodic_file(check_date)
|
||||
if episodic_file.exists():
|
||||
data = self._read_json(episodic_file, [])
|
||||
for i, item in enumerate(data):
|
||||
if item.get("id") == memory_id:
|
||||
del data[i]
|
||||
self._write_json(episodic_file, data)
|
||||
self._log(f"Deleted episodic memory: {memory_id}")
|
||||
return True
|
||||
|
||||
return False
|
||||
|
||||
@@ -1440,7 +1452,7 @@ class FileMemory:
|
||||
sessions_path = self.user_path / "sessions"
|
||||
sessions_path.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
session_file = sessions_path / f"{name}.json"
|
||||
session_file = safe_child(sessions_path, name, ".json")
|
||||
|
||||
# Take snapshots under lock for consistent session data
|
||||
with self._lock:
|
||||
@@ -1479,7 +1491,7 @@ class FileMemory:
|
||||
Session data including conversation history
|
||||
"""
|
||||
sessions_path = self.user_path / "sessions"
|
||||
session_file = sessions_path / f"{name}.json"
|
||||
session_file = safe_child(sessions_path, name, ".json")
|
||||
|
||||
if not session_file.exists():
|
||||
raise FileNotFoundError(f"Session '{name}' not found")
|
||||
@@ -1531,7 +1543,7 @@ class FileMemory:
|
||||
def delete_session(self, name: str) -> bool:
|
||||
"""Delete a saved session."""
|
||||
sessions_path = self.user_path / "sessions"
|
||||
session_file = sessions_path / f"{name}.json"
|
||||
session_file = safe_child(sessions_path, name, ".json")
|
||||
|
||||
if session_file.exists():
|
||||
session_file.unlink()
|
||||
@@ -1687,7 +1699,7 @@ Summary:"""
|
||||
|
||||
from uuid import uuid4
|
||||
checkpoint_id = name or f"checkpoint_{int(time.time())}_{uuid4().hex}"
|
||||
checkpoint_file = checkpoints_path / f"{checkpoint_id}.json"
|
||||
checkpoint_file = safe_child(checkpoints_path, checkpoint_id, ".json")
|
||||
|
||||
checkpoint_data = {
|
||||
"id": checkpoint_id,
|
||||
@@ -1724,7 +1736,7 @@ Summary:"""
|
||||
True if restored successfully
|
||||
"""
|
||||
checkpoints_path = self.user_path / "checkpoints"
|
||||
checkpoint_file = checkpoints_path / f"{checkpoint_id}.json"
|
||||
checkpoint_file = safe_child(checkpoints_path, checkpoint_id, ".json")
|
||||
|
||||
if not checkpoint_file.exists():
|
||||
self._log(f"Checkpoint '{checkpoint_id}' not found", logging.ERROR)
|
||||
@@ -1775,7 +1787,7 @@ Summary:"""
|
||||
def delete_checkpoint(self, checkpoint_id: str) -> bool:
|
||||
"""Delete a checkpoint."""
|
||||
checkpoints_path = self.user_path / "checkpoints"
|
||||
checkpoint_file = checkpoints_path / f"{checkpoint_id}.json"
|
||||
checkpoint_file = safe_child(checkpoints_path, checkpoint_id, ".json")
|
||||
|
||||
if checkpoint_file.exists():
|
||||
checkpoint_file.unlink()
|
||||
|
||||
@@ -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]:
|
||||
|
||||
@@ -35,6 +35,7 @@ import json
|
||||
import logging
|
||||
from praisonaiagents._logging import get_logger
|
||||
from pathlib import Path
|
||||
from ..utils.safe_path import safe_child
|
||||
from typing import Any, Dict, List, Optional
|
||||
from dataclasses import dataclass, field
|
||||
|
||||
@@ -272,7 +273,7 @@ class MCPConfigManager:
|
||||
mcp_dir = self.workspace_path / self.MCP_DIR_NAME.replace("/", os.sep)
|
||||
|
||||
mcp_dir.mkdir(parents=True, exist_ok=True)
|
||||
file_path = mcp_dir / f"{name}.json"
|
||||
file_path = safe_child(mcp_dir, name, ".json")
|
||||
|
||||
# Build config
|
||||
config = MCPConfig(
|
||||
|
||||
@@ -44,6 +44,7 @@ import logging
|
||||
from praisonaiagents._logging import get_logger
|
||||
import subprocess
|
||||
from pathlib import Path
|
||||
from ..utils.safe_path import safe_child
|
||||
from typing import Any, Dict, List, Optional, Literal, Callable
|
||||
from dataclasses import dataclass, field
|
||||
|
||||
@@ -515,9 +516,13 @@ class RulesManager:
|
||||
dir_prefix = rules_dir_name.split("/")[0].strip(".")
|
||||
self._rules[f"gitroot_{dir_prefix}:{rule.name}"] = rule
|
||||
|
||||
# 6. Load subdirectory rules (walk up from cwd to workspace/git root)
|
||||
stop_at = git_root if git_root else self.workspace_path
|
||||
current = Path.cwd()
|
||||
# 6. Load subdirectory rules (walk up from cwd to workspace/git root).
|
||||
# Skipped when cwd is outside that tree, so ancestors such as the
|
||||
# user's home directory are never scanned.
|
||||
stop_at = (git_root if git_root else self.workspace_path).resolve()
|
||||
current = Path.cwd().resolve()
|
||||
if not current.is_relative_to(stop_at):
|
||||
current = stop_at
|
||||
while current != stop_at and current != current.parent:
|
||||
# Check .praisonai/rules/
|
||||
subdir_rules_dir = current / self.RULES_DIR_NAME.replace("/", os.sep)
|
||||
@@ -808,7 +813,7 @@ class RulesManager:
|
||||
rules_dir = self.workspace_path / self.RULES_DIR_NAME.replace("/", os.sep)
|
||||
|
||||
rules_dir.mkdir(parents=True, exist_ok=True)
|
||||
file_path = rules_dir / f"{name}.md"
|
||||
file_path = safe_child(rules_dir, name, ".md")
|
||||
|
||||
# Build frontmatter
|
||||
import yaml
|
||||
|
||||
@@ -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)
|
||||
@@ -231,7 +231,14 @@ def load_plugin(filepath: str) -> Optional[Dict[str, Any]]:
|
||||
except PluginParseError as e:
|
||||
logger.error(f"Invalid plugin header: {e}")
|
||||
return None
|
||||
|
||||
|
||||
# Refuse to exec a plugin built for an incompatible core.
|
||||
from .plugin import core_incompatibility
|
||||
reason = core_incompatibility(metadata.get("dependencies", []))
|
||||
if reason:
|
||||
logger.warning(f"Plugin '{metadata.get('name')}' disabled: {reason}")
|
||||
return None
|
||||
|
||||
# Generate unique module name to avoid conflicts
|
||||
module_name = f"praison_plugin_{path.stem}_{id(path)}"
|
||||
|
||||
|
||||
@@ -13,7 +13,7 @@ import threading
|
||||
from pathlib import Path
|
||||
from typing import Any, Callable, Dict, Iterator, List, Optional, Tuple, TYPE_CHECKING
|
||||
|
||||
from .plugin import Plugin, PluginHook, PluginInfo, FunctionPlugin
|
||||
from .plugin import Plugin, PluginHook, PluginInfo, FunctionPlugin, core_incompatibility
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from ..hooks.registry import HookRegistry
|
||||
@@ -143,6 +143,9 @@ class PluginManager:
|
||||
def __init__(self, disabled: Optional[bool] = None):
|
||||
self._plugins: Dict[str, Plugin] = {}
|
||||
self._enabled: Dict[str, bool] = {}
|
||||
# Names whose on_init is running and that have not been explicitly
|
||||
# enabled/disabled yet; register() applies the default only to these.
|
||||
self._initializing: set = set()
|
||||
self._single_file_plugins: Dict[str, Dict[str, Any]] = {} # WordPress-style plugins
|
||||
# Per-plugin option maps sourced from the unified project config,
|
||||
# delivered to each plugin's on_config hook. {plugin_name: options}.
|
||||
@@ -262,11 +265,10 @@ class PluginManager:
|
||||
if self._enabled.get(name, False) and name in self._plugin_options
|
||||
]
|
||||
for name, plugin, options in targets:
|
||||
on_config = getattr(plugin, "on_config", None)
|
||||
if not callable(on_config):
|
||||
if not callable(getattr(plugin, "on_config", None)):
|
||||
continue
|
||||
try:
|
||||
on_config(options)
|
||||
_ResolvingPlugin(plugin).on_config(options)
|
||||
delivered += 1
|
||||
except Exception as e:
|
||||
logger.warning(f"Plugin '{name}' on_config failed: {e}")
|
||||
@@ -288,26 +290,54 @@ class PluginManager:
|
||||
if info.name in self._plugins:
|
||||
logger.warning(f"Plugin '{info.name}' already registered")
|
||||
return False
|
||||
|
||||
self._plugins[info.name] = plugin
|
||||
self._enabled[info.name] = self._default_enabled_for_registration(info.name)
|
||||
|
||||
# Initialize plugin
|
||||
plugin.on_init({})
|
||||
|
||||
# Make the initialization observable to the hook system
|
||||
# too. ``plugin.on_init`` is the plugin's *own* callback;
|
||||
# ON_INIT is the runtime event other listeners
|
||||
# (metrics/audit plugins, gateways) subscribe to. Without
|
||||
# this emit the enum member was declared but never fired.
|
||||
_emit_plugin_init(plugin)
|
||||
|
||||
logger.info(f"Registered plugin: {info.name} v{info.version}")
|
||||
return True
|
||||
reason = core_incompatibility(info.dependencies)
|
||||
if reason:
|
||||
logger.warning(f"Plugin '{info.name}' disabled: {reason}")
|
||||
return False
|
||||
|
||||
# Reserve the name but keep it out of dispatch until on_init succeeds.
|
||||
self._plugins[info.name] = plugin
|
||||
self._enabled[info.name] = False
|
||||
self._initializing.add(info.name)
|
||||
default_enabled = self._default_enabled_for_registration(info.name)
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to register plugin: {e}")
|
||||
return False
|
||||
|
||||
|
||||
# Initialize outside the lock: an ``async def on_init`` may run on
|
||||
# another thread and call back into this manager.
|
||||
try:
|
||||
_ResolvingPlugin(plugin).on_init({})
|
||||
|
||||
with self._lock:
|
||||
if self._plugins.get(info.name) is not plugin:
|
||||
logger.warning(f"Plugin '{info.name}' was unregistered during on_init")
|
||||
return False
|
||||
# Publish; an explicit enable()/disable() made during on_init is kept.
|
||||
if info.name in self._initializing:
|
||||
self._initializing.discard(info.name)
|
||||
self._enabled[info.name] = default_enabled
|
||||
|
||||
# Make the initialization observable to the hook system
|
||||
# too. ``plugin.on_init`` is the plugin's *own* callback;
|
||||
# ON_INIT is the runtime event other listeners
|
||||
# (metrics/audit plugins, gateways) subscribe to. Without
|
||||
# this emit the enum member was declared but never fired.
|
||||
_emit_plugin_init(plugin)
|
||||
|
||||
logger.info(f"Registered plugin: {info.name} v{info.version}")
|
||||
return True
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to register plugin: {e}")
|
||||
# Roll back so a failed plugin is not left half-registered.
|
||||
with self._lock:
|
||||
if self._plugins.get(info.name) is plugin:
|
||||
del self._plugins[info.name]
|
||||
self._enabled.pop(info.name, None)
|
||||
self._initializing.discard(info.name)
|
||||
return False
|
||||
|
||||
def unregister(self, name: str) -> bool:
|
||||
"""
|
||||
Unregister a plugin.
|
||||
@@ -324,9 +354,10 @@ class PluginManager:
|
||||
return False
|
||||
del self._plugins[name]
|
||||
self._enabled.pop(name, None)
|
||||
self._initializing.discard(name)
|
||||
|
||||
try:
|
||||
plugin.on_shutdown()
|
||||
_ResolvingPlugin(plugin).on_shutdown()
|
||||
except Exception as e:
|
||||
logger.warning(f"Error during plugin shutdown: {e}")
|
||||
|
||||
@@ -342,6 +373,7 @@ class PluginManager:
|
||||
with self._lock:
|
||||
if name in self._enabled:
|
||||
self._enabled[name] = True
|
||||
self._initializing.discard(name)
|
||||
return True
|
||||
return False
|
||||
|
||||
@@ -356,6 +388,7 @@ class PluginManager:
|
||||
if name not in self._enabled:
|
||||
return False
|
||||
self._enabled[name] = False
|
||||
self._initializing.discard(name)
|
||||
try:
|
||||
self.unwire_from_hook_registry(name)
|
||||
except Exception as e:
|
||||
@@ -537,7 +570,8 @@ class PluginManager:
|
||||
continue
|
||||
|
||||
try:
|
||||
method = getattr(plugin, hook.value, None)
|
||||
# Resolving view: an ``async def`` hook is awaited, not dropped.
|
||||
method = getattr(_ResolvingPlugin(plugin), hook.value, None)
|
||||
if method:
|
||||
if hook == PluginHook.BEFORE_TOOL:
|
||||
result = method(args[0], args[1] if len(args) > 1 else {})
|
||||
@@ -839,21 +873,20 @@ class PluginManager:
|
||||
return 0
|
||||
|
||||
loaded = 0
|
||||
# Use lock to ensure thread-safe discovery of multiple entry points.
|
||||
# register() also acquires the lock but that's fine (RLock allows re-entry).
|
||||
with self._lock:
|
||||
for ep in entry_points:
|
||||
try:
|
||||
plugin_cls = ep.load()
|
||||
if callable(plugin_cls):
|
||||
plugin = plugin_cls()
|
||||
if self.register(plugin):
|
||||
loaded += 1
|
||||
logger.info(f"Loaded entry_point plugin: {ep.name}")
|
||||
else:
|
||||
logger.warning(f"Entry point {ep.name} is not callable")
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to load plugin from entry point {ep.name}: {e}")
|
||||
# register() is thread-safe on its own; holding the lock here would
|
||||
# keep it held while a plugin's (possibly async) on_init runs.
|
||||
for ep in entry_points:
|
||||
try:
|
||||
plugin_cls = ep.load()
|
||||
if callable(plugin_cls):
|
||||
plugin = plugin_cls()
|
||||
if self.register(plugin):
|
||||
loaded += 1
|
||||
logger.info(f"Loaded entry_point plugin: {ep.name}")
|
||||
else:
|
||||
logger.warning(f"Entry point {ep.name} is not callable")
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to load plugin from entry point {ep.name}: {e}")
|
||||
|
||||
return loaded
|
||||
|
||||
@@ -1008,6 +1041,44 @@ def _emit_plugin_lifecycle(event: "HookEvent", plugin: Plugin) -> None:
|
||||
logger.debug("Plugin lifecycle hook failed", exc_info=True)
|
||||
|
||||
|
||||
# Deadline for an ``async def`` plugin hook; matches HookRunner's default
|
||||
# per-hook timeout so a hung hook is cancelled and frees its worker thread.
|
||||
_ASYNC_HOOK_TIMEOUT = 60.0
|
||||
|
||||
|
||||
def _resolve_awaitable(value: Any) -> Any:
|
||||
"""Drive a coroutine returned by an ``async def`` plugin hook to completion.
|
||||
|
||||
Hook adapters run on the HookRunner's worker thread (no running loop), so
|
||||
the coroutine runs on a private loop there; from a thread with a running
|
||||
loop the shared bridge offloads it instead of deadlocking. The coroutine is
|
||||
bounded by ``_ASYNC_HOOK_TIMEOUT``. Sync hook return values pass through
|
||||
unchanged.
|
||||
"""
|
||||
if asyncio.iscoroutine(value):
|
||||
from ..utils.async_bridge import run_coroutine_from_any_context
|
||||
return run_coroutine_from_any_context(value, timeout=_ASYNC_HOOK_TIMEOUT)
|
||||
return value
|
||||
|
||||
|
||||
class _ResolvingPlugin:
|
||||
"""View of a plugin whose method calls resolve ``async def`` hooks."""
|
||||
|
||||
__slots__ = ("_plugin",)
|
||||
|
||||
def __init__(self, plugin: Plugin):
|
||||
self._plugin = plugin
|
||||
|
||||
def __getattr__(self, name: str) -> Any:
|
||||
attr = getattr(self._plugin, name)
|
||||
if not callable(attr):
|
||||
return attr
|
||||
|
||||
def call(*args, **kwargs):
|
||||
return _resolve_awaitable(attr(*args, **kwargs))
|
||||
return call
|
||||
|
||||
|
||||
def _write_back(data, attr: str, new_value: dict) -> None:
|
||||
"""Write a plugin's returned dict back onto the hook payload.
|
||||
|
||||
@@ -1044,6 +1115,9 @@ def _adapt_plugin_hooks(plugin: Plugin) -> Iterator[Tuple["HookEvent", Callable]
|
||||
from .plugin import GuardrailBlocked, PluginDecision
|
||||
|
||||
base = Plugin
|
||||
# Adapters call the plugin through this view so ``async def`` hooks are
|
||||
# awaited instead of returning an unawaited coroutine.
|
||||
resolving_plugin = _ResolvingPlugin(plugin)
|
||||
|
||||
def _as_decision(result: Any) -> Optional[HookResult]:
|
||||
"""Coerce a plugin's return value into a denying ``HookResult``.
|
||||
@@ -1081,7 +1155,7 @@ def _adapt_plugin_hooks(plugin: Plugin) -> Iterator[Tuple["HookEvent", Callable]
|
||||
return HookResult.allow()
|
||||
|
||||
if _overrides("before_agent"):
|
||||
def before_agent_hook(data, _p=plugin):
|
||||
def before_agent_hook(data, _p=resolving_plugin):
|
||||
try:
|
||||
new_prompt = _p.before_agent(getattr(data, "prompt", ""),
|
||||
{"agent_name": getattr(data, "agent_name", None)})
|
||||
@@ -1096,7 +1170,7 @@ def _adapt_plugin_hooks(plugin: Plugin) -> Iterator[Tuple["HookEvent", Callable]
|
||||
yield HookEvent.BEFORE_AGENT, before_agent_hook
|
||||
|
||||
if _overrides("after_agent"):
|
||||
def after_agent_hook(data, _p=plugin):
|
||||
def after_agent_hook(data, _p=resolving_plugin):
|
||||
new_resp = _p.after_agent(getattr(data, "response", ""),
|
||||
{"agent_name": getattr(data, "agent_name", None)})
|
||||
if isinstance(new_resp, str) and hasattr(data, "response"):
|
||||
@@ -1105,7 +1179,7 @@ def _adapt_plugin_hooks(plugin: Plugin) -> Iterator[Tuple["HookEvent", Callable]
|
||||
yield HookEvent.AFTER_AGENT, after_agent_hook
|
||||
|
||||
if _overrides("before_llm"):
|
||||
def before_llm_hook(data, _p=plugin):
|
||||
def before_llm_hook(data, _p=resolving_plugin):
|
||||
messages = getattr(data, "messages", [])
|
||||
try:
|
||||
result = _p.before_llm(messages, {"model": getattr(data, "model", "")})
|
||||
@@ -1121,7 +1195,7 @@ def _adapt_plugin_hooks(plugin: Plugin) -> Iterator[Tuple["HookEvent", Callable]
|
||||
yield HookEvent.BEFORE_LLM, before_llm_hook
|
||||
|
||||
if _overrides("after_llm"):
|
||||
def after_llm_hook(data, _p=plugin):
|
||||
def after_llm_hook(data, _p=resolving_plugin):
|
||||
new_resp = _p.after_llm(getattr(data, "response", ""),
|
||||
{"tokens_used": getattr(data, "tokens_used", 0)})
|
||||
if isinstance(new_resp, str) and hasattr(data, "response"):
|
||||
@@ -1130,7 +1204,7 @@ def _adapt_plugin_hooks(plugin: Plugin) -> Iterator[Tuple["HookEvent", Callable]
|
||||
yield HookEvent.AFTER_LLM, after_llm_hook
|
||||
|
||||
if _overrides("before_tool"):
|
||||
def before_tool_hook(data, _p=plugin):
|
||||
def before_tool_hook(data, _p=resolving_plugin):
|
||||
args = getattr(data, "tool_input", {}) or {}
|
||||
try:
|
||||
new_args = _p.before_tool(getattr(data, "tool_name", ""), args)
|
||||
@@ -1146,7 +1220,7 @@ def _adapt_plugin_hooks(plugin: Plugin) -> Iterator[Tuple["HookEvent", Callable]
|
||||
yield HookEvent.BEFORE_TOOL, before_tool_hook
|
||||
|
||||
if _overrides("after_tool"):
|
||||
def after_tool_hook(data, _p=plugin):
|
||||
def after_tool_hook(data, _p=resolving_plugin):
|
||||
try:
|
||||
new_result = _p.after_tool(getattr(data, "tool_name", ""),
|
||||
getattr(data, "tool_output", None))
|
||||
@@ -1161,7 +1235,7 @@ def _adapt_plugin_hooks(plugin: Plugin) -> Iterator[Tuple["HookEvent", Callable]
|
||||
yield HookEvent.AFTER_TOOL, after_tool_hook
|
||||
|
||||
if _overrides("before_tool_definitions"):
|
||||
def before_tool_definitions_hook(data, _p=plugin):
|
||||
def before_tool_definitions_hook(data, _p=resolving_plugin):
|
||||
defs = getattr(data, "tool_definitions", []) or []
|
||||
try:
|
||||
new_defs = _p.before_tool_definitions(defs)
|
||||
@@ -1193,7 +1267,7 @@ def _adapt_plugin_hooks(plugin: Plugin) -> Iterator[Tuple["HookEvent", Callable]
|
||||
return payload
|
||||
|
||||
if _overrides("before_message"):
|
||||
def before_message_hook(data, _p=plugin):
|
||||
def before_message_hook(data, _p=resolving_plugin):
|
||||
try:
|
||||
new = _p.before_message(_message_payload(data))
|
||||
except GuardrailBlocked as e:
|
||||
@@ -1207,7 +1281,7 @@ def _adapt_plugin_hooks(plugin: Plugin) -> Iterator[Tuple["HookEvent", Callable]
|
||||
yield HookEvent.MESSAGE_RECEIVED, before_message_hook
|
||||
|
||||
if _overrides("after_message"):
|
||||
def after_message_hook(data, _p=plugin):
|
||||
def after_message_hook(data, _p=resolving_plugin):
|
||||
try:
|
||||
new = _p.after_message(_message_payload(data))
|
||||
except GuardrailBlocked as e:
|
||||
@@ -1224,14 +1298,14 @@ def _adapt_plugin_hooks(plugin: Plugin) -> Iterator[Tuple["HookEvent", Callable]
|
||||
yield HookEvent.MESSAGE_SENDING, after_message_hook
|
||||
|
||||
if _overrides("message_sent"):
|
||||
def message_sent_hook(data, _p=plugin):
|
||||
def message_sent_hook(data, _p=resolving_plugin):
|
||||
payload = _message_payload(data)
|
||||
_p.message_sent(payload)
|
||||
return HookResult.allow()
|
||||
yield HookEvent.MESSAGE_SENT, message_sent_hook
|
||||
|
||||
if _overrides("message_undelivered"):
|
||||
def message_undelivered_hook(data, _p=plugin):
|
||||
def message_undelivered_hook(data, _p=resolving_plugin):
|
||||
payload = _message_payload(data)
|
||||
payload["error"] = getattr(data, "error", "")
|
||||
payload["notice_delivered"] = getattr(data, "notice_delivered", False)
|
||||
@@ -1240,14 +1314,14 @@ def _adapt_plugin_hooks(plugin: Plugin) -> Iterator[Tuple["HookEvent", Callable]
|
||||
yield HookEvent.MESSAGE_UNDELIVERED, message_undelivered_hook
|
||||
|
||||
if _overrides("on_permission_ask"):
|
||||
def on_permission_ask_hook(data, _p=plugin):
|
||||
def on_permission_ask_hook(data, _p=resolving_plugin):
|
||||
target = getattr(data, "tool_name", "") or getattr(data, "target", "")
|
||||
reason = getattr(data, "reason", "") or "Permission requested"
|
||||
return _permission_result(_p.on_permission_ask(target, reason), reason)
|
||||
yield HookEvent.ON_PERMISSION_ASK, on_permission_ask_hook
|
||||
|
||||
if _overrides("on_config"):
|
||||
def on_config_hook(data, _p=plugin):
|
||||
def on_config_hook(data, _p=resolving_plugin):
|
||||
config = getattr(data, "config", None)
|
||||
if not isinstance(config, dict):
|
||||
config = getattr(data, "extra", {}) or {}
|
||||
@@ -1258,7 +1332,7 @@ def _adapt_plugin_hooks(plugin: Plugin) -> Iterator[Tuple["HookEvent", Callable]
|
||||
yield HookEvent.ON_CONFIG, on_config_hook
|
||||
|
||||
if _overrides("on_auth"):
|
||||
def on_auth_hook(data, _p=plugin):
|
||||
def on_auth_hook(data, _p=resolving_plugin):
|
||||
auth_type = getattr(data, "auth_type", "") or ""
|
||||
credentials = getattr(data, "credentials", None) or {}
|
||||
new_creds = _p.on_auth(auth_type, credentials)
|
||||
@@ -1268,7 +1342,7 @@ def _adapt_plugin_hooks(plugin: Plugin) -> Iterator[Tuple["HookEvent", Callable]
|
||||
yield HookEvent.ON_AUTH, on_auth_hook
|
||||
|
||||
if _overrides("session_start"):
|
||||
def session_start_hook(data, _p=plugin):
|
||||
def session_start_hook(data, _p=resolving_plugin):
|
||||
_p.session_start({
|
||||
"source": getattr(data, "source", None),
|
||||
"session_name": getattr(data, "session_name", None),
|
||||
@@ -1278,7 +1352,7 @@ def _adapt_plugin_hooks(plugin: Plugin) -> Iterator[Tuple["HookEvent", Callable]
|
||||
yield HookEvent.SESSION_START, session_start_hook
|
||||
|
||||
if _overrides("session_end"):
|
||||
def session_end_hook(data, _p=plugin):
|
||||
def session_end_hook(data, _p=resolving_plugin):
|
||||
_p.session_end({
|
||||
"reason": getattr(data, "reason", None),
|
||||
"total_turns": getattr(data, "total_turns", None),
|
||||
@@ -1289,7 +1363,7 @@ def _adapt_plugin_hooks(plugin: Plugin) -> Iterator[Tuple["HookEvent", Callable]
|
||||
yield HookEvent.SESSION_END, session_end_hook
|
||||
|
||||
if _overrides("on_error"):
|
||||
def on_error_hook(data, _p=plugin):
|
||||
def on_error_hook(data, _p=resolving_plugin):
|
||||
_p.on_error(
|
||||
getattr(data, "error_type", "") or "",
|
||||
getattr(data, "error_message", "") or "",
|
||||
@@ -1303,7 +1377,7 @@ def _adapt_plugin_hooks(plugin: Plugin) -> Iterator[Tuple["HookEvent", Callable]
|
||||
yield HookEvent.ON_ERROR, on_error_hook
|
||||
|
||||
if _overrides("cli_backend_execute"):
|
||||
def cli_backend_execute_hook(data, _p=plugin):
|
||||
def cli_backend_execute_hook(data, _p=resolving_plugin):
|
||||
from ..cli_backend.debug import redact_command
|
||||
_p.cli_backend_execute({
|
||||
"agent_name": getattr(data, "agent_name", None),
|
||||
@@ -1361,7 +1435,7 @@ def _adapt_plugin_hooks(plugin: Plugin) -> Iterator[Tuple["HookEvent", Callable]
|
||||
# ``kanban_task_created``) are treated as genuine overrides.
|
||||
if parent is not None and own is parent:
|
||||
continue
|
||||
yield event, _make_observe_hook(method)
|
||||
yield event, _make_observe_hook(getattr(resolving_plugin, event.value))
|
||||
|
||||
|
||||
# Global plugin manager instance
|
||||
|
||||
@@ -111,14 +111,75 @@ class PluginInfo:
|
||||
"plugin_type": self.plugin_type.value,
|
||||
}
|
||||
|
||||
|
||||
def core_incompatibility(dependencies: List[str]) -> Optional[str]:
|
||||
"""Check a plugin's ``dependencies`` against the running core.
|
||||
|
||||
Only a ``praisonaiagents`` entry (e.g. ``"praisonaiagents>=0.9,<2"``) is
|
||||
checked; other entries are ignored. Returns an actionable reason when the
|
||||
installed core does not satisfy it, else ``None`` (permissive: no entry,
|
||||
an unparseable entry, or ``packaging`` unavailable all mean compatible).
|
||||
"""
|
||||
if not dependencies:
|
||||
return None
|
||||
try:
|
||||
from packaging.requirements import InvalidRequirement, Requirement
|
||||
from packaging.utils import canonicalize_name
|
||||
from packaging.version import InvalidVersion, Version
|
||||
from .._version import __version__ as core_version
|
||||
except ImportError:
|
||||
return None
|
||||
if core_version == "unknown":
|
||||
# Wheel installs ship no pyproject.toml; read the installed metadata.
|
||||
try:
|
||||
from importlib.metadata import version
|
||||
core_version = version("praisonaiagents")
|
||||
except Exception:
|
||||
return None
|
||||
|
||||
# Comma-separated manifest headers split "pkg>=1,<2" into ["pkg>=1", "<2"];
|
||||
# rejoin bare specifier fragments onto the preceding requirement.
|
||||
merged: List[str] = []
|
||||
for dep in dependencies:
|
||||
dep = str(dep).strip()
|
||||
if merged and dep[:1] in "<>=!~":
|
||||
merged[-1] = f"{merged[-1]},{dep}"
|
||||
elif dep:
|
||||
merged.append(dep)
|
||||
|
||||
for dep in merged:
|
||||
try:
|
||||
req = Requirement(dep)
|
||||
except InvalidRequirement:
|
||||
continue
|
||||
if canonicalize_name(req.name) != "praisonaiagents" or not req.specifier:
|
||||
continue
|
||||
try:
|
||||
if req.marker and not req.marker.evaluate({"extra": ""}):
|
||||
continue # e.g. "; sys_platform == 'win32'" on another OS
|
||||
except Exception:
|
||||
continue
|
||||
try:
|
||||
ok = req.specifier.contains(Version(core_version), prereleases=True)
|
||||
except InvalidVersion:
|
||||
return None
|
||||
if not ok:
|
||||
return (
|
||||
f"requires praisonaiagents{req.specifier}, installed {core_version}. "
|
||||
"Update the plugin or praisonaiagents."
|
||||
)
|
||||
return None
|
||||
|
||||
|
||||
class Plugin(ABC):
|
||||
"""
|
||||
Base class for PraisonAI plugins.
|
||||
|
||||
Plugins can extend agent functionality by implementing
|
||||
hook methods that are called at various points in the
|
||||
agent lifecycle.
|
||||
|
||||
agent lifecycle. Hook methods may be ``def`` or ``async def``;
|
||||
coroutines returned by async hooks are awaited by the dispatcher.
|
||||
|
||||
Example:
|
||||
class MyPlugin(Plugin):
|
||||
@property
|
||||
|
||||
@@ -115,11 +115,8 @@ class PolicyConfig:
|
||||
rules_dir = os.path.dirname(self.rules_file)
|
||||
os.makedirs(rules_dir, exist_ok=True)
|
||||
|
||||
# Atomic write: tmp file + rename
|
||||
tmp_file = self.rules_file + ".tmp"
|
||||
with open(tmp_file, "w") as f:
|
||||
json.dump({"rules": existing}, f, indent=2)
|
||||
os.replace(tmp_file, self.rules_file)
|
||||
from ..utils.atomic_io import atomic_write_json
|
||||
atomic_write_json(self.rules_file, {"rules": existing}, indent=2)
|
||||
|
||||
logger.debug(f"Saved learned rule: {resource} -> {action}")
|
||||
return True
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -11,6 +11,7 @@ from dataclasses import dataclass
|
||||
from typing import Any, Dict, List, Optional, Protocol, runtime_checkable
|
||||
|
||||
from ..context.budgeter import MODEL_LIMITS as _CANONICAL_LIMITS
|
||||
from ..context.tokens import estimate_tokens_heuristic
|
||||
|
||||
|
||||
# RAG-specific context windows not present in the canonical budgeter table
|
||||
@@ -285,11 +286,28 @@ def estimate_tokens(text: str) -> int:
|
||||
Returns:
|
||||
Estimated token count
|
||||
"""
|
||||
if not text:
|
||||
return 0
|
||||
# Average of 4 characters per token for English
|
||||
# This is a conservative estimate
|
||||
return len(text) // 4 + 1
|
||||
# Delegate to the canonical estimator so dense (CJK/non-ASCII) text is not
|
||||
# under-counted ~5x and over-filled past the budget.
|
||||
return estimate_tokens_heuristic(text)
|
||||
|
||||
|
||||
def _truncate_to_tokens(text: str, max_tokens: int, estimator=None) -> str:
|
||||
"""Return the longest prefix of ``text`` whose estimate fits ``max_tokens``.
|
||||
|
||||
Binary-searches on the (monotonic) estimator instead of assuming
|
||||
~4 chars/token, which overshoots the budget for dense (CJK) text.
|
||||
"""
|
||||
estimator = estimator or estimate_tokens
|
||||
if max_tokens <= 0 or not text:
|
||||
return ""
|
||||
lo, hi = 0, len(text)
|
||||
while lo < hi:
|
||||
mid = (lo + hi + 1) // 2
|
||||
if estimator(text[:mid]) <= max_tokens:
|
||||
lo = mid
|
||||
else:
|
||||
hi = mid - 1
|
||||
return text[:lo]
|
||||
|
||||
|
||||
class DefaultBudgetEnforcer:
|
||||
@@ -340,9 +358,12 @@ class DefaultBudgetEnforcer:
|
||||
# Try to fit partial chunk if there's remaining space
|
||||
remaining = available - used_tokens
|
||||
if remaining > 100: # Only if meaningful space remains
|
||||
# Truncate text to fit
|
||||
chars_to_keep = remaining * 4 # Approximate
|
||||
truncated_text = text[:chars_to_keep]
|
||||
# Truncate text to fit, leaving room for the "..." suffix
|
||||
truncated_text = _truncate_to_tokens(
|
||||
text,
|
||||
remaining - self._token_estimator("..."),
|
||||
self._token_estimator,
|
||||
)
|
||||
if truncated_text:
|
||||
truncated_chunk = {
|
||||
**chunk,
|
||||
|
||||
@@ -12,7 +12,7 @@ No heavy imports at module level.
|
||||
from dataclasses import dataclass, field
|
||||
from typing import Any, Dict, List, Optional, Protocol, runtime_checkable
|
||||
|
||||
from .budget import estimate_tokens
|
||||
from .budget import _truncate_to_tokens, estimate_tokens
|
||||
|
||||
|
||||
@runtime_checkable
|
||||
@@ -219,8 +219,9 @@ class ContextCompressor:
|
||||
# Partial fit
|
||||
remaining = target_tokens - used_tokens
|
||||
if remaining > 50: # Only if meaningful space
|
||||
chars_to_keep = remaining * 4
|
||||
truncated_text = text[:chars_to_keep] + "..."
|
||||
truncated_text = _truncate_to_tokens(
|
||||
text, remaining - estimate_tokens("...")
|
||||
) + "..."
|
||||
result.append({
|
||||
**chunk,
|
||||
'text': truncated_text,
|
||||
|
||||
@@ -8,7 +8,7 @@ No heavy imports - only stdlib.
|
||||
import hashlib
|
||||
from typing import TYPE_CHECKING, Any, Dict, List, Optional, Set, Tuple, Union
|
||||
|
||||
from .budget import estimate_tokens as _budget_estimate_tokens
|
||||
from .budget import _truncate_to_tokens, estimate_tokens as _budget_estimate_tokens
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from ..knowledge.models import SearchResultItem
|
||||
@@ -156,10 +156,10 @@ def truncate_context(
|
||||
# Calculate target character count
|
||||
suffix_tokens = _estimate_tokens(suffix)
|
||||
target_tokens = max_tokens - suffix_tokens
|
||||
target_chars = target_tokens * 4
|
||||
truncated = _truncate_to_tokens(text, target_tokens, _estimate_tokens)
|
||||
target_chars = len(truncated)
|
||||
|
||||
# Truncate at word boundary
|
||||
truncated = text[:target_chars]
|
||||
last_space = truncated.rfind(" ")
|
||||
if last_space > target_chars * 0.8:
|
||||
truncated = truncated[:last_space]
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -119,10 +119,15 @@ class ScheduleRunner:
|
||||
job_name=job.name,
|
||||
)
|
||||
|
||||
if job.should_retire(job.last_run_at):
|
||||
self._store.remove(job.id)
|
||||
else:
|
||||
self._store.update(job)
|
||||
# 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):
|
||||
self._store.remove(job.id)
|
||||
else:
|
||||
self._store.update(job)
|
||||
except OSError as e:
|
||||
logger.error("Failed to persist run state for job %r: %s", job.id, e)
|
||||
|
||||
# ── job-to-job context chaining ───────────────────────────────────
|
||||
|
||||
|
||||
@@ -16,6 +16,7 @@ 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 .hook_emit import emit_schedule_add, emit_schedule_remove
|
||||
from ..utils.atomic_io import atomic_write_json
|
||||
|
||||
logger = get_logger(__name__)
|
||||
|
||||
@@ -41,6 +42,8 @@ class FileScheduleStore:
|
||||
self._lock = threading.RLock()
|
||||
self._jobs: Dict[str, ScheduleJob] = {}
|
||||
self._history: List[RunRecord] = []
|
||||
# Run records whose history write failed; merged into the next save.
|
||||
self._unsaved: List[RunRecord] = []
|
||||
# In-memory record of leases we currently hold, so ``complete`` can
|
||||
# release them. Cross-process leases live in the persisted job
|
||||
# (``lease_until`` / ``lease_owner``) which is the source of truth.
|
||||
@@ -60,7 +63,7 @@ class FileScheduleStore:
|
||||
if job.id in self._jobs:
|
||||
raise ValueError(f"Job '{job.id}' already exists")
|
||||
self._jobs[job.id] = job
|
||||
self._save()
|
||||
self._save_or_rollback(job.id)
|
||||
# Emit SCHEDULE_ADD outside the store locks: a hook must never be able
|
||||
# to deadlock the cross-process file lock by calling back in.
|
||||
emit_schedule_add(job)
|
||||
@@ -120,7 +123,7 @@ class FileScheduleStore:
|
||||
job._lease_until = getattr(existing, "_lease_until", 0.0) or 0.0
|
||||
job._lease_owner = getattr(existing, "_lease_owner", None)
|
||||
self._jobs[job.id] = job
|
||||
self._save()
|
||||
self._save_or_rollback(job.id)
|
||||
|
||||
def remove(self, job_id: str) -> bool:
|
||||
removed = None
|
||||
@@ -128,7 +131,7 @@ class FileScheduleStore:
|
||||
self._reload_locked()
|
||||
if job_id in self._jobs:
|
||||
removed = self._jobs.pop(job_id)
|
||||
self._save()
|
||||
self._save_or_rollback(job_id)
|
||||
if removed is None:
|
||||
return False
|
||||
emit_schedule_remove(removed)
|
||||
@@ -150,7 +153,7 @@ class FileScheduleStore:
|
||||
if principal is not None and job.principal != principal:
|
||||
continue
|
||||
removed = self._jobs.pop(jid)
|
||||
self._save()
|
||||
self._save_or_rollback(jid)
|
||||
break
|
||||
if removed is None:
|
||||
return False
|
||||
@@ -255,7 +258,7 @@ class FileScheduleStore:
|
||||
emit_schedule_remove(job)
|
||||
# Record suppressed occurrences so a missed run is visible in history
|
||||
# rather than silently dropped. Logged outside the file lock (log_run
|
||||
# takes only the thread lock) to keep the critical section short.
|
||||
# takes it itself; flock is not re-entrant across file handles).
|
||||
for job_id, job_name in missed:
|
||||
self.log_run(
|
||||
job_id=job_id,
|
||||
@@ -301,12 +304,26 @@ class FileScheduleStore:
|
||||
delivered=delivered,
|
||||
timestamp=time.time(),
|
||||
)
|
||||
with self._lock:
|
||||
self._history.insert(0, record)
|
||||
# Prune to max_history
|
||||
if len(self._history) > self._max_history:
|
||||
self._history = self._history[:self._max_history]
|
||||
self._save_history()
|
||||
with self._lock, self._file_lock():
|
||||
# Reload under the cross-process lock so concurrent writers (e.g.
|
||||
# the scheduler daemon and a CLI process) don't drop each other's
|
||||
# records (last-writer-wins). Records whose earlier save failed are
|
||||
# carried forward so they are retried, not discarded.
|
||||
pending = ([record] + self._unsaved)[:self._max_history]
|
||||
on_disk = self._read_history()
|
||||
if on_disk is None:
|
||||
# Unreadable history: never overwrite saved records with a
|
||||
# partial view. Keep the new record in memory for a retry.
|
||||
self._history = ([record] + self._history)[:self._max_history]
|
||||
self._unsaved = pending
|
||||
return
|
||||
# Retried records may be older than ones other processes wrote
|
||||
# since; sort so history stays newest-first before trimming.
|
||||
merged = sorted(
|
||||
pending + on_disk, key=lambda r: r.timestamp, reverse=True
|
||||
)
|
||||
self._history = merged[:self._max_history]
|
||||
self._unsaved = [] if self._save_history() else pending
|
||||
|
||||
def get_history(
|
||||
self,
|
||||
@@ -392,6 +409,17 @@ class FileScheduleStore:
|
||||
except Exception as e:
|
||||
logger.warning("Failed to load schedule store from %s: %s", self._path, e)
|
||||
|
||||
def _save_or_rollback(self, job_id: str) -> None:
|
||||
"""Persist jobs or restore on-disk state and raise (file lock held).
|
||||
|
||||
Keeps memory consistent with disk so a failed write never reports
|
||||
success (or fires lifecycle hooks) for a change that will vanish on
|
||||
restart.
|
||||
"""
|
||||
if not self._save():
|
||||
self._reload_locked()
|
||||
raise OSError(f"Failed to persist schedule job '{job_id}'")
|
||||
|
||||
def _save(self) -> bool:
|
||||
"""Persist jobs to disk.
|
||||
|
||||
@@ -401,10 +429,7 @@ class FileScheduleStore:
|
||||
try:
|
||||
os.makedirs(self._dir, exist_ok=True)
|
||||
data = [j.to_dict() for j in self._jobs.values()]
|
||||
tmp = self._path + ".tmp"
|
||||
with open(tmp, "w") as f:
|
||||
json.dump(data, f, indent=2)
|
||||
os.replace(tmp, self._path)
|
||||
atomic_write_json(self._path, data, indent=2)
|
||||
return True
|
||||
except Exception as e:
|
||||
logger.warning("Failed to save schedule store to %s: %s", self._path, e)
|
||||
@@ -412,25 +437,27 @@ class FileScheduleStore:
|
||||
|
||||
def _load_history(self) -> None:
|
||||
"""Load execution history from history.json."""
|
||||
self._history = self._read_history() or []
|
||||
|
||||
def _read_history(self) -> Optional[List[RunRecord]]:
|
||||
"""Read history.json; ``[]`` if absent, ``None`` if unreadable."""
|
||||
if not os.path.exists(self._history_path):
|
||||
return
|
||||
return []
|
||||
try:
|
||||
with open(self._history_path, "r") as f:
|
||||
data = json.load(f)
|
||||
for d in data:
|
||||
record = RunRecord.from_dict(d)
|
||||
self._history.append(record)
|
||||
return [RunRecord.from_dict(d) for d in data]
|
||||
except Exception as e:
|
||||
logger.warning("Failed to load history from %s: %s", self._history_path, e)
|
||||
return None
|
||||
|
||||
def _save_history(self) -> None:
|
||||
"""Save execution history to history.json."""
|
||||
def _save_history(self) -> bool:
|
||||
"""Save execution history to history.json; ``True`` on success."""
|
||||
try:
|
||||
os.makedirs(self._dir, exist_ok=True)
|
||||
data = [r.to_dict() for r in self._history]
|
||||
tmp = self._history_path + ".tmp"
|
||||
with open(tmp, "w") as f:
|
||||
json.dump(data, f, indent=2)
|
||||
os.replace(tmp, self._history_path)
|
||||
atomic_write_json(self._history_path, data, indent=2)
|
||||
return True
|
||||
except Exception as e:
|
||||
logger.warning("Failed to save history to %s: %s", self._history_path, e)
|
||||
return False
|
||||
@@ -23,6 +23,8 @@ logger = logging.getLogger(__name__)
|
||||
|
||||
CACHE_SUBDIR = "remote-skills"
|
||||
_CURRENT_LINK = "current"
|
||||
# Deadline for git calls so a stalled remote falls back to the cached copy.
|
||||
_GIT_TIMEOUT = 120
|
||||
|
||||
|
||||
def get_remote_cache_root() -> Path:
|
||||
@@ -89,7 +91,10 @@ class GitRemoteSkillSource:
|
||||
if self.ref:
|
||||
cmd += ["--branch", self.ref]
|
||||
cmd += [self.url, str(clone_path)]
|
||||
proc = subprocess.run(cmd, capture_output=True, text=True)
|
||||
proc = subprocess.run(
|
||||
cmd, capture_output=True, text=True, timeout=_GIT_TIMEOUT,
|
||||
env={**os.environ, "GIT_TERMINAL_PROMPT": "0"},
|
||||
)
|
||||
if proc.returncode != 0:
|
||||
logger.warning(
|
||||
"Remote skill sync failed for %s: %s; using cached copy.",
|
||||
@@ -123,7 +128,7 @@ class GitRemoteSkillSource:
|
||||
def _resolve_version(self, clone_path: Path) -> str:
|
||||
proc = subprocess.run(
|
||||
["git", "-C", str(clone_path), "rev-parse", "HEAD"],
|
||||
capture_output=True, text=True,
|
||||
capture_output=True, text=True, timeout=_GIT_TIMEOUT,
|
||||
)
|
||||
if proc.returncode == 0 and proc.stdout.strip():
|
||||
return proc.stdout.strip()[:16]
|
||||
|
||||
@@ -120,6 +120,17 @@ class CircuitBreakerException(Exception):
|
||||
super().__init__(f"Circuit breaker open for service '{service_name}', state: {state.value}")
|
||||
|
||||
|
||||
class _NeutralOutcome(Exception):
|
||||
"""Raised by a guarded callable whose outcome says nothing about the service's
|
||||
health (e.g. a caller's bad input): recorded as neither success nor failure,
|
||||
so it can neither open the circuit nor close a half-open one. Carries the
|
||||
real outcome in ``value`` for the caller to unwrap."""
|
||||
|
||||
def __init__(self, value):
|
||||
super().__init__("neutral outcome")
|
||||
self.value = value
|
||||
|
||||
|
||||
class CircuitBreaker:
|
||||
"""Lightweight circuit breaker implementation.
|
||||
|
||||
@@ -226,6 +237,8 @@ class CircuitBreaker:
|
||||
result = func(*args, **kwargs)
|
||||
self._on_success()
|
||||
return result
|
||||
except _NeutralOutcome:
|
||||
raise
|
||||
except Exception as e:
|
||||
self._on_failure()
|
||||
raise
|
||||
@@ -269,6 +282,8 @@ class CircuitBreaker:
|
||||
result = func(*args, **kwargs)
|
||||
self._on_success()
|
||||
return result
|
||||
except _NeutralOutcome:
|
||||
raise
|
||||
except Exception as e:
|
||||
self._on_failure()
|
||||
raise
|
||||
|
||||
@@ -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}")
|
||||
|
||||
@@ -331,8 +331,15 @@ class OpenAPIOperation:
|
||||
# Returned, not raised, for the same reason as the transport errors
|
||||
# below: a raising tool ends the turn.
|
||||
return f"{self.name} failed: {exc}"
|
||||
from .url_safety import is_safe_http_url
|
||||
if not is_safe_http_url(request["url"]):
|
||||
# The host comes from an untrusted spec's `servers:` entry; keep it
|
||||
# off loopback/private/metadata addresses (ALLOW_LOCAL_CRAWL=true
|
||||
# opts out for local development).
|
||||
return (f"{self.name} failed: refusing unsafe URL "
|
||||
f"{request['url']!r}")
|
||||
try:
|
||||
with httpx.Client(timeout=self.timeout) as client:
|
||||
with httpx.Client(timeout=self.timeout, follow_redirects=False) as client:
|
||||
response = client.request(**request)
|
||||
response.raise_for_status()
|
||||
try:
|
||||
@@ -391,7 +398,11 @@ class OpenAPIToolset:
|
||||
text = spec_str
|
||||
else:
|
||||
import httpx
|
||||
with httpx.Client(timeout=30) as client:
|
||||
from .url_safety import is_safe_http_url
|
||||
if not is_safe_http_url(spec_url):
|
||||
raise ValueError(
|
||||
f"OpenAPIToolset refuses unsafe spec_url {spec_url!r}")
|
||||
with httpx.Client(timeout=30, follow_redirects=False) as client:
|
||||
response = client.get(spec_url)
|
||||
response.raise_for_status()
|
||||
text = response.text
|
||||
|
||||
@@ -7,6 +7,7 @@ This module provides a safe interface for executing shell commands with:
|
||||
- Resource limits
|
||||
"""
|
||||
|
||||
import re
|
||||
import subprocess
|
||||
import shlex
|
||||
import logging
|
||||
@@ -89,6 +90,48 @@ def _find_shell_syntax(command: str):
|
||||
return None
|
||||
|
||||
|
||||
# Credential-looking variable names are withheld from model-driven commands:
|
||||
# a gateway serving untrusted chat would otherwise hand `env` / `echo $X` every
|
||||
# LLM key, channel token and its own GATEWAY_AUTH_TOKEN. Operators who need one
|
||||
# inside a command name it in PRAISONAI_SHELL_ENV_ALLOWLIST (comma-separated).
|
||||
_SECRET_NAME_PARTS = frozenset({"TOKEN", "SECRET", "PASSWORD", "PASSWD", "CREDENTIAL", "CREDENTIALS", "APIKEY"})
|
||||
_ENV_REF = re.compile(r"\$(\w+|\{[^}]*\})")
|
||||
_ENV_REF_NT = re.compile(r"\$(\w+|\{[^}]*\})|%(\w+)%")
|
||||
|
||||
|
||||
def _is_secret_name(name: str) -> bool:
|
||||
upper = name.upper()
|
||||
return upper.endswith("_KEY") or any(part in _SECRET_NAME_PARTS for part in upper.split("_"))
|
||||
|
||||
|
||||
def _build_subprocess_env(extra_env: Optional[Dict[str, str]] = None) -> Dict[str, str]:
|
||||
"""Parent environment minus secret-looking names, plus explicit *extra_env*."""
|
||||
allow = {n.strip() for n in os.environ.get("PRAISONAI_SHELL_ENV_ALLOWLIST", "").split(",") if n.strip()}
|
||||
process_env = {k: v for k, v in os.environ.items() if k in allow or not _is_secret_name(k)}
|
||||
if extra_env:
|
||||
process_env.update(extra_env)
|
||||
return process_env
|
||||
|
||||
|
||||
def _expandvars(value: str, env: Dict[str, str]) -> str:
|
||||
"""``os.path.expandvars`` against *env* rather than ``os.environ``.
|
||||
|
||||
On Windows ``%VAR%`` is expanded too and names are case-insensitive,
|
||||
matching ``ntpath.expandvars``.
|
||||
"""
|
||||
windows = platform.system() == 'Windows'
|
||||
if windows:
|
||||
env = {k.upper(): v for k, v in env.items()}
|
||||
|
||||
def _sub(match):
|
||||
name = (match.group(1) or match.group(2)).strip("{}")
|
||||
return env.get(name.upper() if windows else name, match.group(0))
|
||||
|
||||
if windows:
|
||||
return _ENV_REF_NT.sub(_sub, value) if ("$" in value or "%" in value) else value
|
||||
return _ENV_REF.sub(_sub, value) if "$" in value else value
|
||||
|
||||
|
||||
class ShellTools:
|
||||
"""Tools for executing shell commands safely."""
|
||||
|
||||
@@ -187,24 +230,23 @@ class ShellTools:
|
||||
if not command:
|
||||
return {"error": "Empty command", "stdout": "", "stderr": "", "exit_code": 1}
|
||||
|
||||
# Set up process environment (secrets stripped, explicit env kept)
|
||||
process_env = _build_subprocess_env(env)
|
||||
|
||||
# Expand tilde and environment variables in command arguments
|
||||
# (shell=False means the shell won't do this for us)
|
||||
command = [os.path.expanduser(os.path.expandvars(arg)) for arg in command]
|
||||
|
||||
# (shell=False means the shell won't do this for us). Expansion
|
||||
# uses the scrubbed env so `echo $OPENAI_API_KEY` cannot leak it.
|
||||
command = [os.path.expanduser(_expandvars(arg, process_env)) for arg in command]
|
||||
|
||||
# Expand tilde in cwd (subprocess doesn't do this)
|
||||
if cwd:
|
||||
cwd = os.path.expanduser(cwd)
|
||||
cwd = os.path.expandvars(cwd) # Also expand $HOME, $USER, etc.
|
||||
cwd = _expandvars(cwd, process_env) # Also expand $HOME, $USER, etc.
|
||||
if not os.path.isdir(cwd):
|
||||
# Fallback: try home directory, then current working directory
|
||||
fallback = os.path.expanduser("~") if os.path.isdir(os.path.expanduser("~")) else os.getcwd()
|
||||
logging.warning(f"Working directory '{cwd}' does not exist, using '{fallback}'")
|
||||
cwd = fallback
|
||||
|
||||
# Set up process environment
|
||||
process_env = os.environ.copy()
|
||||
if env:
|
||||
process_env.update(env)
|
||||
|
||||
if run_in_background:
|
||||
return self._start_background(command, cwd, process_env)
|
||||
@@ -256,43 +298,38 @@ class ShellTools:
|
||||
result['stderr_artifact'] = stderr_path
|
||||
return result
|
||||
|
||||
except _CommandCancelled:
|
||||
except _CommandCancelled as e:
|
||||
# A cooperative interrupt (/stop, Ctrl-C) fired while the child
|
||||
# was still running. Kill the whole process group so descendants
|
||||
# do not keep mutating state, then report an interrupted outcome.
|
||||
self._kill_process_tree(process)
|
||||
return {
|
||||
'stdout': '',
|
||||
'stderr': 'Command interrupted by user',
|
||||
'exit_code': -1,
|
||||
'success': False,
|
||||
'interrupted': True,
|
||||
'execution_time': time.time() - start_time
|
||||
}
|
||||
result = self._partial_result(
|
||||
e, 'Command interrupted by user', max_output_size, spill, spill_dir
|
||||
)
|
||||
result.update(interrupted=True, execution_time=time.time() - start_time)
|
||||
return result
|
||||
|
||||
except _StreamDrainTimeout:
|
||||
except _StreamDrainTimeout as e:
|
||||
# The direct child already exited; only an inherited pipe kept a
|
||||
# reader blocked. Report a timeout but do NOT kill — the direct
|
||||
# child's PID is gone (and may have been recycled), so killing it
|
||||
# would be unsafe and misleading.
|
||||
return {
|
||||
'stdout': '',
|
||||
'stderr': f'Command output streaming timed out after {timeout} seconds',
|
||||
'exit_code': -1,
|
||||
'success': False,
|
||||
'execution_time': timeout
|
||||
}
|
||||
result = self._partial_result(
|
||||
e, f'Command output streaming timed out after {timeout} seconds',
|
||||
max_output_size, spill, spill_dir,
|
||||
)
|
||||
result.update(timed_out=True, execution_time=timeout)
|
||||
return result
|
||||
|
||||
except subprocess.TimeoutExpired:
|
||||
except subprocess.TimeoutExpired as e:
|
||||
# Kill process on timeout
|
||||
self._kill_process_tree(process)
|
||||
return {
|
||||
'stdout': '',
|
||||
'stderr': f'Command timed out after {timeout} seconds',
|
||||
'exit_code': -1,
|
||||
'success': False,
|
||||
'execution_time': timeout
|
||||
}
|
||||
result = self._partial_result(
|
||||
e, f'Command timed out after {timeout} seconds',
|
||||
max_output_size, spill, spill_dir,
|
||||
)
|
||||
result.update(timed_out=True, execution_time=timeout)
|
||||
return result
|
||||
|
||||
except Exception as e:
|
||||
# Name the exception type and log the traceback. Some exceptions
|
||||
@@ -314,6 +351,28 @@ class ShellTools:
|
||||
'execution_time': 0
|
||||
}
|
||||
|
||||
def _partial_result(self, exc, reason, max_output_size, spill, spill_dir):
|
||||
"""Build a failed result that keeps output captured before ``exc``."""
|
||||
collect = getattr(exc, 'collect_output', None)
|
||||
stdout, stderr = collect() if collect else ('', '')
|
||||
stdout, stdout_path = self._handle_overflow(
|
||||
stdout, "stdout", max_output_size, spill, spill_dir
|
||||
)
|
||||
stderr, stderr_path = self._handle_overflow(
|
||||
stderr, "stderr", max_output_size, spill, spill_dir
|
||||
)
|
||||
result: Dict[str, Union[str, int, bool]] = {
|
||||
'stdout': stdout,
|
||||
'stderr': f"{stderr.rstrip()}\n{reason}" if stderr.strip() else reason,
|
||||
'exit_code': -1,
|
||||
'success': False,
|
||||
}
|
||||
if stdout_path:
|
||||
result['stdout_artifact'] = stdout_path
|
||||
if stderr_path:
|
||||
result['stderr_artifact'] = stderr_path
|
||||
return result
|
||||
|
||||
def _start_background(self, command, cwd, process_env):
|
||||
"""Start *command* detached, writing stdout+stderr to a log file."""
|
||||
import tempfile
|
||||
@@ -560,18 +619,31 @@ class ShellTools:
|
||||
t_out.start()
|
||||
t_err.start()
|
||||
|
||||
def _with_partial(exc, drain=True):
|
||||
# Keep whatever the command already printed: on a timeout/interrupt
|
||||
# the real error usually lives in the output produced so far. The
|
||||
# caller collects only AFTER killing the child, so the pipes close
|
||||
# and readline() returns trailing text written without a newline.
|
||||
def _collect():
|
||||
if drain:
|
||||
for reader in (t_out, t_err):
|
||||
reader.join(timeout=_DRAIN_GRACE_SECONDS)
|
||||
return "".join(stdout_chunks), "".join(stderr_chunks)
|
||||
exc.collect_output = _collect
|
||||
return exc
|
||||
|
||||
deadline = None if timeout is None else time.monotonic() + timeout
|
||||
if cancel_event is None:
|
||||
try:
|
||||
process.wait(timeout=timeout)
|
||||
except subprocess.TimeoutExpired:
|
||||
except subprocess.TimeoutExpired as exc:
|
||||
# The direct child is still running. Disable any further emission
|
||||
# and re-raise TimeoutExpired so the caller's timeout path KILLS
|
||||
# the process before returning — a recorded read error must not
|
||||
# short-circuit that cleanup (otherwise the still-running child
|
||||
# leaks with its pipes open).
|
||||
stop_emitting.set()
|
||||
raise
|
||||
raise _with_partial(exc)
|
||||
else:
|
||||
# Poll in short slices so a cooperative interrupt aborts the child
|
||||
# promptly (within the poll interval) rather than only on timeout.
|
||||
@@ -579,12 +651,12 @@ class ShellTools:
|
||||
while True:
|
||||
if cancel_event.is_set():
|
||||
stop_emitting.set()
|
||||
raise _CommandCancelled()
|
||||
raise _with_partial(_CommandCancelled())
|
||||
if deadline is not None:
|
||||
remaining = deadline - time.monotonic()
|
||||
if remaining <= 0:
|
||||
stop_emitting.set()
|
||||
raise subprocess.TimeoutExpired(process.args, timeout)
|
||||
raise _with_partial(subprocess.TimeoutExpired(process.args, timeout))
|
||||
wait_slice = min(_POLL, remaining)
|
||||
else:
|
||||
wait_slice = _POLL
|
||||
@@ -616,7 +688,7 @@ class ShellTools:
|
||||
# Stop emitting and signal a drain-timeout so the caller reports a
|
||||
# timeout WITHOUT trying to kill an already-exited process.
|
||||
stop_emitting.set()
|
||||
raise _StreamDrainTimeout()
|
||||
raise _with_partial(_StreamDrainTimeout(), drain=False)
|
||||
|
||||
if read_errors:
|
||||
raise read_errors[0]
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -395,6 +395,24 @@ def _html_to_markdown(html: str, base_url: Optional[str] = None) -> tuple:
|
||||
return markdown, unescape(parser.title).strip()
|
||||
|
||||
|
||||
# Raw bytes read per page. Only 50,000 characters are returned, so a few MB is
|
||||
# ample and keeps memory (and HTML parsing) bounded for oversized responses.
|
||||
_MAX_CRAWL_BYTES = 5 * 1024 * 1024
|
||||
|
||||
|
||||
def _read_capped_httpx(response) -> str:
|
||||
"""Read at most ``_MAX_CRAWL_BYTES`` from a streamed httpx response."""
|
||||
buf = bytearray()
|
||||
for chunk in response.iter_bytes(chunk_size=64 * 1024):
|
||||
buf += chunk
|
||||
if len(buf) >= _MAX_CRAWL_BYTES:
|
||||
break
|
||||
try:
|
||||
return bytes(buf[:_MAX_CRAWL_BYTES]).decode(response.charset_encoding or "utf-8", errors="ignore")
|
||||
except LookupError:
|
||||
return bytes(buf[:_MAX_CRAWL_BYTES]).decode("utf-8", errors="ignore")
|
||||
|
||||
|
||||
def _crawl_with_httpx(urls: List[str]) -> List[Dict[str, Any]]:
|
||||
"""Crawl URLs using basic HTTP fetch with httpx or urllib."""
|
||||
import urllib.parse
|
||||
@@ -409,31 +427,49 @@ def _crawl_with_httpx(urls: List[str]) -> List[Dict[str, Any]]:
|
||||
try:
|
||||
import httpx
|
||||
current = url
|
||||
response = None
|
||||
with httpx.Client(follow_redirects=False, timeout=30.0) as client:
|
||||
for _ in range(max_redirects + 1):
|
||||
if not _is_safe_crawl_url(current):
|
||||
raise ValueError("Redirect target failed SSRF validation")
|
||||
response = client.get(current)
|
||||
if response.status_code not in redirect_statuses:
|
||||
break
|
||||
location = response.headers.get("Location")
|
||||
if not location:
|
||||
break
|
||||
# Stream so the body is never buffered past the byte cap.
|
||||
with client.stream("GET", current) as response:
|
||||
location = response.headers.get("Location")
|
||||
if response.status_code not in redirect_statuses or not location:
|
||||
response.raise_for_status()
|
||||
content = _read_capped_httpx(response)
|
||||
base_url = current
|
||||
break
|
||||
current = urllib.parse.urljoin(current, location)
|
||||
else:
|
||||
raise ValueError("Too many redirects")
|
||||
if response is not None:
|
||||
response.raise_for_status()
|
||||
content = response.text
|
||||
base_url = current
|
||||
except ImportError:
|
||||
# Fallback to urllib
|
||||
# Fallback to urllib, walking redirects by hand so every hop
|
||||
# is SSRF-validated before it is requested (as above).
|
||||
import urllib.error
|
||||
import urllib.request
|
||||
if not _is_safe_crawl_url(url):
|
||||
raise ValueError("URL blocked by SSRF policy")
|
||||
with urllib.request.urlopen(url, timeout=30) as response:
|
||||
content = response.read().decode('utf-8', errors='ignore')
|
||||
|
||||
class _NoRedirect(urllib.request.HTTPRedirectHandler):
|
||||
def redirect_request(self, *args, **kwargs):
|
||||
return None # surface 30x as HTTPError
|
||||
|
||||
opener = urllib.request.build_opener(_NoRedirect)
|
||||
current = url
|
||||
for _ in range(max_redirects + 1):
|
||||
if not _is_safe_crawl_url(current):
|
||||
raise ValueError("Redirect target failed SSRF validation")
|
||||
try:
|
||||
with opener.open(current, timeout=30) as response:
|
||||
content = response.read(_MAX_CRAWL_BYTES).decode('utf-8', errors='ignore')
|
||||
break
|
||||
except urllib.error.HTTPError as e:
|
||||
location = e.headers.get("Location")
|
||||
if e.code not in redirect_statuses or not location:
|
||||
raise
|
||||
e.close() # release the redirect response's socket
|
||||
current = urllib.parse.urljoin(current, location)
|
||||
else:
|
||||
raise ValueError("Too many redirects")
|
||||
base_url = current
|
||||
|
||||
# Convert HTML to structured Markdown using the standard library so
|
||||
# headings, lists, links, and code blocks survive on this path.
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -6,11 +6,11 @@ Exposes PraisonAI Agents via the A2A (Agent2Agent) protocol.
|
||||
|
||||
import asyncio
|
||||
from praisonaiagents._logging import get_logger
|
||||
from typing import List, Optional, TYPE_CHECKING
|
||||
from typing import Dict, List, Optional, TYPE_CHECKING
|
||||
|
||||
from praisonaiagents.ui.a2a.types import AgentCard
|
||||
from praisonaiagents.ui.a2a.agent_card import generate_agent_card
|
||||
from praisonaiagents.ui.a2a.task_store import TaskStore
|
||||
from praisonaiagents.ui.a2a.task_store import TaskStore, TERMINAL_STATES
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from praisonaiagents import Agent, AgentTeam
|
||||
@@ -98,6 +98,8 @@ class A2A:
|
||||
|
||||
# Initialize task store
|
||||
self.task_store = TaskStore()
|
||||
# Strong refs to background tasks so they aren't GC'd and can be cancelled
|
||||
self._running: Dict[str, asyncio.Task] = {}
|
||||
|
||||
# Router cache
|
||||
self._router: Optional["APIRouter"] = None
|
||||
@@ -290,12 +292,17 @@ class A2A:
|
||||
self.task_store.add_to_history(task_id, response_msg)
|
||||
self.task_store.add_artifact(task_id, artifact)
|
||||
self.task_store.update_status(task_id, TaskState.COMPLETED)
|
||||
except asyncio.CancelledError:
|
||||
self.task_store.update_status(task_id, TaskState.CANCELLED)
|
||||
raise
|
||||
except Exception as e:
|
||||
logger.error(f"Background agent error: {e}")
|
||||
self.task_store.update_status(task_id, TaskState.FAILED)
|
||||
|
||||
user_input = extract_user_input([message])
|
||||
asyncio.create_task(_background_process(task.id, user_input))
|
||||
bg = asyncio.create_task(_background_process(task.id, user_input))
|
||||
self._running[task.id] = bg
|
||||
bg.add_done_callback(lambda _t, tid=task.id: self._running.pop(tid, None))
|
||||
|
||||
return JSONResponse(content={
|
||||
"jsonrpc": "2.0",
|
||||
@@ -305,6 +312,9 @@ class A2A:
|
||||
|
||||
# Update to working
|
||||
self.task_store.update_status(task.id, TaskState.WORKING)
|
||||
# Keep the task while we wait on the agent: a cancel can make it
|
||||
# terminal (and evictable) before we build the response.
|
||||
self.task_store._pinned.add(task.id)
|
||||
|
||||
try:
|
||||
# Extract user input text
|
||||
@@ -333,6 +343,8 @@ class A2A:
|
||||
except Exception as e:
|
||||
logger.error(f"Agent error in message/send: {e}")
|
||||
updated_task = self.task_store.update_status(task.id, TaskState.FAILED)
|
||||
finally:
|
||||
self.task_store._pinned.discard(task.id)
|
||||
|
||||
# Return JSON-RPC response
|
||||
return JSONResponse(content={
|
||||
@@ -410,9 +422,19 @@ class A2A:
|
||||
if not task_id:
|
||||
return _jsonrpc_error(request_id, -32602, "Invalid params: 'id' required")
|
||||
|
||||
task = self.task_store.cancel_task(task_id)
|
||||
if not task:
|
||||
existing = self.task_store.get_task(task_id)
|
||||
if not existing:
|
||||
return _jsonrpc_error(request_id, -32000, f"Task not found: {task_id}")
|
||||
if existing.status.state in TERMINAL_STATES:
|
||||
return _jsonrpc_error(
|
||||
request_id, -32002,
|
||||
f"Task not cancelable: already {existing.status.state.value}",
|
||||
)
|
||||
|
||||
task = self.task_store.cancel_task(task_id)
|
||||
running = self._running.get(task_id)
|
||||
if running:
|
||||
running.cancel()
|
||||
|
||||
return JSONResponse(content={
|
||||
"jsonrpc": "2.0",
|
||||
|
||||
@@ -6,7 +6,7 @@ In-memory storage for A2A Task lifecycle management.
|
||||
|
||||
import uuid
|
||||
from datetime import datetime, timezone
|
||||
from typing import Dict, List, Optional
|
||||
from typing import Dict, List, Optional, Set
|
||||
|
||||
from praisonaiagents.ui.a2a.types import (
|
||||
Task,
|
||||
@@ -16,6 +16,12 @@ from praisonaiagents.ui.a2a.types import (
|
||||
Artifact,
|
||||
)
|
||||
|
||||
# Terminal states are final: once reached, status can no longer change.
|
||||
TERMINAL_STATES = frozenset({TaskState.COMPLETED, TaskState.FAILED, TaskState.CANCELLED})
|
||||
|
||||
# Upper bound on stored tasks; oldest terminal tasks are evicted beyond this.
|
||||
_MAX_TASKS = 10_000
|
||||
|
||||
|
||||
class TaskStore:
|
||||
"""
|
||||
@@ -28,6 +34,8 @@ class TaskStore:
|
||||
def __init__(self):
|
||||
"""Initialize empty task store."""
|
||||
self._tasks: Dict[str, Task] = {}
|
||||
# Ids of tasks a request is still waiting on; never evicted.
|
||||
self._pinned: Set[str] = set()
|
||||
|
||||
def create_task(self, message: Message, context_id: Optional[str] = None) -> Task:
|
||||
"""
|
||||
@@ -59,8 +67,22 @@ class TaskStore:
|
||||
)
|
||||
|
||||
self._tasks[task_id] = task
|
||||
self._evict()
|
||||
return task
|
||||
|
||||
def _evict(self, keep: Optional[str] = None):
|
||||
"""Drop the oldest terminal tasks (except *keep*) once over the cap."""
|
||||
excess = len(self._tasks) - _MAX_TASKS
|
||||
if excess <= 0:
|
||||
return
|
||||
old = [
|
||||
t.id for t in self._tasks.values()
|
||||
if t.status.state in TERMINAL_STATES
|
||||
and t.id != keep and t.id not in self._pinned
|
||||
]
|
||||
for tid in old[:excess]:
|
||||
del self._tasks[tid]
|
||||
|
||||
def get_task(self, task_id: str) -> Optional[Task]:
|
||||
"""
|
||||
Retrieve a task by its ID.
|
||||
@@ -88,11 +110,14 @@ class TaskStore:
|
||||
message: Optional status message
|
||||
|
||||
Returns:
|
||||
Updated Task if found, None otherwise
|
||||
Updated Task if found, None otherwise. A task already in a
|
||||
terminal state is returned unchanged.
|
||||
"""
|
||||
task = self._tasks.get(task_id)
|
||||
if not task:
|
||||
return None
|
||||
if task.status.state in TERMINAL_STATES:
|
||||
return task
|
||||
|
||||
# Create new status
|
||||
new_status = TaskStatus(
|
||||
@@ -112,6 +137,8 @@ class TaskStore:
|
||||
)
|
||||
|
||||
self._tasks[task_id] = updated
|
||||
if state in TERMINAL_STATES:
|
||||
self._evict(keep=task_id)
|
||||
return updated
|
||||
|
||||
def add_artifact(self, task_id: str, artifact: Artifact) -> Optional[Task]:
|
||||
@@ -123,11 +150,14 @@ class TaskStore:
|
||||
artifact: Artifact to add
|
||||
|
||||
Returns:
|
||||
Updated Task if found, None otherwise
|
||||
Updated Task if found, None otherwise. A task in a terminal
|
||||
state is returned unchanged.
|
||||
"""
|
||||
task = self._tasks.get(task_id)
|
||||
if not task:
|
||||
return None
|
||||
if task.status.state in TERMINAL_STATES:
|
||||
return task
|
||||
|
||||
# Get existing artifacts or create new list
|
||||
artifacts = list(task.artifacts) if task.artifacts else []
|
||||
@@ -155,11 +185,14 @@ class TaskStore:
|
||||
message: Message to add to history
|
||||
|
||||
Returns:
|
||||
Updated Task if found, None otherwise
|
||||
Updated Task if found, None otherwise. A task in a terminal
|
||||
state is returned unchanged.
|
||||
"""
|
||||
task = self._tasks.get(task_id)
|
||||
if not task:
|
||||
return None
|
||||
if task.status.state in TERMINAL_STATES:
|
||||
return task
|
||||
|
||||
# Get existing history or create new list
|
||||
history = list(task.history) if task.history else []
|
||||
|
||||
@@ -120,4 +120,28 @@ async def run_sync_in_executor(func, *args, **kwargs) -> Any:
|
||||
return func(*args, **kwargs)
|
||||
return await loop.run_in_executor(None, wrapper)
|
||||
else:
|
||||
return await loop.run_in_executor(None, func, *args)
|
||||
return await loop.run_in_executor(None, func, *args)
|
||||
|
||||
# Strong refs to fire-and-forget tasks: the event loop only keeps weak refs,
|
||||
# so an unreferenced task can be garbage-collected before it finishes.
|
||||
_background_tasks: "set[asyncio.Task]" = set()
|
||||
|
||||
|
||||
def spawn_background(coro: Awaitable[Any], label: str = "background task") -> "asyncio.Task":
|
||||
"""
|
||||
Schedule *coro* on the running loop without awaiting it.
|
||||
|
||||
Keeps a strong reference until the task finishes and logs any exception
|
||||
with *label*, instead of leaving it to asyncio's "Task exception was never
|
||||
retrieved" warning at GC time. Must be called with a running loop.
|
||||
"""
|
||||
task = asyncio.get_running_loop().create_task(coro)
|
||||
_background_tasks.add(task)
|
||||
|
||||
def _done(t: "asyncio.Task") -> None:
|
||||
_background_tasks.discard(t)
|
||||
if not t.cancelled() and t.exception() is not None:
|
||||
logger.error("Error in %s: %s", label, t.exception(), exc_info=t.exception())
|
||||
|
||||
task.add_done_callback(_done)
|
||||
return task
|
||||
@@ -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
|
||||
@@ -0,0 +1,26 @@
|
||||
"""
|
||||
Resolve a caller-supplied name to a file directly inside a base directory.
|
||||
|
||||
Used wherever a store builds ``base / f"{name}{suffix}"`` from a session,
|
||||
checkpoint, rule, doc or config name, so the name cannot climb out of the
|
||||
store via ``..`` or path separators.
|
||||
"""
|
||||
|
||||
from pathlib import Path
|
||||
|
||||
__all__ = ["safe_child"]
|
||||
|
||||
|
||||
def safe_child(base: Path, name: str, suffix: str = "") -> Path:
|
||||
"""Return ``base / f"{name}{suffix}"``, refusing names that escape ``base``."""
|
||||
if (
|
||||
not isinstance(name, str)
|
||||
or not name.strip()
|
||||
or name in (".", "..")
|
||||
or any(sep in name for sep in ("/", "\\", "\0"))
|
||||
):
|
||||
raise ValueError(f"invalid name: {name!r}")
|
||||
path = Path(base) / f"{name}{suffix}"
|
||||
if path.resolve().parent != Path(base).resolve():
|
||||
raise ValueError(f"name escapes {base}: {name!r}")
|
||||
return 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:
|
||||
|
||||
Loaded 100 of 322 files, more files were not shown because too many files have changed in this diff.
Show more
Reference in new issue
Block a user