zhanggangyong/agent_management
1
0
Fork 1
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity
Files
55780e652bab2a8eec3b1cfdca7715076f214712
agent_management/database.py
T

558 lines
19 KiB
Python
Raw Blame History

"""
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,
pool_pre_ping=True,
pool_recycle=1800,
pool_use_lifo=True,
)
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)
# ============================================================================
# Heicode Integration Models (NEW)
# ============================================================================
class DeploymentStatus(str, enum.Enum):
"""Deployment status for Heicode integration"""
PENDING = "pending"
RUNNING = "running"
STOPPED = "stopped"
FAILED = "failed"
class RiskLevel(str, enum.Enum):
"""Risk level for deployments"""
LOW = "low"
MEDIUM = "medium"
HIGH = "high"
class BillingProvider(str, enum.Enum):
"""Model gateway provider"""
NEWAPI = "newapi"
LITELLM = "litellm"
class Deployment(Base):
"""Deployment model for Heicode integration"""
__tablename__ = "deployments"
id = Column(Integer, primary_key=True, index=True)
deployment_id = Column(String(100), unique=True, nullable=False, index=True)
# Ownership
user_id = Column(String(100), nullable=False, index=True)
binding_scope = Column(String(200), nullable=False, index=True)
correlation_id = Column(String(100), index=True)
# Deployment configuration
orchestration_plan = Column(Text, nullable=False)
risk_level = Column(SQLEnum(RiskLevel), nullable=False)
approval_token = Column(Text)
# Budget
budget_max_usd = Column(Float)
budget_consumed_usd = Column(Float, default=0.0)
budget_alert_threshold_pct = Column(Integer, default=80)
# Model gateway configuration
billing_provider = Column(SQLEnum(BillingProvider), nullable=False)
default_model_id = Column(String(200), nullable=False)
allowed_model_ids = Column(JSON, nullable=False)
secret_ref = Column(String(500))
# Resource grants
resource_grants = Column(JSON, default=[])
# Status
status = Column(SQLEnum(DeploymentStatus), nullable=False, default=DeploymentStatus.PENDING, index=True)
phase = Column(String(100))
error_message = Column(Text)
# Kubernetes resources
namespace = Column(String(100), nullable=False)
configmap_name = Column(String(100))
# Metadata
created_at = Column(DateTime, default=datetime.utcnow, index=True)
updated_at = Column(DateTime, default=datetime.utcnow, onupdate=datetime.utcnow)
stopped_at = Column(DateTime)
# Relationships
agent_instances = relationship("AgentInstance", back_populates="deployment", cascade="all, delete-orphan")
events = relationship("Event", back_populates="deployment", cascade="all, delete-orphan")
class AgentInstance(Base):
"""Agent instance model for Heicode deployments"""
__tablename__ = "agent_instances"
id = Column(Integer, primary_key=True, index=True)
agent_instance_id = Column(String(100), unique=True, nullable=False, index=True)
deployment_id = Column(String(100), ForeignKey("deployments.deployment_id", ondelete="CASCADE"), nullable=False, index=True)
# Agent configuration
role = Column(String(100), nullable=False)
image = Column(String(500), nullable=False)
phase = Column(String(100))
# Kubernetes resources
namespace = Column(String(100), nullable=False)
pod_name = Column(String(100), nullable=False)
service_account = Column(String(100))
# Status
status = Column(SQLEnum(DeploymentStatus), nullable=False, default=DeploymentStatus.PENDING, index=True)
error_message = Column(Text)
# Metadata
created_at = Column(DateTime, default=datetime.utcnow)
updated_at = Column(DateTime, default=datetime.utcnow, onupdate=datetime.utcnow)
# Relationships
deployment = relationship("Deployment", back_populates="agent_instances")
class Event(Base):
"""Event model for deployment events"""
__tablename__ = "events"
id = Column(Integer, primary_key=True, index=True)
event_id = Column(String(100), unique=True, nullable=False, index=True)
deployment_id = Column(String(100), ForeignKey("deployments.deployment_id", ondelete="CASCADE"), nullable=False, index=True)
agent_instance_id = Column(String(100), ForeignKey("agent_instances.agent_instance_id", ondelete="SET NULL"))
# Event details
event_type = Column(String(100), nullable=False, index=True)
correlation_id = Column(String(100))
payload = Column(JSON)
# Metadata
occurred_at = Column(DateTime, default=datetime.utcnow, index=True)
# Relationships
deployment = relationship("Deployment", back_populates="events")
class AuditLog(Base):
"""Audit log model for tracking all operations"""
__tablename__ = "audit_logs"
id = Column(Integer, primary_key=True, index=True)
audit_id = Column(String(100), unique=True, nullable=False, index=True)
# Actor
actor = Column(String(200), nullable=False, index=True)
user_id = Column(String(100), index=True)
binding_scope = Column(String(200), index=True)
# Action
action = Column(String(100), nullable=False, index=True)
resource_type = Column(String(50), nullable=False)
resource_id = Column(String(100))
# Request details
correlation_id = Column(String(100))
request_payload = Column(JSON)
# Result
result = Column(String(50), nullable=False, index=True)
error_code = Column(String(50))
error_message = Column(Text)
# Metadata
occurred_at = Column(DateTime, default=datetime.utcnow, index=True)
ip_address = Column(String(50))
user_agent = Column(Text)
# ============================================================================
# Sub-mode Runtime Internal Models
# ============================================================================
class SwarmStatus(str, enum.Enum):
"""Swarm status enumeration"""
INITIALIZING = "initializing"
RUNNING = "running"
COMPLETED = "completed"
FAILED = "failed"
STOPPED = "stopped"
class SwarmAgentStatus(str, enum.Enum):
"""Swarm agent status enumeration"""
PENDING = "pending"
RUNNING = "running"
COMPLETED = "completed"
FAILED = "failed"
class Swarm(Base):
"""Internal runtime run model used by Heicode sub-mode execution."""
__tablename__ = "swarms"
id = Column(Integer, primary_key=True, index=True)
swarm_id = Column(String(100), unique=True, nullable=False, index=True)
# Task information
task_description = Column(Text, nullable=False)
project_context = Column(JSON) # {repo_url, branch, language, framework}
# Orchestration configuration
orchestration_strategy = Column(String(50), default="sequential") # sequential, parallel, hybrid
max_iterations = Column(Integer, default=3)
timeout_minutes = Column(Integer, default=30)
# Status
status = Column(SQLEnum(SwarmStatus), default=SwarmStatus.INITIALIZING, index=True)
phase = Column(String(50)) # planning, coding, reviewing, testing
progress = Column(Integer, default=0) # 0-100
# Results
artifacts = Column(JSON, default=[]) # Generated code, documents, etc.
error_message = Column(Text)
# Metrics
total_messages = Column(Integer, default=0)
tokens_used = Column(Integer, default=0)
# Callback
callback_url = Column(String(500))
callback_method = Column(String(10), default="POST")
# Ownership
owner_id = Column(String(100), nullable=False, index=True)
# Timestamps
created_at = Column(DateTime, default=datetime.utcnow, index=True)
updated_at = Column(DateTime, default=datetime.utcnow, onupdate=datetime.utcnow)
completed_at = Column(DateTime)
# Relationships
swarm_agents = relationship("SwarmAgent", back_populates="swarm", cascade="all, delete-orphan")
swarm_messages = relationship("SwarmMessage", back_populates="swarm", cascade="all, delete-orphan")
class SwarmAgent(Base):
"""Internal runtime agent model for sub-mode execution."""
__tablename__ = "swarm_agents"
id = Column(Integer, primary_key=True, index=True)
agent_id = Column(String(100), unique=True, nullable=False, index=True)
swarm_id = Column(String(100), ForeignKey("swarms.swarm_id", ondelete="CASCADE"), nullable=False, index=True)
# Agent configuration
role = Column(String(100), nullable=False) # architect, coder, reviewer, tester
template = Column(String(100), nullable=False) # a2a_litellm_agent, code_manager_agent
model = Column(String(200)) # gpt-4, claude-3, etc.
capabilities = Column(JSON, default=[]) # ["design", "coding", "review"]
system_prompt = Column(Text)
# Kubernetes resources
namespace = Column(String(100), nullable=False)
pod_name = Column(String(100), nullable=False)
service_url = Column(String(500))
external_ip = Column(String(100))
# Status
status = Column(SQLEnum(SwarmAgentStatus), default=SwarmAgentStatus.PENDING, index=True)
current_task = Column(Text)
output = Column(Text)
# Timestamps
created_at = Column(DateTime, default=datetime.utcnow)
updated_at = Column(DateTime, default=datetime.utcnow, onupdate=datetime.utcnow)
# Relationships
swarm = relationship("Swarm", back_populates="swarm_agents")
class SwarmMessage(Base):
"""Internal runtime message model for sub-mode execution."""
__tablename__ = "swarm_messages"
id = Column(Integer, primary_key=True, index=True)
message_id = Column(String(100), unique=True, nullable=False, index=True)
swarm_id = Column(String(100), ForeignKey("swarms.swarm_id", ondelete="CASCADE"), nullable=False, index=True)
# Message information
from_agent_id = Column(String(100), index=True) # NULL for orchestrator
to_agent_id = Column(String(100), index=True) # NULL for broadcast
message_type = Column(String(50)) # task, response, broadcast, artifact
content = Column(Text)
message_metadata = Column("metadata", JSON)
# Timestamp
created_at = Column(DateTime, default=datetime.utcnow, index=True)
# Relationships
swarm = relationship("Swarm", back_populates="swarm_messages")
# 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()
Reference in New Issue View Git Blame Copy Permalink