feat(agent): Release agent SDK

.mypy.ini

@@ -5,9 +5,6 @@ exclude = /tests/
 [mypy-dbgpt.app.*]
 follow_imports = skip
 
-[mypy-dbgpt.agent.*]
-follow_imports = skip
-
 [mypy-dbgpt.serve.*]
 follow_imports = skip
 
@@ -89,4 +86,11 @@ ignore_missing_imports = True
 ignore_missing_imports = True
 
 [mypy-fastchat.protocol.api_protocol]
-ignore_missing_imports = True
+ignore_missing_imports = True
+
+# Agent
+[mypy-seaborn.*]
+ignore_missing_imports = True
+
+[mypy-unstructured.*]
+ignore_missing_imports = True

Makefile

@@ -47,7 +47,7 @@ fmt: setup ## Format Python code
 	$(VENV_BIN)/blackdoc examples
 	# TODO: Use flake8 to enforce Python style guide.
 	# https://flake8.pycqa.org/en/latest/
-	$(VENV_BIN)/flake8 dbgpt/core/ dbgpt/rag/ dbgpt/storage/ dbgpt/datasource/ dbgpt/client/
+	$(VENV_BIN)/flake8 dbgpt/core/ dbgpt/rag/ dbgpt/storage/ dbgpt/datasource/ dbgpt/client/ dbgpt/agent/ dbgpt/vis/
 	# TODO: More package checks with flake8.
 
 .PHONY: fmt-check
@@ -56,7 +56,7 @@ fmt-check: setup ## Check Python code formatting and style without making change
 	$(VENV_BIN)/isort --check-only --extend-skip="examples/notebook" examples
 	$(VENV_BIN)/black --check --extend-exclude="examples/notebook" .
 	$(VENV_BIN)/blackdoc --check dbgpt examples
-	$(VENV_BIN)/flake8 dbgpt/core/ dbgpt/rag/ dbgpt/storage/ dbgpt/datasource/ dbgpt/client/
+	$(VENV_BIN)/flake8 dbgpt/core/ dbgpt/rag/ dbgpt/storage/ dbgpt/datasource/ dbgpt/client/ dbgpt/agent/ dbgpt/vis/
 
 .PHONY: pre-commit
 pre-commit: fmt-check test test-doc mypy ## Run formatting and unit tests before committing
@@ -72,7 +72,7 @@ test-doc: $(VENV)/.testenv ## Run doctests
 .PHONY: mypy
 mypy: $(VENV)/.testenv ## Run mypy checks
 	# https://github.com/python/mypy
-	$(VENV_BIN)/mypy --config-file .mypy.ini dbgpt/rag/ dbgpt/datasource/ dbgpt/client/
+	$(VENV_BIN)/mypy --config-file .mypy.ini dbgpt/rag/ dbgpt/datasource/ dbgpt/client/ dbgpt/agent/ dbgpt/vis/
 	# rag depends on core and storage, so we not need to check it again.
 	# $(VENV_BIN)/mypy --config-file .mypy.ini dbgpt/storage/
 	# $(VENV_BIN)/mypy --config-file .mypy.ini dbgpt/core/

dbgpt/_private/config.py

@@ -10,6 +10,7 @@
 if TYPE_CHECKING:
     from auto_gpt_plugin_template import AutoGPTPluginTemplate
 
+    from dbgpt.agent.plugin import CommandRegistry
     from dbgpt.component import SystemApp
     from dbgpt.datasource.manages import ConnectorManager
 
@@ -145,7 +146,7 @@ def __init__(self) -> None:
 
         self.prompt_template_registry = PromptTemplateRegistry()
         ### Related configuration of built-in commands
-        self.command_registry = []  # type: ignore
+        self.command_registry: Optional[CommandRegistry] = None
 
         disabled_command_categories = os.getenv("DISABLED_COMMAND_CATEGORIES")
         if disabled_command_categories:

dbgpt/_private/pydantic.py

@@ -4,6 +4,7 @@
     PYDANTIC_VERSION = 1
     from pydantic import (
         BaseModel,
+        ConfigDict,
         Extra,
         Field,
         NonNegativeFloat,
@@ -20,6 +21,7 @@
     # pydantic 2.x
     from pydantic.v1 import (
         BaseModel,
+        ConfigDict,
         Extra,
         Field,
         NonNegativeFloat,

dbgpt/agent/__init__.py

@@ -1 +1,33 @@
-from .common.schema import PluginStorageType
+"""DB-GPT Multi-Agents Module."""
+
+from .actions.action import Action, ActionOutput  # noqa: F401
+from .core.agent import (  # noqa: F401
+    Agent,
+    AgentContext,
+    AgentGenerateContext,
+    AgentMessage,
+)
+from .core.base_agent import ConversableAgent  # noqa: F401
+from .core.llm.llm import LLMConfig  # noqa: F401
+from .core.schema import PluginStorageType  # noqa: F401
+from .core.user_proxy_agent import UserProxyAgent  # noqa: F401
+from .memory.gpts_memory import GptsMemory  # noqa: F401
+from .resource.resource_api import AgentResource, ResourceType  # noqa: F401
+from .resource.resource_loader import ResourceLoader  # noqa: F401
+
+__ALL__ = [
+    "Agent",
+    "AgentContext",
+    "AgentGenerateContext",
+    "AgentMessage",
+    "ConversableAgent",
+    "Action",
+    "ActionOutput",
+    "LLMConfig",
+    "GptsMemory",
+    "AgentResource",
+    "ResourceType",
+    "ResourceLoader",
+    "PluginStorageType",
+    "UserProxyAgent",
+]

dbgpt/agent/actions/__init__.py

@@ -0,0 +1 @@
+"""Actions of Agent."""

dbgpt/agent/actions/action.py

@@ -1,74 +1,91 @@
+"""Base Action class for defining agent actions."""
 import json
-from abc import ABC
+from abc import ABC, abstractmethod
 from typing import (
     Any,
     Dict,
     Generic,
     List,
     Optional,
-    Tuple,
     Type,
     TypeVar,
     Union,
+    cast,
     get_args,
     get_origin,
 )
 
-from pydantic import BaseModel
-
-from dbgpt.agent.resource.resource_api import AgentResource, ResourceType
-from dbgpt.agent.resource.resource_loader import ResourceLoader
+from dbgpt._private.pydantic import BaseModel
 from dbgpt.util.json_utils import find_json_objects
 
-T = TypeVar("T", None, BaseModel, List[BaseModel])
+from ...vis.base import Vis
+from ..resource.resource_api import AgentResource, ResourceType
+from ..resource.resource_loader import ResourceLoader
+
+T = TypeVar("T", bound=Union[BaseModel, List[BaseModel], None])
+
+JsonMessageType = Union[Dict[str, Any], List[Dict[str, Any]]]
 
 
 class ActionOutput(BaseModel):
+    """Action output model."""
+
     content: str
     is_exe_success: bool = True
-    view: str = None
+    view: Optional[str] = None
     resource_type: Optional[str] = None
     resource_value: Optional[Any] = None
 
-    @staticmethod
-    def from_dict(param: Optional[Dict]):
+    @classmethod
+    def from_dict(
+        cls: Type["ActionOutput"], param: Optional[Dict]
+    ) -> Optional["ActionOutput"]:
+        """Convert dict to ActionOutput object."""
         if not param:
             return None
-        return ActionOutput.parse_obj(param)
+        return cls.parse_obj(param)
 
 
 class Action(ABC, Generic[T]):
+    """Base Action class for defining agent actions."""
+
     def __init__(self):
-        self.resource_loader: ResourceLoader = None
+        """Create an action."""
+        self.resource_loader: Optional[ResourceLoader] = None
 
-    def init_resource_loader(self, resource_loader: ResourceLoader):
+    def init_resource_loader(self, resource_loader: Optional[ResourceLoader]):
+        """Initialize the resource loader."""
         self.resource_loader = resource_loader
 
     @property
     def resource_need(self) -> Optional[ResourceType]:
+        """Return the resource type needed for the action."""
         return None
 
     @property
-    def render_protocal(self):
-        raise NotImplementedError("The run method should be implemented in a subclass.")
+    def render_protocol(self) -> Optional[Vis]:
+        """Return the render protocol."""
+        return None
 
-    def render_prompt(self):
-        if self.render_protocal is None:
+    def render_prompt(self) -> Optional[str]:
+        """Return the render prompt."""
+        if self.render_protocol is None:
             return None
         else:
-            return self.render_protocal.render_prompt()
+            return self.render_protocol.render_prompt()
 
     def _create_example(
         self,
-        model_type: Union[Type[BaseModel], Type[List[BaseModel]]],
-    ) -> Union[Dict[str, Any], List[Dict[str, Any]]]:
+        model_type: Union[Type[BaseModel], List[Type[BaseModel]]],
+    ) -> Optional[Union[Dict[str, Any], List[Dict[str, Any]]]]:
         if model_type is None:
             return None
         origin = get_origin(model_type)
         args = get_args(model_type)
         if origin is None:
             example = {}
-            for field_name, field in model_type.__fields__.items():
+            single_model_type = cast(Type[BaseModel], model_type)
+            for field_name, field in single_model_type.__fields__.items():
                 field_info = field.field_info
                 if field_info.description:
                     example[field_name] = field_info.description
@@ -78,9 +95,11 @@ def _create_example(
                     example[field_name] = ""
             return example
         elif origin is list or origin is List:
-            element_type = args[0]
+            element_type = cast(Type[BaseModel], args[0])
             if issubclass(element_type, BaseModel):
-                return [self._create_example(element_type)]
+                list_example = self._create_example(element_type)
+                typed_list_example = cast(Dict[str, Any], list_example)
+                return [typed_list_example]
             else:
                 raise TypeError("List elements must be BaseModel subclasses")
         else:
@@ -89,40 +108,48 @@ def _create_example(
             )
 
     @property
-    def out_model_type(self) -> T:
+    def out_model_type(self) -> Optional[Union[Type[T], List[Type[T]]]]:
+        """Return the output model type."""
         return None
 
     @property
-    def ai_out_schema(self) -> Union[Dict[str, Any], List[Dict[str, Any]]]:
+    def ai_out_schema(self) -> Optional[str]:
+        """Return the AI output schema."""
         if self.out_model_type is None:
             return None
 
+        json_format_data = json.dumps(
+            self._create_example(self.out_model_type), indent=2, ensure_ascii=False
+        )
         return f"""Please response in the following json format:
-            {json.dumps(self._create_example(self.out_model_type), indent=2, ensure_ascii=False)}
-        Make sure the response is correct json and can be parsed by Python json.loads. 
+            {json_format_data}
+        Make sure the response is correct json and can be parsed by Python json.loads.
         """
 
-    def _ai_mesage_2_json(self, ai_message: str) -> json:
+    def _ai_message_2_json(self, ai_message: str) -> JsonMessageType:
         json_objects = find_json_objects(ai_message)
         json_count = len(json_objects)
         if json_count != 1:
             raise ValueError("Unable to obtain valid output.")
         return json_objects[0]
 
-    def _input_convert(self, ai_message: str, cls: Type[T]) -> Union[T, List[T]]:
-        json_result = self._ai_mesage_2_json(ai_message)
+    def _input_convert(self, ai_message: str, cls: Type[T]) -> T:
+        json_result = self._ai_message_2_json(ai_message)
         if get_origin(cls) == list:
             inner_type = get_args(cls)[0]
-            return [inner_type.parse_obj(item) for item in json_result]
+            typed_cls = cast(Type[BaseModel], inner_type)
+            return [typed_cls.parse_obj(item) for item in json_result]  # type: ignore
         else:
-            return cls.parse_obj(json_result)
+            typed_cls = cast(Type[BaseModel], cls)
+            return typed_cls.parse_obj(json_result)
 
-    async def a_run(
+    @abstractmethod
+    async def run(
         self,
         ai_message: str,
         resource: Optional[AgentResource] = None,
         rely_action_out: Optional[ActionOutput] = None,
         need_vis_render: bool = True,
         **kwargs,
     ) -> ActionOutput:
-        raise NotImplementedError("The run method should be implemented in a subclass.")
+        """Perform the action."""

dbgpt/agent/actions/blank_action.py

@@ -1,40 +1,36 @@
-import json
-import logging
-from typing import Any, Dict, List, Optional, Union
+"""Blank Action for the Agent."""
 
-from pydantic import BaseModel, Field
+import logging
+from typing import Optional
 
-from dbgpt.agent.actions.action import Action, ActionOutput, T
-from dbgpt.agent.common.schema import Status
-from dbgpt.agent.plugin.generator import PluginPromptGenerator
-from dbgpt.agent.resource.resource_api import AgentResource, ResourceType
-from dbgpt.agent.resource.resource_plugin_api import ResourcePluginClient
-from dbgpt.vis.tags.vis_plugin import Vis, VisPlugin
+from ..resource.resource_api import AgentResource
+from .action import Action, ActionOutput
 
 logger = logging.getLogger(__name__)
 
 
 class BlankAction(Action):
-    def __init__(self, **kwargs):
-        super().__init__(**kwargs)
-
-    @property
-    def resource_need(self) -> Optional[ResourceType]:
-        return None
+    """Blank action class."""
 
-    @property
-    def render_protocal(self) -> Optional[Vis]:
-        return None
+    def __init__(self):
+        """Create a blank action."""
+        super().__init__()
 
     @property
-    def ai_out_schema(self) -> Union[Dict[str, Any], List[Dict[str, Any]]]:
+    def ai_out_schema(self) -> Optional[str]:
+        """Return the AI output schema."""
         return None
 
-    async def a_run(
+    async def run(
         self,
         ai_message: str,
         resource: Optional[AgentResource] = None,
         rely_action_out: Optional[ActionOutput] = None,
         need_vis_render: bool = True,
+        **kwargs,
     ) -> ActionOutput:
+        """Perform the action.
+
+        Just return the AI message.
+        """
         return ActionOutput(is_exe_success=True, content=ai_message, view=ai_message)