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

import os
from pathlib import Path
from typing import List, Literal, Union, Optional, Dict, Any, Type
from datetime import datetime
import json
from autogen_agentchat.task import MaxMessageTermination, TextMentionTermination, StopMessageTermination
import yaml
import logging
from packaging import version

from ..datamodel import (
    TeamConfig, AgentConfig, ModelConfig, ToolConfig,
    TeamTypes, AgentTypes, ModelTypes, ToolTypes,
    ComponentType, ComponentConfig, ComponentConfigInput, TerminationConfig, TerminationTypes, Response
)
from autogen_agentchat.agents import AssistantAgent
from autogen_agentchat.teams import RoundRobinGroupChat, SelectorGroupChat
from autogen_ext.models import OpenAIChatCompletionClient
from autogen_core.components.tools import FunctionTool

logger = logging.getLogger(__name__)

# Type definitions for supported components
TeamComponent = Union[RoundRobinGroupChat, SelectorGroupChat]
AgentComponent = Union[AssistantAgent]  # Will grow with more agent types
# Will grow with more model types
ModelComponent = Union[OpenAIChatCompletionClient]
ToolComponent = Union[FunctionTool]  # Will grow with more tool types
TerminationComponent = Union[MaxMessageTermination,
                             StopMessageTermination, TextMentionTermination]

# Config type definitions

Component = Union[TeamComponent, AgentComponent, ModelComponent, ToolComponent]


ReturnType = Literal['object', 'dict', 'config']
Component = Union[RoundRobinGroupChat, SelectorGroupChat,
                  AssistantAgent, OpenAIChatCompletionClient, FunctionTool]


class ComponentFactory:
    """Creates and manages agent components with versioned configuration loading"""

    SUPPORTED_VERSIONS = {
        ComponentType.TEAM: ["1.0.0"],
        ComponentType.AGENT: ["1.0.0"],
        ComponentType.MODEL: ["1.0.0"],
        ComponentType.TOOL: ["1.0.0"],
        ComponentType.TERMINATION: ["1.0.0"]
    }

    def __init__(self):
        self._model_cache: Dict[str, OpenAIChatCompletionClient] = {}
        self._tool_cache: Dict[str, FunctionTool] = {}
        self._last_cache_clear = datetime.now()

    async def load(self, component: ComponentConfigInput, return_type: ReturnType = 'object') -> Union[Component, dict, ComponentConfig]:
        """
        Universal loader for any component type

        Args:
            component: Component configuration (file path, dict, or ComponentConfig)
            return_type: Type of return value ('object', 'dict', or 'config')

        Returns:
            Component instance, config dict, or ComponentConfig based on return_type

        Raises:
            ValueError: If component type is unknown or version unsupported
        """
        try:
            # Load and validate config
            if isinstance(component, (str, Path)):
                component_dict = await self._load_from_file(component)
                config = self._dict_to_config(component_dict)
            elif isinstance(component, dict):
                config = self._dict_to_config(component)
            else:
                config = component

            # Validate version
            if not self._is_version_supported(config.component_type, config.version):
                raise ValueError(
                    f"Unsupported version {config.version} for "
                    f"component type {config.component_type}. "
                    f"Supported versions: {self.SUPPORTED_VERSIONS[config.component_type]}"
                )

            # Return early if dict or config requested
            if return_type == 'dict':
                return config.model_dump()
            elif return_type == 'config':
                return config

            # Otherwise create and return component instance
            handlers = {
                ComponentType.TEAM: self.load_team,
                ComponentType.AGENT: self.load_agent,
                ComponentType.MODEL: self.load_model,
                ComponentType.TOOL: self.load_tool,
                ComponentType.TERMINATION: self.load_termination
            }

            handler = handlers.get(config.component_type)
            if not handler:
                raise ValueError(
                    f"Unknown component type: {config.component_type}")

            return await handler(config)

        except Exception as e:
            logger.error(f"Failed to load component: {str(e)}")
            raise

    async def load_directory(self, directory: Union[str, Path], check_exists: bool = False, return_type: ReturnType = 'object') -> List[Union[Component, dict, ComponentConfig]]:
        """
        Import all component configurations from a directory.
        """
        components = []
        try:
            directory = Path(directory)
            # Using Path.iterdir() instead of os.listdir
            for path in list(directory.glob("*")):
                if path.suffix.lower().endswith(('.json', '.yaml', '.yml')):
                    try:
                        component = await self.load(path, return_type)
                        components.append(component)
                    except Exception as e:
                        logger.info(
                            f"Failed to load component: {str(e)}, {path}")

            return components
        except Exception as e:
            logger.info(f"Failed to load directory: {str(e)}")
            return components

    def _dict_to_config(self, config_dict: dict) -> ComponentConfig:
        """Convert dictionary to appropriate config type based on component_type"""
        if "component_type" not in config_dict:
            raise ValueError("component_type is required in configuration")

        config_types = {
            ComponentType.TEAM: TeamConfig,
            ComponentType.AGENT: AgentConfig,
            ComponentType.MODEL: ModelConfig,
            ComponentType.TOOL: ToolConfig,
            ComponentType.TERMINATION: TerminationConfig  # Add mapping for termination
        }

        component_type = ComponentType(config_dict["component_type"])
        config_class = config_types.get(component_type)

        if not config_class:
            raise ValueError(f"Unknown component type: {component_type}")

        return config_class(**config_dict)

    async def load_termination(self, config: TerminationConfig) -> TerminationComponent:
        """Create termination condition instance from configuration."""
        try:
            if config.termination_type == TerminationTypes.MAX_MESSAGES:
                return MaxMessageTermination(max_messages=config.max_messages)
            elif config.termination_type == TerminationTypes.STOP_MESSAGE:
                return StopMessageTermination()
            elif config.termination_type == TerminationTypes.TEXT_MENTION:
                if not config.text:
                    raise ValueError(
                        "text parameter required for TextMentionTermination")
                return TextMentionTermination(text=config.text)
            else:
                raise ValueError(
                    f"Unsupported termination type: {config.termination_type}")
        except Exception as e:
            logger.error(f"Failed to create termination condition: {str(e)}")
            raise ValueError(
                f"Termination condition creation failed: {str(e)}")

    async def load_team(self, config: TeamConfig) -> TeamComponent:
        """Create team instance from configuration."""

        default_selector_prompt = """You are in a role play game. The following roles are available:
{roles}.
Read the following conversation. Then select the next role from {participants} to play. Only return the role.

{history}

Read the above conversation. Then select the next role from {participants} to play. Only return the role.
"""
        try:
            # Load participants (agents)
            participants = []
            for participant in config.participants:
                agent = await self.load(participant)
                participants.append(agent)

            # Load model client if specified
            model_client = None
            if config.model_client:
                model_client = await self.load(config.model_client)

            # Load termination condition if specified
            termination = None
            if config.termination_condition:
                # Now we can use the universal load() method since termination is a proper component
                termination = await self.load(config.termination_condition)

            # Create team based on type
            if config.team_type == TeamTypes.ROUND_ROBIN:
                return RoundRobinGroupChat(
                    participants=participants,
                    termination_condition=termination
                )
            elif config.team_type == TeamTypes.SELECTOR:
                if not model_client:
                    raise ValueError(
                        "SelectorGroupChat requires a model_client")
                selector_prompt = config.selector_prompt if config.selector_prompt else default_selector_prompt
                return SelectorGroupChat(
                    participants=participants,
                    model_client=model_client,
                    termination_condition=termination,
                    selector_prompt=selector_prompt
                )
            else:
                raise ValueError(f"Unsupported team type: {config.team_type}")

        except Exception as e:
            logger.error(f"Failed to create team {config.name}: {str(e)}")
            raise ValueError(f"Team creation failed: {str(e)}")

    async def load_agent(self, config: AgentConfig) -> AgentComponent:
        """Create agent instance from configuration."""
        try:
            # Load model client if specified
            model_client = None
            if config.model_client:
                model_client = await self.load(config.model_client)
            system_message = config.system_message if config.system_message else "You are a helpful assistant"
            # Load tools if specified
            tools = []
            if config.tools:
                for tool_config in config.tools:
                    tool = await self.load(tool_config)
                    tools.append(tool)

            if config.agent_type == AgentTypes.ASSISTANT:
                return AssistantAgent(
                    name=config.name,
                    model_client=model_client,
                    tools=tools,
                    system_message=system_message
                )
            else:
                raise ValueError(
                    f"Unsupported agent type: {config.agent_type}")

        except Exception as e:
            logger.error(f"Failed to create agent {config.name}: {str(e)}")
            raise ValueError(f"Agent creation failed: {str(e)}")

    async def load_model(self, config: ModelConfig) -> ModelComponent:
        """Create model instance from configuration."""
        try:
            # Check cache first
            cache_key = str(config.model_dump())
            if cache_key in self._model_cache:
                logger.debug(f"Using cached model for {config.model}")
                return self._model_cache[cache_key]

            if config.model_type == ModelTypes.OPENAI:
                model = OpenAIChatCompletionClient(
                    model=config.model,
                    api_key=config.api_key,
                    base_url=config.base_url
                )
                self._model_cache[cache_key] = model
                return model
            else:
                raise ValueError(
                    f"Unsupported model type: {config.model_type}")

        except Exception as e:
            logger.error(f"Failed to create model {config.model}: {str(e)}")
            raise ValueError(f"Model creation failed: {str(e)}")

    async def load_tool(self, config: ToolConfig) -> ToolComponent:
        """Create tool instance from configuration."""
        try:
            # Validate required fields
            if not all([config.name, config.description, config.content, config.tool_type]):
                raise ValueError("Tool configuration missing required fields")

            # Check cache first
            cache_key = str(config.model_dump())
            if cache_key in self._tool_cache:
                logger.debug(f"Using cached tool '{config.name}'")
                return self._tool_cache[cache_key]

            if config.tool_type == ToolTypes.PYTHON_FUNCTION:
                tool = FunctionTool(
                    name=config.name,
                    description=config.description,
                    func=self._func_from_string(config.content)
                )
                self._tool_cache[cache_key] = tool
                return tool
            else:
                raise ValueError(f"Unsupported tool type: {config.tool_type}")

        except Exception as e:
            logger.error(f"Failed to create tool '{config.name}': {str(e)}")
            raise

    # Helper methods remain largely the same
    async def _load_from_file(self, path: Union[str, Path]) -> dict:
        """Load configuration from JSON or YAML file."""
        path = Path(path)
        if not path.exists():
            raise FileNotFoundError(f"Config file not found: {path}")

        try:
            with open(path) as f:
                if path.suffix == '.json':
                    return json.load(f)
                elif path.suffix in ('.yml', '.yaml'):
                    return yaml.safe_load(f)
                else:
                    raise ValueError(f"Unsupported file format: {path.suffix}")
        except Exception as e:
            raise ValueError(f"Failed to load file {path}: {str(e)}")

    def _func_from_string(self, content: str) -> callable:
        """Convert function string to callable."""
        try:
            namespace = {}
            exec(content, namespace)
            for item in namespace.values():
                if callable(item) and not isinstance(item, type):
                    return item
            raise ValueError("No function found in provided code")
        except Exception as e:
            raise ValueError(f"Failed to create function: {str(e)}")

    def _is_version_supported(self, component_type: ComponentType, ver: str) -> bool:
        """Check if version is supported for component type."""
        try:
            v = version.parse(ver)
            return ver in self.SUPPORTED_VERSIONS[component_type]
        except version.InvalidVersion:
            return False

    async def cleanup(self) -> None:
        """Cleanup resources and clear caches."""
        for model in self._model_cache.values():
            if hasattr(model, 'cleanup'):
                await model.cleanup()

        for tool in self._tool_cache.values():
            if hasattr(tool, 'cleanup'):
                await tool.cleanup()

        self._model_cache.clear()
        self._tool_cache.clear()
        self._last_cache_clear = datetime.now()
        logger.info("Cleared all component caches")