models/workflow.py

"""
Workflow data models
"""
from pydantic import BaseModel, Field
from typing import Optional, Any
from enum import Enum
import uuid


class NodeType(str, Enum):
    """Types of nodes available in Tritan"""
    TRIGGER = "trigger"          # Start of workflow
    ACTION = "action"            # Generic action
    CONDITION = "condition"      # If/else branching
    LOOP = "loop"               # Iterate over data
    LLM = "llm"                 # AI/LLM node
    HTTP = "http"               # HTTP request
    CODE = "code"               # Custom code execution
    TRANSFORM = "transform"     # Data transformation


class Position(BaseModel):
    """Node position on canvas"""
    x: float
    y: float


class NodeData(BaseModel):
    """Node configuration data"""
    label: str
    description: Optional[str] = None
    config: dict[str, Any] = Field(default_factory=dict)
    # For LLM nodes
    provider: Optional[str] = None  # groq, openrouter, cerebras, gemini
    model: Optional[str] = None
    prompt: Optional[str] = None
    temperature: Optional[float] = 0.7
    # For HTTP nodes
    url: Optional[str] = None
    method: Optional[str] = "GET"
    headers: Optional[dict[str, str]] = None
    body: Optional[Any] = None
    # For Code nodes
    code: Optional[str] = None
    # For Condition nodes
    condition: Optional[str] = None
    # For Transform nodes
    transform: Optional[str] = None
    input: Optional[str] = None


class Node(BaseModel):
    """A node in the workflow"""
    id: str = Field(default_factory=lambda: str(uuid.uuid4()))
    type: NodeType
    position: Position
    data: NodeData


class Edge(BaseModel):
    """Connection between nodes"""
    id: str = Field(default_factory=lambda: str(uuid.uuid4()))
    source: str          # Source node ID
    target: str          # Target node ID
    sourceHandle: Optional[str] = None  # For condition nodes (true/false)
    targetHandle: Optional[str] = None


class Workflow(BaseModel):
    """A complete workflow definition"""
    id: str = Field(default_factory=lambda: str(uuid.uuid4()))
    name: str
    description: Optional[str] = None
    nodes: list[Node] = Field(default_factory=list)
    edges: list[Edge] = Field(default_factory=list)
    created_at: Optional[str] = None
    updated_at: Optional[str] = None