簡介

文章主要介紹了如何擴展一個自定義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了