46 lines
1.5 KiB
Python
46 lines
1.5 KiB
Python
|
"""
|
|||
|
|
执行预算：全局默认 + Agent.budget_config 合并，供 WorkflowEngine 与任务入口使用。
|
|||
|
|
"""
|
|||
|
|
from typing import Dict, Optional
|
|||
|
|
|
|||
|
|
from sqlalchemy.orm import Session
|
|||
|
|
|
|||
|
|
from app.core.config import settings
|
|||
|
|
from app.models.agent import Agent
|
|||
|
|
from app.models.execution import Execution
|
|||
|
|
|
|||
|
|
|
|||
|
|
def merge_budget_for_execution(db: Session, execution: Optional[Execution]) -> Dict[str, int]:
|
|||
|
|
"""
|
|||
|
|
返回合并后的整数预算（至少为 1）。
|
|||
|
|
keys: max_steps, max_llm_invocations, max_tool_calls
|
|||
|
|
"""
|
|||
|
|
out: Dict[str, int] = {
|
|||
|
|
"max_steps": max(1, int(getattr(settings, "WORKFLOW_MAX_STEPS_PER_RUN", 2000) or 2000)),
|
|||
|
|
"max_llm_invocations": max(
|
|||
|
|
1, int(getattr(settings, "WORKFLOW_MAX_LLM_INVOCATIONS_PER_RUN", 200) or 200)
|
|||
|
|
),
|
|||
|
|
"max_tool_calls": max(
|
|||
|
|
1, int(getattr(settings, "WORKFLOW_MAX_TOOL_CALLS_PER_RUN", 500) or 500)
|
|||
|
|
),
|
|||
|
|
}
|
|||
|
|
if not execution or not execution.agent_id:
|
|||
|
|
return out
|
|||
|
|
ag = db.query(Agent).filter(Agent.id == execution.agent_id).first()
|
|||
|
|
if not ag or not isinstance(ag.budget_config, dict):
|
|||
|
|
return out
|
|||
|
|
bc = ag.budget_config
|
|||
|
|
mapping = {
|
|||
|
|
"max_steps": "max_steps",
|
|||
|
|
"max_llm_invocations": "max_llm_invocations",
|
|||
|
|
"max_tool_calls": "max_tool_calls",
|
|||
|
|
}
|
|||
|
|
for json_key, out_key in mapping.items():
|
|||
|
|
if json_key not in bc or bc[json_key] is None:
|
|||
|
|
continue
|
|||
|
|
try:
|
|||
|
|
out[out_key] = max(1, int(bc[json_key]))
|
|||
|
|
except (TypeError, ValueError):
|
|||
|
|
continue
|
|||
|
|
return out
|