runner_service/models.py

from __future__ import annotations

from enum import Enum
from typing import Any, Literal

from pydantic import BaseModel, Field


class TaskStatus(str, Enum):
    running = "running"
    success = "success"
    failed = "failed"
    needs_human = "needs_human"


class Citation(BaseModel):
    url: str
    title: str | None = None


class Screenshot(BaseModel):
    url: str
    ts: str


class TaskResultPayload(BaseModel):
    answer_text: str
    citations: list[Citation] = Field(default_factory=list)
    screenshots: list[Screenshot] = Field(default_factory=list)
    raw_html_url: str | None = None
    log_url: str | None = None


class RunnerTaskResult(BaseModel):
    task_id: str
    status: TaskStatus
    reason: str | None = None
    handoff_url: str | None = None
    result: TaskResultPayload | None = None


class TaskCreate(BaseModel):
    task_id: str
    site: str
    question_id: str
    question_text: str
    profile_id: str
    locale: str | None = None


class TaskCreateRequest(BaseModel):
    run_id: str
    tasks: list[TaskCreate]


class TaskCreateResponse(BaseModel):
    accepted: list[dict[str, str]]
    rejected: list[dict[str, str]]


class ResumeRequest(BaseModel):
    resume_token: str


class ResumeResponse(BaseModel):
    ok: Literal[True] = True