autogen/python/packages/autogen-studio/autogenstudio/database/config_manager.py

import logging
from typing import Optional, Union, Dict, Any, List
from pathlib import Path
from loguru import logger
from ..datamodel import (
    Model, Team, Agent, Tool,
    Response, ComponentTypes, LinkTypes,
    ComponentConfigInput
)

from .component_factory import ComponentFactory
from .db_manager import DatabaseManager


class ConfigurationManager:
    """Manages persistence and relationships of components using ComponentFactory for validation"""

    DEFAULT_UNIQUENESS_FIELDS = {
        ComponentTypes.MODEL: ['model_type', 'model'],
        ComponentTypes.TOOL: ['name'],
        ComponentTypes.AGENT: ['agent_type', 'name'],
        ComponentTypes.TEAM: ['team_type', 'name']
    }

    def __init__(self, db_manager: DatabaseManager, uniqueness_fields: Dict[ComponentTypes, List[str]] = None):
        self.db_manager = db_manager
        self.component_factory = ComponentFactory()
        self.uniqueness_fields = uniqueness_fields or self.DEFAULT_UNIQUENESS_FIELDS

    async def import_component(self, component_config: ComponentConfigInput, user_id: str, check_exists: bool = False) -> Response:
        """
        Import a component configuration, validate it, and store the resulting component.

        Args:
            component_config: Configuration for the component (file path, dict, or ComponentConfig)
            user_id: User ID to associate with imported component
            check_exists: Whether to check for existing components before storing (default: False)

        Returns:
            Response containing import results or error
        """
        try:
            # Get validated config as dict
            config = await self.component_factory.load(component_config, return_type='dict')

            # Get component type
            component_type = self._determine_component_type(config)
            if not component_type:
                raise ValueError(
                    f"Unable to determine component type from config")

            # Check existence if requested
            if check_exists:
                existing = self._check_exists(component_type, config, user_id)
                if existing:
                    return Response(
                        message=self._format_exists_message(
                            component_type, config),
                        status=True,
                        data={"id": existing.id}
                    )

            # Route to appropriate storage method
            if component_type == ComponentTypes.TEAM:
                return await self._store_team(config, user_id, check_exists)
            elif component_type == ComponentTypes.AGENT:
                return await self._store_agent(config, user_id, check_exists)
            elif component_type == ComponentTypes.MODEL:
                return await self._store_model(config, user_id)
            elif component_type == ComponentTypes.TOOL:
                return await self._store_tool(config, user_id)
            else:
                raise ValueError(
                    f"Unsupported component type: {component_type}")

        except Exception as e:
            logger.error(f"Failed to import component: {str(e)}")
            return Response(message=str(e), status=False)

    async def import_directory(self, directory: Union[str, Path], user_id: str, check_exists: bool = False) -> Response:
        """
        Import all component configurations from a directory.

        Args:
            directory: Path to directory containing configuration files
            user_id: User ID to associate with imported components
            check_exists: Whether to check for existing components before storing (default: False)

        Returns:
            Response containing import results for all files
        """
        try:
            configs = await self.component_factory.load_directory(directory, return_type='dict')

            results = []
            for config in configs:
                result = await self.import_component(config, user_id, check_exists)
                results.append({
                    "component": self._get_component_type(config),
                    "status": result.status,
                    "message": result.message,
                    "id": result.data.get("id") if result.status else None
                })

            return Response(
                message="Directory import complete",
                status=True,
                data=results
            )

        except Exception as e:
            logger.error(f"Failed to import directory: {str(e)}")
            return Response(message=str(e), status=False)

    async def _store_team(self, config: dict, user_id: str, check_exists: bool = False) -> Response:
        """Store team component and manage its relationships with agents"""
        try:
            # Store the team
            team_db = Team(
                user_id=user_id,
                config=config
            )
            team_result = self.db_manager.upsert(team_db)
            if not team_result.status:
                return team_result

            team_id = team_result.data["id"]

            # Handle participants (agents)
            for participant in config.get("participants", []):
                if check_exists:
                    # Check for existing agent
                    agent_type = self._determine_component_type(participant)
                    existing_agent = self._check_exists(
                        agent_type, participant, user_id)
                    if existing_agent:
                        # Link existing agent
                        self.db_manager.link(
                            LinkTypes.TEAM_AGENT,
                            team_id,
                            existing_agent.id
                        )
                        logger.info(
                            f"Linked existing agent to team: {existing_agent}")
                        continue

                # Store and link new agent
                agent_result = await self._store_agent(participant, user_id, check_exists)
                if agent_result.status:
                    self.db_manager.link(
                        LinkTypes.TEAM_AGENT,
                        team_id,
                        agent_result.data["id"]
                    )

            return team_result

        except Exception as e:
            logger.error(f"Failed to store team: {str(e)}")
            return Response(message=str(e), status=False)

    async def _store_agent(self, config: dict, user_id: str, check_exists: bool = False) -> Response:
        """Store agent component and manage its relationships with tools and model"""
        try:
            # Store the agent
            agent_db = Agent(
                user_id=user_id,
                config=config
            )
            agent_result = self.db_manager.upsert(agent_db)
            if not agent_result.status:
                return agent_result

            agent_id = agent_result.data["id"]

            # Handle model client
            if "model_client" in config:
                if check_exists:
                    # Check for existing model
                    model_type = self._determine_component_type(
                        config["model_client"])
                    existing_model = self._check_exists(
                        model_type, config["model_client"], user_id)
                    if existing_model:
                        # Link existing model
                        self.db_manager.link(
                            LinkTypes.AGENT_MODEL,
                            agent_id,
                            existing_model.id
                        )
                        logger.info(
                            f"Linked existing model to agent: {existing_model.config.model_type}")
                    else:
                        # Store and link new model
                        model_result = await self._store_model(config["model_client"], user_id)
                        if model_result.status:
                            self.db_manager.link(
                                LinkTypes.AGENT_MODEL,
                                agent_id,
                                model_result.data["id"]
                            )
                else:
                    # Store and link new model without checking
                    model_result = await self._store_model(config["model_client"], user_id)
                    if model_result.status:
                        self.db_manager.link(
                            LinkTypes.AGENT_MODEL,
                            agent_id,
                            model_result.data["id"]
                        )

            # Handle tools
            for tool_config in config.get("tools", []):
                if check_exists:
                    # Check for existing tool
                    tool_type = self._determine_component_type(tool_config)
                    existing_tool = self._check_exists(
                        tool_type, tool_config, user_id)
                    if existing_tool:
                        # Link existing tool
                        self.db_manager.link(
                            LinkTypes.AGENT_TOOL,
                            agent_id,
                            existing_tool.id
                        )
                        logger.info(
                            f"Linked existing tool to agent: {existing_tool.config.name}")
                        continue

                # Store and link new tool
                tool_result = await self._store_tool(tool_config, user_id)
                if tool_result.status:
                    self.db_manager.link(
                        LinkTypes.AGENT_TOOL,
                        agent_id,
                        tool_result.data["id"]
                    )

            return agent_result

        except Exception as e:
            logger.error(f"Failed to store agent: {str(e)}")
            return Response(message=str(e), status=False)

    async def _store_model(self, config: dict, user_id: str) -> Response:
        """Store model component (leaf node - no relationships)"""
        try:
            model_db = Model(
                user_id=user_id,
                config=config
            )
            return self.db_manager.upsert(model_db)

        except Exception as e:
            logger.error(f"Failed to store model: {str(e)}")
            return Response(message=str(e), status=False)

    async def _store_tool(self, config: dict, user_id: str) -> Response:
        """Store tool component (leaf node - no relationships)"""
        try:
            tool_db = Tool(
                user_id=user_id,
                config=config
            )
            return self.db_manager.upsert(tool_db)

        except Exception as e:
            logger.error(f"Failed to store tool: {str(e)}")
            return Response(message=str(e), status=False)

    def _check_exists(self, component_type: ComponentTypes, config: dict, user_id: str) -> Optional[Union[Model, Tool, Agent, Team]]:
        """Check if component exists based on configured uniqueness fields."""
        fields = self.uniqueness_fields.get(component_type, [])
        if not fields:
            return None

        component_class = {
            ComponentTypes.MODEL: Model,
            ComponentTypes.TOOL: Tool,
            ComponentTypes.AGENT: Agent,
            ComponentTypes.TEAM: Team
        }.get(component_type)

        components = self.db_manager.get(
            component_class, {"user_id": user_id}).data

        for component in components:
            matches = all(
                component.config.get(field) == config.get(field)
                for field in fields
            )
            if matches:
                return component

        return None

    def _format_exists_message(self, component_type: ComponentTypes, config: dict) -> str:
        """Format existence message with identifying fields."""
        fields = self.uniqueness_fields.get(component_type, [])
        field_values = [f"{field}='{config.get(field)}'" for field in fields]
        return f"{component_type.value} with {' and '.join(field_values)} already exists"

    def _determine_component_type(self, config: dict) -> Optional[ComponentTypes]:
        """Determine component type from configuration dictionary"""
        if "team_type" in config:
            return ComponentTypes.TEAM
        elif "agent_type" in config:
            return ComponentTypes.AGENT
        elif "model_type" in config:
            return ComponentTypes.MODEL
        elif "tool_type" in config:
            return ComponentTypes.TOOL
        return None

    def _get_component_type(self, config: dict) -> str:
        """Helper to get component type string from config"""
        component_type = self._determine_component_type(config)
        return component_type.value if component_type else "unknown"

    async def cleanup(self):
        """Cleanup resources"""
        await self.component_factory.cleanup()