Merge remote-tracking branch 'origin/dev' into dev
This commit is contained in:
mjalmoudy94
@@ -0,0 +1,475 @@
|
||||
"""Repositories for the AI-workforce engine's own tables.
|
||||
|
||||
Four repositories, one per engine table, all following the house rule: the
|
||||
ORM row never leaves this module — the domain dataclasses in
|
||||
``app.domain.workforce.pipeline`` / ``engine`` are what the application layer
|
||||
holds.
|
||||
|
||||
The interesting one is :class:`PipelineStateRepository`. Its ``advance`` is an
|
||||
optimistic compare-and-set: the ``UPDATE`` carries the caller's expected stage
|
||||
and version in its ``WHERE`` clause, so two workers moving the same task
|
||||
cannot both win — the loser sees zero updated rows and gets a
|
||||
:class:`StaleStageError` instead of silently overwriting the winner.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from sqlalchemy import func, select, update
|
||||
from sqlalchemy.ext.asyncio import AsyncSession
|
||||
|
||||
from app.domain.base import utcnow
|
||||
from app.domain.workforce.engine import (
|
||||
EmployeeProfile,
|
||||
EngineTurn,
|
||||
GateName,
|
||||
GateResult,
|
||||
)
|
||||
from app.domain.workforce.models import RoleKey
|
||||
from app.domain.workforce.pipeline import (
|
||||
TRANSITIONS,
|
||||
Mode,
|
||||
PipelineKey,
|
||||
Risk,
|
||||
Stage,
|
||||
TaskKind,
|
||||
TaskPipelineState,
|
||||
)
|
||||
from app.infrastructure.db.models_engine import (
|
||||
EmployeeProfileRow,
|
||||
EngineTurnRow,
|
||||
GateResultRow,
|
||||
TaskPipelineStateRow,
|
||||
)
|
||||
|
||||
|
||||
class StaleStageError(Exception):
|
||||
"""Another worker moved the task first; the caller must re-read it."""
|
||||
|
||||
|
||||
class UnknownTaskError(LookupError):
|
||||
"""The task has no pipeline state row (it was never registered)."""
|
||||
|
||||
|
||||
class InvalidTransitionError(ValueError):
|
||||
"""The requested move is not in the stage transition table."""
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# task_pipeline_state
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
#: Columns ``patch`` may touch. An allow-list, so a typo can never become SQL.
|
||||
_PATCHABLE = frozenset(
|
||||
{
|
||||
"pipeline",
|
||||
"kind",
|
||||
"risk",
|
||||
"stage",
|
||||
"cycle",
|
||||
"max_cycles",
|
||||
"branch",
|
||||
"last_diff_hash",
|
||||
"same_diff_count",
|
||||
"retry_count",
|
||||
"waiver_reason",
|
||||
"parent_task_id",
|
||||
"depends_on",
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
def _state(row: TaskPipelineStateRow) -> TaskPipelineState:
|
||||
return TaskPipelineState(
|
||||
task_id=row.task_id,
|
||||
organization_id=row.organization_id,
|
||||
project_id=row.project_id,
|
||||
pipeline=PipelineKey(row.pipeline),
|
||||
kind=TaskKind(row.kind),
|
||||
risk=Risk(row.risk),
|
||||
stage=Stage(row.stage),
|
||||
cycle=row.cycle,
|
||||
max_cycles=row.max_cycles,
|
||||
branch=row.branch,
|
||||
last_diff_hash=row.last_diff_hash,
|
||||
same_diff_count=row.same_diff_count,
|
||||
retry_count=row.retry_count,
|
||||
waiver_reason=row.waiver_reason,
|
||||
parent_task_id=row.parent_task_id,
|
||||
depends_on=list(row.depends_on or []),
|
||||
version=row.version,
|
||||
updated_at=row.updated_at,
|
||||
)
|
||||
|
||||
|
||||
def _legal_targets(stage: Stage) -> frozenset[Stage]:
|
||||
"""The stage moves the domain allows from here."""
|
||||
return TRANSITIONS.get(stage, frozenset())
|
||||
|
||||
|
||||
class PipelineStateRepository:
|
||||
"""1:1 with a task: pipeline, stage, cycle guards and the version lock."""
|
||||
|
||||
def __init__(self, session: AsyncSession) -> None:
|
||||
self.session = session
|
||||
|
||||
async def get(self, task_id: str) -> TaskPipelineState | None:
|
||||
row = await self.session.get(TaskPipelineStateRow, task_id)
|
||||
return _state(row) if row else None
|
||||
|
||||
async def get_or_raise(self, task_id: str) -> TaskPipelineState:
|
||||
state = await self.get(task_id)
|
||||
if state is None:
|
||||
raise UnknownTaskError(task_id)
|
||||
return state
|
||||
|
||||
async def list_for_org(
|
||||
self, organization_id: str, *, limit: int = 500
|
||||
) -> list[TaskPipelineState]:
|
||||
stmt = (
|
||||
select(TaskPipelineStateRow)
|
||||
.where(TaskPipelineStateRow.organization_id == organization_id)
|
||||
.order_by(TaskPipelineStateRow.updated_at.desc())
|
||||
.limit(min(limit, 1000))
|
||||
)
|
||||
return [_state(r) for r in (await self.session.execute(stmt)).scalars()]
|
||||
|
||||
async def list_at_stages(
|
||||
self, stages: tuple[Stage, ...], *, limit: int = 50
|
||||
) -> list[TaskPipelineState]:
|
||||
"""The scheduler's hot query: tasks currently at a turn stage."""
|
||||
if not stages:
|
||||
return []
|
||||
stmt = (
|
||||
select(TaskPipelineStateRow)
|
||||
.where(TaskPipelineStateRow.stage.in_([str(s) for s in stages]))
|
||||
.order_by(TaskPipelineStateRow.updated_at.asc())
|
||||
.limit(min(limit, 500))
|
||||
)
|
||||
return [_state(r) for r in (await self.session.execute(stmt)).scalars()]
|
||||
|
||||
async def list_dependents(self, task_id: str) -> list[TaskPipelineState]:
|
||||
"""Tasks that declare ``task_id`` in ``depends_on`` and are not closed.
|
||||
|
||||
JSON columns have no portable containment operator (SQLite differs from
|
||||
PostgreSQL), so the match happens in Python. The set is small — it is
|
||||
the tasks waiting on one specific task — so this is not a scan of the
|
||||
table.
|
||||
"""
|
||||
stmt = select(TaskPipelineStateRow).where(
|
||||
TaskPipelineStateRow.stage.notin_(
|
||||
[str(Stage.COMPLETED), str(Stage.CANCELLED), str(Stage.FAILED)]
|
||||
)
|
||||
)
|
||||
dependents: list[TaskPipelineState] = []
|
||||
for row in (await self.session.execute(stmt)).scalars():
|
||||
if task_id in list(row.depends_on or []):
|
||||
dependents.append(_state(row))
|
||||
return dependents
|
||||
|
||||
async def create(self, state: TaskPipelineState) -> TaskPipelineState:
|
||||
row = TaskPipelineStateRow(
|
||||
task_id=state.task_id,
|
||||
organization_id=state.organization_id,
|
||||
project_id=state.project_id,
|
||||
pipeline=str(state.pipeline),
|
||||
kind=str(state.kind),
|
||||
risk=str(state.risk),
|
||||
stage=str(state.stage),
|
||||
cycle=state.cycle,
|
||||
max_cycles=state.max_cycles,
|
||||
branch=state.branch,
|
||||
last_diff_hash=state.last_diff_hash,
|
||||
same_diff_count=state.same_diff_count,
|
||||
retry_count=state.retry_count,
|
||||
waiver_reason=state.waiver_reason,
|
||||
parent_task_id=state.parent_task_id,
|
||||
depends_on=list(state.depends_on),
|
||||
version=state.version,
|
||||
)
|
||||
self.session.add(row)
|
||||
await self.session.flush()
|
||||
return _state(row)
|
||||
|
||||
async def advance(
|
||||
self, *, task_id: str, expected_from: Stage, to: Stage
|
||||
) -> TaskPipelineState:
|
||||
"""Move a task one stage forward, but only if nobody moved it yet.
|
||||
|
||||
The row is read for its current version, then written with both the
|
||||
expected stage and that version in the ``WHERE``. If another worker
|
||||
got there first, zero rows match and :class:`StaleStageError` is
|
||||
raised — the caller re-reads and retries (or gives up cleanly).
|
||||
"""
|
||||
row = await self.session.get(TaskPipelineStateRow, task_id)
|
||||
if row is None:
|
||||
raise UnknownTaskError(task_id)
|
||||
if Stage(row.stage) is not Stage(expected_from):
|
||||
raise StaleStageError(
|
||||
f"task {task_id} is at {row.stage}, not {expected_from.value}"
|
||||
)
|
||||
if Stage(to) not in _legal_targets(Stage(row.stage)):
|
||||
raise InvalidTransitionError(f"{row.stage} -> {to.value} is not allowed")
|
||||
expected_version = row.version
|
||||
# synchronize_session="fetch" keeps the identity-mapped object truthful
|
||||
# (SQLite has no RETURNING for this), so the check below reads what the
|
||||
# database actually did rather than what we hoped it would do.
|
||||
await self.session.execute(
|
||||
update(TaskPipelineStateRow)
|
||||
.where(
|
||||
TaskPipelineStateRow.task_id == task_id,
|
||||
TaskPipelineStateRow.stage == str(expected_from),
|
||||
TaskPipelineStateRow.version == expected_version,
|
||||
)
|
||||
.values(stage=str(to), version=expected_version + 1, updated_at=utcnow())
|
||||
.execution_options(synchronize_session="fetch")
|
||||
)
|
||||
await self.session.flush()
|
||||
# A version that did not move means the WHERE matched nothing: another
|
||||
# worker wrote this row first and this move never happened.
|
||||
if row.version == expected_version or Stage(row.stage) is not Stage(to):
|
||||
raise StaleStageError(f"task {task_id} was moved by another worker")
|
||||
return _state(row)
|
||||
|
||||
async def patch(
|
||||
self, *, task_id: str, expected_version: int, **fields: object
|
||||
) -> TaskPipelineState:
|
||||
"""Change counters without moving the stage, under the version lock."""
|
||||
unknown = set(fields) - _PATCHABLE
|
||||
if unknown:
|
||||
raise ValueError(f"non-patchable columns: {sorted(unknown)}")
|
||||
row = await self.session.get(TaskPipelineStateRow, task_id)
|
||||
if row is None:
|
||||
raise UnknownTaskError(task_id)
|
||||
if not fields:
|
||||
return _state(row)
|
||||
before = row.version
|
||||
await self.session.execute(
|
||||
update(TaskPipelineStateRow)
|
||||
.where(
|
||||
TaskPipelineStateRow.task_id == task_id,
|
||||
TaskPipelineStateRow.version == expected_version,
|
||||
)
|
||||
.values(**fields, version=expected_version + 1, updated_at=utcnow())
|
||||
.execution_options(synchronize_session="fetch")
|
||||
)
|
||||
await self.session.flush()
|
||||
if row.version == before:
|
||||
raise StaleStageError(f"task {task_id} was moved by another worker")
|
||||
return _state(row)
|
||||
|
||||
async def count_at(self, stage: Stage) -> int:
|
||||
stmt = (
|
||||
select(func.count())
|
||||
.select_from(TaskPipelineStateRow)
|
||||
.where(TaskPipelineStateRow.stage == str(stage))
|
||||
)
|
||||
return int((await self.session.execute(stmt)).scalar_one())
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# employee_profiles
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def _profile(row: EmployeeProfileRow) -> EmployeeProfile:
|
||||
return EmployeeProfile(
|
||||
employee_id=row.employee_id,
|
||||
ai_model_id=row.ai_model_id,
|
||||
effort=row.effort,
|
||||
roles=[r.strip() for r in (row.roles or "").split(",") if r.strip()],
|
||||
hired_by=row.hired_by,
|
||||
busy_task_id=row.busy_task_id,
|
||||
updated_at=row.updated_at,
|
||||
)
|
||||
|
||||
|
||||
class EmployeeProfileRepository:
|
||||
"""The engine's view of an employee: granted model, roles, current task."""
|
||||
|
||||
def __init__(self, session: AsyncSession) -> None:
|
||||
self.session = session
|
||||
|
||||
async def get(self, employee_id: str) -> EmployeeProfile | None:
|
||||
row = await self.session.get(EmployeeProfileRow, employee_id)
|
||||
return _profile(row) if row else None
|
||||
|
||||
async def list_free_with_role(self, role: str, *, limit: int = 10) -> list[str]:
|
||||
"""Employee ids that hold ``role`` and are not busy, least busy first."""
|
||||
stmt = select(EmployeeProfileRow).where(
|
||||
EmployeeProfileRow.busy_task_id.is_(None),
|
||||
EmployeeProfileRow.roles.like(f"%{role}%"),
|
||||
)
|
||||
return [r.employee_id for r in (await self.session.execute(stmt)).scalars()][:limit]
|
||||
|
||||
async def upsert(self, profile: EmployeeProfile) -> EmployeeProfile:
|
||||
row = await self.session.get(EmployeeProfileRow, profile.employee_id)
|
||||
if row is None:
|
||||
row = EmployeeProfileRow(
|
||||
employee_id=profile.employee_id,
|
||||
ai_model_id=profile.ai_model_id,
|
||||
effort=profile.effort,
|
||||
roles=",".join(profile.role_list()),
|
||||
hired_by=profile.hired_by,
|
||||
busy_task_id=profile.busy_task_id,
|
||||
)
|
||||
self.session.add(row)
|
||||
else:
|
||||
row.ai_model_id = profile.ai_model_id
|
||||
row.effort = profile.effort
|
||||
row.roles = ",".join(profile.role_list())
|
||||
row.hired_by = profile.hired_by
|
||||
row.busy_task_id = profile.busy_task_id
|
||||
await self.session.flush()
|
||||
return _profile(row)
|
||||
|
||||
async def set_busy(self, *, employee_id: str, task_id: str | None) -> None:
|
||||
row = await self.session.get(EmployeeProfileRow, employee_id)
|
||||
if row is None:
|
||||
return
|
||||
row.busy_task_id = task_id
|
||||
await self.session.flush()
|
||||
|
||||
async def set_model(self, *, employee_id: str, ai_model_id: str | None) -> None:
|
||||
row = await self.session.get(EmployeeProfileRow, employee_id)
|
||||
if row is None:
|
||||
return
|
||||
row.ai_model_id = ai_model_id
|
||||
await self.session.flush()
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# engine_turns
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def _turn(row: EngineTurnRow) -> EngineTurn:
|
||||
return EngineTurn(
|
||||
id=row.id,
|
||||
task_id=row.task_id,
|
||||
run_id=row.run_id,
|
||||
employee_id=row.employee_id,
|
||||
role=RoleKey(row.role),
|
||||
mode=Mode(row.mode),
|
||||
stage=Stage(row.stage) if row.stage else Stage.PENDING_IMPLEMENTATION,
|
||||
prompt_chars=row.prompt_chars,
|
||||
reply_excerpt=row.reply_excerpt,
|
||||
tags=dict(row.tags or {}),
|
||||
tool_calls=row.tool_calls,
|
||||
ok=row.ok,
|
||||
error_kind=row.error_kind,
|
||||
started_at=row.started_at,
|
||||
finished_at=row.finished_at,
|
||||
)
|
||||
|
||||
|
||||
class EngineTurnRepository:
|
||||
"""One row per agent turn: the audit trail of who thought what."""
|
||||
|
||||
def __init__(self, session: AsyncSession) -> None:
|
||||
self.session = session
|
||||
|
||||
async def create(self, turn: EngineTurn) -> EngineTurn:
|
||||
row = EngineTurnRow(
|
||||
id=turn.id,
|
||||
task_id=turn.task_id,
|
||||
run_id=turn.run_id,
|
||||
employee_id=turn.employee_id,
|
||||
role=str(turn.role),
|
||||
mode=str(turn.mode),
|
||||
stage=str(turn.stage),
|
||||
prompt_chars=turn.prompt_chars,
|
||||
reply_excerpt=turn.reply_excerpt,
|
||||
tags=dict(turn.tags),
|
||||
tool_calls=turn.tool_calls,
|
||||
ok=turn.ok,
|
||||
error_kind=turn.error_kind,
|
||||
started_at=turn.started_at,
|
||||
finished_at=turn.finished_at,
|
||||
)
|
||||
self.session.add(row)
|
||||
await self.session.flush()
|
||||
return _turn(row)
|
||||
|
||||
async def list_for_task(self, task_id: str, *, limit: int = 200) -> list[EngineTurn]:
|
||||
stmt = (
|
||||
select(EngineTurnRow)
|
||||
.where(EngineTurnRow.task_id == task_id)
|
||||
.order_by(EngineTurnRow.started_at.asc())
|
||||
.limit(min(limit, 1000))
|
||||
)
|
||||
return [_turn(r) for r in (await self.session.execute(stmt)).scalars()]
|
||||
|
||||
async def list_for_run(self, run_id: str) -> list[EngineTurn]:
|
||||
stmt = (
|
||||
select(EngineTurnRow)
|
||||
.where(EngineTurnRow.run_id == run_id)
|
||||
.order_by(EngineTurnRow.started_at.asc())
|
||||
)
|
||||
return [_turn(r) for r in (await self.session.execute(stmt)).scalars()]
|
||||
|
||||
async def finish(
|
||||
self, turn_id: str, *, ok: bool, tool_calls: int, error_kind: str = ""
|
||||
) -> None:
|
||||
row = await self.session.get(EngineTurnRow, turn_id)
|
||||
if row is None:
|
||||
return
|
||||
row.ok = ok
|
||||
row.tool_calls = tool_calls
|
||||
row.error_kind = error_kind
|
||||
row.finished_at = utcnow()
|
||||
await self.session.flush()
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# gate_results
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def _gate(row: GateResultRow) -> GateResult:
|
||||
return GateResult(
|
||||
id=row.id,
|
||||
task_id=row.task_id,
|
||||
gate=GateName(row.gate),
|
||||
ok=row.ok,
|
||||
detail=row.detail,
|
||||
created_at=row.created_at,
|
||||
)
|
||||
|
||||
|
||||
class GateResultRepository:
|
||||
"""One row per gate verdict, with the detail that justifies it."""
|
||||
|
||||
def __init__(self, session: AsyncSession) -> None:
|
||||
self.session = session
|
||||
|
||||
async def create(self, result: GateResult) -> GateResult:
|
||||
row = GateResultRow(
|
||||
id=result.id,
|
||||
task_id=result.task_id,
|
||||
gate=str(result.gate),
|
||||
ok=result.ok,
|
||||
detail=result.detail,
|
||||
created_at=result.created_at,
|
||||
)
|
||||
self.session.add(row)
|
||||
await self.session.flush()
|
||||
return _gate(row)
|
||||
|
||||
async def list_for_task(self, task_id: str, *, limit: int = 200) -> list[GateResult]:
|
||||
stmt = (
|
||||
select(GateResultRow)
|
||||
.where(GateResultRow.task_id == task_id)
|
||||
.order_by(GateResultRow.created_at.asc())
|
||||
.limit(min(limit, 1000))
|
||||
)
|
||||
return [_gate(r) for r in (await self.session.execute(stmt)).scalars()]
|
||||
|
||||
async def latest_for(self, task_id: str, gate: str) -> GateResult | None:
|
||||
stmt = (
|
||||
select(GateResultRow)
|
||||
.where(GateResultRow.task_id == task_id, GateResultRow.gate == gate)
|
||||
.order_by(GateResultRow.created_at.desc())
|
||||
.limit(1)
|
||||
)
|
||||
row = (await self.session.execute(stmt)).scalar_one_or_none()
|
||||
return _gate(row) if row else None
|
||||
Reference in New Issue
Block a user