state_schema.py

"""
State Schema for UI Regression Testing Workflow
Defines the structure of data passed between agents.
"""

from typing import TypedDict, Dict, List, Any, Optional


class WorkflowState(TypedDict, total=False):
    """Complete state for the UI regression testing workflow."""
    
    # Configuration
    figma_file_key: str
    figma_access_token: str
    website_url: str
    execution_id: str
    hf_token: Optional[str]
    
    # Screenshot paths
    figma_screenshots: Dict[str, str]  # {"desktop": "path/to/file.png", "mobile": "..."}
    website_screenshots: Dict[str, str]
    
    # Screenshot metadata (for proper comparison)
    figma_dimensions: Dict[str, Dict[str, int]]  # {"desktop": {"width": 1440, "height": 1649}}
    website_dimensions: Dict[str, Dict[str, int]]
    
    # Comparison results
    comparison_images: Dict[str, str]  # {"desktop": "path/to/diff.png"}
    visual_differences: List[Dict[str, Any]]
    similarity_scores: Dict[str, float]  # {"desktop": 85.5, "mobile": 78.2}
    overall_score: float
    
    # Workflow control
    user_approval: bool
    status: str
    error_message: Optional[str]
    
    # Logs
    logs: List[str]


def create_initial_state(
    figma_file_key: str,
    figma_access_token: str,
    website_url: str,
    execution_id: str,
    hf_token: str = ""
) -> WorkflowState:
    """Create initial state for a new workflow run."""
    return WorkflowState(
        figma_file_key=figma_file_key,
        figma_access_token=figma_access_token,
        website_url=website_url,
        execution_id=execution_id,
        hf_token=hf_token,
        figma_screenshots={},
        website_screenshots={},
        figma_dimensions={},
        website_dimensions={},
        comparison_images={},
        visual_differences=[],
        similarity_scores={},
        overall_score=0.0,
        user_approval=False,
        status="initialized",
        error_message=None,
        logs=[]
    )