4e9145db46
Some checks failed
CI / test (push) Failing after 32s
- Remove unused imports (json, datetime, timedelta, Any, Optional, IntelligentEval) - Remove unused variable (estimated_sessions) - Remove whitespace from blank line - Organize import blocks All ruff checks passing.
277 lines
8.2 KiB
Python
277 lines
8.2 KiB
Python
"""Decision logic for intelligent eval workers (决策逻辑).
|
|
|
|
Worker analyzes current situation and decides what to do:
|
|
- execute_session: Execute a new session
|
|
- wait: Wait for next tick
|
|
- start_analysis: Start analysis (all sessions completed)
|
|
"""
|
|
|
|
from datetime import datetime, timedelta
|
|
from enum import Enum
|
|
from typing import Optional
|
|
|
|
from sqlmodel import Session, select
|
|
|
|
from agenteval.intelligent_eval.models import IntelligentEvalStatus
|
|
from agenteval.storage.db import (
|
|
IntelligentEvalDB,
|
|
IntelligentEvalSessionDB,
|
|
as_utc,
|
|
utc_now,
|
|
)
|
|
|
|
|
|
class DecisionType(str, Enum):
|
|
"""Decision types for worker."""
|
|
|
|
EXECUTE_SESSION = "execute_session"
|
|
WAIT = "wait"
|
|
START_ANALYSIS = "start_analysis"
|
|
|
|
|
|
class Decision:
|
|
"""A decision made by a worker."""
|
|
|
|
def __init__(self, decision_type: DecisionType, reason: str, context: dict):
|
|
self.decision_type = decision_type
|
|
self.reason = reason
|
|
self.context = context
|
|
|
|
def to_dict(self) -> dict:
|
|
return {
|
|
"decision_type": self.decision_type.value,
|
|
"reason": self.reason,
|
|
"context": self.context,
|
|
}
|
|
|
|
|
|
def _parse_time_slot(time_slot: str) -> Optional[tuple[int, int]]:
|
|
"""Parse time slot string (e.g., "8-10h") to (start_hour, end_hour).
|
|
|
|
Returns:
|
|
(start_hour, end_hour) tuple, or None if invalid format
|
|
"""
|
|
try:
|
|
parts = time_slot.replace("h", "").split("-")
|
|
if len(parts) != 2:
|
|
return None
|
|
return (int(parts[0]), int(parts[1]))
|
|
except (ValueError, AttributeError):
|
|
return None
|
|
|
|
|
|
def _get_current_slot(time_distribution: list[dict], current_offset: timedelta) -> Optional[dict]:
|
|
"""Get current time slot based on offset.
|
|
|
|
Returns:
|
|
Current slot dict, or None if not in any slot
|
|
"""
|
|
current_hours = current_offset.total_seconds() / 3600
|
|
|
|
for slot in time_distribution:
|
|
time_slot = slot.get("time_slot", "")
|
|
parsed = _parse_time_slot(time_slot)
|
|
if parsed is None:
|
|
continue
|
|
|
|
start_hour, end_hour = parsed
|
|
if start_hour <= current_hours < end_hour:
|
|
return slot
|
|
|
|
return None
|
|
|
|
|
|
def _count_sessions_in_slot(
|
|
eval_id: str, slot: dict, eval_started_at: datetime, session: Session
|
|
) -> int:
|
|
"""Count sessions created in a time slot.
|
|
|
|
Args:
|
|
eval_id: Evaluation ID
|
|
slot: Time slot dict (e.g., {"time_slot": "8-10h", "sessions": 2})
|
|
eval_started_at: When the eval started
|
|
session: Database session
|
|
|
|
Returns:
|
|
Number of sessions created in this slot
|
|
"""
|
|
time_slot = slot.get("time_slot", "")
|
|
parsed = _parse_time_slot(time_slot)
|
|
if parsed is None:
|
|
return 0
|
|
|
|
start_hour, end_hour = parsed
|
|
|
|
# Convert to naive datetime for SQLite comparison
|
|
slot_start = (as_utc(eval_started_at) + timedelta(hours=start_hour)).replace(tzinfo=None)
|
|
slot_end = (as_utc(eval_started_at) + timedelta(hours=end_hour)).replace(tzinfo=None)
|
|
|
|
# Count sessions created within slot time range
|
|
sessions = session.exec(
|
|
select(IntelligentEvalSessionDB).where(
|
|
IntelligentEvalSessionDB.eval_id == eval_id,
|
|
IntelligentEvalSessionDB.created_at >= slot_start,
|
|
IntelligentEvalSessionDB.created_at < slot_end,
|
|
)
|
|
).all()
|
|
|
|
return len(sessions)
|
|
|
|
|
|
def _has_high_severity_issues(eval_id: str, session: Session) -> bool:
|
|
"""Check if any completed session has high severity issues.
|
|
|
|
Returns:
|
|
True if any session's verdict contains high severity issue
|
|
"""
|
|
completed_sessions = session.exec(
|
|
select(IntelligentEvalSessionDB).where(
|
|
IntelligentEvalSessionDB.eval_id == eval_id,
|
|
IntelligentEvalSessionDB.status == "completed",
|
|
)
|
|
).all()
|
|
|
|
for sess in completed_sessions:
|
|
verdict = sess.get_verdict()
|
|
if verdict and verdict.get("severity") == "high":
|
|
return True
|
|
|
|
return False
|
|
|
|
|
|
def make_decision(eval_db: IntelligentEvalDB, session: Session) -> Decision:
|
|
"""Make a decision based on current evaluation state.
|
|
|
|
Args:
|
|
eval_db: Evaluation database record
|
|
session: Database session
|
|
|
|
Returns:
|
|
Decision object
|
|
"""
|
|
# Check if eval is still executing
|
|
if eval_db.status != IntelligentEvalStatus.EXECUTING.value:
|
|
return Decision(
|
|
DecisionType.WAIT,
|
|
f"评估状态为 {eval_db.status}，不在执行中",
|
|
{"status": eval_db.status},
|
|
)
|
|
|
|
# Check if eval has plan and started_at
|
|
if not eval_db.plan or not eval_db.started_at:
|
|
return Decision(
|
|
DecisionType.WAIT,
|
|
"评估缺少计划或未开始",
|
|
{"has_plan": bool(eval_db.plan), "has_started_at": bool(eval_db.started_at)},
|
|
)
|
|
|
|
plan = eval_db.get_plan()
|
|
time_distribution = plan.get("time_distribution", [])
|
|
estimated_sessions = plan.get("estimated_sessions", 0)
|
|
|
|
# Calculate current offset
|
|
current_offset = utc_now() - as_utc(eval_db.started_at)
|
|
current_hours = current_offset.total_seconds() / 3600
|
|
|
|
# Get all sessions
|
|
all_sessions = session.exec(
|
|
select(IntelligentEvalSessionDB).where(IntelligentEvalSessionDB.eval_id == eval_db.id)
|
|
).all()
|
|
|
|
completed_sessions = [s for s in all_sessions if s.status == "completed"]
|
|
|
|
# Check if all sessions completed and estimated reached
|
|
if len(completed_sessions) >= estimated_sessions:
|
|
return Decision(
|
|
DecisionType.START_ANALYSIS,
|
|
f"所有 {estimated_sessions} 个会话已完成，开始分析",
|
|
{
|
|
"completed_sessions": len(completed_sessions),
|
|
"estimated_sessions": estimated_sessions,
|
|
},
|
|
)
|
|
|
|
# Get current time slot
|
|
current_slot = _get_current_slot(time_distribution, current_offset)
|
|
|
|
if current_slot is None:
|
|
return Decision(
|
|
DecisionType.WAIT,
|
|
f"当前时间偏移 {current_hours:.1f}h 不在任何时段内",
|
|
{"current_offset_hours": current_hours},
|
|
)
|
|
|
|
# Check if current slot has deficit
|
|
slot_name = current_slot.get("time_slot", "")
|
|
expected_sessions = current_slot.get("sessions", 0)
|
|
current_sessions = _count_sessions_in_slot(eval_db.id, current_slot, eval_db.started_at, session)
|
|
|
|
deficit = expected_sessions - current_sessions
|
|
|
|
if deficit > 0:
|
|
return Decision(
|
|
DecisionType.EXECUTE_SESSION,
|
|
f"时段 {slot_name} 欠账 {deficit} 个会话",
|
|
{
|
|
"current_slot": slot_name,
|
|
"expected_sessions": expected_sessions,
|
|
"current_sessions": current_sessions,
|
|
"deficit": deficit,
|
|
},
|
|
)
|
|
|
|
# Check if any high severity issues found
|
|
if _has_high_severity_issues(eval_db.id, session):
|
|
return Decision(
|
|
DecisionType.EXECUTE_SESSION,
|
|
"发现高严重度问题，需要深入挖掘",
|
|
{"has_high_severity": True},
|
|
)
|
|
|
|
# No deficit, no high severity issues, wait
|
|
return Decision(
|
|
DecisionType.WAIT,
|
|
f"时段 {slot_name} 无欠账，等待下一时段",
|
|
{
|
|
"current_slot": slot_name,
|
|
"expected_sessions": expected_sessions,
|
|
"current_sessions": current_sessions,
|
|
},
|
|
)
|
|
|
|
|
|
def is_eval_completed(eval_db: IntelligentEvalDB, session: Session) -> bool:
|
|
"""Check if evaluation is completed.
|
|
|
|
An eval is completed when:
|
|
1. All sessions are completed
|
|
2. Report is submitted
|
|
|
|
Returns:
|
|
True if eval is completed
|
|
"""
|
|
if eval_db.status != IntelligentEvalStatus.EXECUTING.value:
|
|
return False
|
|
|
|
if not eval_db.plan:
|
|
return False
|
|
|
|
plan = eval_db.get_plan()
|
|
estimated_sessions = plan.get("estimated_sessions", 0)
|
|
|
|
# Check if all sessions completed
|
|
all_sessions = session.exec(
|
|
select(IntelligentEvalSessionDB).where(IntelligentEvalSessionDB.eval_id == eval_db.id)
|
|
).all()
|
|
|
|
completed_sessions = [s for s in all_sessions if s.status == "completed"]
|
|
|
|
if len(completed_sessions) < estimated_sessions:
|
|
return False
|
|
|
|
# Check if report submitted
|
|
if not eval_db.report:
|
|
return False
|
|
|
|
return True
|