DB-GPT/dbgpt/agent/resource/manage.py

"""Resource manager."""

import logging
from collections import defaultdict
from typing import Any, Dict, List, Optional, Type, Union, cast

from dbgpt._private.pydantic import BaseModel, ConfigDict, model_validator
from dbgpt.component import BaseComponent, ComponentType, SystemApp
from dbgpt.util.parameter_utils import ParameterDescription

from .base import AgentResource, Resource, ResourceParameters, ResourceType
from .pack import ResourcePack
from .tool.pack import ToolResourceType, _is_function_tool, _to_tool_list

logger = logging.getLogger(__name__)


class RegisterResource(BaseModel):
    """Register resource model."""

    model_config = ConfigDict(arbitrary_types_allowed=True)

    name: Optional[str] = None
    resource_type: ResourceType
    resource_type_alias: Optional[str] = None
    resource_cls: Type[Resource]
    resource_instance: Optional[Resource] = None
    is_class: bool = True

    @property
    def key(self) -> str:
        """Return the key."""
        full_cls = f"{self.resource_cls.__module__}.{self.resource_cls.__qualname__}"
        name = self.name or full_cls
        resource_type_alias = self.resource_type_alias or self.resource_type.value
        return f"{resource_type_alias}:{name}"

    @property
    def type_unique_key(self) -> str:
        """Return the key."""
        resource_type_alias = self.resource_type_alias or self.resource_type.value
        return resource_type_alias

    @model_validator(mode="before")
    @classmethod
    def pre_fill(cls, values):
        """Pre-fill the model."""
        if not isinstance(values, dict):
            return values
        resource_instance = values.get("resource_instance")
        if resource_instance is not None:
            values["name"] = values["name"] or resource_instance.name
            values["is_class"] = False
            if not isinstance(resource_instance, Resource):
                raise ValueError(
                    f"resource_instance must be a Resource instance, not "
                    f"{type(resource_instance)}"
                )
        if not values.get("resource_type"):
            values["resource_type"] = values["resource_cls"].type()
        if not values.get("resource_type_alias"):
            values["resource_type_alias"] = values["resource_cls"].type_alias()
        return values

    def get_parameter_class(
        self,
        **kwargs,
    ) -> Type[ResourceParameters]:
        """Return the parameter description."""
        if self.is_class:
            return self.resource_cls.resource_parameters_class(**kwargs)
        return self.resource_instance.prefer_resource_parameters_class(**kwargs)  # type: ignore # noqa


class ResourceManager(BaseComponent):
    """Resource manager.

    To manage the resources.
    """

    name = ComponentType.RESOURCE_MANAGER

    def __init__(self, system_app: SystemApp):
        """Create a new AgentManager."""
        super().__init__(system_app)
        self.system_app = system_app
        self._resources: Dict[str, RegisterResource] = {}
        self._type_to_resources: Dict[str, List[RegisterResource]] = defaultdict(list)

    def init_app(self, system_app: SystemApp):
        """Initialize the AgentManager."""
        self.system_app = system_app

    def after_start(self):
        """Register all resources."""
        # TODO: Register some internal resources
        pass

    def register_resource(
        self,
        resource_cls: Optional[Type[Resource]] = None,
        resource_instance: Optional[Union[Resource, ToolResourceType]] = None,
        resource_type: Optional[ResourceType] = None,
        resource_type_alias: Optional[str] = None,
        ignore_duplicate: bool = False,
    ):
        """Register a resource."""
        if resource_instance and _is_function_tool(resource_instance):
            resource_instance = _to_tool_list(resource_instance)[0]  # type: ignore

        if resource_cls is None and resource_instance is None:
            raise ValueError("Resource class or instance must be provided.")
        name: Optional[str] = None
        if resource_instance is not None:
            resource_cls = resource_cls or type(resource_instance)  # type: ignore
            name = resource_instance.name  # type: ignore
        resource = RegisterResource(
            name=name,
            resource_cls=resource_cls,
            resource_instance=resource_instance,
            resource_type=resource_type,
            resource_type_alias=resource_type_alias,
        )
        if resource.key in self._resources:
            if ignore_duplicate:
                return
            else:
                raise ValueError(f"Resource {resource.key} already exists.")
        self._resources[resource.key] = resource
        self._type_to_resources[resource.type_unique_key].append(resource)

    def get_supported_resources_type(self) -> List[str]:
        """Get supported resources types."""
        unique_types = []
        for key, resource in self._resources.items():
            if resource.type_unique_key not in unique_types:
                unique_types.append(resource.type_unique_key)
        return unique_types

    def get_supported_resources(
        self, version: Optional[str] = None, type: Optional[str] = None, **kwargs
    ) -> Dict[str, Union[List[ParameterDescription], List[str]]]:
        """Return the resources."""
        results: Dict[str, Union[List[ParameterDescription], List[str]]] = defaultdict(
            list
        )
        for key, resource in self._resources.items():
            if type and type != resource.type_unique_key:
                continue
            parameter_class = resource.get_parameter_class(**kwargs)
            resource_type = resource.type_unique_key
            configs: Any = parameter_class.to_configurations(
                parameter_class,
                version=version,
            )
            if (
                version == "v1"
                and isinstance(configs, list)
                and len(configs) > 0
                and isinstance(configs[0], ParameterDescription)
            ):
                # v1, not compatible with class
                # set_configs_name = set(results[resource_type])
                set_configs: list = []
                if not resource.is_class:
                    for r in self._type_to_resources[resource_type]:
                        if not r.is_class:
                            set_configs.append(
                                {
                                    "label": r.resource_instance.name,  # type: ignore
                                    "key": r.resource_instance.name,  # type: ignore
                                    "description": r.resource_instance.description,  # type: ignore # noqa
                                }
                            )  # type: ignore
                configs = set_configs
            results[resource_type] = configs

        return results

    def build_resource_by_type(
        self,
        type_unique_key: str,
        agent_resource: AgentResource,
        version: Optional[str] = None,
    ) -> Resource:
        """Return the resource by type."""
        item = self._type_to_resources.get(type_unique_key)
        if not item:
            raise ValueError(f"Resource type {type_unique_key} not found.")
        inst_items = [i for i in item if not i.is_class]
        if inst_items:
            if version == "v1":
                for i in inst_items:
                    if (
                        i.resource_instance
                        and i.resource_instance.name == agent_resource.value
                    ):
                        return i.resource_instance
                raise ValueError(
                    f"Resource {agent_resource.value} not found in {type_unique_key}"
                )
            return cast(Resource, inst_items[0].resource_instance)
        elif len(inst_items) > 1:
            raise ValueError(
                f"Multiple instances of resource {type_unique_key} found, "
                f"please specify the resource name."
            )
        else:
            single_item = item[0]
            try:
                parameter_cls = single_item.get_parameter_class()
                param = parameter_cls.from_dict(agent_resource.to_dict())
                resource_inst = single_item.resource_cls(**param.to_dict())
                return resource_inst
            except Exception as e:
                logger.warning(f"Failed to build resource {single_item.key}: {str(e)}")
                raise ValueError(
                    f"Failed to build resource {single_item.key}: {str(e)}"
                )

    def build_resource(
        self,
        agent_resources: Optional[List[AgentResource]] = None,
        version: Optional[str] = None,
    ) -> Optional[Resource]:
        """Build a resource.

        If there is only one resource, return the resource instance, otherwise return a
        ResourcePack.

        Args:
            agent_resources: The agent resources.
            version: The resource version.

        Returns:
            Optional[Resource]: The resource instance.
        """
        if not agent_resources:
            return None
        dependencies: List[Resource] = []
        for resource in agent_resources:
            resource_inst = self.build_resource_by_type(
                resource.type, resource, version=version
            )
            dependencies.append(resource_inst)
        if len(dependencies) == 1:
            return dependencies[0]
        else:
            return ResourcePack(dependencies)


_SYSTEM_APP: Optional[SystemApp] = None


def initialize_resource(system_app: SystemApp):
    """Initialize the resource manager."""
    global _SYSTEM_APP
    _SYSTEM_APP = system_app
    resource_manager = ResourceManager(system_app)
    system_app.register_instance(resource_manager)


def get_resource_manager(system_app: Optional[SystemApp] = None) -> ResourceManager:
    """Return the resource manager."""
    if not _SYSTEM_APP:
        if not system_app:
            system_app = SystemApp()
        initialize_resource(system_app)
    app = system_app or _SYSTEM_APP
    return ResourceManager.get_instance(cast(SystemApp, app))