Refactor schemas: split Domain/Experiment into separate files, rename I/O functions, use dataclass objects for hierarchical relationships, and improve documentation

Author: kohankhaki

src/schemas/PIPELINE_SCHEMAS.md

@@ -25,7 +25,7 @@ Each stage follows a consistent pattern:
 
 **Important:** All stage implementations must follow this pattern to ensure the pipeline is clean, consistent, and maintainable. This enables interoperability between different implementations, resumability of failed runs, and clear traceability through the pipeline.
 
-**Note:** The dataclasses, save functions (`save_<stage>_output(data, metadata, output_path)`), and load functions (`load_<stage>_output(file_path) -> <OutputDataclass>`) for each stage will be provided and must be used. Do not implement custom serialization or data structures - use the standardized schemas to ensure consistency across the pipeline. Dataclasses provide type safety, validation, and clear structure. JSON is the serialization format.
+**Note:** The dataclasses, save functions (`save_<stage>(data, metadata, output_path)`), and load functions (`load_<stage>(file_path) -> <OutputDataclass>`) for each stage will be provided and must be used. Do not implement custom serialization or data structures - use the standardized schemas to ensure consistency across the pipeline. Dataclasses provide type safety, validation, and clear structure. JSON is the serialization format.
 
 **Iteration Note:** Some stages operate on subsets (one area, capability, or task at a time) and require an outer orchestrator/loop script to iterate over all items:
 - **Stage 2 (Capability Generation)**: Operates on one area at a time - orchestrator loops over all areas from Stage 1
@@ -342,7 +342,7 @@ This stage creates two files:
 #### Output 1: `experiment.json`
 
 **Stage Output:** Experiment dataclass + PipelineMetadata
-**Save Function:** `save_experiment_output(experiment: Experiment, metadata: PipelineMetadata, output_path: Path)`
+**Save Function:** `save_experiment(experiment: Experiment, metadata: PipelineMetadata, output_path: Path)`
 
 **File Path:** `<output_dir>/<experiment_id>/experiment.json`
 
@@ -373,7 +373,7 @@ This stage creates two files:
 #### Output 2: `domain.json`
 
 **Stage Output:** Domain dataclass object + PipelineMetadata
-**Save Function:** `save_domain_output(domain: Domain, metadata: PipelineMetadata, output_path: Path)`
+**Save Function:** `save_domain(domain: Domain, metadata: PipelineMetadata, output_path: Path)`
 
 **File Path:** `<output_dir>/<experiment_id>/domain/domain.json`
 
@@ -411,7 +411,7 @@ This stage creates two files:
 ### Output: `areas.json`
 
 **Stage Output:** List[Area] dataclasses + PipelineMetadata
-**Save Function:** `save_areas_output(areas: List[Area], metadata: PipelineMetadata, output_path: Path)`
+**Save Function:** `save_areas(areas: List[Area], metadata: PipelineMetadata, output_path: Path)`
 
 **File Path:** `<output_dir>/<experiment_id>/areas/<tag>/areas.json`
 ```json
@@ -456,7 +456,7 @@ This stage creates two files:
 ### Output: `capabilities.json` (one per area)
 
 **Stage Output:** List[Capability] dataclasses + PipelineMetadata
-**Save Function:** `save_capabilities_output(capabilities: List[Capability], metadata: PipelineMetadata, output_path: Path)`
+**Save Function:** `save_capabilities(capabilities: List[Capability], metadata: PipelineMetadata, output_path: Path)`
 
 **File Path:** `<output_dir>/<experiment_id>/capabilities/<cap_tag>/<area_id>/capabilities.json`
 
@@ -504,7 +504,7 @@ This stage creates two files:
 ### Output: `tasks.json` (one per capability)
 
 **Stage Output:** List[Task] dataclasses + PipelineMetadata
-**Save Function:** `save_tasks_output(tasks: List[Task], metadata: PipelineMetadata, output_path: Path)`
+**Save Function:** `save_tasks(tasks: List[Task], metadata: PipelineMetadata, output_path: Path)`
 
 **File Path:** `<output_dir>/<experiment_id>/tasks/<task_tag>/<area_id>/<capability_id>/tasks.json`
 
@@ -563,7 +563,7 @@ This stage creates two files:
 ### Output: `solution.json` (one per task)
 
 **Stage Output:** TaskSolution dataclass + PipelineMetadata
-**Save Function:** `save_solution_output(task_solution: TaskSolution, metadata: PipelineMetadata, output_path: Path)`
+**Save Function:** `save_solution(task_solution: TaskSolution, metadata: PipelineMetadata, output_path: Path)`
 
 **File Path:** `<output_dir>/<experiment_id>/solutions/<solution_tag>/<area_id>/<capability_id>/<task_id>/solution.json`
 
@@ -634,7 +634,7 @@ This stage creates two files:
 ### Output: `validation.json` (one per task)
 
 **Stage Output:** ValidationResult dataclass + PipelineMetadata
-**Save Function:** `save_validation_output(validation_result: ValidationResult, metadata: PipelineMetadata, output_path: Path)`
+**Save Function:** `save_validation(validation_result: ValidationResult, metadata: PipelineMetadata, output_path: Path)`
 
 **File Path:** `<output_dir>/<experiment_id>/validation/<validation_tag>/<area_id>/<capability_id>/<task_id>/validation.json`
 

src/schemas/README.md

@@ -6,7 +6,8 @@ This directory contains standardized schemas for all ACE pipeline stages, ensuri
 
 - **[`PIPELINE_SCHEMAS.md`](PIPELINE_SCHEMAS.md)** - Complete documentation of input/output formats for each stage
 - **Python Dataclasses** - Type-safe data structures for each stage:
-  - [`experiment_schemas.py`](experiment_schemas.py) - Experiment and Domain (Stage 0)
+  - [`experiment_schemas.py`](experiment_schemas.py) - Experiment (Stage 0)
+  - [`domain_schemas.py`](domain_schemas.py) - Domain (Stage 0)
   - [`metadata_schemas.py`](metadata_schemas.py) - Common metadata (PipelineMetadata)
   - [`area_schemas.py`](area_schemas.py) - Area generation (Stage 1)
   - [`capability_schemas.py`](capability_schemas.py) - Capability generation (Stage 2)
@@ -22,8 +23,8 @@ This directory contains standardized schemas for all ACE pipeline stages, ensuri
 
 ```python
 from src.schemas import (
-    Experiment,
     Domain,
+    Experiment,
     PipelineMetadata,
     Area,
     Capability,
@@ -33,12 +34,12 @@ from src.schemas import (
 )
 
 # Create area
+domain = Domain(name="Personal Finance", domain_id="domain_000")
 area = Area(
     name="Cash Flow & Budget Management",
     area_id="area_000",
     description="Design and monitor budgets...",
-    domain="personal finance",
-    domain_id="domain_000",
+    domain=domain,
     # generation_metadata is optional
 )
 
@@ -54,8 +55,8 @@ area = Area.from_dict(data)
 ```python
 from pathlib import Path
 from src.schemas import (
-    save_areas_output,
-    load_areas_output,
+    save_areas,
+    load_areas,
     PipelineMetadata,
     Area,
 )
@@ -68,10 +69,10 @@ metadata = PipelineMetadata(
     timestamp="2025-11-06T12:00:00Z",
     output_stage_tag="_20251009_122040"
 )
-save_areas_output(areas, metadata, Path("output/areas.json"))
+save_areas(areas, metadata, Path("output/areas.json"))
 
 # Load areas
-areas, metadata = load_areas_output(Path("output/areas.json"))
+areas, metadata = load_areas(Path("output/areas.json"))
 ```
 
 ## Pipeline Stages

src/schemas/__init__.py

@@ -6,22 +6,23 @@
 
 from src.schemas.area_schemas import Area
 from src.schemas.capability_schemas import Capability
-from src.schemas.experiment_schemas import Domain, Experiment
+from src.schemas.domain_schemas import Domain
+from src.schemas.experiment_schemas import Experiment
 from src.schemas.io_utils import (
-    load_areas_output,
-    load_capabilities_output,
-    load_domain_output,
-    load_experiment_output,
-    load_solution_output,
-    load_tasks_output,
-    load_validation_output,
-    save_areas_output,
-    save_capabilities_output,
-    save_domain_output,
-    save_experiment_output,
-    save_solution_output,
-    save_tasks_output,
-    save_validation_output,
+    load_areas,
+    load_capabilities,
+    load_domain,
+    load_experiment,
+    load_solution,
+    load_tasks,
+    load_validation,
+    save_areas,
+    save_capabilities,
+    save_domain,
+    save_experiment,
+    save_solution,
+    save_tasks,
+    save_validation,
 )
 from src.schemas.metadata_schemas import PipelineMetadata
 from src.schemas.solution_schemas import TaskSolution
@@ -46,19 +47,19 @@
     # Validation schemas
     "ValidationResult",
     # I/O functions - Save
-    "save_experiment_output",
-    "save_domain_output",
-    "save_areas_output",
-    "save_capabilities_output",
-    "save_tasks_output",
-    "save_solution_output",
-    "save_validation_output",
+    "save_experiment",
+    "save_domain",
+    "save_areas",
+    "save_capabilities",
+    "save_tasks",
+    "save_solution",
+    "save_validation",
     # I/O functions - Load
-    "load_experiment_output",
-    "load_domain_output",
-    "load_areas_output",
-    "load_capabilities_output",
-    "load_tasks_output",
-    "load_solution_output",
-    "load_validation_output",
+    "load_experiment",
+    "load_domain",
+    "load_areas",
+    "load_capabilities",
+    "load_tasks",
+    "load_solution",
+    "load_validation",
 ]

src/schemas/area_schemas.py

@@ -1,8 +1,14 @@
-"""Schemas for area generation stage."""
+"""Schemas for area generation stage (Stage 1).
+
+Defines Area dataclass representing a domain area. Areas are high-level categories
+within a domain (e.g., "Budgeting" within "Personal Finance").
+"""
 
 from dataclasses import dataclass, field
 from typing import Dict, Optional
 
+from src.schemas.domain_schemas import Domain
+
 
 @dataclass
 class Area:
@@ -11,18 +17,18 @@ class Area:
     name: str
     area_id: str
     description: Optional[str] = None
-    domain: str = ""
-    domain_id: str = ""
+    domain: Optional[Domain] = None
     generation_metadata: Optional[Dict] = field(default_factory=dict)
 
     def to_dict(self):
         """Convert to dictionary."""
         result = {
             "name": self.name,
             "area_id": self.area_id,
-            "domain": self.domain,
-            "domain_id": self.domain_id,
         }
+        if self.domain is not None:
+            result["domain"] = self.domain.name
+            result["domain_id"] = self.domain.domain_id
         if self.description is not None:
             result["description"] = self.description
         if self.generation_metadata:
@@ -32,11 +38,17 @@ def to_dict(self):
     @classmethod
     def from_dict(cls, data: dict):
         """Create from dictionary."""
+        domain = None
+        if "domain" in data and "domain_id" in data:
+            domain = Domain(
+                name=data["domain"],
+                domain_id=data["domain_id"],
+                description=None,
+            )
         return cls(
             name=data["name"],
             area_id=data["area_id"],
             description=data.get("description"),
-            domain=data.get("domain", ""),
-            domain_id=data.get("domain_id", ""),
+            domain=domain,
             generation_metadata=data.get("generation_metadata", {}),
         )

src/schemas/capability_schemas.py

@@ -1,8 +1,15 @@
-"""Schemas for capability generation stage."""
+"""Schemas for capability generation stage (Stage 2).
+
+Defines Capability dataclass representing a capability within an area. Capabilities
+are specific skills or abilities (e.g., "Budget Creation" within "Budgeting" area).
+"""
 
 from dataclasses import dataclass, field
 from typing import Dict, Optional
 
+from src.schemas.area_schemas import Area
+from src.schemas.domain_schemas import Domain
+
 
 @dataclass
 class Capability:
@@ -11,22 +18,21 @@ class Capability:
     name: str
     capability_id: str
     description: Optional[str] = None
-    area: str = ""
-    area_id: str = ""
-    domain: str = ""
-    domain_id: str = ""
+    area: Optional[Area] = None
     generation_metadata: Optional[Dict] = field(default_factory=dict)
 
     def to_dict(self):
         """Convert to dictionary."""
         result = {
             "name": self.name,
             "capability_id": self.capability_id,
-            "area": self.area,
-            "area_id": self.area_id,
-            "domain": self.domain,
-            "domain_id": self.domain_id,
         }
+        if self.area is not None:
+            result["area"] = self.area.name
+            result["area_id"] = self.area.area_id
+            if self.area.domain is not None:
+                result["domain"] = self.area.domain.name
+                result["domain_id"] = self.area.domain.domain_id
         if self.description is not None:
             result["description"] = self.description
         if self.generation_metadata:
@@ -36,13 +42,25 @@ def to_dict(self):
     @classmethod
     def from_dict(cls, data: dict):
         """Create from dictionary."""
+        area = None
+        if "area" in data and "area_id" in data:
+            domain = None
+            if "domain" in data and "domain_id" in data:
+                domain = Domain(
+                    name=data["domain"],
+                    domain_id=data["domain_id"],
+                    description=None,
+                )
+            area = Area(
+                name=data["area"],
+                area_id=data["area_id"],
+                description=None,
+                domain=domain,
+            )
         return cls(
             name=data["name"],
             capability_id=data["capability_id"],
             description=data.get("description"),
-            area=data.get("area", ""),
-            area_id=data.get("area_id", ""),
-            domain=data.get("domain", ""),
-            domain_id=data.get("domain_id", ""),
+            area=area,
             generation_metadata=data.get("generation_metadata", {}),
         )

src/schemas/domain_schemas.py

@@ -0,0 +1,35 @@
+"""Schemas for domain (Stage 0).
+
+Defines Domain dataclass representing the domain being evaluated in the experiment.
+"""
+
+from dataclasses import dataclass
+from typing import Optional
+
+
+@dataclass
+class Domain:
+    """Represents a domain."""
+
+    name: str
+    domain_id: str
+    description: Optional[str] = None
+
+    def to_dict(self):
+        """Convert to dictionary."""
+        result = {
+            "name": self.name,
+            "domain_id": self.domain_id,
+        }
+        if self.description is not None:
+            result["description"] = self.description
+        return result
+
+    @classmethod
+    def from_dict(cls, data: dict):
+        """Create from dictionary."""
+        return cls(
+            name=data["name"],
+            domain_id=data["domain_id"],
+            description=data.get("description"),
+        )