wugren
49 lines
2.0 KiB
Python
49 lines
2.0 KiB
Python
import copy
|
|
|
|
from aios_kernel import CustomAIAgent, AgentMsg, AgentMsgType, AgentPrompt
|
|
from aios_kernel.compute_task import ComputeTaskResultCode
|
|
from knowledge.data.writer import split_text
|
|
|
|
|
|
class TextSummaryAgent(CustomAIAgent):
|
|
def __init__(self):
|
|
super().__init__("TextSummary", "Text Summary", 128000)
|
|
|
|
async def _process_msg(self, msg: AgentMsg, workspace=None) -> AgentMsg:
|
|
if msg.msg_type is not AgentMsgType.TYPE_MSG:
|
|
return AgentMsg.create_error_resp(msg, "only support msg type")
|
|
|
|
if msg.body_mime is not None and msg.body_mime != "text/plain":
|
|
return AgentMsg.create_error_resp(msg, "only support text/plain mime type")
|
|
|
|
chunks = split_text(msg.body, separators=["\n\n", "\n"], chunk_size=4000, chunk_overlap=200, length_function=len)
|
|
|
|
prompt = AgentPrompt()
|
|
prompt.system_message = "Your job is to generate a summary based on the input."
|
|
if len(chunks) == 1:
|
|
prompt.append(AgentPrompt(chunks[0]))
|
|
resp = await self.do_llm_complection(prompt)
|
|
if resp.result_code != ComputeTaskResultCode.OK:
|
|
return msg.create_error_resp(resp.error_str)
|
|
return msg.create_resp_msg(resp.result_str)
|
|
|
|
segments = []
|
|
for i, chunk in enumerate(chunks):
|
|
seg_prompt = copy.deepcopy(prompt)
|
|
seg_prompt.append(AgentPrompt(chunk))
|
|
resp = await self.do_llm_complection(seg_prompt)
|
|
if resp.result_code != ComputeTaskResultCode.OK:
|
|
return msg.create_error_resp(resp.error_str)
|
|
segments.append(resp.result_str)
|
|
|
|
segments_str = "\n".join(segments)
|
|
prompt.append(AgentPrompt(f"以下文本分段之后的各段摘要，请合并生成一个完整摘要：\n{segments_str}"))
|
|
resp = await self.do_llm_complection(prompt)
|
|
if resp.result_code != ComputeTaskResultCode.OK:
|
|
return msg.create_error_resp(resp.error_str)
|
|
return msg.create_resp_msg(resp.result_str)
|
|
|
|
|
|
def init():
|
|
return TextSummaryAgent()
|