hxf/backend/th_agenter/services/agent_config.py

"""Agent configuration service."""

from typing import List, Dict, Any, Optional
from sqlalchemy.orm import Session
from sqlalchemy import and_, select, update

from ..models.agent_config import AgentConfig
from utils.util_exceptions import ValidationError, NotFoundError
from loguru import logger


class AgentConfigService:
    """Service for managing agent configurations."""
    
    def __init__(self, db: Session):
        self.db = db
    
    async def create_config(self, config_data: Dict[str, Any]) -> AgentConfig:
        """Create a new agent configuration."""
        try:
            # Validate required fields
            if not config_data.get("name"):
                raise ValidationError("Configuration name is required")
            
            # Check if name already exists
            stmt = select(AgentConfig).where(AgentConfig.name == config_data["name"])
            existing = (await self.db.execute(stmt)).scalar_one_or_none()
            if existing:
                raise ValidationError(f"Configuration with name '{config_data['name']}' already exists")
            
            # Create new configuration
            config = AgentConfig(
                name=config_data["name"],
                description=config_data.get("description", ""),
                enabled_tools=config_data.get("enabled_tools", ["calculator", "weather", "search", "datetime", "file"]),
                max_iterations=config_data.get("max_iterations", 10),
                temperature=config_data.get("temperature", 0.1),
                system_message=config_data.get("system_message", "You are a helpful AI assistant with access to various tools. Use the available tools to help answer user questions accurately. Always explain your reasoning and the tools you're using."),
                verbose=config_data.get("verbose", True),
                is_active=config_data.get("is_active", True),
                is_default=config_data.get("is_default", False)
            )
            
            # If this is set as default, unset other defaults
            if config.is_default:
                stmt = update(AgentConfig).where(AgentConfig.is_default == True).values({"is_default": False})
                await self.db.execute(stmt)
            
            self.db.add(config)
            await self.db.commit()
            await self.db.refresh(config)
            
            logger.info(f"Created agent configuration: {config.name}")
            return config
            
        except Exception as e:
            await self.db.rollback()
            logger.error(f"Error creating agent configuration: {str(e)}")
            raise
    
    async def get_config(self, config_id: int) -> Optional[AgentConfig]:
        """Get agent configuration by ID."""
        stmt = select(AgentConfig).where(AgentConfig.id == config_id)
        return (await self.db.execute(stmt)).scalar_one_or_none()
    
    async def get_config_by_name(self, name: str) -> Optional[AgentConfig]:
        """Get agent configuration by name."""
        stmt = select(AgentConfig).where(AgentConfig.name == name)
        return (await self.db.execute(stmt)).scalar_one_or_none()
    
    async def get_default_config(self) -> Optional[AgentConfig]:
        """Get default agent configuration."""
        stmt = select(AgentConfig).where(and_(AgentConfig.is_default == True, AgentConfig.is_active == True))
        return (await self.db.execute(stmt)).scalar_one_or_none()
    
    def list_configs(self, active_only: bool = True) -> List[AgentConfig]:
        """List all agent configurations."""
        stmt = select(AgentConfig)
        if active_only:
            stmt = stmt.where(AgentConfig.is_active == True)
        stmt = stmt.order_by(AgentConfig.created_at.desc())
        return self.db.execute(stmt).scalars().all()
    
    async def update_config(self, config_id: int, config_data: Dict[str, Any]) -> AgentConfig:
        """Update agent configuration."""
        try:
            config = await self.get_config(config_id)
            if not config:
                raise NotFoundError(f"Agent configuration with ID {config_id} not found")
            
            # Check if name change conflicts with existing
            if "name" in config_data and config_data["name"] != config.name:
                stmt = select(AgentConfig).where(
                    and_(
                        AgentConfig.name == config_data["name"],
                        AgentConfig.id != config_id
                    )
                )
                existing = (await self.db.execute(stmt)).scalar_one_or_none()
                if existing:
                    raise ValidationError(f"Configuration with name '{config_data['name']}' already exists")
            
            # Update fields
            for key, value in config_data.items():
                if hasattr(config, key):
                    setattr(config, key, value)
            
            # If this is set as default, unset other defaults
            if config_data.get("is_default", False):
                stmt = update(AgentConfig).where(
                    and_(
                        AgentConfig.is_default == True,
                        AgentConfig.id != config_id
                    )
                ).values({"is_default": False})
                await self.db.execute(stmt)
            
            await self.db.commit()
            await self.db.refresh(config)
            
            logger.info(f"Updated agent configuration: {config.name}")
            return config
            
        except Exception as e:
            await self.db.rollback()
            logger.error(f"Error updating agent configuration: {str(e)}")
            raise
    
    async def delete_config(self, config_id: int) -> bool:
        """Delete agent configuration (soft delete by setting is_active=False)."""
        try:
            config = await self.get_config(config_id)
            if not config:
                raise NotFoundError(f"Agent configuration with ID {config_id} not found")
            
            # Don't allow deleting the default configuration
            if config.is_default:
                raise ValidationError("Cannot delete the default configuration")
            
            config.is_active = False
            self.db.commit()
            
            logger.info(f"Deleted agent configuration: {config.name}")
            return True
            
        except Exception as e:
            await self.db.rollback()
            logger.error(f"Error deleting agent configuration: {str(e)}")
            raise
    
    async def set_default_config(self, config_id: int) -> AgentConfig:
        """Set a configuration as default."""
        try:
            config = await self.get_config(config_id)
            if not config:
                raise NotFoundError(f"Agent configuration with ID {config_id} not found")
            
            if not config.is_active:
                raise ValidationError("Cannot set inactive configuration as default")
            
            # Unset other defaults
            stmt = update(AgentConfig).where(AgentConfig.is_default == True).values({"is_default": False})
            await self.db.execute(stmt)
            
            # Set this as default
            config.is_default = True
            await self.db.commit()
            await self.db.refresh(config)
            
            logger.info(f"Set default agent configuration: {config.name}")
            return config
            
        except Exception as e:
            await self.db.rollback()
            logger.error(f"Error setting default agent configuration: {str(e)}")
            raise
    
    async def get_config_dict(self, config_id: Optional[int] = None) -> Dict[str, Any]:
        """Get configuration as dictionary. If no ID provided, returns default config."""
        if config_id:
            config = await self.get_config(config_id)
        else:
            config = await self.get_default_config()
        
        if not config:
            # Return default values if no configuration found
            return {
                "enabled_tools": ["calculator", "weather", "search", "datetime", "file", "generate_image"],
                "max_iterations": 10,
                "temperature": 0.1,
                "system_message": "You are a helpful AI assistant with access to various tools. Use the available tools to help answer user questions accurately. Always explain your reasoning and the tools you're using.",
                "verbose": True
            }
        
        return {
            "enabled_tools": config.enabled_tools,
            "max_iterations": config.max_iterations,
            "temperature": config.temperature,
            "system_message": config.system_message,
            "verbose": config.verbose
        }