Files
agent_management/api/swarm/models.py
T

142 lines
5.5 KiB
Python

"""Pydantic models for sub-mode runtime compatibility requests and responses."""
from typing import List, Optional, Dict, Any
from pydantic import BaseModel, ConfigDict, Field, model_validator
from datetime import datetime
class AgentConfig(BaseModel):
"""Agent configuration for sub-mode runtime execution."""
role: str = Field(..., description="Agent role (architect/coder/reviewer/tester)")
template: str = Field(default="a2a_litellm_agent", description="Agent template name")
model: str = Field(default="gpt-4", description="Model name")
capabilities: List[str] = Field(default=[], description="Agent capabilities")
system_prompt: Optional[str] = Field(None, description="Custom system prompt")
replicas: int = Field(default=1, description="Number of agent replicas")
@model_validator(mode="before")
@classmethod
def normalize_sub_mode_agent(cls, data: Any) -> Any:
"""Accept Manager sub-mode agent fields when the compatibility API is used."""
if not isinstance(data, dict):
return data
data = dict(data)
if not data.get("role"):
data["role"] = data.get("role_template") or data.get("target_role") or "worker"
if not data.get("model"):
data["model"] = data.get("default_model_id") or data.get("model_ref") or "gpt-4"
return data
class OrchestrationConfig(BaseModel):
"""Orchestration configuration"""
strategy: str = Field(default="sequential", description="Orchestration strategy (sequential/parallel/hybrid)")
max_iterations: int = Field(default=3, description="Maximum iterations")
timeout_minutes: int = Field(default=30, description="Timeout in minutes")
class ProjectContext(BaseModel):
"""Project context information"""
model_config = ConfigDict(extra="allow")
repo_url: Optional[str] = Field(None, description="Repository URL")
branch: Optional[str] = Field(None, description="Git branch")
language: Optional[str] = Field(None, description="Programming language")
framework: Optional[str] = Field(None, description="Framework")
intent_id: Optional[str] = None
template_hint: Optional[str] = None
sub_mode: Optional[str] = None
agile_context: Optional[Dict[str, Any]] = None
correlation_id: Optional[str] = None
class CallbackConfig(BaseModel):
"""Callback configuration"""
url: str = Field(..., description="Callback URL")
method: str = Field(default="POST", description="HTTP method")
signing_secret_ref: Optional[str] = Field(None, description="Azure Key Vault reference for callback HMAC signing")
subscribed_events: Optional[List[str]] = Field(None, description="Runtime event types subscribed by Manager")
class SwarmCreateRequest(BaseModel):
"""Compatibility request model for creating a sub-mode runtime run."""
task_description: str = Field(..., description="Task description")
project_context: Optional[ProjectContext] = Field(None, description="Project context")
agents: List[AgentConfig] = Field(..., description="Agent configurations")
orchestration: OrchestrationConfig = Field(default_factory=OrchestrationConfig, description="Orchestration config")
callback: Optional[CallbackConfig] = Field(None, description="Callback configuration")
owner_id: str = Field(default="default", description="Owner ID")
class SwarmAgentInfo(BaseModel):
"""Sub-mode runtime agent information."""
agent_id: str
role: str
status: str
namespace: str
service_url: Optional[str] = None
current_task: Optional[str] = None
output: Optional[str] = None
class SwarmCreateResponse(BaseModel):
"""Compatibility response model for sub-mode runtime creation."""
deployment_id: Optional[str] = None
swarm_id: str
status: str
agents: List[SwarmAgentInfo]
created_at: datetime
estimated_ready_at: Optional[datetime] = None
class SwarmMetrics(BaseModel):
"""Sub-mode runtime metrics."""
total_messages: int
tokens_used: int
elapsed_seconds: int
class SwarmStatusResponse(BaseModel):
"""Compatibility response model for sub-mode runtime status."""
deployment_id: Optional[str] = None
swarm_id: str
status: str
phase: Optional[str] = None
progress: int
agents: List[SwarmAgentInfo]
metrics: SwarmMetrics
artifacts: List[Dict[str, Any]] = []
error_message: Optional[str] = None
created_at: Optional[datetime] = None
updated_at: Optional[datetime] = None
class SwarmStopRequest(BaseModel):
"""Compatibility request model for stopping a sub-mode runtime run."""
reason: Optional[str] = Field(None, description="Reason for stopping")
cleanup: bool = Field(default=True, description="Whether to cleanup K8s resources")
class SwarmStopResponse(BaseModel):
"""Compatibility response model for stopping a sub-mode runtime run."""
deployment_id: Optional[str] = None
swarm_id: str
status: str
stopped_at: datetime
class ApprovalDecisionRequest(BaseModel):
"""Manager approval decision for a paused high-risk Runtime action."""
approval_id: str
decision: str = Field(..., description="approved or rejected")
manager_deployment_id: Optional[str] = None
runtime_deployment_id: Optional[str] = None
operation: Optional[str] = None
resource_id: Optional[str] = None
resource_type: Optional[str] = None
target_role: Optional[str] = None
requires_credential: bool = False
credential_ref: Optional[str] = None
lease_id: Optional[str] = None
lease_expires_at: Optional[int] = None