pytrace/server/app/models/graph.py

from pydantic import BaseModel
from typing import List, Dict, Any, Optional

from pytrace.model.enums import NodeType, DimensionMode

class NodeSchema(BaseModel):
    """Data structure for a node in a graph request."""
    id: str
    type: str # This corresponds to the NodeSpec's 'type' field
    name: str
    params: Optional[Dict[str, Any]] = {}

class EdgeSchema(BaseModel):
    """Data structure for an edge in a graph request."""
    source_id: str
    source_port: str
    target_id: str
    target_port: str

class GraphExecuteRequest(BaseModel):
    """Request model for graph execution."""
    nodes: List[NodeSchema]
    edges: List[EdgeSchema]
    # initial_inputs could be provided to kickstart the graph
    initial_inputs: Optional[Dict[str, Any]] = {}