eveAI/common/utils/cache/config_cache.py

from typing import Dict, Any, Optional
from pathlib import Path
import yaml
from packaging import version
import os
from flask import current_app

from common.utils.cache.base import CacheHandler, CacheKey
from config.type_defs import agent_types, task_types, tool_types, specialist_types, retriever_types, prompt_types


def is_major_minor(version: str) -> bool:
    parts = version.strip('.').split('.')
    return len(parts) == 2 and all(part.isdigit() for part in parts)


class BaseConfigCacheHandler(CacheHandler[Dict[str, Any]]):
    """Base handler for configuration caching"""

    def __init__(self, region, config_type: str):
        """
        Args:
            region: Cache region
            config_type: Type of configuration (agents, tasks, etc.)
        """
        super().__init__(region, f'config_{config_type}')
        self.config_type = config_type
        self._types_module = None  # Set by subclasses
        self._config_dir = None  # Set by subclasses
        self.version_tree_cache = None
        self.configure_keys('type_name', 'version')

    def _to_cache_data(self, instance: Dict[str, Any]) -> Dict[str, Any]:
        """Convert the data to a cacheable format"""
        # For configuration data, we can just return the dictionary as is
        # since it's already in a serializable format
        return instance

    def _from_cache_data(self, data: Dict[str, Any], **kwargs) -> Dict[str, Any]:
        """Convert cached data back to usable format"""
        # Similarly, we can return the data directly since it's already
        # in the format we need
        return data

    def _should_cache(self, value: Dict[str, Any]) -> bool:
        """
        Validate if the value should be cached

        Args:
            value: The value to be cached

        Returns:
            bool: True if the value should be cached
        """
        return isinstance(value, dict)  # Cache all dictionaries

    def set_version_tree_cache(self, cache):
        """Set the version tree cache dependency."""
        self.version_tree_cache = cache

    def _load_specific_config(self, type_name: str, version_str: str) -> Dict[str, Any]:
        """
        Load a specific configuration version

        Args:
            type_name: Type name
            version_str: Version string

        Returns:
            Configuration data
        """
        current_app.logger.debug(f"Loading specific configuration for {type_name}, version: {version_str} - no cache")
        version_tree = self.version_tree_cache.get_versions(type_name)
        versions = version_tree['versions']
        current_app.logger.debug(f"Loaded specific versions for {type_name}, versions: {versions}")

        if version_str == 'latest':
            version_str = version_tree['latest_version']

        if version_str not in versions:
            raise ValueError(f"Version {version_str} not found for {type_name}")

        file_path = versions[version_str]['file_path']
        current_app.logger.debug(f'Trying to load configuration from {file_path}')

        try:
            with open(file_path) as f:
                config = yaml.safe_load(f)
                current_app.logger.debug(f"Loaded config for {type_name}{version_str}: {config}")
                return config
        except Exception as e:
            raise ValueError(f"Error loading config from {file_path}: {e}")

    def get_config(self, type_name: str, version: Optional[str] = None) -> Dict[str, Any]:
        """
        Get configuration for a specific type and version
        If version not specified, returns latest

        Args:
            type_name: Configuration type name
            version: Optional specific version to retrieve

        Returns:
            Configuration data
        """
        current_app.logger.debug(f"Trying to retrieve config for {self.config_type}, type name: {type_name}, "
                                 f"version: {version}")
        if version is None:
            version_str = self.version_tree_cache.get_latest_version(type_name)
        elif is_major_minor(version):
            version_str = self.version_tree_cache.get_latest_patch_version(type_name, version)
        else:
            version_str = version

        result = self.get(
            lambda type_name, version: self._load_specific_config(type_name, version),
            type_name=type_name,
            version=version_str
        )
        return result


class BaseConfigVersionTreeCacheHandler(CacheHandler[Dict[str, Any]]):
    """Base handler for configuration version tree caching"""

    def __init__(self, region, config_type: str):
        """
        Args:
            region: Cache region
            config_type: Type of configuration (agents, tasks, etc.)
        """
        super().__init__(region, f'config_{config_type}_version_tree')
        self.config_type = config_type
        self._types_module = None  # Set by subclasses
        self._config_dir = None  # Set by subclasses
        self.configure_keys('type_name')

    def _load_version_tree(self, type_name: str) -> Dict[str, Any]:
        """
        Load version tree for a specific type without loading full configurations

        Args:
            type_name: Name of configuration type

        Returns:
            Dict containing available versions and their metadata
        """
        current_app.logger.debug(f"Loading version tree for {type_name} - no cache")
        type_path = Path(self._config_dir) / type_name
        if not type_path.exists():
            raise ValueError(f"No configuration found for type {type_name}")

        version_files = list(type_path.glob('*.yaml'))
        if not version_files:
            raise ValueError(f"No versions found for type {type_name}")

        versions = {}
        latest_version = None
        latest_version_obj = None

        for file_path in version_files:
            ver = file_path.stem  # Get version from filename
            try:
                ver_obj = version.parse(ver)
                # Only load minimal metadata for version tree
                with open(file_path) as f:
                    yaml_data = yaml.safe_load(f)
                    metadata = yaml_data.get('metadata', {})
                    versions[ver] = {
                        'metadata': metadata,
                        'file_path': str(file_path)
                    }

                # Track latest version
                if latest_version_obj is None or ver_obj > latest_version_obj:
                    latest_version = ver
                    latest_version_obj = ver_obj

            except Exception as e:
                current_app.logger.error(f"Error loading version {ver}: {e}")
                continue

        current_app.logger.debug(f"Loaded versions for {type_name}: {versions}")
        current_app.logger.debug(f"Latest version for {type_name}: {latest_version}")
        return {
            'versions': versions,
            'latest_version': latest_version
        }

    def _to_cache_data(self, instance: Dict[str, Any]) -> Dict[str, Any]:
        """Convert the data to a cacheable format"""
        # For configuration data, we can just return the dictionary as is
        # since it's already in a serializable format
        return instance

    def _from_cache_data(self, data: Dict[str, Any], **kwargs) -> Dict[str, Any]:
        """Convert cached data back to usable format"""
        # Similarly, we can return the data directly since it's already
        # in the format we need
        return data

    def _should_cache(self, value: Dict[str, Any]) -> bool:
        """
        Validate if the value should be cached

        Args:
            value: The value to be cached

        Returns:
            bool: True if the value should be cached
        """
        return isinstance(value, dict)  # Cache all dictionaries

    def get_versions(self, type_name: str) -> Dict[str, Any]:
        """
        Get version tree for a type

        Args:
            type_name: Type to get versions for

        Returns:
            Dict with version information
        """
        current_app.logger.debug(f"Trying to get version tree for {self.config_type}, {type_name}")
        return self.get(
            lambda type_name: self._load_version_tree(type_name),
            type_name=type_name
        )

    def get_latest_version(self, type_name: str) -> str:
        """
        Get the latest version for a given type name.

        Args:
            type_name: Name of the configuration type

        Returns:
            Latest version string

        Raises:
            ValueError: If type not found or no versions available
        """
        version_tree = self.get_versions(type_name)
        if not version_tree or 'latest_version' not in version_tree:
            raise ValueError(f"No versions found for {type_name}")

        return version_tree['latest_version']

    def get_latest_patch_version(self, type_name: str, major_minor: str) -> str:
        """
        Get the latest patch version for a given major.minor version.

        Args:
            type_name: Name of the configuration type
            major_minor: Major.minor version (e.g. "1.0")

        Returns:
            Latest patch version string (e.g. "1.0.3")

        Raises:
            ValueError: If type not found or no matching versions
        """
        version_tree = self.get_versions(type_name)
        if not version_tree or 'versions' not in version_tree:
            raise ValueError(f"No versions found for {type_name}")

        # Filter versions that match the major.minor prefix
        matching_versions = [
            ver for ver in version_tree['versions'].keys()
            if ver.startswith(major_minor + '.')
        ]

        if not matching_versions:
            raise ValueError(f"No versions found for {type_name} with prefix {major_minor}")

        # Return highest matching version
        latest_patch = max(matching_versions, key=version.parse)
        return latest_patch


class BaseConfigTypesCacheHandler(CacheHandler[Dict[str, Any]]):
    """Base handler for configuration types caching"""

    def __init__(self, region, config_type: str):
        """
        Args:
            region: Cache region
            config_type: Type of configuration (agents, tasks, etc.)
        """
        super().__init__(region, f'config_{config_type}_types')
        self.config_type = config_type
        self._types_module = None  # Set by subclasses
        self._config_dir = None  # Set by subclasses
        self.configure_keys()

    def _to_cache_data(self, instance: Dict[str, Any]) -> Dict[str, Any]:
        """Convert the data to a cacheable format"""
        # For configuration data, we can just return the dictionary as is
        # since it's already in a serializable format
        return instance

    def _from_cache_data(self, data: Dict[str, Any], **kwargs) -> Dict[str, Any]:
        """Convert cached data back to usable format"""
        # Similarly, we can return the data directly since it's already
        # in the format we need
        return data

    def _should_cache(self, value: Dict[str, Any]) -> bool:
        """
        Validate if the value should be cached

        Args:
            value: The value to be cached

        Returns:
            bool: True if the value should be cached
        """
        return isinstance(value, dict)  # Cache all dictionaries

    def _load_type_definitions(self) -> Dict[str, Dict[str, str]]:
        """Load type definitions from the corresponding type_defs module"""
        current_app.logger.debug(f"Loading type definitions for {self.config_type} - no cache")
        if not self._types_module:
            raise ValueError("_types_module must be set by subclass")

        type_definitions = {
            type_id: {
                'name': info['name'],
                'description': info['description']
            }
            for type_id, info in self._types_module.items()
        }

        current_app.logger.debug(f"Loaded type definitions: {type_definitions}")

        return type_definitions

    def get_types(self) -> Dict[str, Dict[str, str]]:
        """Get dictionary of available types with name and description"""
        current_app.logger.debug(f"Trying to retrieve type definitions for {self.config_type}")
        result = self.get(
            lambda type_name: self._load_type_definitions(),
            type_name=f'{self.config_type}_types',
        )
        return result


def create_config_cache_handlers(config_type: str, config_dir: str, types_module: dict) -> tuple:
    """
    Factory function to dynamically create the 3 cache handler classes for a given configuration type.
    The following cache names are created:
    - <config_type>_config_cache
    - <config_type>_version_tree_cache
    - <config_type>_types_cache


    Args:
        config_type: The configuration type (e.g., 'agents', 'tasks').
        config_dir: The directory where configuration files are stored.
        types_module: The types module defining the available types for this config.

    Returns:
        A tuple of dynamically created classes for config, version tree, and types handlers.
    """

    class ConfigCacheHandler(BaseConfigCacheHandler):
        handler_name = f"{config_type}_config_cache"

        def __init__(self, region):
            super().__init__(region, config_type)
            self._types_module = types_module
            self._config_dir = config_dir

    class VersionTreeCacheHandler(BaseConfigVersionTreeCacheHandler):
        handler_name = f"{config_type}_version_tree_cache"

        def __init__(self, region):
            super().__init__(region, config_type)
            self._types_module = types_module
            self._config_dir = config_dir

    class TypesCacheHandler(BaseConfigTypesCacheHandler):
        handler_name = f"{config_type}_types_cache"

        def __init__(self, region):
            super().__init__(region, config_type)
            self._types_module = types_module
            self._config_dir = config_dir

    return ConfigCacheHandler, VersionTreeCacheHandler, TypesCacheHandler


AgentConfigCacheHandler, AgentConfigVersionTreeCacheHandler, AgentConfigTypesCacheHandler = (
    create_config_cache_handlers(
        config_type='agents',
        config_dir='config/agents',
        types_module=agent_types.AGENT_TYPES
    ))


TaskConfigCacheHandler, TaskConfigVersionTreeCacheHandler, TaskConfigTypesCacheHandler = (
    create_config_cache_handlers(
        config_type='tasks',
        config_dir='config/tasks',
        types_module=task_types.TASK_TYPES
    ))


ToolConfigCacheHandler, ToolConfigVersionTreeCacheHandler, ToolConfigTypesCacheHandler = (
    create_config_cache_handlers(
        config_type='tools',
        config_dir='config/tools',
        types_module=tool_types.TOOL_TYPES
    ))


SpecialistConfigCacheHandler, SpecialistConfigVersionTreeCacheHandler, SpecialistConfigTypesCacheHandler = (
    create_config_cache_handlers(
        config_type='specialists',
        config_dir='config/specialists',
        types_module=specialist_types.SPECIALIST_TYPES
    ))


RetrieverConfigCacheHandler, RetrieverConfigVersionTreeCacheHandler, RetrieverConfigTypesCacheHandler = (
    create_config_cache_handlers(
        config_type='retrievers',
        config_dir='config/retrievers',
        types_module=retriever_types.RETRIEVER_TYPES
    ))


PromptConfigCacheHandler, PromptConfigVersionTreeCacheHandler, PromptConfigTypesCacheHandler = (
    create_config_cache_handlers(
        config_type='prompts',
        config_dir='config/prompts',
        types_module=prompt_types.PROMPT_TYPES

    ))


def register_config_cache_handlers(cache_manager) -> None:
    cache_manager.register_handler(AgentConfigCacheHandler, 'eveai_config')
    cache_manager.register_handler(AgentConfigTypesCacheHandler, 'eveai_config')
    cache_manager.register_handler(AgentConfigVersionTreeCacheHandler, 'eveai_config')
    cache_manager.register_handler(TaskConfigCacheHandler, 'eveai_config')
    cache_manager.register_handler(TaskConfigTypesCacheHandler, 'eveai_config')
    cache_manager.register_handler(TaskConfigVersionTreeCacheHandler, 'eveai_config')
    cache_manager.register_handler(ToolConfigCacheHandler, 'eveai_config')
    cache_manager.register_handler(ToolConfigTypesCacheHandler, 'eveai_config')
    cache_manager.register_handler(ToolConfigVersionTreeCacheHandler, 'eveai_config')
    cache_manager.register_handler(SpecialistConfigCacheHandler, 'eveai_config')
    cache_manager.register_handler(SpecialistConfigTypesCacheHandler, 'eveai_config')
    cache_manager.register_handler(SpecialistConfigVersionTreeCacheHandler, 'eveai_config')
    cache_manager.register_handler(RetrieverConfigCacheHandler, 'eveai_config')
    cache_manager.register_handler(RetrieverConfigTypesCacheHandler, 'eveai_config')
    cache_manager.register_handler(RetrieverConfigVersionTreeCacheHandler, 'eveai_config')
    cache_manager.register_handler(PromptConfigCacheHandler, 'eveai_config')
    cache_manager.register_handler(PromptConfigVersionTreeCacheHandler, 'eveai_config')
    cache_manager.register_handler(PromptConfigTypesCacheHandler, 'eveai_config')

    cache_manager.agents_config_cache.set_version_tree_cache(cache_manager.agents_version_tree_cache)
    cache_manager.tasks_config_cache.set_version_tree_cache(cache_manager.tasks_version_tree_cache)
    cache_manager.tools_config_cache.set_version_tree_cache(cache_manager.tools_version_tree_cache)
    cache_manager.specialists_config_cache.set_version_tree_cache(cache_manager.specialists_version_tree_cache)
    cache_manager.retrievers_config_cache.set_version_tree_cache(cache_manager.retrievers_version_tree_cache)
    cache_manager.prompts_config_cache.set_version_tree_cache(cache_manager.prompts_version_tree_cache)