state_schema.py

"""
State Schema for LangGraph UI Regression Testing Workflow
Updated to be fully serializable for LangGraph persistence.
"""

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

# Use TypedDict for all state components to ensure serializability
class ViewportDict(TypedDict):
    name: str
    width: int
    height: int

class VisualDifferenceDict(TypedDict):
    category: str
    severity: str
    title: str
    description: str
    viewport: str
    location: Optional[Dict[str, int]]

class VisualDifference:
    """Class representing a visual difference between design and implementation."""
    def __init__(
        self,
        category: str,
        severity: str,
        issue_id: str,
        title: str,
        description: str,
        design_value: str,
        website_value: str,
        viewport: str,
        confidence: float = 1.0,
        detection_method: str = "manual",
        location: Optional[Dict[str, int]] = None
    ):
        self.category = category
        self.severity = severity
        self.issue_id = issue_id
        self.title = title
        self.description = description
        self.design_value = design_value
        self.website_value = website_value
        self.viewport = viewport
        self.confidence = confidence
        self.detection_method = detection_method
        self.location = location

    def to_dict(self) -> VisualDifferenceDict:
        """Convert to dictionary for serialization."""
        return {
            "category": self.category,
            "severity": self.severity,
            "title": self.title,
            "description": self.description,
            "viewport": self.viewport,
            "location": self.location
        }

class WorkflowState(TypedDict):
    # Configuration
    figma_file_key: str
    figma_access_token: str
    website_url: str
    hf_token: str
    execution_id: str
    
    # Viewports and Categories
    viewports: List[ViewportDict]
    test_plan: Dict[str, Any]
    
    # Agent Results
    figma_screenshots: Dict[str, str]
    website_screenshots: Dict[str, str]
    visual_differences: Annotated[List[VisualDifferenceDict], operator.add]
    
    # Analysis Results
    hf_analysis: Optional[Dict[str, Any]]
    similarity_score: float
    
    # Status and Metadata
    status: str
    error_message: str
    user_approval: bool

def create_initial_state(
    figma_file_key: str,
    figma_access_token: str,
    website_url: str,
    hf_token: str,
    execution_id: str
) -> WorkflowState:
    """Create initial workflow state."""
    return {
        "figma_file_key": figma_file_key,
        "figma_access_token": figma_access_token,
        "website_url": website_url,
        "hf_token": hf_token,
        "execution_id": execution_id,
        "viewports": [],
        "test_plan": {},
        "figma_screenshots": {},
        "website_screenshots": {},
        "visual_differences": [],
        "hf_analysis": None,
        "similarity_score": 0.0,
        "status": "initialized",
        "error_message": "",
        "user_approval": False
    }