| from __future__ import annotations |
|
|
| from typing import Any, Literal |
|
|
| from pydantic import BaseModel, Field, model_validator |
|
|
|
|
| class GraderSpec(BaseModel): |
| type: Literal["exact", "contains", "numeric", "regex", "llm_judge"] = "exact" |
| tolerance: float = 0.0 |
| pattern: str | None = None |
| rubric: str | None = None |
| case_sensitive: bool = False |
|
|
|
|
| class TaskRecord(BaseModel): |
| task_id: str |
| prompt: str |
| domain: str = "general" |
| reference_answer: str | None = None |
| grader: GraderSpec = Field(default_factory=GraderSpec) |
| split: Literal["train", "validation", "test"] | None = None |
| tags: list[str] = Field(default_factory=list) |
| metadata: dict[str, Any] = Field(default_factory=dict) |
|
|
|
|
| class WorkerResult(BaseModel): |
| worker_id: str |
| requested_model: str |
| served_model: str | None = None |
| response: str |
| quality: float |
| cost_usd: float | None = None |
| latency_ms: float |
| prompt_tokens: int | None = None |
| completion_tokens: int | None = None |
| generation_id: str | None = None |
| error: str | None = None |
| utility: float = 0.0 |
|
|
|
|
| class RewardRecord(BaseModel): |
| task_id: str |
| prompt: str |
| domain: str |
| split: Literal["train", "validation", "test"] | None = None |
| tags: list[str] = Field(default_factory=list) |
| metadata: dict[str, Any] = Field(default_factory=dict) |
| repetitions: int = Field(default=1, ge=1) |
| worker_ids: list[str] |
| rewards: list[float] |
| results: list[WorkerResult] = Field(default_factory=list) |
|
|
| @model_validator(mode="after") |
| def validate_vector(self) -> RewardRecord: |
| if not self.worker_ids: |
| raise ValueError("worker_ids cannot be empty") |
| if len(self.worker_ids) != len(self.rewards): |
| raise ValueError("worker_ids and rewards must have equal length") |
| if len(set(self.worker_ids)) != len(self.worker_ids): |
| raise ValueError("worker_ids must be unique") |
| return self |
|
|
|
|
| def route_text( |
| prompt: str, |
| domain: str = "general", |
| tags: list[str] | None = None, |
| worker_descriptions: dict[str, str] | None = None, |
| ) -> str: |
| """Create the router input. Do not include answers or worker outcomes.""" |
| tag_text = ", ".join(tags or []) or "none" |
| experts = "" |
| if worker_descriptions: |
| experts = "\nAvailable workers:\n" + "\n".join( |
| f"- {worker_id}: {description}" |
| for worker_id, description in worker_descriptions.items() |
| ) |
| return ( |
| "Select the best worker for this task.\n" |
| f"Domain: {domain}\n" |
| f"Tags: {tag_text}{experts}\n" |
| f"Task:\n{prompt}" |
| ) |
|
|
|
|
