151 lines
4.8 KiB
Python
151 lines
4.8 KiB
Python
"""Small-step workflow executor helpers for generation harness adoption."""
|
|
|
|
from collections.abc import Awaitable, Callable
|
|
from dataclasses import dataclass
|
|
from typing import TYPE_CHECKING, Any
|
|
|
|
from sqlalchemy.ext.asyncio import AsyncSession
|
|
|
|
from app.services.harness.artifacts import AssetCompletionResult
|
|
from app.services.harness.plans import WorkflowPlan
|
|
from app.services.harness.trace import TraceRecorder
|
|
from app.services.harness.types import ArtifactKind, WorkflowStep
|
|
|
|
if TYPE_CHECKING:
|
|
from app.db.models import GenerationJob
|
|
|
|
AssetTask = Callable[[], Awaitable[AssetCompletionResult]]
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class AssetPlanRunResult:
|
|
"""Result of executing asset-producing tasks from one workflow plan."""
|
|
|
|
task_results: tuple[AssetCompletionResult, ...]
|
|
executed_task_keys: tuple[str, ...]
|
|
ignored_task_keys: tuple[str, ...]
|
|
|
|
@property
|
|
def result_assets(self) -> tuple[str, ...]:
|
|
"""Assets returned by executed task handlers."""
|
|
|
|
return tuple(result.asset for result in self.task_results)
|
|
|
|
def to_metadata(self, plan: WorkflowPlan) -> dict[str, Any]:
|
|
"""Return internal executor coverage metadata for admin-only analytics."""
|
|
|
|
return {
|
|
"plan_mode": plan.mode.value,
|
|
"planned_task_count": len(plan.tasks),
|
|
"executed_task_count": len(self.executed_task_keys),
|
|
"ignored_task_count": len(self.ignored_task_keys),
|
|
"result_count": len(self.task_results),
|
|
"executed_task_keys": list(self.executed_task_keys),
|
|
"ignored_task_keys": list(self.ignored_task_keys),
|
|
"result_assets": list(self.result_assets),
|
|
}
|
|
|
|
|
|
async def record_workflow_plan(
|
|
db: AsyncSession,
|
|
*,
|
|
job: "GenerationJob | None",
|
|
plan: WorkflowPlan,
|
|
) -> None:
|
|
"""Persist a workflow plan snapshot for a tracked job."""
|
|
|
|
await TraceRecorder(db).record_step(
|
|
job=job,
|
|
event_type="workflow_planned",
|
|
status="succeeded",
|
|
message="Workflow plan selected for this generation request.",
|
|
metadata={"plan": plan.to_snapshot()},
|
|
step=WorkflowStep.REQUEST_ACCEPTANCE,
|
|
artifact=ArtifactKind.NONE,
|
|
blocks_main_result=True,
|
|
)
|
|
|
|
|
|
async def record_evaluation_result(
|
|
db: AsyncSession,
|
|
*,
|
|
job: "GenerationJob | None",
|
|
story_id: int | None = None,
|
|
metadata: dict[str, Any],
|
|
status: str,
|
|
artifact: ArtifactKind | str = ArtifactKind.STORY_TEXT,
|
|
) -> None:
|
|
"""Persist a deterministic evaluation result for a tracked job."""
|
|
|
|
await TraceRecorder(db).record_step(
|
|
job=job,
|
|
story_id=story_id,
|
|
event_type="evaluation_completed",
|
|
status=status,
|
|
message="Generated content evaluation completed.",
|
|
metadata=metadata,
|
|
step=WorkflowStep.EVALUATION,
|
|
artifact=artifact,
|
|
blocks_main_result=status != "succeeded",
|
|
)
|
|
|
|
|
|
async def record_executor_result(
|
|
db: AsyncSession,
|
|
*,
|
|
job: "GenerationJob | None",
|
|
plan: WorkflowPlan,
|
|
result: AssetPlanRunResult,
|
|
) -> None:
|
|
"""Persist internal executor coverage metadata for a tracked job."""
|
|
|
|
await TraceRecorder(db).record_step(
|
|
job=job,
|
|
event_type="executor_completed",
|
|
status="succeeded",
|
|
message="Workflow executor completed planned asset tasks.",
|
|
metadata=result.to_metadata(plan),
|
|
step=WorkflowStep.UNKNOWN,
|
|
artifact=ArtifactKind.NONE,
|
|
blocks_main_result=False,
|
|
)
|
|
|
|
|
|
async def run_asset_plan(
|
|
plan: WorkflowPlan,
|
|
*,
|
|
image_task: AssetTask | None = None,
|
|
audio_task: AssetTask | None = None,
|
|
) -> AssetPlanRunResult:
|
|
"""Execute asset-producing tasks in the order declared by a workflow plan."""
|
|
|
|
if plan.mode.value not in {"asset_generation", "asset_retry"}:
|
|
raise ValueError("run_asset_plan only supports asset workflow plans")
|
|
|
|
task_results: list[AssetCompletionResult] = []
|
|
executed_task_keys: list[str] = []
|
|
ignored_task_keys: list[str] = []
|
|
|
|
for task in plan.tasks:
|
|
if task.key == "complete_image_asset":
|
|
if image_task is None:
|
|
raise ValueError("Asset workflow plan requires an image task handler")
|
|
task_results.append(await image_task())
|
|
executed_task_keys.append(task.key)
|
|
continue
|
|
|
|
if task.key == "complete_audio_asset":
|
|
if audio_task is None:
|
|
raise ValueError("Asset workflow plan requires an audio task handler")
|
|
task_results.append(await audio_task())
|
|
executed_task_keys.append(task.key)
|
|
continue
|
|
|
|
ignored_task_keys.append(task.key)
|
|
|
|
return AssetPlanRunResult(
|
|
task_results=tuple(task_results),
|
|
executed_task_keys=tuple(executed_task_keys),
|
|
ignored_task_keys=tuple(ignored_task_keys),
|
|
)
|