a2a-source-edit: write agent.py

agent.py

@@ -854,9 +854,21 @@ def _prioritize(analysis: AnalyzeRevenueLeakageOutput, policy: PolicyConfig, max
         low_risk_bonus = 15 if risk == RiskLevel.low else 5 if risk == RiskLevel.medium else 0
         blocker_penalty = min(60, len(blockers) * 18)
         score = max(0.0, min(100.0, value_score * 0.45 + likelihood * 35 + low_risk_bonus - blocker_penalty))
-        accounts.append(PrioritizedAccount(rank=1, customer_token=token, priority_score=round(score, 2), recoverable_amount=min(amount, policy.maximum_account_amount), currency=findings[0].currency, leakage_types=sorted({f.leakage_type for f in findings}, key=lambda x: x.value), likelihood=likelihood, customer_risk=risk, policy_blockers=blockers))
-    accounts.sort(key=lambda item: (-item.priority_score, -item.recoverable_amount, item.customer_token))
-    return [item.model_copy(update={"rank": idx + 1}) for idx, item in enumerate(accounts[:max_accounts])]
+        leakage_types = sorted({finding.leakage_type for finding in findings}, key=_leakage_type_sort_key)
+        accounts.append(PrioritizedAccount(rank=1, customer_token=token, priority_score=round(score, 2), recoverable_amount=min(amount, policy.maximum_account_amount), currency=findings[0].currency, leakage_types=leakage_types, likelihood=likelihood, customer_risk=risk, policy_blockers=blockers))
+    accounts.sort(key=_prioritized_account_sort_key)
+    ranked_accounts: list[PrioritizedAccount] = []
+    for index, account in enumerate(accounts[:max_accounts]):
+        ranked_accounts.append(account.model_copy(update={"rank": index + 1}))
+    return ranked_accounts
+
+
+def _leakage_type_sort_key(leakage_type: LeakageType) -> str:
+    return leakage_type.value
+
+
+def _prioritized_account_sort_key(account: PrioritizedAccount) -> tuple[float, float, str]:
+    return (-account.priority_score, -account.recoverable_amount, account.customer_token)
 
 
 def _propose(analysis: AnalyzeRevenueLeakageOutput, prioritized: PrioritizeAccountsOutput, max_plans: int) -> list[RecoveryPlan]: