Fix/mypy schema types

src/schemas/area_schemas.py

@@ -4,8 +4,10 @@
 within a domain.
 """
 
+from __future__ import annotations
+
 from dataclasses import dataclass, field
-from typing import Dict, Optional
+from typing import Any, Dict, Optional
 
 from src.schemas.domain_schemas import Domain
 
@@ -18,11 +20,11 @@ class Area:
     area_id: str
     domain: Domain
     area_description: str
-    generation_metadata: Optional[Dict] = field(default_factory=dict)
+    generation_metadata: Optional[Dict[str, Any]] = field(default_factory=dict)
 
-    def to_dict(self):
+    def to_dict(self) -> Dict[str, Any]:
         """Convert to dictionary."""
-        result = {
+        result: Dict[str, Any] = {
             "area_name": self.area_name,
             "area_id": self.area_id,
             "area_description": self.area_description,
@@ -33,7 +35,7 @@ def to_dict(self):
         return result
 
     @classmethod
-    def from_dict(cls, data: dict):
+    def from_dict(cls, data: Dict[str, Any]) -> Area:
         """Create from dictionary."""
         domain = Domain.from_dict(data)
         return cls(

src/schemas/capability_schemas.py

@@ -4,8 +4,10 @@
 are specific skills or abilities.
 """
 
+from __future__ import annotations
+
 from dataclasses import dataclass, field
-from typing import Dict, Optional
+from typing import Any, Dict, Optional
 
 from src.schemas.area_schemas import Area
 
@@ -18,11 +20,11 @@ class Capability:
     capability_id: str
     area: Area
     capability_description: str
-    generation_metadata: Optional[Dict] = field(default_factory=dict)
+    generation_metadata: Optional[Dict[str, Any]] = field(default_factory=dict)
 
-    def to_dict(self):
+    def to_dict(self) -> Dict[str, Any]:
         """Convert to dictionary."""
-        result = {
+        result: Dict[str, Any] = {
             "capability_name": self.capability_name,
             "capability_id": self.capability_id,
             "capability_description": self.capability_description,
@@ -33,7 +35,7 @@ def to_dict(self):
         return result
 
     @classmethod
-    def from_dict(cls, data: dict):
+    def from_dict(cls, data: Dict[str, Any]) -> Capability:
         """Create from dictionary."""
         area = Area.from_dict(data)
         return cls(

src/schemas/domain_schemas.py

@@ -3,8 +3,10 @@
 Defines Domain dataclass for domain.
 """
 
+from __future__ import annotations
+
 from dataclasses import dataclass
-from typing import Optional
+from typing import Any, Dict, Optional
 
 
 @dataclass
@@ -15,9 +17,9 @@ class Domain:
     domain_id: str
     domain_description: Optional[str] = None
 
-    def to_dict(self):
+    def to_dict(self) -> Dict[str, Any]:
         """Convert to dictionary."""
-        result = {
+        result: Dict[str, Any] = {
             "domain_name": self.domain_name,
             "domain_id": self.domain_id,
         }
@@ -26,7 +28,7 @@ def to_dict(self):
         return result
 
     @classmethod
-    def from_dict(cls, data: dict):
+    def from_dict(cls, data: Dict[str, Any]) -> Domain:
         """Create from dictionary."""
         return cls(
             domain_name=data["domain_name"],

src/schemas/experiment_schemas.py

@@ -3,6 +3,8 @@
 Defines Experiment dataclass containing experiment configuration and metadata.
 """
 
+from __future__ import annotations
+
 from dataclasses import dataclass, field
 from typing import Any, Dict, Optional
 
@@ -17,9 +19,9 @@ class Experiment:
     pipeline_type: Optional[str] = None
     configuration: Dict[str, Any] = field(default_factory=dict)
 
-    def to_dict(self):
+    def to_dict(self) -> Dict[str, Any]:
         """Convert to dictionary."""
-        result = {
+        result: Dict[str, Any] = {
             "experiment_id": self.experiment_id,
             "domain": self.domain,
             "domain_id": self.domain_id,
@@ -30,7 +32,7 @@ def to_dict(self):
         return result
 
     @classmethod
-    def from_dict(cls, data: dict):
+    def from_dict(cls, data: Dict[str, Any]) -> Experiment:
         """Create from dictionary."""
         return cls(
             experiment_id=data["experiment_id"],

src/schemas/metadata_schemas.py

@@ -9,9 +9,11 @@
 in the data objects themselves).
 """
 
+from __future__ import annotations
+
 from dataclasses import dataclass
 from datetime import datetime
-from typing import Optional
+from typing import Any, Dict, Optional
 
 
 @dataclass
@@ -42,17 +44,17 @@ class PipelineMetadata:
     output_stage_tag: Optional[str] = None
     resume: bool = False
 
-    def __post_init__(self):
+    def __post_init__(self) -> None:
         """Set default timestamp if not provided.
 
         Automatically generates a UTC timestamp in ISO 8601 format if not set.
         """
         if not self.timestamp:
             self.timestamp = datetime.utcnow().isoformat() + "Z"
 
-    def to_dict(self):
+    def to_dict(self) -> Dict[str, Any]:
         """Convert metadata to dictionary for JSON serialization."""
-        result = {
+        result: Dict[str, Any] = {
             "experiment_id": self.experiment_id,
             "output_base_dir": self.output_base_dir,
             "timestamp": self.timestamp,
@@ -65,7 +67,7 @@ def to_dict(self):
         return result
 
     @classmethod
-    def from_dict(cls, data: dict):
+    def from_dict(cls, data: Dict[str, Any]) -> PipelineMetadata:
         """Create PipelineMetadata from dictionary (e.g., loaded from JSON)."""
         return cls(
             experiment_id=data["experiment_id"],

src/schemas/solution_schemas.py

@@ -4,8 +4,10 @@
 reasoning, and optional numerical answer.
 """
 
+from __future__ import annotations
+
 from dataclasses import dataclass, field
-from typing import Dict, Optional
+from typing import Any, Dict, Optional
 
 from src.schemas.task_schemas import Task
 
@@ -18,7 +20,7 @@ class TaskSolution:
     solution: str
     reasoning: str
     numerical_answer: Optional[str] = None
-    generation_metadata: Optional[Dict] = field(default_factory=dict)
+    generation_metadata: Optional[Dict[str, Any]] = field(default_factory=dict)
 
     @property
     def task_id(self) -> str:
@@ -30,12 +32,12 @@ def task_statement(self) -> str:
         """Get task statement from the task object for convenience."""
         return self.task.task_statement
 
-    def to_dict(self):
+    def to_dict(self) -> Dict[str, Any]:
         """Convert to dictionary.
 
         Flattens the task object fields into the result for JSON serialization.
         """
-        result = self.task.to_dict()
+        result: Dict[str, Any] = self.task.to_dict()
         result["solution"] = self.solution
         result["reasoning"] = self.reasoning
         if self.numerical_answer is not None:
@@ -45,7 +47,7 @@ def to_dict(self):
         return result
 
     @classmethod
-    def from_dict(cls, data: dict):
+    def from_dict(cls, data: Dict[str, Any]) -> TaskSolution:
         """Create from dictionary."""
         task = Task.from_dict(data)
         return cls(

src/schemas/task_schemas.py

@@ -4,8 +4,10 @@
 that test a capability.
 """
 
+from __future__ import annotations
+
 from dataclasses import dataclass, field
-from typing import Dict, List, Optional
+from typing import Any, Dict, List, Optional
 
 from src.schemas.capability_schemas import Capability
 
@@ -24,11 +26,11 @@ class Task:
     choices: Optional[List[Dict[str, str]]] = (
         None  # [{"label": "A", "solution": "..."}]
     )
-    generation_metadata: Optional[Dict] = field(default_factory=dict)
+    generation_metadata: Optional[Dict[str, Any]] = field(default_factory=dict)
 
-    def to_dict(self):
+    def to_dict(self) -> Dict[str, Any]:
         """Convert to dictionary."""
-        result = {
+        result: Dict[str, Any] = {
             "task_id": self.task_id,
             "task_statement": self.task_statement,
             "task_type": self.task_type,
@@ -43,7 +45,7 @@ def to_dict(self):
         return result
 
     @classmethod
-    def from_dict(cls, data: dict):
+    def from_dict(cls, data: Dict[str, Any]) -> Task:
         """Create from dictionary."""
         capability = Capability.from_dict(data)
         return cls(

src/schemas/validation_schemas.py

@@ -4,8 +4,10 @@
 verification status, feedback, and optional score.
 """
 
+from __future__ import annotations
+
 from dataclasses import dataclass, field
-from typing import Dict, Optional
+from typing import Any, Dict, Optional
 
 from src.schemas.solution_schemas import TaskSolution
 
@@ -18,7 +20,7 @@ class ValidationResult:
     verification: bool
     feedback: str
     score: Optional[float] = None
-    generation_metadata: Optional[Dict] = field(default_factory=dict)
+    generation_metadata: Optional[Dict[str, Any]] = field(default_factory=dict)
 
     @property
     def task_id(self) -> str:
@@ -30,12 +32,12 @@ def task_statement(self) -> str:
         """Get task statement from the task_solution for convenience."""
         return self.task_solution.task_statement
 
-    def to_dict(self):
+    def to_dict(self) -> Dict[str, Any]:
         """Convert to dictionary.
 
         Flattens the task_solution fields into the result for JSON serialization.
         """
-        result = self.task_solution.to_dict()
+        result: Dict[str, Any] = self.task_solution.to_dict()
         result["verification"] = self.verification
         result["feedback"] = self.feedback
         if self.score is not None:
@@ -45,7 +47,7 @@ def to_dict(self):
         return result
 
     @classmethod
-    def from_dict(cls, data: dict):
+    def from_dict(cls, data: Dict[str, Any]) -> ValidationResult:
         """Create from dictionary."""
         task_solution = TaskSolution.from_dict(data)
         return cls(