""" Database models and session management for Agent Manager. Supports both SQLite (development) and PostgreSQL (production). """ from datetime import datetime from typing import Optional, Dict, Any from sqlalchemy import ( create_engine, Column, Integer, String, DateTime, Boolean, JSON, Float, ForeignKey, Text, Enum as SQLEnum ) from sqlalchemy.ext.declarative import declarative_base from sqlalchemy.orm import sessionmaker, relationship, Session import enum import os # Database URL from environment or default to SQLite DATABASE_URL = os.getenv("DATABASE_URL", "sqlite:///./agent_manager.db") engine = create_engine( DATABASE_URL, connect_args={"check_same_thread": False} if "sqlite" in DATABASE_URL else {} ) SessionLocal = sessionmaker(autocommit=False, autoflush=False, bind=engine) Base = declarative_base() class AgentType(str, enum.Enum): """Agent type enumeration""" PLATFORM = "platform" CUSTOM = "custom" class AgentStatus(str, enum.Enum): """Agent status enumeration""" PENDING = "pending" RUNNING = "running" STOPPED = "stopped" FAILED = "failed" SCALING = "scaling" class Template(Base): """Template model for agent templates""" __tablename__ = "templates" id = Column(Integer, primary_key=True, index=True) name = Column(String(100), unique=True, nullable=False, index=True) display_name = Column(String(200), nullable=False) description = Column(Text) agent_type = Column(SQLEnum(AgentType), nullable=False, index=True) # Framework configuration (NEW) agent_framework = Column(String(50), default="langchain") # "langchain", "mcp", "a2a" # Image configuration image = Column(String(500), nullable=False) port = Column(Integer, nullable=True) # Environment variable requirements (JSON format) # {"required": {"KEY": "description"}, "optional": {"KEY": "description"}} env_requirements = Column(JSON, default={}) # Tools configuration (NEW) - JSON format for tool definitions tools_config = Column(JSON, default={}) # Model configuration defaults (NEW) default_model_provider = Column(String(100)) # e.g., "openai", "azure-openai" default_model_name = Column(String(200)) # e.g., "gpt-4", "claude-3" # Resource configuration (for platform agents, fixed by admin) cpu_request = Column(String(20)) # e.g., "100m" cpu_limit = Column(String(20)) # e.g., "500m" memory_request = Column(String(20)) # e.g., "128Mi" memory_limit = Column(String(20)) # e.g., "512Mi" # Scaling configuration defaults min_replicas = Column(Integer, default=1) max_replicas = Column(Integer, default=3) target_cpu_utilization = Column(Integer, default=80) # Metadata is_active = Column(Boolean, default=True) created_at = Column(DateTime, default=datetime.utcnow) updated_at = Column(DateTime, default=datetime.utcnow, onupdate=datetime.utcnow) created_by = Column(String(100)) # Relationships agents = relationship("Agent", back_populates="template") class Agent(Base): """Agent instance model""" __tablename__ = "agents" id = Column(Integer, primary_key=True, index=True) name = Column(String(100), unique=True, nullable=False, index=True) display_name = Column(String(200)) # Template reference template_id = Column(Integer, ForeignKey("templates.id"), nullable=False) template = relationship("Template", back_populates="agents") # Ownership and organization owner_id = Column(String(100), nullable=False, index=True) channel_id = Column(String(100), index=True) tenant_id = Column(String(100), index=True) # Agent configuration agent_type = Column(SQLEnum(AgentType), nullable=False, index=True) status = Column(SQLEnum(AgentStatus), default=AgentStatus.PENDING, index=True) # Framework configuration (NEW) agent_framework = Column(String(50), default="langchain") # "langchain", "mcp", "a2a" # Environment variables (encrypted in production) environment_vars = Column(JSON, default={}) # Tools configuration (NEW) - Instance-level tools override tools_config = Column(JSON, default={}) tool_endpoint = Column(String(500)) # External tool endpoint URL tool_api_key = Column(String(500)) # Encrypted tool API key # Model configuration (NEW) - Instance-level model settings model_provider = Column(String(100)) # e.g., "openai", "azure-openai" model_name = Column(String(200)) # e.g., "gpt-4" model_endpoint = Column(String(500)) # Model API endpoint model_api_key = Column(String(500)) # Encrypted model API key # Storage configuration (NEW) - For agents that need storage storage_connection_string = Column(String(1000)) # Encrypted storage connection storage_account_name = Column(String(200)) # Resource configuration (for custom agents) cpu_request = Column(String(20)) cpu_limit = Column(String(20)) memory_request = Column(String(20)) memory_limit = Column(String(20)) # Scaling configuration min_replicas = Column(Integer, default=1) max_replicas = Column(Integer, default=3) target_cpu_utilization = Column(Integer, default=80) current_replicas = Column(Integer, default=0) # Kubernetes resources deployment_name = Column(String(100)) service_name = Column(String(100)) service_url = Column(String(500)) namespace = Column(String(100), default="ai-agents") # Metadata created_at = Column(DateTime, default=datetime.utcnow) updated_at = Column(DateTime, default=datetime.utcnow, onupdate=datetime.utcnow) last_accessed_at = Column(DateTime, nullable=True) # Relationships metrics = relationship("AgentMetric", back_populates="agent", cascade="all, delete-orphan") class Quota(Base): """Resource quota model""" __tablename__ = "quotas" id = Column(Integer, primary_key=True, index=True) # Quota owner (hierarchical: admin -> channel -> tenant) owner_type = Column(String(20), nullable=False) # "admin", "channel", "tenant" owner_id = Column(String(100), nullable=False, index=True) channel_id = Column(String(100), index=True) # For platform agents: quota is Pod count platform_pod_quota = Column(Integer, default=0) platform_pod_used = Column(Integer, default=0) # For custom agents: quota is CPU/Memory totals custom_cpu_quota = Column(Float, default=0.0) # in cores custom_cpu_used = Column(Float, default=0.0) custom_memory_quota = Column(Float, default=0.0) # in GB custom_memory_used = Column(Float, default=0.0) # Metadata created_at = Column(DateTime, default=datetime.utcnow) updated_at = Column(DateTime, default=datetime.utcnow, onupdate=datetime.utcnow) # Unique constraint __table_args__ = ( # Unique per owner # UniqueConstraint('owner_type', 'owner_id', name='uq_quota_owner'), ) class AgentMetric(Base): """Agent metrics tracking""" __tablename__ = "agent_metrics" id = Column(Integer, primary_key=True, index=True) agent_id = Column(Integer, ForeignKey("agents.id"), nullable=False, index=True) agent = relationship("Agent", back_populates="metrics") # Timestamp timestamp = Column(DateTime, default=datetime.utcnow, index=True) # Resource metrics cpu_usage = Column(Float) # in cores memory_usage = Column(Float) # in MB network_rx_bytes = Column(Integer, default=0) network_tx_bytes = Column(Integer, default=0) # Replica count replica_count = Column(Integer, default=0) # Request metrics request_count = Column(Integer, default=0) error_count = Column(Integer, default=0) # Database initialization def init_db(): """Initialize database tables""" Base.metadata.create_all(bind=engine) def get_db() -> Session: """Get database session (FastAPI dependency)""" db = SessionLocal() try: yield db finally: db.close() def parse_resource_string(resource_str: str) -> float: """Parse Kubernetes resource string to float Examples: "100m" -> 0.1 (cores) "2" -> 2.0 (cores) "128Mi" -> 128.0 (MB) "1Gi" -> 1024.0 (MB) """ if not resource_str: return 0.0 resource_str = resource_str.strip() # CPU resources if resource_str.endswith('m'): return float(resource_str[:-1]) / 1000.0 # Memory resources if resource_str.endswith('Mi'): return float(resource_str[:-2]) elif resource_str.endswith('Gi'): return float(resource_str[:-2]) * 1024.0 elif resource_str.endswith('Ki'): return float(resource_str[:-2]) / 1024.0 # Plain number try: return float(resource_str) except ValueError: return 0.0 # Initialize database on import init_db()