src/models/agent_models.py

"""
Agent Models - Pydantic Models for Agents
"""
from pydantic import BaseModel, Field
from typing import Optional, List, Dict, Any
from enum import Enum


class AgentType(str, Enum):
    SKILL = "skill"
    LONG_CHAIN = "long_chain"
    SCRIPT = "script"


class TaskStatus(str, Enum):
    PENDING = "pending"
    RUNNING = "running"
    COMPLETED = "completed"
    FAILED = "failed"


class SkillRequest(BaseModel):
    """Request model for Skill Agent"""
    prompt: str = Field(..., description="The task prompt")
    skill_type: str = Field(..., description="Type of skill to execute")
    context: Optional[Dict[str, Any]] = Field(default=None, description="Additional context")
    tools: Optional[List[str]] = Field(default=None, description="Available tools")


class SkillResponse(BaseModel):
    """Response model for Skill Agent"""
    status: TaskStatus
    result: Optional[Dict[str, Any]] = None
    error: Optional[str] = None
    execution_time: float = 0.0


class ChainRequest(BaseModel):
    """Request model for Long Chain Agent"""
    workflow: str = Field(..., description="Workflow name or definition")
    input_data: Dict[str, Any] = Field(..., description="Input data for workflow")
    steps: Optional[List[str]] = Field(default=None, description="Specific steps to execute")


class ChainResponse(BaseModel):
    """Response model for Long Chain Agent"""
    status: TaskStatus
    workflow: str
    steps_completed: List[Dict[str, Any]] = []
    final_result: Optional[Dict[str, Any]] = None
    error: Optional[str] = None
    total_time: float = 0.0


class ScriptRequest(BaseModel):
    """Request model for Script Agent"""
    language: str = Field(default="python", description="Programming language")
    prompt: str = Field(..., description="Code generation prompt")
    template: Optional[str] = Field(default=None, description="Code template to use")
    optimize: bool = Field(default=False, description="Whether to optimize the code")


class ScriptResponse(BaseModel):
    """Response model for Script Agent"""
    status: TaskStatus
    code: Optional[str] = None
    language: str
    explanation: Optional[str] = None
    error: Optional[str] = None