65 lines
1.8 KiB
Python
65 lines
1.8 KiB
Python
"""Trace recording helpers for generation harness workflows."""
|
|
|
|
from typing import TYPE_CHECKING, Any
|
|
|
|
from sqlalchemy.ext.asyncio import AsyncSession
|
|
|
|
from app.services.generation_jobs import record_generation_event
|
|
from app.services.harness.types import (
|
|
ArtifactKind,
|
|
FailureCategory,
|
|
WorkflowStep,
|
|
normalize_trace_metadata,
|
|
)
|
|
|
|
if TYPE_CHECKING:
|
|
from app.db.models import GenerationJob
|
|
|
|
|
|
class TraceRecorder:
|
|
"""Append workflow events with standard harness trace metadata."""
|
|
|
|
def __init__(self, db: AsyncSession):
|
|
self.db = db
|
|
|
|
async def record_step(
|
|
self,
|
|
*,
|
|
job: "GenerationJob | None",
|
|
event_type: str,
|
|
status: str,
|
|
story_id: int | None = None,
|
|
message: str | None = None,
|
|
metadata: dict[str, Any] | None = None,
|
|
step: WorkflowStep | str | None = None,
|
|
artifact: ArtifactKind | str | None = None,
|
|
failure_category: FailureCategory | str | None = None,
|
|
retryable: bool | None = None,
|
|
blocks_main_result: bool | None = None,
|
|
commit: bool = True,
|
|
):
|
|
"""Append a workflow event when the caller is running under a tracked job."""
|
|
|
|
if job is None:
|
|
return None
|
|
|
|
return await record_generation_event(
|
|
self.db,
|
|
job=job,
|
|
story_id=story_id,
|
|
event_type=event_type,
|
|
status=status,
|
|
message=message,
|
|
metadata=normalize_trace_metadata(
|
|
event_type,
|
|
metadata,
|
|
step=step,
|
|
artifact=artifact,
|
|
failure_category=failure_category,
|
|
retryable=retryable,
|
|
blocks_main_result=blocks_main_result,
|
|
),
|
|
commit=commit,
|
|
)
|
|
|