mirror of
https://github.com/MervinPraison/PraisonAI.git
synced 2026-11-07 17:33:33 -10:00
Compare commits
7
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
2ead7ca927 | ||
|
|
ec0efc2b69 | ||
|
|
5a19ab07b2 | ||
|
|
a5dca259b0 | ||
|
|
59fb20af8e | ||
|
|
f9e058f939 | ||
|
|
e0bdf9648b |
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
|
||||
|
||||
@@ -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
|
||||
@@ -1917,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
|
||||
@@ -4119,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.
|
||||
"""
|
||||
@@ -4737,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
|
||||
|
||||
@@ -4811,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():
|
||||
@@ -4844,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",
|
||||
@@ -4858,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.
|
||||
@@ -5242,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
|
||||
@@ -5294,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():
|
||||
@@ -5328,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",
|
||||
@@ -5342,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.
|
||||
@@ -6684,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]'")
|
||||
|
||||
@@ -400,6 +400,26 @@ class AutonomyResult:
|
||||
"""
|
||||
return self.output or ""
|
||||
|
||||
|
||||
def _timeout_result(timeout_seconds, iterations, stage, actions, start_time, started_at,
|
||||
turn_timed_out=False):
|
||||
"""Return a ``timeout`` AutonomyResult once a run is past ``timeout_seconds``
|
||||
or the last turn was cut by ``max_execution_time`` (``turn_timed_out``), else None."""
|
||||
import time
|
||||
if not turn_timed_out and (not timeout_seconds or (time.time() - start_time) <= timeout_seconds):
|
||||
return None
|
||||
return AutonomyResult(
|
||||
success=False,
|
||||
output="Task timed out",
|
||||
completion_reason="timeout",
|
||||
iterations=iterations,
|
||||
stage=stage,
|
||||
actions=actions,
|
||||
duration_seconds=time.time() - start_time,
|
||||
started_at=started_at,
|
||||
)
|
||||
|
||||
|
||||
class DoomLoopTracker:
|
||||
"""Tracks actions to detect doom loops with graduated recovery.
|
||||
|
||||
|
||||
@@ -6,6 +6,7 @@ and response processing. Extracted from agent.py for maintainability.
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import contextlib
|
||||
import contextvars
|
||||
import os
|
||||
import re
|
||||
@@ -105,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))
|
||||
@@ -129,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:
|
||||
@@ -165,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):
|
||||
@@ -1988,7 +2030,8 @@ Your Goal: {self.goal}"""
|
||||
llm_kwargs, tools_ran = self._track_tool_runs(llm_kwargs)
|
||||
while True:
|
||||
try:
|
||||
return llm.get_response(**llm_kwargs)
|
||||
with self._llm_usage_sink(llm):
|
||||
return llm.get_response(**llm_kwargs)
|
||||
except (ToolExecutionError, BudgetExceededError, InterruptedError):
|
||||
raise
|
||||
except Exception as e:
|
||||
@@ -2007,7 +2050,8 @@ Your Goal: {self.goal}"""
|
||||
llm_kwargs, tools_ran = self._track_tool_runs(llm_kwargs)
|
||||
while True:
|
||||
try:
|
||||
return await llm.get_response_async(**llm_kwargs)
|
||||
with self._llm_usage_sink(llm):
|
||||
return await llm.get_response_async(**llm_kwargs)
|
||||
except (ToolExecutionError, BudgetExceededError, InterruptedError):
|
||||
raise
|
||||
except Exception as e:
|
||||
@@ -2182,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()
|
||||
|
||||
@@ -2381,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,
|
||||
@@ -2403,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
|
||||
@@ -3098,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
|
||||
@@ -3635,6 +3687,8 @@ Your Goal: {self.goal}"""
|
||||
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):
|
||||
@@ -3992,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"
|
||||
@@ -4000,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"
|
||||
@@ -4484,6 +4538,8 @@ Output MUST be JSON with 'reflection' and 'satisfactory'.
|
||||
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):
|
||||
@@ -4722,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"
|
||||
@@ -5236,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:
|
||||
|
||||
@@ -163,10 +163,14 @@ class ImageAgent(Agent):
|
||||
# Gemini provider doesn't support response_format parameter
|
||||
# Apply this filter if provider is explicitly 'gemini' or as fallback for gemini models
|
||||
config.pop('response_format', None)
|
||||
elif custom_llm_provider == "openai" and config.get('response_format') == "url":
|
||||
# "url" is OpenAI's default for DALL-E and is rejected by models that
|
||||
# don't accept response_format (e.g. gpt-image), so don't send it
|
||||
config.pop('response_format', None)
|
||||
elif custom_llm_provider == "openai":
|
||||
if config.get('response_format') == "url":
|
||||
# "url" is OpenAI's default for DALL-E and is rejected by models that
|
||||
# don't accept response_format (e.g. gpt-image), so don't send it
|
||||
config.pop('response_format', None)
|
||||
if 'dall-e-3' not in model_name.lower():
|
||||
# Only dall-e-3 accepts 'style'; other OpenAI image models reject it
|
||||
config.pop('style', None)
|
||||
from rich.progress import Progress, SpinnerColumn, TextColumn
|
||||
with Progress(
|
||||
SpinnerColumn(),
|
||||
|
||||
@@ -1744,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:
|
||||
@@ -1770,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
|
||||
|
||||
@@ -1803,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
|
||||
@@ -2192,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,
|
||||
@@ -2202,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
|
||||
@@ -2210,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,
|
||||
)
|
||||
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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},
|
||||
}
|
||||
|
||||
|
||||
|
||||
@@ -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
|
||||
@@ -344,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.
|
||||
@@ -3988,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}")
|
||||
@@ -4514,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}")
|
||||
@@ -4861,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
|
||||
@@ -6856,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
|
||||
|
||||
|
||||
@@ -199,6 +199,17 @@ class RAG:
|
||||
|
||||
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,
|
||||
@@ -395,7 +406,7 @@ class RAG:
|
||||
"num_results": len(results),
|
||||
"num_citations": len(citations),
|
||||
"retrieval_strategy": self.config.retrieval_strategy.value,
|
||||
"reranked": eff_rerank,
|
||||
"reranked": self._reranked(eff_rerank),
|
||||
"top_k": eff_top_k,
|
||||
},
|
||||
)
|
||||
@@ -495,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),
|
||||
},
|
||||
)
|
||||
|
||||
|
||||
@@ -332,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
|
||||
|
||||
@@ -0,0 +1,199 @@
|
||||
"""Budget and cost accounting on the custom-LLM (``provider/model``) path (#6029).
|
||||
|
||||
That path hands the whole turn, tool loop included, to ``LLM.get_response*``.
|
||||
Every call it makes must reach the Agent's cost/token counters so
|
||||
``ExecutionConfig.max_budget`` and the autonomy spend caps can fire.
|
||||
``litellm.completion`` is stubbed; nothing leaves the machine.
|
||||
"""
|
||||
import asyncio
|
||||
|
||||
import pytest
|
||||
|
||||
pytest.importorskip("litellm")
|
||||
|
||||
from praisonaiagents import Agent # noqa: E402
|
||||
from praisonaiagents.config.feature_configs import ExecutionConfig # noqa: E402
|
||||
from praisonaiagents.errors import BudgetExceededError # noqa: E402
|
||||
from praisonaiagents.utils.cost_utils import calculate_llm_cost # noqa: E402
|
||||
|
||||
MODEL = "anthropic/claude-3-5-sonnet-20240620"
|
||||
USAGE = {"prompt_tokens": 1500, "completion_tokens": 1500, "total_tokens": 3000}
|
||||
|
||||
|
||||
def _litellm():
|
||||
"""The litellm module LLM will import now.
|
||||
|
||||
Other tests evict ``litellm`` from sys.modules, so a module-level binding
|
||||
can go stale and stubbing it would let the real client run.
|
||||
"""
|
||||
import litellm
|
||||
return litellm
|
||||
|
||||
|
||||
def _text(model):
|
||||
return _litellm().ModelResponse(model=model, usage=USAGE, choices=[{
|
||||
"index": 0, "finish_reason": "stop",
|
||||
"message": {"role": "assistant", "content": "ok"}}])
|
||||
|
||||
|
||||
def _tool_call(model, n):
|
||||
return _litellm().ModelResponse(model=model, usage=USAGE, choices=[{
|
||||
"index": 0, "finish_reason": "tool_calls",
|
||||
"message": {"role": "assistant", "content": None, "tool_calls": [{
|
||||
"id": f"call_{n}", "type": "function",
|
||||
"function": {"name": "step", "arguments": "{}"}}]}}])
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def calls(monkeypatch):
|
||||
"""Stub litellm; answer with a tool call while ``tool_rounds`` remain."""
|
||||
monkeypatch.setenv("ANTHROPIC_API_KEY", "sk-test")
|
||||
# Agent caches its env-derived default model per process; keep the fake
|
||||
# key's Anthropic default from leaking into later tests.
|
||||
monkeypatch.setattr(Agent, "_default_model", None)
|
||||
monkeypatch.setattr(Agent, "_default_model_checked", False)
|
||||
state = {"n": 0, "tool_rounds": 0}
|
||||
|
||||
def fake(*a, **k):
|
||||
state["n"] += 1
|
||||
if k.get("tools") and state["tool_rounds"] > 0:
|
||||
state["tool_rounds"] -= 1
|
||||
return _tool_call(k.get("model"), state["n"])
|
||||
return _text(k.get("model"))
|
||||
|
||||
async def afake(*a, **k):
|
||||
return fake(*a, **k)
|
||||
|
||||
litellm = _litellm()
|
||||
monkeypatch.setattr(litellm, "completion", fake)
|
||||
monkeypatch.setattr(litellm, "acompletion", afake)
|
||||
return state
|
||||
|
||||
|
||||
def _agent(**kw):
|
||||
return Agent(name="p", instructions="x", llm=MODEL, output="silent", **kw)
|
||||
|
||||
|
||||
def step() -> str:
|
||||
"""One step of work."""
|
||||
return "done"
|
||||
|
||||
|
||||
def test_sync_chat_counts_cost_and_tokens(calls):
|
||||
agent = _agent()
|
||||
assert agent.chat("hi") == "ok"
|
||||
summary = agent.cost_summary
|
||||
assert summary["llm_calls"] == 1
|
||||
assert summary["tokens_in"] == 1500 and summary["tokens_out"] == 1500
|
||||
assert agent.total_cost > 0
|
||||
|
||||
|
||||
def test_sync_budget_raises_between_turns(calls):
|
||||
agent = _agent(execution=ExecutionConfig(max_budget=0.05))
|
||||
with pytest.raises(BudgetExceededError):
|
||||
for i in range(5):
|
||||
agent.chat(f"hello {i}")
|
||||
assert calls["n"] < 5
|
||||
assert agent.total_cost >= 0.05
|
||||
|
||||
|
||||
def test_async_budget_raises_between_turns(calls):
|
||||
agent = _agent(execution=ExecutionConfig(max_budget=0.05))
|
||||
|
||||
async def run():
|
||||
for i in range(5):
|
||||
await agent.achat(f"hello {i}")
|
||||
|
||||
with pytest.raises(BudgetExceededError):
|
||||
asyncio.run(run())
|
||||
assert calls["n"] < 5
|
||||
assert agent.total_cost >= 0.05
|
||||
|
||||
|
||||
def test_budget_trips_mid_tool_loop(calls):
|
||||
calls["tool_rounds"] = 10 # the model would keep calling tools for 11 calls
|
||||
agent = _agent(execution=ExecutionConfig(max_budget=1.0))
|
||||
agent.chat("price one call")
|
||||
one_call = agent.total_cost
|
||||
calls["n"] = 0
|
||||
# A cap between one and two calls must stop the loop at the second call.
|
||||
agent = _agent(tools=[step], execution=ExecutionConfig(max_budget=one_call * 1.5))
|
||||
with pytest.raises(BudgetExceededError):
|
||||
agent.chat("go")
|
||||
assert calls["n"] == 2
|
||||
|
||||
|
||||
def test_warn_mode_does_not_raise(calls):
|
||||
agent = _agent(execution=ExecutionConfig(max_budget=0.001, on_budget_exceeded="warn"))
|
||||
assert agent.chat("hi") == "ok"
|
||||
assert agent.total_cost > 0.001
|
||||
|
||||
|
||||
def test_autonomy_budget_exhausted_on_provider_path(calls):
|
||||
agent = _agent(autonomy={"max_budget_usd": 0.001})
|
||||
result = agent.run_autonomous("go", max_iterations=3)
|
||||
assert result.completion_reason == "budget_exhausted"
|
||||
assert calls["n"] == 1
|
||||
|
||||
|
||||
def test_sink_ignores_other_llm_instances(calls):
|
||||
"""A call made by a different LLM object is not billed to this agent."""
|
||||
from praisonaiagents.llm.llm import LLM
|
||||
|
||||
agent = _agent()
|
||||
other = LLM(model=MODEL)
|
||||
with agent._llm_usage_sink():
|
||||
other._track_token_usage(_text(MODEL), MODEL)
|
||||
assert agent.cost_summary["llm_calls"] == 0
|
||||
|
||||
|
||||
class _Unavailable(Exception):
|
||||
status_code = 503
|
||||
|
||||
|
||||
def test_fallback_model_spend_is_counted(calls, monkeypatch):
|
||||
"""A turn answered by a ``fallback_models`` entry is still billed to the agent."""
|
||||
backup = "anthropic/claude-3-haiku-20240307"
|
||||
|
||||
def fake(*a, **k):
|
||||
calls["n"] += 1
|
||||
if k.get("model") != backup:
|
||||
raise _Unavailable("503 overloaded")
|
||||
return _text(backup)
|
||||
|
||||
async def afake(*a, **k):
|
||||
return fake(*a, **k)
|
||||
|
||||
monkeypatch.setattr(_litellm(), "completion", fake)
|
||||
monkeypatch.setattr(_litellm(), "acompletion", afake)
|
||||
for run in (lambda a: a.chat("hi"), lambda a: asyncio.run(a.achat("hi"))):
|
||||
agent = _agent(retry=False) # fail over at once, no backoff sleeps
|
||||
agent.fallback_models = [backup]
|
||||
assert run(agent) == "ok"
|
||||
assert agent.cost_summary["llm_calls"] == 1
|
||||
# Priced as the model that answered, not the (pricier) primary.
|
||||
assert agent.total_cost == pytest.approx(calculate_llm_cost(1500, 1500, backup))
|
||||
assert agent.total_cost != pytest.approx(calculate_llm_cost(1500, 1500, MODEL))
|
||||
|
||||
|
||||
def test_retry_after_exhaustion_sends_no_request(calls):
|
||||
agent = _agent(execution=ExecutionConfig(max_budget=0.001))
|
||||
with pytest.raises(BudgetExceededError):
|
||||
agent.chat("first")
|
||||
sent = calls["n"]
|
||||
with pytest.raises(BudgetExceededError):
|
||||
agent.chat("retry")
|
||||
with pytest.raises(BudgetExceededError):
|
||||
asyncio.run(agent.achat("retry async"))
|
||||
assert calls["n"] == sent
|
||||
|
||||
|
||||
def test_autonomy_budget_error_keeps_budget_exhausted(calls):
|
||||
"""A max_budget stop inside the turn still reports the autonomy cap result."""
|
||||
agent = _agent(execution=ExecutionConfig(max_budget=0.001),
|
||||
autonomy={"max_budget_usd": 0.001})
|
||||
result = agent.run_autonomous("go", max_iterations=3)
|
||||
assert result.completion_reason == "budget_exhausted"
|
||||
result = asyncio.run(_agent(execution=ExecutionConfig(max_budget=0.001),
|
||||
autonomy={"max_budget_usd": 0.001}).run_autonomous_async("go", max_iterations=3))
|
||||
assert result.completion_reason == "budget_exhausted"
|
||||
@@ -0,0 +1,126 @@
|
||||
"""``ExecutionConfig.max_execution_time`` and ``run_autonomous(timeout_seconds=)`` (#6029).
|
||||
|
||||
Both caps become a deadline on the turn's cancel token, so the LLM tool loop
|
||||
stops at its next boundary instead of running to completion.
|
||||
``litellm.completion`` is stubbed; nothing leaves the machine.
|
||||
"""
|
||||
import asyncio
|
||||
import time
|
||||
|
||||
import pytest
|
||||
|
||||
pytest.importorskip("litellm")
|
||||
|
||||
from praisonaiagents import Agent # noqa: E402
|
||||
from praisonaiagents.config.feature_configs import ExecutionConfig # noqa: E402
|
||||
from praisonaiagents.escalation.loop_guard import LoopGuardConfig # noqa: E402
|
||||
|
||||
MODEL = "anthropic/claude-3-5-sonnet-20240620"
|
||||
USAGE = {"prompt_tokens": 1, "completion_tokens": 1, "total_tokens": 2}
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def calls(monkeypatch):
|
||||
"""Stub litellm so the model asks for ``slow_step`` on every call."""
|
||||
# Import here, not at module level: other tests evict ``litellm`` from
|
||||
# sys.modules, and LLM then imports a fresh module object at call time.
|
||||
import litellm
|
||||
from litellm import ModelResponse
|
||||
|
||||
monkeypatch.setenv("ANTHROPIC_API_KEY", "sk-test")
|
||||
# Agent caches its env-derived default model per process; keep the fake
|
||||
# key's Anthropic default from leaking into later tests.
|
||||
monkeypatch.setattr(Agent, "_default_model", None)
|
||||
monkeypatch.setattr(Agent, "_default_model_checked", False)
|
||||
state = {"n": 0}
|
||||
|
||||
def fake(*a, **k):
|
||||
state["n"] += 1
|
||||
if not k.get("tools"):
|
||||
return ModelResponse(model=k.get("model"), usage=USAGE, choices=[{
|
||||
"index": 0, "finish_reason": "stop",
|
||||
"message": {"role": "assistant", "content": "done"}}])
|
||||
return ModelResponse(model=k.get("model"), usage=USAGE, choices=[{
|
||||
"index": 0, "finish_reason": "tool_calls",
|
||||
"message": {"role": "assistant", "content": None, "tool_calls": [{
|
||||
"id": f"call_{state['n']}", "type": "function",
|
||||
"function": {"name": "slow_step", "arguments": "{}"}}]}}])
|
||||
|
||||
async def afake(*a, **k):
|
||||
return fake(*a, **k)
|
||||
|
||||
monkeypatch.setattr(litellm, "completion", fake)
|
||||
monkeypatch.setattr(litellm, "acompletion", afake)
|
||||
return state
|
||||
|
||||
|
||||
def slow_step() -> str:
|
||||
"""A slow step."""
|
||||
time.sleep(0.4)
|
||||
return "ok"
|
||||
|
||||
|
||||
def _agent(**kw):
|
||||
return Agent(name="t", instructions="x", llm=MODEL, tools=[slow_step],
|
||||
output="silent", **kw)
|
||||
|
||||
|
||||
def test_max_execution_time_stops_tool_loop(calls):
|
||||
agent = _agent(execution=ExecutionConfig(max_execution_time=1, max_iter=50))
|
||||
t0 = time.time()
|
||||
out = agent.chat("go")
|
||||
elapsed = time.time() - t0
|
||||
assert elapsed < 3
|
||||
assert "timeout" in (out or "")
|
||||
assert agent.last_stop_reason == "timeout"
|
||||
assert calls["n"] < 10
|
||||
|
||||
|
||||
def test_max_execution_time_stops_tool_loop_async(calls):
|
||||
agent = _agent(execution=ExecutionConfig(max_execution_time=1, max_iter=50))
|
||||
t0 = time.time()
|
||||
out = asyncio.run(agent.achat("go"))
|
||||
assert time.time() - t0 < 3
|
||||
assert "timeout" in (out or "")
|
||||
assert agent.last_stop_reason == "timeout"
|
||||
|
||||
|
||||
def test_unset_max_execution_time_creates_no_deadline():
|
||||
agent = Agent(name="t", instructions="x", output="silent")
|
||||
assert agent._turn_cancel_token(None, explicit=False) is None
|
||||
|
||||
|
||||
def test_loop_guard_turn_limit_follows_max_execution_time():
|
||||
agent = Agent(name="t", instructions="x", output="silent",
|
||||
execution=ExecutionConfig(max_execution_time=600))
|
||||
assert agent._ensure_loop_guard().config.max_time_per_turn == 600
|
||||
default = Agent(name="d", instructions="x", output="silent")
|
||||
assert default._ensure_loop_guard().config.max_time_per_turn == LoopGuardConfig.max_time_per_turn
|
||||
|
||||
|
||||
def test_run_autonomous_timeout_bounds_the_running_turn(calls):
|
||||
agent = _agent(autonomy=True, execution=ExecutionConfig(max_iter=50))
|
||||
t0 = time.time()
|
||||
result = agent.run_autonomous("go", max_iterations=1, timeout_seconds=1)
|
||||
assert time.time() - t0 < 3
|
||||
assert result.completion_reason == "timeout"
|
||||
assert not result.success
|
||||
|
||||
|
||||
def test_run_autonomous_async_timeout_bounds_the_running_turn(calls):
|
||||
agent = _agent(autonomy=True, execution=ExecutionConfig(max_iter=50))
|
||||
t0 = time.time()
|
||||
result = asyncio.run(agent.run_autonomous_async("go", max_iterations=1, timeout_seconds=1))
|
||||
assert time.time() - t0 < 3
|
||||
assert result.completion_reason == "timeout"
|
||||
|
||||
|
||||
def test_run_autonomous_reports_turn_max_execution_time(calls):
|
||||
"""A turn cut by max_execution_time is a failed timeout, not a goal."""
|
||||
agent = _agent(autonomy=True, execution=ExecutionConfig(max_execution_time=1, max_iter=50))
|
||||
result = agent.run_autonomous("go", max_iterations=3)
|
||||
assert result.completion_reason == "timeout"
|
||||
assert not result.success
|
||||
result = asyncio.run(agent.run_autonomous_async("go", max_iterations=3))
|
||||
assert result.completion_reason == "timeout"
|
||||
assert not result.success
|
||||
@@ -0,0 +1,176 @@
|
||||
"""Config options that were accepted but never read must either work or fail loud.
|
||||
|
||||
Covers GuardrailConfig.on_fail, ToolConfig.output_max_lines/output_direction and
|
||||
rerank=True on the Agent's RAG pipeline (issue #6042). No network calls.
|
||||
"""
|
||||
import pytest
|
||||
|
||||
from praisonaiagents import Agent
|
||||
from praisonaiagents.config import GuardrailConfig, ToolConfig
|
||||
from praisonaiagents.config.presets import GUARDRAIL_PRESETS
|
||||
|
||||
|
||||
def _ok(task_output):
|
||||
return True, task_output
|
||||
|
||||
|
||||
# --- GuardrailConfig.on_fail ------------------------------------------------
|
||||
|
||||
@pytest.mark.parametrize("on_fail", ["skip", "raise"])
|
||||
def test_unsupported_on_fail_fails_loud(on_fail):
|
||||
with pytest.raises(ValueError, match="on_fail"):
|
||||
Agent(name="g", instructions="x", llm="gpt-4o-mini",
|
||||
guardrails=GuardrailConfig(validator=_ok, on_fail=on_fail))
|
||||
|
||||
|
||||
def test_retry_on_fail_still_accepted():
|
||||
a = Agent(name="g", instructions="x", llm="gpt-4o-mini",
|
||||
guardrails=GuardrailConfig(validator=_ok, max_retries=2, on_fail="retry"))
|
||||
assert a.max_guardrail_retries == 2
|
||||
|
||||
|
||||
def test_presets_carry_no_unimplemented_on_fail():
|
||||
assert all("on_fail" not in p for p in GUARDRAIL_PRESETS.values())
|
||||
a = Agent(name="g", instructions="x", llm="gpt-4o-mini",
|
||||
guardrails=["strict", {"validator": _ok}])
|
||||
assert a.max_guardrail_retries == 5
|
||||
assert a.guardrail is _ok
|
||||
|
||||
|
||||
# --- ToolConfig.output_max_lines / output_direction --------------------------
|
||||
|
||||
def many_lines() -> str:
|
||||
"""Return 200 short lines."""
|
||||
return "\n".join(f"line {i}" for i in range(200))
|
||||
|
||||
|
||||
def _run_tool(cfg):
|
||||
a = Agent(name="t", instructions="x", llm="gpt-4o-mini", tools=[many_lines],
|
||||
tool_config=cfg, output="silent")
|
||||
return str(a.execute_tool("many_lines", {}))
|
||||
|
||||
|
||||
def test_default_tool_config_does_not_cap_lines():
|
||||
out = _run_tool(ToolConfig())
|
||||
assert out.splitlines()[-1] == "line 199"
|
||||
assert out.count("\n") == 199
|
||||
|
||||
|
||||
def test_output_max_lines_head():
|
||||
lines = _run_tool(ToolConfig(output_max_lines=5, output_direction="head")).splitlines()
|
||||
assert lines[:5] == [f"line {i}" for i in range(5)]
|
||||
assert "line 5" not in lines and "line 199" not in lines
|
||||
assert "200 lines total" in lines[5]
|
||||
|
||||
|
||||
def test_output_max_lines_tail():
|
||||
lines = _run_tool(ToolConfig(output_max_lines=5, output_direction="tail")).splitlines()
|
||||
assert lines[-5:] == [f"line {i}" for i in range(195, 200)]
|
||||
assert "line 0" not in lines
|
||||
|
||||
|
||||
def test_output_max_lines_both():
|
||||
lines = _run_tool(ToolConfig(output_max_lines=4)).splitlines()
|
||||
assert lines[:2] == ["line 0", "line 1"]
|
||||
assert lines[-2:] == ["line 198", "line 199"]
|
||||
assert len(lines) == 5 # 4 kept lines + marker
|
||||
|
||||
|
||||
def test_output_direction_applies_to_char_cap():
|
||||
out = _run_tool(ToolConfig(output_limit=50, output_direction="head"))
|
||||
assert out.startswith("line 0\nline 1")
|
||||
assert "line 199" not in out
|
||||
|
||||
|
||||
def short_lines() -> str:
|
||||
"""Return ten one-character lines."""
|
||||
return "\n".join("abcdefghij")
|
||||
|
||||
|
||||
def big_dict() -> dict:
|
||||
"""Return a dict with one large field."""
|
||||
return {"body": "START" + "x" * 45000 + "END"}
|
||||
|
||||
|
||||
def test_line_cap_applies_to_short_lines():
|
||||
a = Agent(name="t", instructions="x", llm="gpt-4o-mini", tools=[short_lines],
|
||||
tool_config=ToolConfig(output_max_lines=5, output_direction="head"), output="silent")
|
||||
lines = str(a.execute_tool("short_lines", {})).splitlines()
|
||||
assert lines[:5] == list("abcde")
|
||||
assert "j" not in lines
|
||||
|
||||
|
||||
def test_output_direction_applies_to_dict_results():
|
||||
a = Agent(name="t", instructions="x", llm="gpt-4o-mini", tools=[big_dict],
|
||||
tool_config=ToolConfig(output_limit=40000, output_direction="head"), output="silent")
|
||||
out = str(a.execute_tool("big_dict", {}))
|
||||
assert "START" in out and "END" not in out
|
||||
|
||||
|
||||
# --- rerank=True on Agent.rag -----------------------------------------------
|
||||
|
||||
class StubKnowledge:
|
||||
"""Minimal Knowledge stand-in: search() returns hits worst-first."""
|
||||
|
||||
def add(self, *a, **k):
|
||||
pass
|
||||
|
||||
def search(self, query, **kw):
|
||||
return {"results": [
|
||||
{"memory": "bananas are yellow", "score": 0.9, "metadata": {"source": "a"}},
|
||||
{"memory": "paris is the capital of france", "score": 0.8, "metadata": {"source": "b"}},
|
||||
]}
|
||||
|
||||
|
||||
def _rag(retrieval_config=None):
|
||||
a = Agent(name="r", instructions="x", llm="gpt-4o-mini", knowledge=StubKnowledge(),
|
||||
retrieval_config=retrieval_config)
|
||||
return a.rag
|
||||
|
||||
|
||||
def test_rerank_true_reranks_and_reports_it():
|
||||
pack = _rag({"rerank": True}).retrieve("capital of france")
|
||||
assert pack.metadata["reranked"] is True
|
||||
assert pack.context.index("paris") < pack.context.index("bananas")
|
||||
|
||||
|
||||
def test_rerank_false_reports_not_reranked():
|
||||
rag = _rag()
|
||||
assert rag.reranker is None
|
||||
assert rag.retrieve("capital of france").metadata["reranked"] is False
|
||||
|
||||
|
||||
def test_rerank_override_without_reranker_reports_false():
|
||||
assert _rag().retrieve("q", rerank=True).metadata["reranked"] is False
|
||||
|
||||
|
||||
def test_rerank_model_leaves_reranking_to_backend():
|
||||
rag = _rag({"rerank": True, "rerank_model": "cohere/rerank-v3"})
|
||||
assert rag.reranker is None
|
||||
|
||||
|
||||
def test_rerank_model_backend_reranking_reported():
|
||||
class BackendRerankKnowledge(StubKnowledge):
|
||||
config = {"reranker": {"default_rerank": True, "provider": "cohere"}}
|
||||
|
||||
a = Agent(name="r", instructions="x", llm="gpt-4o-mini", knowledge=BackendRerankKnowledge(),
|
||||
retrieval_config={"rerank": True, "rerank_model": "cohere/rerank-v3"})
|
||||
assert a.rag.reranker is None
|
||||
assert a.rag.retrieve("q").metadata["reranked"] is True
|
||||
assert a.rag.retrieve("q", rerank=False).metadata["reranked"] is False
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_rerank_true_async():
|
||||
pack = await _rag({"rerank": True}).aretrieve("capital of france")
|
||||
assert pack.metadata["reranked"] is True
|
||||
assert pack.context.index("paris") < pack.context.index("bananas")
|
||||
|
||||
|
||||
def test_simple_reranker_accepts_search_result_items():
|
||||
from praisonaiagents.knowledge.models import SearchResultItem
|
||||
from praisonaiagents.rag.retriever import SimpleReranker
|
||||
hits = [SearchResultItem(text="bananas are yellow", score=0.9),
|
||||
SearchResultItem(text="paris is the capital of france", score=0.8)]
|
||||
ranked = SimpleReranker().rerank("capital of france", hits)
|
||||
assert ranked[0]["text"].startswith("paris")
|
||||
@@ -0,0 +1,31 @@
|
||||
"""OpenAI image calls must only forward style to models that accept it (dall-e-3)."""
|
||||
|
||||
import pytest
|
||||
|
||||
from praisonaiagents import ImageAgent
|
||||
|
||||
|
||||
@pytest.mark.parametrize("llm,kwargs,expected", [
|
||||
("openai/dall-e-2", {}, None),
|
||||
("openai/gpt-image-1", {}, None),
|
||||
("openai/gpt-image-1", {"style": "vivid"}, None),
|
||||
("openai/dall-e-3", {}, "natural"),
|
||||
("openai/dall-e-3", {"style": "vivid"}, "vivid"),
|
||||
])
|
||||
def test_style_forwarding(llm, kwargs, expected):
|
||||
agent = ImageAgent(llm=llm, verbose=False)
|
||||
received = {}
|
||||
|
||||
def fake_image_generation(**call_kwargs):
|
||||
received.update(call_kwargs)
|
||||
return {"data": [{"url": "https://example.com/image.png"}]}
|
||||
|
||||
agent._litellm = fake_image_generation
|
||||
try:
|
||||
agent.generate("A small green square icon", **kwargs)
|
||||
if expected is None:
|
||||
assert "style" not in received
|
||||
else:
|
||||
assert received["style"] == expected
|
||||
finally:
|
||||
agent.close()
|
||||
@@ -14,6 +14,7 @@ _LAZY_EXPORTS = {
|
||||
"praisonai_bot.gateway.admin",
|
||||
"provision_gateway_config",
|
||||
),
|
||||
"edit_gateway_channel": ("praisonai_bot.gateway.admin", "edit_gateway_channel"),
|
||||
}
|
||||
|
||||
|
||||
@@ -34,5 +35,6 @@ __all__ = [
|
||||
"GatewayRepairResult",
|
||||
"repair_gateway_config",
|
||||
"provision_gateway_config",
|
||||
"edit_gateway_channel",
|
||||
"__version__",
|
||||
]
|
||||
@@ -1553,8 +1553,16 @@ def gateway_test(
|
||||
raise typer.Exit(1)
|
||||
|
||||
|
||||
@app.command("channels")
|
||||
channels_app = typer.Typer(
|
||||
help="List, probe and edit channels in gateway.yaml",
|
||||
invoke_without_command=True,
|
||||
)
|
||||
app.add_typer(channels_app, name="channels")
|
||||
|
||||
|
||||
@channels_app.callback()
|
||||
def gateway_channels(
|
||||
ctx: typer.Context,
|
||||
config: str = typer.Option("gateway.yaml", "--config", "-c", help="Path to gateway.yaml"),
|
||||
json_output: bool = typer.Option(False, "--json", help="Output JSON"),
|
||||
probe: bool = typer.Option(False, "--probe", help="Probe each channel's credentials"),
|
||||
@@ -1574,6 +1582,9 @@ def gateway_channels(
|
||||
import os
|
||||
import yaml
|
||||
|
||||
if ctx.invoked_subcommand is not None:
|
||||
return
|
||||
|
||||
if available:
|
||||
try:
|
||||
from praisonai_bot.bots._registry import list_platforms
|
||||
@@ -1662,6 +1673,125 @@ def gateway_channels(
|
||||
print(f"{name:<20} {platform:<12} {has_token:<12}")
|
||||
|
||||
|
||||
def _parse_channel_settings(pairs: Optional[List[str]]) -> dict:
|
||||
"""Parse repeated ``--set key=value`` flags; values are YAML-typed."""
|
||||
import yaml
|
||||
|
||||
settings = {}
|
||||
for pair in pairs or []:
|
||||
key, sep, value = pair.partition("=")
|
||||
if not sep or not key.strip():
|
||||
print(f"Error: --set expects key=value, got {pair!r}")
|
||||
raise typer.Exit(1)
|
||||
settings[key.strip()] = yaml.safe_load(value) if value else ""
|
||||
return settings
|
||||
|
||||
|
||||
def _edit_channel(
|
||||
ctx: typer.Context, config: Optional[str], name: str, settings=None, *, remove=False, must_exist=True
|
||||
) -> None:
|
||||
"""Run one validated channel edit, mapping failures to CLI exits (#6079)."""
|
||||
import os
|
||||
import yaml
|
||||
|
||||
from praisonai_bot.gateway.admin import _editable_channels, edit_gateway_channel
|
||||
|
||||
# ``gateway channels --config X <cmd>`` applies to the subcommand too.
|
||||
config = _resolve_doctor_config(config or ctx.parent.params.get("config") or "gateway.yaml")
|
||||
if not os.path.exists(config):
|
||||
print(f"Error: Config file not found: {config}")
|
||||
raise typer.Exit(1)
|
||||
with open(config) as fh:
|
||||
exists = name in _editable_channels(yaml.safe_load(fh) or {})
|
||||
if must_exist and not exists:
|
||||
print(f"Error: channel {name!r} is not configured in {config}")
|
||||
raise typer.Exit(1)
|
||||
if not must_exist and exists:
|
||||
print(f"Error: channel {name!r} already exists in {config}; use 'update'")
|
||||
raise typer.Exit(1)
|
||||
try:
|
||||
edit_gateway_channel(config, name, settings, remove=remove)
|
||||
except ValueError as exc:
|
||||
print(f"Error: {exc}\nNo changes written to {config}.")
|
||||
raise typer.Exit(1) from exc
|
||||
|
||||
|
||||
@channels_app.command("add")
|
||||
def gateway_channels_add(
|
||||
ctx: typer.Context,
|
||||
name: str = typer.Argument(..., help="Channel name, e.g. 'telegram'"),
|
||||
platform: Optional[str] = typer.Option(None, "--platform", help="Platform (default: the channel name)"),
|
||||
token: Optional[str] = typer.Option(
|
||||
None, "--token",
|
||||
help="Channel token (default: ${VAR} reference to the platform's credential env var)",
|
||||
),
|
||||
set_: Optional[List[str]] = typer.Option(None, "--set", help="Extra setting key=value (repeatable)"),
|
||||
config: Optional[str] = typer.Option(None, "--config", "-c", help="Path to gateway.yaml"),
|
||||
):
|
||||
"""Add one channel to gateway.yaml without touching the others.
|
||||
|
||||
Examples:
|
||||
praisonai gateway channels add telegram
|
||||
praisonai gateway channels add support --platform discord --set group_policy=mention_only
|
||||
"""
|
||||
from praisonai_bot.bots._config_schema import GatewayConfigSchema
|
||||
|
||||
platform = platform or name
|
||||
settings = {"platform": platform}
|
||||
settings["token"] = token if token is not None else GatewayConfigSchema._credential_env_ref(platform)
|
||||
settings.update(_parse_channel_settings(set_))
|
||||
_edit_channel(ctx, config, name, settings, must_exist=False)
|
||||
print(f"Added channel {name!r}.")
|
||||
|
||||
|
||||
@channels_app.command("update")
|
||||
def gateway_channels_update(
|
||||
ctx: typer.Context,
|
||||
name: str = typer.Argument(..., help="Channel name to update"),
|
||||
token: Optional[str] = typer.Option(None, "--token", help="New channel token"),
|
||||
set_: Optional[List[str]] = typer.Option(None, "--set", help="Setting key=value (repeatable)"),
|
||||
config: Optional[str] = typer.Option(None, "--config", "-c", help="Path to gateway.yaml"),
|
||||
):
|
||||
"""Update settings of one configured channel.
|
||||
|
||||
Examples:
|
||||
praisonai gateway channels update telegram --set group_policy=mention_only
|
||||
"""
|
||||
settings = _parse_channel_settings(set_)
|
||||
if token is not None:
|
||||
settings["token"] = token
|
||||
if not settings:
|
||||
print("Error: nothing to update; pass --set key=value, --token, or both")
|
||||
raise typer.Exit(1)
|
||||
_edit_channel(ctx, config, name, settings)
|
||||
print(f"Updated channel {name!r}.")
|
||||
|
||||
|
||||
@channels_app.command("disable")
|
||||
def gateway_channels_disable(
|
||||
ctx: typer.Context,
|
||||
name: str = typer.Argument(..., help="Channel name to disable"),
|
||||
config: Optional[str] = typer.Option(None, "--config", "-c", help="Path to gateway.yaml"),
|
||||
):
|
||||
"""Stop serving a channel but keep its config (sets ``enabled: false``).
|
||||
|
||||
Re-enable with ``praisonai gateway channels update <name> --set enabled=true``.
|
||||
"""
|
||||
_edit_channel(ctx, config, name, {"enabled": False})
|
||||
print(f"Disabled channel {name!r}.")
|
||||
|
||||
|
||||
@channels_app.command("remove")
|
||||
def gateway_channels_remove(
|
||||
ctx: typer.Context,
|
||||
name: str = typer.Argument(..., help="Channel name to remove"),
|
||||
config: Optional[str] = typer.Option(None, "--config", "-c", help="Path to gateway.yaml"),
|
||||
):
|
||||
"""Delete a channel block from gateway.yaml."""
|
||||
_edit_channel(ctx, config, name, remove=True)
|
||||
print(f"Removed channel {name!r}.")
|
||||
|
||||
|
||||
def _resolve_gateway_rest_url(
|
||||
url: Optional[str],
|
||||
host: Optional[str] = None,
|
||||
|
||||
@@ -15,6 +15,7 @@ one canonical flow (DRY).
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import contextlib
|
||||
import os
|
||||
import sys
|
||||
from dataclasses import dataclass, field
|
||||
@@ -201,7 +202,6 @@ def _repair_config_version(config_path: str):
|
||||
an interruption, full disk, or I/O error during ``doctor --fix`` can never
|
||||
leave the config truncated or half-written.
|
||||
"""
|
||||
import tempfile
|
||||
import yaml
|
||||
|
||||
with open(config_path) as fh:
|
||||
@@ -210,6 +210,14 @@ def _repair_config_version(config_path: str):
|
||||
from praisonaiagents.gateway.config import migrate_config_with_doctor
|
||||
|
||||
migrated, reasons = migrate_config_with_doctor(cfg)
|
||||
_atomic_write_yaml(config_path, migrated)
|
||||
return reasons
|
||||
|
||||
|
||||
def _atomic_write_yaml(config_path: str, data: dict) -> None:
|
||||
"""Serialise ``data`` to a temp file and ``os.replace`` it over the config."""
|
||||
import tempfile
|
||||
import yaml
|
||||
|
||||
directory = os.path.dirname(os.path.abspath(config_path))
|
||||
fd, tmp_path = tempfile.mkstemp(
|
||||
@@ -217,9 +225,7 @@ def _repair_config_version(config_path: str):
|
||||
)
|
||||
try:
|
||||
with os.fdopen(fd, "w") as fh:
|
||||
yaml.safe_dump(
|
||||
migrated, fh, default_flow_style=False, sort_keys=False
|
||||
)
|
||||
yaml.safe_dump(data, fh, default_flow_style=False, sort_keys=False)
|
||||
fh.flush()
|
||||
os.fsync(fh.fileno())
|
||||
os.replace(tmp_path, config_path)
|
||||
@@ -230,7 +236,127 @@ def _repair_config_version(config_path: str):
|
||||
pass
|
||||
raise
|
||||
|
||||
return reasons
|
||||
|
||||
def edit_gateway_channel(
|
||||
config_path: PathLike,
|
||||
name: str,
|
||||
settings: Optional[Dict] = None,
|
||||
*,
|
||||
remove: bool = False,
|
||||
) -> dict:
|
||||
"""Add, update or remove one ``channels:`` entry in place (Issue #6079).
|
||||
|
||||
A surgical read-modify-write: ``settings`` is merged into
|
||||
``channels[name]`` (created if absent, ``platform`` defaulting to ``name``)
|
||||
or, with ``remove=True``, the block is deleted. Other channels and
|
||||
top-level keys are untouched. The edited document is validated against
|
||||
``GatewayConfigSchema`` before the atomic write, so an invalid mutation is
|
||||
rejected rather than persisted. Disable a channel with
|
||||
``settings={"enabled": False}``.
|
||||
|
||||
Returns:
|
||||
The resulting channel block (``{}`` when removed).
|
||||
|
||||
Raises:
|
||||
FileNotFoundError: If ``config_path`` does not exist.
|
||||
KeyError: If ``remove=True`` and the channel is not configured.
|
||||
ValueError: If the edited config fails validation.
|
||||
"""
|
||||
import yaml
|
||||
|
||||
config_path = os.fspath(config_path)
|
||||
if not os.path.isfile(config_path):
|
||||
raise FileNotFoundError(config_path)
|
||||
# Hold a sidecar lock across read-modify-write so concurrent edits of
|
||||
# different channels cannot overwrite each other.
|
||||
with _config_edit_lock(config_path):
|
||||
with open(config_path) as fh:
|
||||
cfg = yaml.safe_load(fh) or {}
|
||||
return _apply_channel_edit(config_path, cfg, name, settings, remove)
|
||||
|
||||
|
||||
def _apply_channel_edit(config_path, cfg, name, settings, remove) -> dict:
|
||||
from praisonai_bot.bots._config_schema import validate_gateway_config
|
||||
|
||||
channels = _editable_channels(cfg)
|
||||
if remove:
|
||||
if name not in channels:
|
||||
raise KeyError(f"channel {name!r} is not configured in {config_path}")
|
||||
channels.pop(name)
|
||||
block: dict = {}
|
||||
else:
|
||||
block = dict(channels.get(name) or {})
|
||||
block.setdefault("platform", name)
|
||||
block.update(settings or {})
|
||||
channels[name] = block
|
||||
cfg["channels"] = channels
|
||||
|
||||
# Check the document itself: the validator may otherwise satisfy its
|
||||
# "at least one channel" rule from env credentials, which are not saved.
|
||||
if not any(
|
||||
not isinstance(b, dict) or b.get("enabled", True) is not False
|
||||
for b in channels.values()
|
||||
):
|
||||
raise ValueError("at least one enabled channel must remain configured")
|
||||
# Validate an env-resolved copy (as the runtime does; unset ``${VAR}``
|
||||
# refs are left as-is) while the saved file keeps the references.
|
||||
validate_gateway_config(cfg)
|
||||
_atomic_write_yaml(config_path, cfg)
|
||||
return block
|
||||
|
||||
|
||||
def _editable_channels(cfg: dict) -> dict:
|
||||
"""Return ``cfg``'s ``channels`` map, moving legacy declarations into it.
|
||||
|
||||
Single-bot (top-level ``platform``) and BotOS (``platforms``) files are
|
||||
only honoured by the loader while ``channels`` is empty, so they are
|
||||
converted in place before an edit adds a ``channels:`` block. ``${VAR}``
|
||||
references are copied verbatim.
|
||||
"""
|
||||
channels = cfg.get("channels") or {}
|
||||
if channels:
|
||||
return channels
|
||||
platform = cfg.get("platform")
|
||||
if platform:
|
||||
# Mirrors the loader's single-bot migration, which wins over BotOS.
|
||||
from praisonai_bot.bots._config_schema import GatewayConfigSchema
|
||||
|
||||
# Resolve a ``${VAR}`` platform first, as the loader does.
|
||||
raw = str(cfg.pop("platform"))
|
||||
platform = str(_resolve_env_token(raw) or raw).lower().strip()
|
||||
block = {"platform": platform}
|
||||
token = cfg.pop("token", None)
|
||||
block["token"] = token or GatewayConfigSchema._credential_env_ref(platform)
|
||||
if platform == "slack" and cfg.get("app_token") is not None:
|
||||
block["app_token"] = cfg.pop("app_token")
|
||||
for key in ("unknown_user_policy", "owner_user_id"):
|
||||
if cfg.get(key) is not None:
|
||||
block[key] = cfg[key]
|
||||
channels[platform] = block
|
||||
elif cfg.get("platforms"):
|
||||
for pname, pcfg in cfg.pop("platforms").items():
|
||||
channels[pname] = {"platform": pname, **(pcfg or {})}
|
||||
return channels
|
||||
|
||||
|
||||
@contextlib.contextmanager
|
||||
def _config_edit_lock(config_path: str):
|
||||
"""Exclusive OS advisory lock on ``<config>.lock`` (no-op if unsupported)."""
|
||||
with open(config_path + ".lock", "a+") as fh:
|
||||
try:
|
||||
import fcntl
|
||||
except ImportError: # pragma: no cover - Windows
|
||||
fcntl = None
|
||||
if fcntl is not None:
|
||||
fcntl.flock(fh.fileno(), fcntl.LOCK_EX)
|
||||
else: # pragma: no cover - Windows
|
||||
try:
|
||||
import msvcrt
|
||||
|
||||
msvcrt.locking(fh.fileno(), msvcrt.LK_LOCK, 1)
|
||||
except Exception:
|
||||
pass
|
||||
yield # released when ``fh`` closes
|
||||
|
||||
|
||||
def _repair_gateway_secret(dry_run: bool = False, config_path: str = ""):
|
||||
|
||||
@@ -0,0 +1,167 @@
|
||||
"""Per-channel config lifecycle: add / update / disable / remove (Issue #6079)."""
|
||||
|
||||
import pytest
|
||||
import yaml
|
||||
from typer.testing import CliRunner
|
||||
|
||||
from praisonai_bot.cli.commands.gateway import app
|
||||
from praisonai_bot.gateway.admin import edit_gateway_channel
|
||||
|
||||
BASE = {
|
||||
"agents": {"assistant": {"instructions": "hi"}},
|
||||
"channels": {
|
||||
"telegram": {"platform": "telegram", "token": "${TELEGRAM_BOT_TOKEN}"},
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def config(tmp_path):
|
||||
path = tmp_path / "gateway.yaml"
|
||||
path.write_text(yaml.safe_dump(BASE, sort_keys=False))
|
||||
return path
|
||||
|
||||
|
||||
def _load(path):
|
||||
return yaml.safe_load(path.read_text())
|
||||
|
||||
|
||||
def _run(*args):
|
||||
return CliRunner().invoke(app, ["channels", *args])
|
||||
|
||||
|
||||
def test_add_channel_keeps_others_and_defaults_env_token(config):
|
||||
result = _run("add", "discord", "--config", str(config), "--set", "group_policy=mention_only")
|
||||
assert result.exit_code == 0, result.output
|
||||
channels = _load(config)["channels"]
|
||||
assert channels["telegram"] == BASE["channels"]["telegram"]
|
||||
assert channels["discord"] == {
|
||||
"platform": "discord",
|
||||
"token": "${DISCORD_BOT_TOKEN}",
|
||||
"group_policy": "mention_only",
|
||||
}
|
||||
assert _load(config)["agents"] == BASE["agents"]
|
||||
|
||||
|
||||
def test_add_existing_channel_is_rejected(config):
|
||||
result = _run("add", "telegram", "--config", str(config))
|
||||
assert result.exit_code == 1
|
||||
assert _load(config) == BASE
|
||||
|
||||
|
||||
def test_update_and_disable_channel(config):
|
||||
assert _run("update", "telegram", "--config", str(config), "--token", "${NEW_TOKEN}").exit_code == 0
|
||||
assert _run("add", "discord", "--config", str(config)).exit_code == 0
|
||||
assert _run("disable", "telegram", "--config", str(config)).exit_code == 0
|
||||
telegram = _load(config)["channels"]["telegram"]
|
||||
assert telegram["token"] == "${NEW_TOKEN}"
|
||||
assert telegram["enabled"] is False
|
||||
|
||||
|
||||
def test_invalid_update_is_not_persisted(config):
|
||||
result = _run("update", "telegram", "--config", str(config), "--set", "allowed_users=[1, 2]")
|
||||
assert result.exit_code == 1
|
||||
assert _load(config) == BASE
|
||||
|
||||
|
||||
def test_unknown_platform_is_rejected(config):
|
||||
assert _run("add", "bogus", "--config", str(config)).exit_code == 1
|
||||
assert _load(config) == BASE
|
||||
|
||||
|
||||
def test_remove_channel(config):
|
||||
assert _run("add", "discord", "--config", str(config)).exit_code == 0
|
||||
assert _run("remove", "discord", "--config", str(config)).exit_code == 0
|
||||
assert "discord" not in _load(config)["channels"]
|
||||
assert _run("remove", "discord", "--config", str(config)).exit_code == 1
|
||||
|
||||
|
||||
def test_python_helper_matches_cli(config):
|
||||
block = edit_gateway_channel(config, "discord", {"token": "${D}"})
|
||||
assert block == {"platform": "discord", "token": "${D}"}
|
||||
with pytest.raises(KeyError):
|
||||
edit_gateway_channel(config, "nope", remove=True)
|
||||
with pytest.raises(ValueError):
|
||||
edit_gateway_channel(config, "bogus", {})
|
||||
assert "bogus" not in _load(config)["channels"]
|
||||
|
||||
|
||||
def test_list_still_works_without_subcommand(config):
|
||||
result = _run("--config", str(config), "--json")
|
||||
assert result.exit_code == 0
|
||||
assert '"telegram"' in result.output
|
||||
|
||||
|
||||
def test_parent_config_option_applies_to_subcommand(config, tmp_path, monkeypatch):
|
||||
monkeypatch.chdir(tmp_path)
|
||||
other = tmp_path / "other.yaml"
|
||||
other.write_text(yaml.safe_dump(BASE, sort_keys=False))
|
||||
result = _run("--config", str(other), "update", "telegram", "--set", "group_policy=mention_only")
|
||||
assert result.exit_code == 0, result.output
|
||||
assert _load(other)["channels"]["telegram"]["group_policy"] == "mention_only"
|
||||
assert _load(config) == BASE
|
||||
|
||||
|
||||
def test_single_bot_platform_is_converted_not_hidden(tmp_path):
|
||||
path = tmp_path / "bot.yaml"
|
||||
path.write_text(yaml.safe_dump({"agent": {"instructions": "hi"}, "platform": "telegram", "token": "${T}"}))
|
||||
edit_gateway_channel(path, "discord", {"token": "${D}"})
|
||||
cfg = _load(path)
|
||||
assert "platform" not in cfg and "token" not in cfg
|
||||
assert cfg["channels"]["telegram"] == {"platform": "telegram", "token": "${T}"}
|
||||
assert cfg["channels"]["discord"]["token"] == "${D}"
|
||||
|
||||
|
||||
def test_botos_platforms_are_converted(tmp_path):
|
||||
path = tmp_path / "bot.yaml"
|
||||
path.write_text(yaml.safe_dump({"agent": {"instructions": "hi"}, "platforms": {"telegram": {"token": "${T}"}}}))
|
||||
result = CliRunner().invoke(app, ["channels", "update", "telegram", "--config", str(path), "--set", "group_policy=mention_only"])
|
||||
assert result.exit_code == 0, result.output
|
||||
cfg = _load(path)
|
||||
assert "platforms" not in cfg
|
||||
assert cfg["channels"]["telegram"] == {"platform": "telegram", "token": "${T}", "group_policy": "mention_only"}
|
||||
|
||||
|
||||
def test_env_placeholder_in_typed_field_does_not_block_edit(config, monkeypatch):
|
||||
cfg = _load(config)
|
||||
cfg["channels"]["telegram"]["max_inbound_media_bytes"] = "${MEDIA_LIMIT}"
|
||||
config.write_text(yaml.safe_dump(cfg, sort_keys=False))
|
||||
monkeypatch.setenv("MEDIA_LIMIT", "1024")
|
||||
edit_gateway_channel(config, "telegram", {"token": "${NEW}"})
|
||||
saved = _load(config)["channels"]["telegram"]
|
||||
assert saved["max_inbound_media_bytes"] == "${MEDIA_LIMIT}"
|
||||
assert saved["token"] == "${NEW}"
|
||||
|
||||
|
||||
def test_removing_last_channel_is_rejected_even_with_env_token(config, monkeypatch):
|
||||
monkeypatch.setenv("TELEGRAM_BOT_TOKEN", "x")
|
||||
with pytest.raises(ValueError):
|
||||
edit_gateway_channel(config, "telegram", remove=True)
|
||||
with pytest.raises(ValueError):
|
||||
edit_gateway_channel(config, "telegram", {"enabled": False})
|
||||
assert _load(config) == BASE
|
||||
|
||||
|
||||
def test_env_platform_reference_is_resolved_before_conversion(tmp_path, monkeypatch):
|
||||
monkeypatch.setenv("BOT_PLATFORM", "telegram")
|
||||
path = tmp_path / "bot.yaml"
|
||||
path.write_text(yaml.safe_dump({"agent": {"instructions": "hi"}, "platform": "${BOT_PLATFORM}", "token": "${T}"}))
|
||||
edit_gateway_channel(path, "discord", {"token": "${D}"})
|
||||
cfg = _load(path)
|
||||
assert cfg["channels"]["telegram"] == {"platform": "telegram", "token": "${T}"}
|
||||
assert cfg["channels"]["discord"]["token"] == "${D}"
|
||||
|
||||
|
||||
def test_concurrent_edits_are_not_lost(config):
|
||||
from concurrent.futures import ThreadPoolExecutor
|
||||
|
||||
names = [f"ch{i}" for i in range(8)]
|
||||
with ThreadPoolExecutor(len(names)) as pool:
|
||||
list(pool.map(lambda n: edit_gateway_channel(config, n, {"platform": "telegram", "token": "${X}"}), names))
|
||||
assert set(_load(config)["channels"]) == {"telegram", *names}
|
||||
|
||||
|
||||
def test_missing_config_raises_without_lock_file(tmp_path):
|
||||
with pytest.raises(FileNotFoundError):
|
||||
edit_gateway_channel(tmp_path / "nope.yaml", "discord", {})
|
||||
assert not list(tmp_path.iterdir())
|
||||
@@ -219,7 +219,7 @@ TS-only members: `previousResult`?, `onToken`?, `signal`?, `onEvent`?, `options`
|
||||
|
||||
### `Agent.chat`
|
||||
|
||||
- Python: `src/praisonai-agents/praisonaiagents/agent/chat_mixin.py:3478`
|
||||
- Python: `src/praisonai-agents/praisonaiagents/agent/chat_mixin.py:3530`
|
||||
- TypeScript: `src/praisonai-ts/src/agent/simple.ts:3592`
|
||||
- Counts: 17 python params: 7 exact, 9 camelCase, 1 alias, 0 flattened, 0 missing; 0 mismatches; 0 waived; 2 TS-only of 19
|
||||
|
||||
@@ -279,7 +279,7 @@ TS-only members: `condition`?, `contextPolicy`?, `maxContextTokens`?, `maxContex
|
||||
|
||||
### `LLM.__init__`
|
||||
|
||||
- Python: `src/praisonai-agents/praisonaiagents/llm/llm.py:673`
|
||||
- Python: `src/praisonai-agents/praisonaiagents/llm/llm.py:678`
|
||||
- TypeScript: `src/praisonai-ts/src/llm/index.ts:236` (ctor `src/praisonai-ts/src/llm/index.ts:416`)
|
||||
- Counts: 25 python params: 8 exact, 16 camelCase, 1 alias, 0 flattened, 0 missing; 1 mismatches; 1 waived; 1 TS-only of 26
|
||||
|
||||
|
||||
@@ -1797,7 +1797,7 @@
|
||||
}
|
||||
],
|
||||
"python": {
|
||||
"location": "src/praisonai-agents/praisonaiagents/agent/chat_mixin.py:3478",
|
||||
"location": "src/praisonai-agents/praisonaiagents/agent/chat_mixin.py:3530",
|
||||
"resolved_class": "ChatMixin"
|
||||
},
|
||||
"ts_only": [
|
||||
@@ -4659,7 +4659,7 @@
|
||||
}
|
||||
],
|
||||
"python": {
|
||||
"location": "src/praisonai-agents/praisonaiagents/llm/llm.py:673",
|
||||
"location": "src/praisonai-agents/praisonaiagents/llm/llm.py:678",
|
||||
"resolved_class": "LLM"
|
||||
},
|
||||
"ts_only": [
|
||||
|
||||
@@ -60,9 +60,11 @@ def _project_root(tmp_path, monkeypatch):
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def mock_endpoint(monkeypatch):
|
||||
def mock_endpoint(monkeypatch, tmp_path):
|
||||
server = HTTPServer(("127.0.0.1", 0), _MockOpenAI)
|
||||
threading.Thread(target=server.serve_forever, daemon=True).start()
|
||||
# output_file writes are contained to the project root (#5770).
|
||||
monkeypatch.setenv("PRAISONAI_PROJECT_ROOT", str(tmp_path))
|
||||
monkeypatch.setenv("OPENAI_API_KEY", "placeholder")
|
||||
monkeypatch.setenv("OPENAI_BASE_URL", f"http://127.0.0.1:{server.server_port}/v1")
|
||||
monkeypatch.setenv("PRAISONAI_TELEMETRY_DISABLED", "true")
|
||||
|
||||
Reference in new issue
Block a user