105 lines
3.3 KiB
Python
105 lines
3.3 KiB
Python
import logging
|
|
import time
|
|
import uuid
|
|
|
|
from yuxi.channel.cron.types import (
|
|
CronAgentExecutionPhase,
|
|
CronJob,
|
|
CronPhaseRecord,
|
|
CronRunDiagnostic,
|
|
CronRunOutcome,
|
|
)
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
def create_run_id(job_id: str, started_at_ms: int) -> str:
|
|
return f"cron:{job_id}:{started_at_ms}:{uuid.uuid4().hex}"
|
|
|
|
|
|
def create_diagnostic(job: CronJob, started_at_ms: int) -> CronRunDiagnostic:
|
|
return CronRunDiagnostic(
|
|
run_id=create_run_id(job.id, started_at_ms),
|
|
job_id=job.id,
|
|
started_at_ms=started_at_ms,
|
|
)
|
|
|
|
|
|
def record_phase(
|
|
diagnostic: CronRunDiagnostic,
|
|
phase: CronAgentExecutionPhase,
|
|
) -> CronPhaseRecord:
|
|
now_ms = int(time.time() * 1000)
|
|
record = CronPhaseRecord(
|
|
phase=phase,
|
|
started_at_ms=now_ms,
|
|
)
|
|
diagnostic.phases.append(record)
|
|
return record
|
|
|
|
|
|
def finish_phase(record: CronPhaseRecord, error: str | None = None) -> None:
|
|
record.finished_at_ms = int(time.time() * 1000)
|
|
if error:
|
|
record.error = error
|
|
|
|
|
|
def finalize_diagnostic(
|
|
diagnostic: CronRunDiagnostic,
|
|
outcome: CronRunOutcome,
|
|
error: str | None = None,
|
|
token_usage: dict | None = None,
|
|
started_mono: float = 0.0,
|
|
) -> None:
|
|
diagnostic.finished_at_ms = int(time.time() * 1000)
|
|
diagnostic.total_duration_ms = _mono_duration_ms(started_mono) if started_mono > 0 else diagnostic.finished_at_ms - diagnostic.started_at_ms
|
|
diagnostic.outcome = outcome
|
|
if error:
|
|
diagnostic.error = error
|
|
if token_usage:
|
|
diagnostic.token_usage = token_usage
|
|
|
|
|
|
def _mono_duration_ms(started_mono: float) -> int:
|
|
return int((time.monotonic() - started_mono) * 1000)
|
|
|
|
|
|
class PhaseTracker:
|
|
def __init__(self, job: CronJob):
|
|
self.job = job
|
|
started_at = int(time.time() * 1000)
|
|
self._started_mono = time.monotonic()
|
|
self.diagnostic = create_diagnostic(job, started_at)
|
|
self._init_phase = record_phase(self.diagnostic, CronAgentExecutionPhase.INITIALIZING)
|
|
finish_phase(self._init_phase)
|
|
|
|
def start_planning(self) -> CronPhaseRecord:
|
|
return record_phase(self.diagnostic, CronAgentExecutionPhase.PLANNING)
|
|
|
|
def finish_planning(self, record: CronPhaseRecord, error: str | None = None) -> None:
|
|
finish_phase(record, error)
|
|
|
|
def start_executing(self) -> CronPhaseRecord:
|
|
return record_phase(self.diagnostic, CronAgentExecutionPhase.EXECUTING)
|
|
|
|
def finish_executing(self, record: CronPhaseRecord, error: str | None = None) -> None:
|
|
finish_phase(record, error)
|
|
|
|
def start_summarizing(self) -> CronPhaseRecord:
|
|
return record_phase(self.diagnostic, CronAgentExecutionPhase.SUMMARIZING)
|
|
|
|
def finish_summarizing(self, record: CronPhaseRecord, error: str | None = None) -> None:
|
|
finish_phase(record, error)
|
|
|
|
def start_delivering(self) -> CronPhaseRecord:
|
|
return record_phase(self.diagnostic, CronAgentExecutionPhase.DELIVERING)
|
|
|
|
def finish_delivering(self, record: CronPhaseRecord, error: str | None = None) -> None:
|
|
finish_phase(record, error)
|
|
|
|
def close(
|
|
self, outcome: CronRunOutcome, error: str | None = None, token_usage: dict | None = None
|
|
) -> CronRunDiagnostic:
|
|
finalize_diagnostic(self.diagnostic, outcome, error, token_usage, self._started_mono)
|
|
return self.diagnostic
|