feat(engine): workforce engine core — tag protocol, pipelines, optimistic state
Phase 1 of TODO-agent-2: pure logic plus the tables it writes to. - domain/workforce/pipeline.py: Stage/PipelineKey/Outcome, the four pipelines, TRANSITIONS, next_stage, STAGE_ROLE (role+mode), resolve_pipeline with the no-reviewer fallback, task-kind defaults, stage -> coarse tasks.status, and the TaskPipelineState record (with the optimistic version). - application/engine/protocol.py: the [[TAG attr="v"]]body[[/TAG]] parser and its self-closing form, the per-role/mode allow-list, literal bodies (a bracket inside a body is not a tag), typed readers (ASSIGN/CREATE_TASK/ team changes/KNOWLEDGE/EDIT_FILE blocks) and the per-reply limits (3 tasks, 3 team changes, 3 knowledge items). - application/engine/state.py: advance() with the optimistic compare-and-set and park_dependents() so a failed dependency never deadlocks its children. - db/models_engine.py + migration a2b3c4d5e6f2: task_pipeline_state, employee_profiles, engine_turns, gate_results as side tables (tasks and employees stay locked). Linear head; alembic check reports no drift. - repositories/engine.py + 58 new tests (test_engine_protocol, test_engine_pipeline) covering every pipeline's happy path, illegal transitions, exactly-one-winner concurrent advance, dependency parking and all parser edge cases. Suite: 594 passed, ruff clean, mypy clean.
This commit is contained in:
mjalmoudy94
@@ -5,6 +5,170 @@ Use ``from app.core.database import Base`` and the helpers in
|
||||
``app.infrastructure.db.models`` here: that module imports this one (import cycle).
|
||||
Foreign keys to existing tables are written as plain strings, e.g.
|
||||
``ForeignKey("organizations.id", ondelete="CASCADE")``. See docs/plan/PLAN-v2.md §6.
|
||||
|
||||
Four tables, deliberately split so ``tasks`` and ``employees`` stay untouched
|
||||
(they are locked):
|
||||
|
||||
* ``task_pipeline_state`` — the engine's own view of a task: pipeline, stage,
|
||||
cycle counters, branch and the optimistic ``version``. One row per task.
|
||||
* ``employee_profiles`` — which granted model an employee runs on, and how
|
||||
busy it is. Model identity never leaves the server.
|
||||
* ``engine_turns`` — one row per agent turn: the audit trail of who thought
|
||||
what, with a bounded reply excerpt.
|
||||
* ``gate_results`` — one row per gate verdict (tests, diff guard, no-progress,
|
||||
green), with the detail that justifies it.
|
||||
|
||||
Both side tables cascade from their parent, so deleting a task or an employee
|
||||
never leaves an orphaned engine row behind.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from datetime import datetime
|
||||
|
||||
from sqlalchemy import (
|
||||
JSON,
|
||||
Boolean,
|
||||
DateTime,
|
||||
ForeignKey,
|
||||
Index,
|
||||
Integer,
|
||||
String,
|
||||
Text,
|
||||
)
|
||||
from sqlalchemy.orm import Mapped, mapped_column
|
||||
|
||||
from app.core.database import Base
|
||||
from app.domain.base import utcnow
|
||||
from app.infrastructure.db.columns import created, pk
|
||||
|
||||
#: The engine's default limits, mirroring the settings block in ``config.py``.
|
||||
#: They are column defaults (not settings) so that a database written by an
|
||||
#: older deployment still carries a sane value.
|
||||
DEFAULT_MAX_CYCLES = 6
|
||||
|
||||
|
||||
class TaskPipelineStateRow(Base):
|
||||
__tablename__ = "task_pipeline_state"
|
||||
|
||||
task_id: Mapped[str] = mapped_column(
|
||||
ForeignKey("tasks.id", ondelete="CASCADE"), primary_key=True
|
||||
)
|
||||
#: Denormalised from the task so the tenant filter stays a single-table
|
||||
#: WHERE: the scheduler's hot query never joins back to ``tasks``.
|
||||
organization_id: Mapped[str] = mapped_column(
|
||||
ForeignKey("organizations.id", ondelete="CASCADE"), nullable=False, index=True
|
||||
)
|
||||
project_id: Mapped[str | None] = mapped_column(
|
||||
ForeignKey("projects.id", ondelete="CASCADE"), nullable=True
|
||||
)
|
||||
pipeline: Mapped[str] = mapped_column(
|
||||
String(32), default="plan_code_test", nullable=False
|
||||
)
|
||||
kind: Mapped[str] = mapped_column(String(16), default="feature", nullable=False)
|
||||
risk: Mapped[str] = mapped_column(String(8), default="medium", nullable=False)
|
||||
#: Indexed: every scheduler tick asks "which tasks are at a turn stage".
|
||||
stage: Mapped[str] = mapped_column(
|
||||
String(32), default="pending_manager_review", nullable=False, index=True
|
||||
)
|
||||
cycle: Mapped[int] = mapped_column(Integer, default=0, nullable=False)
|
||||
max_cycles: Mapped[int] = mapped_column(
|
||||
Integer, default=DEFAULT_MAX_CYCLES, nullable=False
|
||||
)
|
||||
branch: Mapped[str] = mapped_column(String(120), default="", nullable=False)
|
||||
last_diff_hash: Mapped[str] = mapped_column(String(64), default="", nullable=False)
|
||||
same_diff_count: Mapped[int] = mapped_column(Integer, default=0, nullable=False)
|
||||
retry_count: Mapped[int] = mapped_column(Integer, default=0, nullable=False)
|
||||
#: A waived test pass must always be explainable: the reason is stored,
|
||||
#: never inferred (PLAN-v2 §3, green gate).
|
||||
waiver_reason: Mapped[str] = mapped_column(Text, default="", nullable=False)
|
||||
parent_task_id: Mapped[str | None] = mapped_column(
|
||||
ForeignKey("tasks.id", ondelete="SET NULL"), nullable=True
|
||||
)
|
||||
depends_on: Mapped[list] = mapped_column(JSON, default=list, nullable=False)
|
||||
#: Optimistic lock. Every stage move bumps it; a stale write is refused.
|
||||
version: Mapped[int] = mapped_column(Integer, default=0, nullable=False)
|
||||
updated_at: Mapped[datetime] = mapped_column(
|
||||
DateTime(timezone=True), default=utcnow, onupdate=utcnow, nullable=False
|
||||
)
|
||||
|
||||
|
||||
class EmployeeProfileRow(Base):
|
||||
__tablename__ = "employee_profiles"
|
||||
__table_args__ = (Index("ix_employee_profile_busy", "busy_task_id"),)
|
||||
|
||||
employee_id: Mapped[str] = mapped_column(
|
||||
ForeignKey("employees.id", ondelete="CASCADE"), primary_key=True
|
||||
)
|
||||
#: A granted model id (``org_model_grants.ai_model_id``). Nullable until the
|
||||
#: owner grants something: an employee exists but is inactive, which is a
|
||||
#: state the UI shows instead of hiding the employee.
|
||||
ai_model_id: Mapped[str | None] = mapped_column(String(32), nullable=True)
|
||||
effort: Mapped[str] = mapped_column(String(16), default="medium", nullable=False)
|
||||
roles: Mapped[str] = mapped_column(String(255), default="", nullable=False)
|
||||
hired_by: Mapped[str] = mapped_column(String(8), default="user", nullable=False)
|
||||
busy_task_id: Mapped[str | None] = mapped_column(
|
||||
ForeignKey("tasks.id", ondelete="SET NULL"), nullable=True
|
||||
)
|
||||
updated_at: Mapped[datetime] = mapped_column(
|
||||
DateTime(timezone=True), default=utcnow, onupdate=utcnow, nullable=False
|
||||
)
|
||||
|
||||
|
||||
class EngineTurnRow(Base):
|
||||
__tablename__ = "engine_turns"
|
||||
__table_args__ = (
|
||||
Index("ix_engine_turn_task_stage", "task_id", "stage"),
|
||||
Index("ix_engine_turn_run", "run_id"),
|
||||
)
|
||||
|
||||
id: Mapped[str] = pk()
|
||||
task_id: Mapped[str] = mapped_column(
|
||||
ForeignKey("tasks.id", ondelete="CASCADE"), nullable=False
|
||||
)
|
||||
run_id: Mapped[str] = mapped_column(
|
||||
ForeignKey("runs.id", ondelete="CASCADE"), nullable=False
|
||||
)
|
||||
employee_id: Mapped[str | None] = mapped_column(
|
||||
ForeignKey("employees.id", ondelete="SET NULL"), nullable=True
|
||||
)
|
||||
role: Mapped[str] = mapped_column(String(24), default="developer", nullable=False)
|
||||
mode: Mapped[str] = mapped_column(String(16), default="work", nullable=False)
|
||||
stage: Mapped[str] = mapped_column(String(32), default="", nullable=False)
|
||||
prompt_chars: Mapped[int] = mapped_column(Integer, default=0, nullable=False)
|
||||
#: Bounded on write: a reply is unbounded, a row is not.
|
||||
reply_excerpt: Mapped[str] = mapped_column(Text, default="", nullable=False)
|
||||
tags: Mapped[dict] = mapped_column(JSON, default=dict, nullable=False)
|
||||
tool_calls: Mapped[int] = mapped_column(Integer, default=0, nullable=False)
|
||||
ok: Mapped[bool] = mapped_column(Boolean, default=True, nullable=False)
|
||||
error_kind: Mapped[str] = mapped_column(String(16), default="", nullable=False)
|
||||
started_at: Mapped[datetime] = mapped_column(
|
||||
DateTime(timezone=True), default=utcnow, nullable=False
|
||||
)
|
||||
finished_at: Mapped[datetime | None] = mapped_column(
|
||||
DateTime(timezone=True), nullable=True
|
||||
)
|
||||
created_at: Mapped[datetime] = created()
|
||||
|
||||
|
||||
class GateResultRow(Base):
|
||||
__tablename__ = "gate_results"
|
||||
__table_args__ = (Index("ix_gate_result_task_gate", "task_id", "gate"),)
|
||||
|
||||
id: Mapped[str] = pk()
|
||||
task_id: Mapped[str] = mapped_column(
|
||||
ForeignKey("tasks.id", ondelete="CASCADE"), nullable=False
|
||||
)
|
||||
gate: Mapped[str] = mapped_column(String(16), default="tests", nullable=False)
|
||||
ok: Mapped[bool] = mapped_column(Boolean, default=True, nullable=False)
|
||||
detail: Mapped[str] = mapped_column(Text, default="", nullable=False)
|
||||
created_at: Mapped[datetime] = created()
|
||||
|
||||
|
||||
__all__ = [
|
||||
"DEFAULT_MAX_CYCLES",
|
||||
"EmployeeProfileRow",
|
||||
"EngineTurnRow",
|
||||
"GateResultRow",
|
||||
"TaskPipelineStateRow",
|
||||
]
|
||||
|
||||
@@ -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