简介
文章主要介绍了如何扩展一个自定义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了
