简介

文章主要介绍了如何扩展一个自定义Agent，这里是用官方提供的总结摘要的Agent做了个示例，先给大家看下显示效果

代码目录

博主将代码放在core目录了，后续经过对源码的解读感觉放在dbgpt_serve.agent.agents.expand目录下可能更合适，大家自行把控即可

代码详情

summarizer_action.py

from typing import Optional

from pydantic import BaseModel, Field

from dbgpt.vis import Vis

from dbgpt.agent import Action, ActionOutput, AgentResource, ResourceType

from dbgpt.agent.util import cmp_string_equal

NOT_RELATED_MESSAGE = "Did not find the information you want."

# The parameter object that the Action that the current Agent needs to execute needs to output.

class SummaryActionInput(BaseModel):

    summary: str = Field(

        ...,

        description="The summary content",

    )

class SummaryAction(Action[SummaryActionInput]):

    def __init__(self, **kwargs):

        super().__init__(**kwargs)

    @property

    def resource_need(self) -> Optional[ResourceType]:

        # The resource type that the current Agent needs to use

        # here we do not need to use resources, just return None

        return None

    @property

    def render_protocol(self) -> Optional[Vis]:

        # The visualization rendering protocol that the current Agent needs to use

        # here we do not need to use visualization rendering, just return None

        return None

    @property

    def out_model_type(self):

        return SummaryActionInput

    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.

        The entry point for actual execution of Action. Action execution will be

        automatically initiated after model inference.

        """

        try:

            # Parse the input message

            param: SummaryActionInput = self._input_convert(ai_message, SummaryActionInput)

        except Exception:

            return ActionOutput(

                is_exe_success=False,

                content="The requested correctly structured answer could not be found, "

                        f"ai message: {ai_message}",

            )

        # Check if the summary content is not related to user questions

        if param.summary and cmp_string_equal(

                param.summary,

                NOT_RELATED_MESSAGE,

                ignore_case=True,

                ignore_punctuation=True,

                ignore_whitespace=True,

        ):

            return ActionOutput(

                is_exe_success=False,

                content="the provided text content is not related to user questions at all."

                        f"ai message: {ai_message}",

            )

        else:

            return ActionOutput(

                is_exe_success=True,

                content=param.summary,

            )

summarizer_agent.py

from typing import Optional

from pydantic import BaseModel, Field

from dbgpt.vis import Vis

from dbgpt.agent import Action, ActionOutput, AgentResource, ResourceType

from dbgpt.agent.util import cmp_string_equal

NOT_RELATED_MESSAGE = "Did not find the information you want."

# The parameter object that the Action that the current Agent needs to execute needs to output.

class SummaryActionInput(BaseModel):

    summary: str = Field(

        ...,

        description="The summary content",

    )

class SummaryAction(Action[SummaryActionInput]):

    def __init__(self, **kwargs):

        super().__init__(**kwargs)

    @property

    def resource_need(self) -> Optional[ResourceType]:

        # The resource type that the current Agent needs to use

        # here we do not need to use resources, just return None

        return None

    @property

    def render_protocol(self) -> Optional[Vis]:

        # The visualization rendering protocol that the current Agent needs to use

        # here we do not need to use visualization rendering, just return None

        return None

    @property

    def out_model_type(self):

        return SummaryActionInput

    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.

        The entry point for actual execution of Action. Action execution will be

        automatically initiated after model inference.

        """

        try:

            # Parse the input message

            param: SummaryActionInput = self._input_convert(ai_message, SummaryActionInput)

        except Exception:

            return ActionOutput(

                is_exe_success=False,

                content="The requested correctly structured answer could not be found, "

                        f"ai message: {ai_message}",

            )

        # Check if the summary content is not related to user questions

        if param.summary and cmp_string_equal(

                param.summary,

                NOT_RELATED_MESSAGE,

                ignore_case=True,

                ignore_punctuation=True,

                ignore_whitespace=True,

        ):

            return ActionOutput(

                is_exe_success=False,

                content="the provided text content is not related to user questions at all."

                        f"ai message: {ai_message}",

            )

        else:

            return ActionOutput(

                is_exe_success=True,

                content=param.summary,

            )

这样重启项目就能看到自定义的agent了