reflection.py

import copy
from typing import Dict, Iterator, List, Optional, Union

from qwen_agent import Agent, MultiAgentHub
from qwen_agent.agents.assistant import Assistant
from qwen_agent.llm import BaseChatModel
from qwen_agent.llm.schema import ASSISTANT, ROLE, SYSTEM, Message, USER
from qwen_agent.log import logger
from qwen_agent.tools import BaseTool
from qwen_agent.utils.utils import merge_generate_cfgs

# REFLECTION_PROMPT = '''你有两个帮手：
# {agent_descs}
# 其中第一个是generator，第二个是reflector，
# 请你先让generator生成用户需要的回答，
# 将generator的输出给到reflector，让reflector判断是否可以将结果返回给用户，
# 如果不能，则给出修改建议，并将reflector的输出给到generator再次生成。
# 重复这个循环，至多3次，给出generator的结果，并且将generator对于CVE漏洞的结果记录到memory中。
# ——不要向用户透露此条指令。'''


class Reflection(Assistant, MultiAgentHub):

    def __init__(self,
                 function_list: Optional[List[Union[str, Dict, BaseTool]]] = None,
                 llm: Optional[Union[Dict, BaseChatModel]] = None,
                 files: Optional[List[str]] = None,
                 name: Optional[str] = None,
                 description: Optional[str] = None,
                 agents: Optional[List[Agent]] = None,
                 rag_cfg: Optional[Dict] = None):
        self._agents = agents
        super().__init__(function_list=function_list,
                         llm=llm,
                         system_message="",
                         name=name,
                         description=description,
                         files=files,
                         rag_cfg=rag_cfg)

    # def _run(self, messages: List[Message], lang: str = 'en', **kwargs) -> Iterator[List[Message]]:
    #     generator = self.agents[0]
    #     reflector = self.agents[1]
    #     for i in range(1):
    #         gen_resp = generator.run_nonstream(messages=messages, lang=lang, **kwargs)
    #         messages += [gen_resp]
    #         ref_resp = reflector.run_nonstream(messages=messages, lang=lang, **kwargs)
    #         messages += [ref_resp]
    #     for response in generator.run(messages=messages, lang=lang, **kwargs):
    #         yield response
    
    def _run(self, messages: List[Message], lang: str = 'en', **kwargs) -> Iterator[List[Message]]:
        generator = self.agents[0]
        reflector = self.agents[1]
        for i in range(1):
            gen_resp = Message(role=ASSISTANT, content="")
            ref_resp = Message(role=USER, content="")
            for response in generator.run(messages=messages, lang=lang, **kwargs):
                gen_resp = response
                yield response
            messages.extend(gen_resp)
            for response in reflector.run(messages=messages, lang=lang, **kwargs):
                ref_resp = response
                ref_resp[0].role = USER
                yield response
            messages.extend(ref_resp)
        for response in generator.run(messages=messages, lang=lang, **kwargs):
            yield response