460 lines
16 KiB
Python
460 lines
16 KiB
Python
"""Repository for intelligent evaluation entities."""
|
|
|
|
import json
|
|
from collections.abc import Sequence
|
|
from dataclasses import dataclass
|
|
from enum import Enum
|
|
from typing import Optional
|
|
|
|
from sqlalchemy import func, literal
|
|
from sqlalchemy import insert as sql_insert
|
|
from sqlalchemy import update as sql_update
|
|
from sqlmodel import Session, select
|
|
|
|
from agenteval.intelligent_eval.models import (
|
|
IntelligentEval,
|
|
IntelligentEvalMessage,
|
|
IntelligentEvalSession,
|
|
IntelligentEvalSessionStatus,
|
|
IntelligentEvalStatus,
|
|
)
|
|
from agenteval.storage.db import (
|
|
IntelligentEvalDB,
|
|
IntelligentEvalMessageDB,
|
|
IntelligentEvalSessionDB,
|
|
get_session,
|
|
new_uuid,
|
|
utc_now,
|
|
)
|
|
|
|
|
|
class CompareAndSetStatus(str, Enum):
|
|
"""Outcome of a conditional intelligent-evaluation write."""
|
|
|
|
APPLIED = "applied"
|
|
NOT_FOUND = "not_found"
|
|
CONFLICT = "conflict"
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class CompareAndSetResult:
|
|
"""Typed result for a status update guarded by an expected status."""
|
|
|
|
status: CompareAndSetStatus
|
|
evaluation: Optional[IntelligentEval] = None
|
|
|
|
@property
|
|
def applied(self) -> bool:
|
|
return self.status is CompareAndSetStatus.APPLIED
|
|
|
|
|
|
class IntelligentEvalRepository:
|
|
"""CRUD + lifecycle operations for intelligent evaluations."""
|
|
|
|
def __init__(self, session: Optional[Session] = None):
|
|
self.session = session or get_session()
|
|
|
|
def _to_db(self, obj: IntelligentEval) -> IntelligentEvalDB:
|
|
db = IntelligentEvalDB(
|
|
id=obj.id,
|
|
name=obj.name,
|
|
target_id=obj.target_id,
|
|
status=obj.status.value if isinstance(obj.status, IntelligentEvalStatus) else obj.status,
|
|
goal=obj.goal,
|
|
intent=obj.intent,
|
|
role_description=obj.role_description,
|
|
time_window_hours=obj.time_window_hours,
|
|
plan_feedback=obj.plan_feedback,
|
|
created_at=obj.created_at,
|
|
updated_at=obj.updated_at,
|
|
started_at=obj.started_at,
|
|
completed_at=obj.completed_at,
|
|
)
|
|
db.set_seeds(obj.seeds)
|
|
if obj.plan is not None:
|
|
db.set_plan(obj.plan)
|
|
if obj.report is not None:
|
|
db.set_report(obj.report)
|
|
return db
|
|
|
|
def _from_db(self, db: IntelligentEvalDB) -> IntelligentEval:
|
|
return IntelligentEval(
|
|
id=db.id,
|
|
name=db.name,
|
|
target_id=db.target_id,
|
|
status=IntelligentEvalStatus(db.status),
|
|
goal=db.goal,
|
|
seeds=db.get_seeds(),
|
|
intent=db.intent,
|
|
role_description=db.role_description,
|
|
plan=db.get_plan(),
|
|
plan_feedback=db.plan_feedback,
|
|
time_window_hours=db.time_window_hours,
|
|
report=db.get_report(),
|
|
created_at=db.created_at,
|
|
updated_at=db.updated_at,
|
|
started_at=db.started_at,
|
|
completed_at=db.completed_at,
|
|
)
|
|
|
|
def list_all(self) -> list[IntelligentEval]:
|
|
statement = select(IntelligentEvalDB).order_by(IntelligentEvalDB.created_at.desc())
|
|
return [self._from_db(r) for r in self.session.exec(statement).all()]
|
|
|
|
def get(self, eval_id: str) -> Optional[IntelligentEval]:
|
|
db = self.session.get(IntelligentEvalDB, eval_id)
|
|
return self._from_db(db) if db else None
|
|
|
|
def get_with_sessions(
|
|
self,
|
|
eval_id: str,
|
|
) -> tuple[IntelligentEval, list[IntelligentEvalSession]] | None:
|
|
"""Load an evaluation and all session summaries in one DB snapshot."""
|
|
|
|
statement = (
|
|
select(IntelligentEvalDB, IntelligentEvalSessionDB)
|
|
.join(
|
|
IntelligentEvalSessionDB,
|
|
IntelligentEvalSessionDB.eval_id == IntelligentEvalDB.id,
|
|
isouter=True,
|
|
)
|
|
.where(IntelligentEvalDB.id == eval_id)
|
|
.order_by(IntelligentEvalSessionDB.created_at.asc())
|
|
)
|
|
rows = self.session.exec(statement).all()
|
|
if not rows:
|
|
return None
|
|
evaluation = self._from_db(rows[0][0])
|
|
session_repo = IntelligentEvalSessionRepository(self.session)
|
|
sessions = [session_repo._from_db(session_db) for _, session_db in rows if session_db is not None]
|
|
return evaluation, sessions
|
|
|
|
def _create(self, obj: IntelligentEval) -> IntelligentEval:
|
|
db = self._to_db(obj)
|
|
self.session.add(db)
|
|
self.session.commit()
|
|
self.session.refresh(db)
|
|
return self._from_db(db)
|
|
|
|
def _compare_and_set_status(
|
|
self,
|
|
eval_id: str,
|
|
*,
|
|
expected_status: IntelligentEvalStatus,
|
|
new_status: IntelligentEvalStatus,
|
|
) -> CompareAndSetResult:
|
|
"""Atomically transition status only when the expected state still holds."""
|
|
now = utc_now()
|
|
values = {
|
|
"status": new_status.value,
|
|
"updated_at": now,
|
|
}
|
|
if new_status is IntelligentEvalStatus.EXECUTING:
|
|
values["started_at"] = func.coalesce(IntelligentEvalDB.started_at, now)
|
|
if new_status in {
|
|
IntelligentEvalStatus.COMPLETED,
|
|
IntelligentEvalStatus.CANCELLED,
|
|
IntelligentEvalStatus.FAILED,
|
|
}:
|
|
values["completed_at"] = now
|
|
|
|
return self._compare_and_set_fields(
|
|
eval_id,
|
|
expected_status=expected_status,
|
|
values=values,
|
|
)
|
|
|
|
def _compare_and_set_fields(
|
|
self,
|
|
eval_id: str,
|
|
*,
|
|
expected_status: IntelligentEvalStatus,
|
|
values: dict[str, object],
|
|
) -> CompareAndSetResult:
|
|
"""Apply one conditional SQL update and classify its outcome."""
|
|
try:
|
|
result = self.session.exec(
|
|
sql_update(IntelligentEvalDB)
|
|
.where(
|
|
IntelligentEvalDB.id == eval_id,
|
|
IntelligentEvalDB.status == expected_status.value,
|
|
)
|
|
.values(**values)
|
|
)
|
|
if result.rowcount != 1:
|
|
self.session.rollback()
|
|
current = self.session.get(IntelligentEvalDB, eval_id)
|
|
status = CompareAndSetStatus.NOT_FOUND if current is None else CompareAndSetStatus.CONFLICT
|
|
return CompareAndSetResult(status=status)
|
|
self.session.commit()
|
|
except Exception:
|
|
self.session.rollback()
|
|
raise
|
|
|
|
db = self.session.get(IntelligentEvalDB, eval_id)
|
|
if db is None:
|
|
return CompareAndSetResult(status=CompareAndSetStatus.NOT_FOUND)
|
|
return CompareAndSetResult(status=CompareAndSetStatus.APPLIED, evaluation=self._from_db(db))
|
|
|
|
def _submit_plan_if_planning(self, eval_id: str, plan: dict) -> CompareAndSetResult:
|
|
now = utc_now()
|
|
return self._compare_and_set_fields(
|
|
eval_id,
|
|
expected_status=IntelligentEvalStatus.PLANNING,
|
|
values={
|
|
"status": IntelligentEvalStatus.PENDING_APPROVAL.value,
|
|
"plan": json.dumps(plan, ensure_ascii=False),
|
|
"plan_feedback": None,
|
|
"updated_at": now,
|
|
},
|
|
)
|
|
|
|
def _reject_plan_if_pending(self, eval_id: str, feedback: str) -> CompareAndSetResult:
|
|
now = utc_now()
|
|
return self._compare_and_set_fields(
|
|
eval_id,
|
|
expected_status=IntelligentEvalStatus.PENDING_APPROVAL,
|
|
values={
|
|
"status": IntelligentEvalStatus.PLANNING.value,
|
|
"plan_feedback": feedback,
|
|
"updated_at": now,
|
|
},
|
|
)
|
|
|
|
def _submit_report_if_executing(self, eval_id: str, report: dict) -> CompareAndSetResult:
|
|
now = utc_now()
|
|
return self._compare_and_set_fields(
|
|
eval_id,
|
|
expected_status=IntelligentEvalStatus.EXECUTING,
|
|
values={
|
|
"status": IntelligentEvalStatus.COMPLETED.value,
|
|
"report": json.dumps(report, ensure_ascii=False),
|
|
"updated_at": now,
|
|
"completed_at": now,
|
|
},
|
|
)
|
|
|
|
class IntelligentEvalSessionRepository:
|
|
"""CRUD for intelligent eval sessions."""
|
|
|
|
def __init__(self, session: Optional[Session] = None):
|
|
self.session = session or get_session()
|
|
|
|
def _from_db(self, db: IntelligentEvalSessionDB) -> IntelligentEvalSession:
|
|
return IntelligentEvalSession(
|
|
id=db.id,
|
|
eval_id=db.eval_id,
|
|
target_id=db.target_id,
|
|
persona=db.get_persona(),
|
|
goal=db.goal,
|
|
dimension=db.dimension,
|
|
status=IntelligentEvalSessionStatus(db.status),
|
|
verdict=db.get_verdict(),
|
|
turn_count=db.turn_count,
|
|
created_at=db.created_at,
|
|
closed_at=db.closed_at,
|
|
)
|
|
|
|
def list_by_eval(self, eval_id: str) -> list[IntelligentEvalSession]:
|
|
statement = (
|
|
select(IntelligentEvalSessionDB)
|
|
.where(IntelligentEvalSessionDB.eval_id == eval_id)
|
|
.order_by(IntelligentEvalSessionDB.created_at.asc())
|
|
)
|
|
return [self._from_db(r) for r in self.session.exec(statement).all()]
|
|
|
|
def list_by_evals(self, eval_ids: Sequence[str]) -> dict[str, list[IntelligentEvalSession]]:
|
|
"""Load sessions for several evaluations with one query.
|
|
|
|
The result contains an empty list for every requested evaluation. This
|
|
keeps read-model assembly deterministic while avoiding one count query
|
|
per row in the evaluation list.
|
|
"""
|
|
|
|
grouped = {eval_id: [] for eval_id in eval_ids}
|
|
if not grouped:
|
|
return grouped
|
|
statement = (
|
|
select(IntelligentEvalSessionDB)
|
|
.where(IntelligentEvalSessionDB.eval_id.in_(tuple(grouped)))
|
|
.order_by(IntelligentEvalSessionDB.created_at.asc())
|
|
)
|
|
for row in self.session.exec(statement).all():
|
|
grouped.setdefault(row.eval_id, []).append(self._from_db(row))
|
|
return grouped
|
|
|
|
def get(self, session_id: str) -> Optional[IntelligentEvalSession]:
|
|
db = self.session.get(IntelligentEvalSessionDB, session_id)
|
|
return self._from_db(db) if db else None
|
|
|
|
def _create(self, obj: IntelligentEvalSession) -> IntelligentEvalSession:
|
|
db = IntelligentEvalSessionDB(
|
|
id=obj.id,
|
|
eval_id=obj.eval_id,
|
|
target_id=obj.target_id,
|
|
goal=obj.goal,
|
|
dimension=obj.dimension,
|
|
status=obj.status.value if isinstance(obj.status, IntelligentEvalSessionStatus) else obj.status,
|
|
turn_count=obj.turn_count,
|
|
)
|
|
db.set_persona(obj.persona)
|
|
if obj.verdict is not None:
|
|
db.set_verdict(obj.verdict)
|
|
self.session.add(db)
|
|
self.session.commit()
|
|
self.session.refresh(db)
|
|
return self._from_db(db)
|
|
|
|
def _create_if_executing(
|
|
self,
|
|
obj: IntelligentEvalSession,
|
|
) -> tuple[CompareAndSetStatus, Optional[IntelligentEvalSession]]:
|
|
"""Insert a session only while its parent evaluation is executing."""
|
|
session_id = obj.id or new_uuid()
|
|
now = obj.created_at or utc_now()
|
|
statement = sql_insert(IntelligentEvalSessionDB).from_select(
|
|
[
|
|
"id",
|
|
"eval_id",
|
|
"target_id",
|
|
"persona",
|
|
"goal",
|
|
"dimension",
|
|
"status",
|
|
"turn_count",
|
|
"created_at",
|
|
],
|
|
select(
|
|
literal(session_id),
|
|
IntelligentEvalDB.id,
|
|
IntelligentEvalDB.target_id,
|
|
literal(json.dumps(obj.persona, ensure_ascii=False)),
|
|
literal(obj.goal),
|
|
literal(obj.dimension),
|
|
literal(IntelligentEvalSessionStatus.RUNNING.value),
|
|
literal(0),
|
|
literal(now),
|
|
).where(
|
|
IntelligentEvalDB.id == obj.eval_id,
|
|
IntelligentEvalDB.status == IntelligentEvalStatus.EXECUTING.value,
|
|
),
|
|
)
|
|
try:
|
|
result = self.session.exec(statement)
|
|
if result.rowcount != 1:
|
|
self.session.rollback()
|
|
evaluation = self.session.get(IntelligentEvalDB, obj.eval_id)
|
|
status = CompareAndSetStatus.NOT_FOUND if evaluation is None else CompareAndSetStatus.CONFLICT
|
|
return status, None
|
|
self.session.commit()
|
|
except Exception:
|
|
self.session.rollback()
|
|
raise
|
|
|
|
db = self.session.get(IntelligentEvalSessionDB, session_id)
|
|
if db is None:
|
|
return CompareAndSetStatus.NOT_FOUND, None
|
|
return CompareAndSetStatus.APPLIED, self._from_db(db)
|
|
|
|
def _close_if_running(
|
|
self,
|
|
session_id: str,
|
|
verdict: dict,
|
|
status: IntelligentEvalSessionStatus = IntelligentEvalSessionStatus.COMPLETED,
|
|
) -> tuple[CompareAndSetStatus, Optional[IntelligentEvalSession]]:
|
|
"""Close exactly one running session without overwriting a race winner."""
|
|
now = utc_now()
|
|
try:
|
|
result = self.session.exec(
|
|
sql_update(IntelligentEvalSessionDB)
|
|
.where(
|
|
IntelligentEvalSessionDB.id == session_id,
|
|
IntelligentEvalSessionDB.status == IntelligentEvalSessionStatus.RUNNING.value,
|
|
)
|
|
.values(
|
|
status=status.value,
|
|
verdict=json.dumps(verdict, ensure_ascii=False),
|
|
closed_at=now,
|
|
)
|
|
)
|
|
if result.rowcount != 1:
|
|
self.session.rollback()
|
|
current = self.session.get(IntelligentEvalSessionDB, session_id)
|
|
outcome = CompareAndSetStatus.NOT_FOUND if current is None else CompareAndSetStatus.CONFLICT
|
|
return outcome, None
|
|
self.session.commit()
|
|
except Exception:
|
|
self.session.rollback()
|
|
raise
|
|
|
|
db = self.session.get(IntelligentEvalSessionDB, session_id)
|
|
if db is None:
|
|
return CompareAndSetStatus.NOT_FOUND, None
|
|
return CompareAndSetStatus.APPLIED, self._from_db(db)
|
|
|
|
class IntelligentEvalMessageRepository:
|
|
"""CRUD for intelligent eval session messages."""
|
|
|
|
def __init__(self, session: Optional[Session] = None):
|
|
self.session = session or get_session()
|
|
|
|
def _from_db(self, db: IntelligentEvalMessageDB) -> IntelligentEvalMessage:
|
|
return IntelligentEvalMessage(
|
|
id=db.id,
|
|
session_id=db.session_id,
|
|
role=db.role,
|
|
content=db.content,
|
|
latency_ms=db.latency_ms,
|
|
created_at=db.created_at,
|
|
)
|
|
|
|
def list_by_session(self, session_id: str) -> list[IntelligentEvalMessage]:
|
|
statement = (
|
|
select(IntelligentEvalMessageDB)
|
|
.where(IntelligentEvalMessageDB.session_id == session_id)
|
|
.order_by(IntelligentEvalMessageDB.created_at.asc())
|
|
)
|
|
return [self._from_db(r) for r in self.session.exec(statement).all()]
|
|
|
|
def _create(self, obj: IntelligentEvalMessage) -> IntelligentEvalMessage:
|
|
db = IntelligentEvalMessageDB(
|
|
id=obj.id,
|
|
session_id=obj.session_id,
|
|
role=obj.role,
|
|
content=obj.content,
|
|
latency_ms=obj.latency_ms,
|
|
created_at=obj.created_at,
|
|
)
|
|
self.session.add(db)
|
|
self.session.commit()
|
|
self.session.refresh(db)
|
|
return self._from_db(db)
|
|
|
|
def _create_user_and_increment(self, obj: IntelligentEvalMessage) -> CompareAndSetStatus:
|
|
"""Persist the sent user message and consume one turn atomically."""
|
|
session_db = self.session.get(IntelligentEvalSessionDB, obj.session_id)
|
|
if session_db is None:
|
|
return CompareAndSetStatus.NOT_FOUND
|
|
if session_db.status != IntelligentEvalSessionStatus.RUNNING.value:
|
|
return CompareAndSetStatus.CONFLICT
|
|
|
|
message_db = IntelligentEvalMessageDB(
|
|
id=obj.id,
|
|
session_id=obj.session_id,
|
|
role="user",
|
|
content=obj.content,
|
|
latency_ms=None,
|
|
created_at=obj.created_at,
|
|
)
|
|
try:
|
|
session_db.turn_count += 1
|
|
self.session.add(session_db)
|
|
self.session.add(message_db)
|
|
self.session.commit()
|
|
self.session.refresh(message_db)
|
|
except Exception:
|
|
self.session.rollback()
|
|
raise
|
|
obj.id = message_db.id
|
|
return CompareAndSetStatus.APPLIED
|