solar_eval/models/project.py

from typing import Any

from solar_eval.models.base import CustomBaseModel
from solar_eval.models.enums import EvaluatorType, PipelineType


class EvaluatorCategory(CustomBaseModel):
    name: str
    weight: float


class EvaluatorConfig(CustomBaseModel):
    type: EvaluatorType
    categories: list[EvaluatorCategory] | None = None
    rules: list[str] | None = None
    metrics: list[str] | None = None
    components: list[dict[str, Any]] | None = None


class PipelineStepConfig(CustomBaseModel):
    name: str
    type: str  # "llm" or "rule"
    model: str | None = None
    prompt_key: str | None = None
    config: dict[str, Any] | None = None


class PipelineConfig(CustomBaseModel):
    text_split_size: int | None = None
    steps: list[PipelineStepConfig] | None = None


class TaskConfig(CustomBaseModel):
    name: str
    dataset_path: str
    input_fields: list[str]
    golden_field: str | None = None
    golden_fields: dict[str, str] | None = None
    pipeline: PipelineType
    pipeline_config: PipelineConfig | None = None
    evaluator: EvaluatorConfig


class DatasetConfig(CustomBaseModel):
    source: str = "huggingface"
    repo: str


class ProjectConfig(CustomBaseModel):
    name: str
    display_name: str
    description: str = ""
    dataset: DatasetConfig
    tasks: list[TaskConfig]


class ProjectCreate(CustomBaseModel):
    name: str
    display_name: str
    description: str = ""
    config: dict  # Raw YAML-parsed config


class ProjectResponse(CustomBaseModel):
    id: str
    name: str
    display_name: str
    description: str
    config: dict
    created_at: str


class ProjectListResponse(CustomBaseModel):
    id: str
    name: str
    display_name: str
    description: str
    task_count: int
    created_at: str