feat: improve typing throughout the codebase

Author: rayokamoto

AGENTS.md

@@ -96,9 +96,13 @@ Keep it current when commands or conventions change.
 
 ### Typing
 - Python 3.12 syntax (`list[int]`, `dict[str, float]`, `str | None`).
-- Add return types to public functions and methods.
+- Add type hints to all function signatures (parameters and return types), not just public ones.
 - Prefer `BaseSettings` and `BaseModel` type annotations for config/DTOs.
 - `ty` runs in CI; avoid `Any` unless required and explain why in code.
+- Use `NewType` to distinguish domain identifiers that share an underlying primitive (e.g., `MinerUID`, `BlockNumber`, `EnvironmentId`, `Hotkey`). This prevents accidental misuse such as passing a `MinerUID` where a `BlockNumber` is expected.
+- Centralise shared newtypes, type aliases, enums, and `TypedDict` definitions in `kinitro/types.py`. Import from there rather than re-defining types locally.
+- When introducing a new domain concept that is fundamentally a `str`, `int`, or other primitive, create a `NewType` for it in `kinitro/types.py` and use it consistently across signatures, models, and data structures.
+- Prefer `TypedDict` or `dataclasses.dataclass` over plain `dict` for structured data with known keys.
 
 ### Naming
 - `snake_case` for functions, variables, modules.

environments/_template/env.py

@@ -75,9 +75,9 @@ async def list_environments(self) -> list[str]:
     async def evaluate(
         self,
         task_id: int,
+        base_url: str,
         seed: int | None = None,
         model: str | None = None,
-        base_url: str | None = None,
         env_id: str = "myenv/v0",  # TODO: Change default env_id
         max_timesteps: int = 500,
         action_timeout: float = 0.5,
@@ -108,9 +108,6 @@ async def evaluate(
                 error=f"Invalid env_id: {env_id}. Must start with 'myenv/'",
             )
 
-        if base_url is None:
-            raise ValueError("base_url (miner endpoint) is required")
-
         seed = seed if seed is not None else task_id
         start_time = time.time()
 

environments/metaworld/env.py

@@ -98,9 +98,9 @@ async def list_environments(self) -> list[str]:
     async def evaluate(
         self,
         task_id: int,
+        base_url: str,
         seed: int | None = None,
         model: str | None = None,
-        base_url: str | None = None,
         env_id: str = "metaworld/pick-place-v3",
         max_timesteps: int = 500,
         action_timeout: float = 0.5,
@@ -148,9 +148,6 @@ async def evaluate(
                 error=f"Invalid env_id for MetaWorld container: {env_id}. Must start with 'metaworld/'",
             )
 
-        if base_url is None:
-            raise ValueError("base_url (miner endpoint) is required")
-
         if seed is None:
             seed = task_id
 

kinitro/api/routes/health.py

@@ -17,7 +17,7 @@
 
 
 @router.get("/health", response_model=HealthResponse)
-async def health_check(session: AsyncSession = Depends(get_session)):
+async def health_check(session: AsyncSession = Depends(get_session)) -> HealthResponse:
     """Health check endpoint."""
     try:
         await session.execute(text("SELECT 1"))
@@ -35,7 +35,7 @@ async def health_check(session: AsyncSession = Depends(get_session)):
 async def get_status(
     session: AsyncSession = Depends(get_session),
     storage: Storage = Depends(get_storage),
-):
+) -> StatusResponse:
     """Get current backend status."""
     # Get current/latest cycles
     current_cycle = await storage.get_running_cycle(session)

kinitro/api/routes/miners.py

@@ -7,6 +7,7 @@
 from kinitro.backend.models import EnvironmentInfo, MinerInfo
 from kinitro.backend.storage import Storage
 from kinitro.environments import get_all_environment_ids
+from kinitro.types import EnvironmentId, EnvStatsEntry
 
 router = APIRouter(prefix="/v1", tags=["Miners & Environments"])
 
@@ -15,7 +16,7 @@
 async def list_miners(
     session: AsyncSession = Depends(get_session),
     storage: Storage = Depends(get_storage),
-):
+) -> list[MinerInfo]:
     """List all miners that have been evaluated."""
     # Get latest cycle's scores to get miner info
     cycle = await storage.get_latest_cycle(session, completed_only=True)
@@ -50,29 +51,32 @@ async def list_miners(
 async def list_environments(
     session: AsyncSession = Depends(get_session),
     storage: Storage = Depends(get_storage),
-):
+) -> list[EnvironmentInfo]:
     """List all evaluation environments."""
     env_ids = get_all_environment_ids()
 
     # Get latest cycle for stats
     cycle = await storage.get_latest_cycle(session, completed_only=True)
 
-    env_stats: dict[str, dict] = {env_id: {"count": 0, "total_sr": 0.0} for env_id in env_ids}
+    env_stats: dict[EnvironmentId, EnvStatsEntry] = {
+        EnvironmentId(env_id): EnvStatsEntry(count=0, total_sr=0.0) for env_id in env_ids
+    }
 
     if cycle:
         scores = await storage.get_scores_for_cycle(session, cycle.id)
         for s in scores:
-            if s.env_id in env_stats:
-                env_stats[s.env_id]["count"] += 1
-                env_stats[s.env_id]["total_sr"] += s.success_rate
+            eid = EnvironmentId(s.env_id)
+            if eid in env_stats:
+                env_stats[eid]["count"] += 1
+                env_stats[eid]["total_sr"] += s.success_rate
 
     result = []
     for env_id in env_ids:
         parts = env_id.split("/")
         env_name = parts[0] if parts else env_id
         task_name = parts[1] if len(parts) > 1 else ""
 
-        stats = env_stats[env_id]
+        stats = env_stats[EnvironmentId(env_id)]
         avg_sr = stats["total_sr"] / stats["count"] if stats["count"] > 0 else None
 
         result.append(

kinitro/api/routes/scores.py

@@ -12,6 +12,7 @@
     ScoresResponse,
 )
 from kinitro.backend.storage import Storage
+from kinitro.types import EnvironmentId, Hotkey, MinerUID
 
 router = APIRouter(prefix="/v1/scores", tags=["Scores"])
 
@@ -22,9 +23,9 @@ def _build_scores_response(
     """Build a ScoresResponse from a cycle ORM object and its scores."""
     scores = [
         MinerScore(
-            uid=s.uid,
-            hotkey=s.hotkey,
-            env_id=s.env_id,
+            uid=MinerUID(s.uid),
+            hotkey=Hotkey(s.hotkey),
+            env_id=EnvironmentId(s.env_id),
             success_rate=s.success_rate,
             mean_reward=s.mean_reward,
             episodes_completed=s.episodes_completed,
@@ -33,11 +34,12 @@ def _build_scores_response(
         for s in scores_orm
     ]
 
-    miner_summary: dict[int, dict[str, float]] = {}
+    miner_summary: dict[MinerUID, dict[EnvironmentId, float]] = {}
     for s in scores_orm:
-        if s.uid not in miner_summary:
-            miner_summary[s.uid] = {}
-        miner_summary[s.uid][s.env_id] = s.success_rate
+        uid = MinerUID(s.uid)
+        if uid not in miner_summary:
+            miner_summary[uid] = {}
+        miner_summary[uid][EnvironmentId(s.env_id)] = s.success_rate
 
     return ScoresResponse(
         cycle=EvaluationCycle.model_validate(cycle),
@@ -50,7 +52,7 @@ def _build_scores_response(
 async def get_latest_scores(
     session: AsyncSession = Depends(get_session),
     storage: Storage = Depends(get_storage),
-):
+) -> ScoresResponse:
     """Get scores from the most recent completed evaluation cycle."""
     cycle = await storage.get_latest_cycle(session, completed_only=True)
     if cycle is None:
@@ -65,7 +67,7 @@ async def get_scores_for_cycle(
     cycle_id: int,
     session: AsyncSession = Depends(get_session),
     storage: Storage = Depends(get_storage),
-):
+) -> ScoresResponse:
     """Get scores for a specific evaluation cycle."""
     cycle = await storage.get_cycle(session, cycle_id)
     if cycle is None:

kinitro/api/routes/tasks.py

@@ -13,6 +13,7 @@
     TaskSubmitResponse,
 )
 from kinitro.backend.storage import Storage
+from kinitro.types import EnvironmentId, Hotkey, MinerUID, Seed, TaskUUID
 
 router = APIRouter(prefix="/v1/tasks", tags=["Tasks"])
 
@@ -23,7 +24,7 @@ async def fetch_tasks(
     session: AsyncSession = Depends(get_session),
     storage: Storage = Depends(get_storage),
     _auth: None = Depends(verify_api_key),
-):
+) -> TaskFetchResponse:
     """
     Fetch tasks from the task pool.
 
@@ -46,15 +47,15 @@ async def fetch_tasks(
 
     tasks = [
         Task(
-            task_uuid=t.task_uuid,
+            task_uuid=TaskUUID(t.task_uuid),
             cycle_id=t.cycle_id,
-            miner_uid=t.miner_uid,
-            miner_hotkey=t.miner_hotkey,
+            miner_uid=MinerUID(t.miner_uid),
+            miner_hotkey=Hotkey(t.miner_hotkey),
             miner_endpoint=t.miner_endpoint,
             miner_repo=t.miner_repo,
             miner_revision=t.miner_revision,
-            env_id=t.env_id,
-            seed=t.seed,
+            env_id=EnvironmentId(t.env_id),
+            seed=Seed(t.seed),
             status=t.status,
             created_at=t.created_at,
         )
@@ -73,7 +74,7 @@ async def submit_tasks(
     session: AsyncSession = Depends(get_session),
     storage: Storage = Depends(get_storage),
     _auth: None = Depends(verify_api_key),
-):
+) -> TaskSubmitResponse:
     """
     Submit results for completed tasks.
 
@@ -117,7 +118,7 @@ async def get_task_stats(
     cycle_id: int | None = None,
     session: AsyncSession = Depends(get_session),
     storage: Storage = Depends(get_storage),
-):
+) -> TaskPoolStats:
     """
     Get task pool statistics.
 
@@ -129,14 +130,4 @@ async def get_task_stats(
         running_cycle = await storage.get_running_cycle(session)
         cycle_id = running_cycle.id if running_cycle else None
 
-    stats = await storage.get_task_pool_stats(session, cycle_id=cycle_id)
-
-    return TaskPoolStats(
-        total_tasks=stats["total_tasks"],
-        pending_tasks=stats["pending_tasks"],
-        assigned_tasks=stats["assigned_tasks"],
-        completed_tasks=stats["completed_tasks"],
-        failed_tasks=stats["failed_tasks"],
-        active_executors=stats["active_executors"],
-        current_cycle_id=stats["current_cycle_id"],
-    )
+    return await storage.get_task_pool_stats(session, cycle_id=cycle_id)

kinitro/backend/models.py

@@ -10,6 +10,8 @@
 from sqlalchemy.dialects.postgresql import JSONB
 from sqlalchemy.orm import DeclarativeBase, Mapped, mapped_column, relationship
 
+from kinitro.types import EnvironmentId, Hotkey, MinerUID, Seed, TaskUUID
+
 
 def generate_task_uuid() -> str:
     """Generate a unique task UUID."""
@@ -198,9 +200,9 @@ class HealthResponse(BaseModel):
 class MinerScore(BaseModel):
     """Score for one miner on one environment."""
 
-    uid: int
-    hotkey: str
-    env_id: str
+    uid: MinerUID
+    hotkey: Hotkey
+    env_id: EnvironmentId
     success_rate: float
     mean_reward: float
     episodes_completed: int
@@ -214,7 +216,7 @@ class EvaluationCycle(BaseModel):
     block_number: int
     started_at: datetime
     completed_at: datetime | None
-    status: str
+    status: EvaluationCycleStatus
     n_miners: int | None
     n_environments: int | None
     duration_seconds: float | None
@@ -230,7 +232,7 @@ class ScoresResponse(BaseModel):
     scores: list[MinerScore]
 
     # Aggregated by miner
-    miner_summary: dict[int, dict[str, float]] = Field(
+    miner_summary: dict[MinerUID, dict[EnvironmentId, float]] = Field(
         default_factory=dict,
         description="Aggregated scores per miner: {uid: {env_id: success_rate}}",
     )
@@ -239,7 +241,7 @@ class ScoresResponse(BaseModel):
 class WeightsU16(BaseModel):
     """Weights in u16 format for chain submission."""
 
-    uids: list[int]
+    uids: list[MinerUID]
     values: list[int]
 
 
@@ -249,7 +251,7 @@ class WeightsResponse(BaseModel):
     cycle_id: int
     block_number: int
     timestamp: datetime
-    weights: dict[int, float] = Field(description="Normalized weights: {uid: weight}")
+    weights: dict[MinerUID, float] = Field(description="Normalized weights: {uid: weight}")
     weights_u16: WeightsU16 = Field(description="Weights formatted for chain submission")
     metadata: dict[str, Any] = Field(default_factory=dict)
 
@@ -261,24 +263,24 @@ class StatusResponse(BaseModel):
     last_completed_cycle: EvaluationCycle | None
     total_cycles: int
     total_miners_evaluated: int
-    environments: list[str]
+    environments: list[EnvironmentId]
     is_evaluating: bool
 
 
 class MinerInfo(BaseModel):
     """Information about a miner."""
 
-    uid: int
-    hotkey: str
+    uid: MinerUID
+    hotkey: Hotkey
     last_evaluated_block: int | None
     avg_success_rate: float | None
-    environments_evaluated: list[str]
+    environments_evaluated: list[EnvironmentId]
 
 
 class EnvironmentInfo(BaseModel):
     """Information about an evaluation environment."""
 
-    env_id: str
+    env_id: EnvironmentId
     env_name: str
     task_name: str
     n_evaluations: int
@@ -293,16 +295,16 @@ class EnvironmentInfo(BaseModel):
 class Task(BaseModel):
     """A single evaluation task from the task pool."""
 
-    task_uuid: str  # Unique identifier for API calls
+    task_uuid: TaskUUID  # Unique identifier for API calls
     cycle_id: int
-    miner_uid: int
-    miner_hotkey: str
+    miner_uid: MinerUID
+    miner_hotkey: Hotkey
     miner_endpoint: str
     miner_repo: str | None = None  # HuggingFace repo for verification
     miner_revision: str | None = None  # HuggingFace revision for verification
-    env_id: str
-    seed: int  # Deterministic seed for reproducibility
-    status: str
+    env_id: EnvironmentId
+    seed: Seed  # Deterministic seed for reproducibility
+    status: TaskStatus
     created_at: datetime
 
     class Config:
@@ -314,7 +316,9 @@ class TaskFetchRequest(BaseModel):
 
     executor_id: str = Field(description="Unique identifier for the executor")
     batch_size: int = Field(default=10, ge=1, le=100, description="Number of tasks to fetch")
-    env_ids: list[str] | None = Field(default=None, description="Filter by environment IDs")
+    env_ids: list[EnvironmentId] | None = Field(
+        default=None, description="Filter by environment IDs"
+    )
 
 
 class TaskFetchResponse(BaseModel):
@@ -327,7 +331,7 @@ class TaskFetchResponse(BaseModel):
 class TaskResult(BaseModel):
     """Result of a single task execution."""
 
-    task_uuid: str = Field(description="UUID of the task")
+    task_uuid: TaskUUID = Field(description="UUID of the task")
     success: bool
     score: float = Field(default=0.0)
     total_reward: float = Field(default=0.0)