68 lines
2.4 KiB
Python
68 lines
2.4 KiB
Python
"""Token estimates and tier limits used by the executable role planner."""
|
|
from __future__ import annotations
|
|
|
|
import os
|
|
from typing import Any
|
|
|
|
import yaml
|
|
|
|
ROOT = os.environ.get("CLAUDE_PROJECT_DIR") or os.path.dirname(
|
|
os.path.dirname(os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))))
|
|
)
|
|
KPI_PATH = os.path.join(ROOT, "org-os", "06-agent-work", "agent-operating-kpi.yaml")
|
|
|
|
TIER_ROLE_LIMITS = {"light": 2, "standard": 5, "heavy": 12}
|
|
TIER_ROLE_TOKENS = {
|
|
"light": {"input": 9000, "output": 2500},
|
|
"standard": {"input": 18000, "output": 4500},
|
|
"heavy": {"input": 32000, "output": 8000},
|
|
}
|
|
ROLE_TYPE_MULTIPLIER = {
|
|
"coordinator": 0.65,
|
|
"worker": 1.0,
|
|
"recommender": 0.9,
|
|
"reviewer": 0.85,
|
|
"auditor": 0.9,
|
|
"decider": 1.05,
|
|
}
|
|
|
|
|
|
def workflow_budget(tier: str, override: int | None = None) -> int:
|
|
if override is not None:
|
|
return max(0, int(override))
|
|
try:
|
|
data = yaml.safe_load(open(KPI_PATH, encoding="utf-8")) or {}
|
|
budgets = data["agent-operating-kpi"]["token-budgets"]["per-wave"]
|
|
return int(budgets[tier])
|
|
except Exception:
|
|
return {"light": 150000, "standard": 500000, "heavy": 2000000}.get(tier, 500000)
|
|
|
|
|
|
def estimate_role(role: dict[str, Any], tier: str, stage: str | None = None) -> dict[str, int]:
|
|
base = TIER_ROLE_TOKENS.get(tier, TIER_ROLE_TOKENS["standard"])
|
|
multiplier = ROLE_TYPE_MULTIPLIER.get(str(role.get("role-type") or "worker"), 1.0)
|
|
if stage in {"verification", "acceptance"}:
|
|
multiplier *= 0.8
|
|
input_tokens = int(base["input"] * multiplier)
|
|
output_tokens = int(base["output"] * multiplier)
|
|
return {"input": input_tokens, "output": output_tokens, "total": input_tokens + output_tokens}
|
|
|
|
|
|
def estimate_plan(selected_roles: list[dict[str, Any]], tier: str, stage: str | None = None) -> dict[str, int]:
|
|
input_tokens = output_tokens = 0
|
|
for role in selected_roles:
|
|
estimate = estimate_role(role, tier, stage)
|
|
input_tokens += estimate["input"]
|
|
output_tokens += estimate["output"]
|
|
synthesis = 0 if len(selected_roles) <= 1 else (3500 if tier == "light" else 7000) * (len(selected_roles) - 1)
|
|
return {
|
|
"input": input_tokens,
|
|
"output": output_tokens,
|
|
"synthesis": synthesis,
|
|
"total": input_tokens + output_tokens + synthesis,
|
|
}
|
|
|
|
|
|
def max_selected_roles(tier: str) -> int:
|
|
return TIER_ROLE_LIMITS.get(tier, TIER_ROLE_LIMITS["standard"])
|