"""Resources for the agent.""" import dataclasses import json from abc import ABC, abstractmethod from enum import Enum from typing import Any, Dict, Generic, List, Optional, Tuple, Type, TypeVar, cast from pydantic import field_validator from dbgpt._private.pydantic import BaseModel, model_to_dict from dbgpt.core import Chunk from dbgpt.util.parameter_utils import BaseParameters, _get_parameter_descriptions P = TypeVar("P", bound="ResourceParameters") T = TypeVar("T", bound="Resource") class ResourceType(str, Enum): """Resource type enumeration.""" DB = "database" Knowledge = "knowledge" Internet = "internet" Tool = "tool" Plugin = "plugin" TextFile = "text_file" ExcelFile = "excel_file" ImageFile = "image_file" AWELFlow = "awel_flow" # Resource type for resource pack Pack = "pack" @dataclasses.dataclass class ResourceParameters(BaseParameters): """Resource parameters class. It defines the parameters for building a resource. """ name: str = dataclasses.field(metadata={"help": "Resource name", "tags": "fixed"}) @classmethod def _resource_version(cls) -> str: """Return the resource version.""" return "v2" @classmethod def to_configurations( cls, parameters: Type["ResourceParameters"], version: Optional[str] = None, ) -> Any: """Convert the parameters to configurations.""" return _get_parameter_descriptions(parameters) class Resource(ABC, Generic[P]): """Resource for the agent.""" @classmethod @abstractmethod def type(cls) -> ResourceType: """Return the resource type.""" @classmethod def type_alias(cls) -> str: """Return the resource type alias.""" return cls.type().value @property @abstractmethod def name(self) -> str: """Return the resource name.""" @classmethod def resource_parameters_class(cls, **kwargs) -> Type[P]: """Return the parameters class.""" return ResourceParameters def prefer_resource_parameters_class(self, **kwargs) -> Type[P]: """Return the parameters class. You can override this method to return a different parameters class. It will be used to initialize the resource with parameters. """ return self.resource_parameters_class(**kwargs) def initialize_with_parameters(self, resource_parameters: P): """Initialize the resource with parameters.""" pass def preload_resource(self): """Preload the resource.""" pass @classmethod def from_resource( cls: Type[T], resource: Optional["Resource"], expected_type: Optional[ResourceType] = None, ) -> List[T]: """Create a resource from another resource. Another resource can be a pack or a single resource, if it is a pack, it will return all resources which type is the same as the current resource. Args: resource(Resource): The resource. expected_type(ResourceType): The expected resource type. Returns: List[Resource]: The resources. """ if not resource: return [] typed_resources = [] for r in resource.get_resource_by_type(expected_type or cls.type()): typed_resources.append(cast(T, r)) return typed_resources async def get_resources_info( self, *, lang: str = "en", prompt_type: str = "default", question: Optional[str] = None, resources: Optional[List] = None, **kwargs, ): """Get prompts for multiple resources at the same time. Args: lang(str): The language. prompt_type(str): The prompt type. question(str): The question. resource_name(str): The resource name, just for the pack, it will be used to select specific resource in the pack. """ @abstractmethod async def get_prompt( self, *, lang: str = "en", prompt_type: str = "default", question: Optional[str] = None, resource_name: Optional[str] = None, **kwargs, ) -> Tuple[str, Optional[Dict]]: """Get the prompt. Args: lang(str): The language. prompt_type(str): The prompt type. question(str): The question. resource_name(str): The resource name, just for the pack, it will be used to select specific resource in the pack. """ async def get_resources( self, lang: str = "en", prompt_type: str = "default", question: Optional[str] = None, resource_name: Optional[str] = None, ) -> Tuple[Optional[List[Chunk]], str, Optional[Dict]]: """Get the resources.""" raise NotImplementedError def execute(self, *args, resource_name: Optional[str] = None, **kwargs) -> Any: """Execute the resource.""" raise NotImplementedError async def async_execute( self, *args, resource_name: Optional[str] = None, **kwargs ) -> Any: """Execute the resource asynchronously.""" raise NotImplementedError @property def is_async(self) -> bool: """Return whether the resource is asynchronous.""" return False @property def is_pack(self) -> bool: """Return whether the resource is a pack.""" return False @property def sub_resources(self) -> List["Resource"]: """Return the resources.""" if not self.is_pack: raise ValueError("The resource is not a pack, no sub-resources.") return [] def get_resource_by_type(self, resource_type: ResourceType) -> List["Resource"]: """Get resources by type. If the resource is a pack, it will search the sub-resources. Otherwise, it will return itself if the type matches. Args: resource_type(ResourceType): The resource type. Returns: List[Resource]: The resources. """ if not self.is_pack: if self.type() == resource_type: return [self] else: return [] resources = [] for resource in self.sub_resources: if resource.type() == resource_type: resources.append(resource) return resources class AgentResource(BaseModel): """Agent resource class.""" type: str value: str name: Optional[str] = None is_dynamic: bool = ( False # Is the current resource predefined or dynamically passed in? ) context: Optional[dict] = None def resource_prompt_template(self, **kwargs) -> str: """Get the resource prompt template.""" return "{data_type} --{data_introduce}" @field_validator("value", mode="before") def parse_value(cls, value): """Parse value.""" return str(value) @staticmethod def from_dict(d: Dict[str, Any]) -> Optional["AgentResource"]: """Create an AgentResource object from a dictionary.""" if d is None: return None return AgentResource( type=d.get("type"), name=d.get("name"), introduce=d.get("introduce"), value=d.get("value", None), is_dynamic=d.get("is_dynamic", False), context=d.get("context", None), ) @staticmethod def from_json_list_str(d: Optional[str]) -> Optional[List["AgentResource"]]: """Create a list of AgentResource objects from a json string.""" if d is None: return None try: json_array = json.loads(d) except Exception: raise ValueError(f"Illegal AgentResource json string!{d}") if not isinstance(json_array, list): raise ValueError(f"Illegal AgentResource json string!{d}") json_list = [] for item in json_array: r = AgentResource.from_dict(item) if r: json_list.append(r) return json_list def to_dict(self) -> Dict[str, Any]: """Convert the AgentResource object to a dictionary.""" temp = model_to_dict(self) for field, value in temp.items(): if isinstance(value, Enum): temp[field] = value.value return temp