from pydantic import BaseModel
from typing import List
from .general import SampleIndex
from .output import AgentOutput, TaskOutput

class RegisterRequest(BaseModel):
    name: str
    address: str
    concurrency: int
    indices: list
    
    
class TaskConfigRequest(BaseModel):
    tag: str
    question: str
    assignment: dict
    task: dict


class StartSampleRequest(BaseModel):
    name: str
    index: SampleIndex


class InteractRequest(BaseModel):
    session_id: int
    agent_response: AgentOutput


class CancelRequest(BaseModel):
    session_id: int


class HeartbeatRequest(BaseModel):
    name: str
    address: str


class CalculateOverallRequest(BaseModel):
    name: str
    results: List[TaskOutput]


class WorkerStartSampleRequest(BaseModel):
    index: SampleIndex
    session_id: int


class SampleStatusRequest(BaseModel):
    session_id: int
