272 lines
9.3 KiB
Python
272 lines
9.3 KiB
Python
"""
|
|
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 - PostgreSQL (hardcoded)
|
|
DATABASE_URL = "postgresql://taiji:By%40123456.@taijipda.postgres.database.azure.com:5432/taijiagnet"
|
|
|
|
engine = create_engine(DATABASE_URL)
|
|
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 (nullable for backward compatibility with non-template agents)
|
|
template_id = Column(Integer, ForeignKey("templates.id"), nullable=True)
|
|
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")
|
|
|
|
# Access information (NEW) - LoadBalancer and DNS details
|
|
external_ip = Column(String(100), nullable=True) # LoadBalancer external IP
|
|
domain = Column(String(500), nullable=True) # Full DNS domain (e.g., agent-name.taijiagnet.com)
|
|
ip_url = Column(String(500), nullable=True) # HTTP URL via IP (http://x.x.x.x:80)
|
|
domain_url = Column(String(500), nullable=True) # HTTP URL via domain
|
|
recommended_url = Column(String(500), nullable=True) # Recommended access URL
|
|
|
|
# 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()
|