optimizer = AgentOptimizer(max_actions_per_step=3, llm_config = llm_config)
# ------------ 解决一个问题的代码 ------------
# ......
# -------------------------------------------------
history = assistant.chat_messages_for_summary(UserProxy)
optimizer.record_one_conversation(history, is_satisfied=result)

register_for_llm, register_for_exector = optimizer.step()
for item in register_for_llm:
    assistant.update_function_signature(**item)
if len(register_for_exector.keys()) > 0:
    user_proxy.register_function(function_map=register_for_exector)

optimizer = AgentOptimizer(max_actions_per_step=3, llm_config = llm_config)
for i in range(EPOCH):
    is_correct = user_proxy.initiate_chat(assistant, message = problem)
    history = assistant.chat_messages_for_summary(user_proxy)
    optimizer.record_one_conversation(history, is_satisfied=is_correct)
    register_for_llm, register_for_exector = optimizer.step()
    for item in register_for_llm:
        assistant.update_function_signature(**item)
    if len(register_for_exector.keys()) > 0:
        user_proxy.register_function(function_map=register_for_exector)

ADD_FUNC = {
    "type": "function",
    "function": {
        "name": "add_function",
        "description": "在对话的上下文中添加一个函数。必须声明所需的 Python 包。函数的名称必须与您生成的代码中的函数名称相同。",
        "parameters": {
            "type": "object",
            "properties": {
                "name": {"type": "string", "description": "代码实现中函数的名称。"},
                "description": {"type": "string", "description": "函数的简短描述。"},
                "arguments": {
                    "type": "string",
                    "description": '以字符串形式编码的参数的 JSON 模式。请注意，JSON 模式仅支持特定类型，包括字符串、整数、对象、数组、布尔值（不支持浮点数类型）。例如：{ "url": { "type": "string", "description": "URL" }}。在使用数组类型时，请避免出现错误 "array schema missing items"。',
                },
                "packages": {
                    "type": "string",
                    "description": "函数导入的包名称列表，需要在调用函数之前使用 pip 安装。这解决了 ModuleNotFoundError。它应该是字符串，而不是列表。",
                },
                "code": {
                    "type": "string",
                    "description": "Python 中的实现。不要包含函数声明。",
                },
            },
            "required": ["name", "description", "arguments", "packages", "code"],
        },
    },
}

REVISE_FUNC = {
    "type": "function",
    "function": {
        "name": "revise_function",
        "description": "在对话的上下文中修改函数。必须声明必要的Python包。函数的名称必须与您生成的代码中的函数名称相同。",
        "parameters": {
            "type": "object",
            "properties": {
                "name": {"type": "string", "description": "代码实现中函数的名称。"},
                "description": {"type": "string", "description": "函数的简短描述。"},
                "arguments": {
                    "type": "string",
                    "description": '以字符串形式编码的参数的JSON模式。请注意，JSON模式仅支持特定类型，包括字符串、整数、对象、数组、布尔值（不支持浮点类型）。例如：{ "url": { "type": "string", "description": "URL" }}。在使用数组类型时，请避免出现错误\'array schema missing items\'。',
                },
                "packages": {
                    "type": "string",
                    "description": "函数导入的包名称列表，需要在调用函数之前使用pip安装。这解决了ModuleNotFoundError。它应该是字符串，而不是列表。",
                },
                "code": {
                    "type": "string",
                    "description": "Python中的实现。不要包含函数声明。",
                },
            },
            "required": ["name", "description", "arguments", "packages", "code"],
        },
    },
}

REMOVE_FUNC = {
    "type": "function",
    "function": {
        "name": "remove_function",
        "description": "在对话的上下文中删除一个函数。一旦删除一个函数，助手将不会在未来的对话中使用该函数。",
        "parameters": {
            "type": "object",
            "properties": {
                "name": {"type": "string", "description": "代码实现中函数的名称。"}
            },
            "required": ["name"],
        },
    },
}

llm_config = LLMConfig(
    config_list=[{
        "model": "gpt-4",
        "api_key": "<azure_api_key>",
        "base_url": "<azure api base url>",
        "api_type": "azure",
        "api_version": "2024-02-01"
    }],
    temperature=0,
)

import json
from autogenstudio import AutoGenWorkFlowManager, AgentWorkFlowConfig

# 加载一个以 JSON 格式存储的 agent 规范
agent_spec = json.load(open('agent_spec.json'))

# 从 agent 规范创建一个 AutoGen 工作流配置
agent_work_flow_config = FlowConfig(**agent_spec)

# 从配置创建一个工作流
agent_work_flow = AutoGenWorkFlowManager(agent_work_flow_config)

# 在一个任务上运行工作流
task_query = "巴黎铁塔的高度是多少？"
agent_work_flow.run(message=task_query)

config_file_or_env = '/home/elpis_ubuntu/LLM/autogen/OAI_CONFIG_LIST'  # 修改路径
default_llm_config = {
    'temperature': 0
}

from autogen.agentchat.contrib.agent_builder import AgentBuilder

builder = AgentBuilder(config_file_or_env=config_file_or_env, builder_model='gpt-4-1106-preview', agent_model='gpt-4-1106-preview')

building_task = "在 arxiv 上找一篇与编程相关的论文，并分析其在某个领域的应用。例如，在 arxiv 上找一篇关于 GPT-4 的最新论文，并找出它在软件领域的潜在应用。"

agent_list, agent_configs = builder.build(building_task, default_llm_config, coding=True)

import autogen

def start_task(execution_task: str, agent_list: list, llm_config: dict):
    config_list = autogen.config_list_from_json(config_file_or_env, filter_dict={"model": ["gpt-4-1106-preview"]})

    group_chat = autogen.GroupChat(agents=agent_list, messages=[], max_round=12)
    manager = autogen.GroupChatManager(
        groupchat=group_chat, llm_config={"config_list": config_list, **llm_config}
    )
    agent_list[0].initiate_chat(manager, message=execution_task)

start_task(
    execution_task="在 arxiv 上找一篇关于 gpt-4 的最新论文，并找出它在软件领域的潜在应用。",
    agent_list=agent_list,
    llm_config=default_llm_config
)

builder.clear_all_agents(recycle_endpoint=True)

saved_path = builder.save()

new_builder = AgentBuilder(config_file_or_env=config_file_or_env)
agent_list, agent_config = new_builder.load(saved_path)
start_task(...)  # 跳过 build()

# 转换为 OpenAI Assistant API。
agent_list, agent_config = new_builder.build(building_task, default_llm_config, use_oai_assistant=True)
...

critic = autogen.AssistantAgent(
    name="critic",
    llm_config={"config_list": config_list},
    system_message="""你是一个乐于助人的助理。你会为评估不同任务提供评价标准。这些标准应该是可区分的、可量化的，并且不重复。
将评估标准转化为一个字典，其中键是标准的名称。
每个键的值是一个字典，格式如下：{"description": 标准描述, "accepted_values": 该键的可接受输入}
确保键是用于评估给定任务的标准。"accepted_values" 包括每个键的可接受输入，这些输入应该是细粒度的，最好是多级别的。"description" 包括标准的描述。
只返回字典。"""
)

quantifier = autogen.AssistantAgent(
    name="quantifier",
    llm_config={"config_list": config_list},
    system_message = """You are a helpful assistant. You quantify the output of different tasks based on the given criteria.
    The criterion is given in a dictionary format where each key is a distinct criteria.
    The value of each key is a dictionary as follows {"description": criteria description , "accepted_values": possible accepted inputs for this key}
    You are going to quantify each of the criteria for a given task based on the task description.
    Return a dictionary where the keys are the criteria and the values are the assessed performance based on accepted values for each criteria.
    Return only the dictionary."""

)

from autogen import config_list_from_json
from autogen.agentchat.contrib.gpt_assistant_agent import GPTAssistantAgent
from autogen import UserProxyAgent

config_list = config_list_from_json("OAI_CONFIG_LIST")

# 使用助手API创建新助手
gpt_assistant = GPTAssistantAgent(
    name="assistant",
    llm_config={
        "config_list": config_list,
        "assistant_id": None
    })

user_proxy = UserProxyAgent(name="user_proxy",
    code_execution_config={
        "work_dir": "coding"
    },
    human_input_mode="NEVER")

user_proxy.initiate_chat(gpt_assistant, message="打印Hello World")

# 使用 Assistant API 创建新的助手
gpt_assistant = GPTAssistantAgent(
    name="assistant",
    llm_config={
        "config_list": config_list,
        "assistant_id": None,
        "tools": [
            {
                "type": "code_interpreter"
            }
        ],
    })

user_proxy = UserProxyAgent(name="user_proxy",
    code_execution_config={
        "work_dir": "coding"
    },
    human_input_mode="NEVER")

user_proxy.initiate_chat(gpt_assistant, message="Print hello world")

from autogen.agentchat.contrib.multimodal_conversable_agent import MultimodalConversableAgent  # 用于GPT-4V
from autogen.agentchat.contrib.llava_agent import LLaVAAgent  # 用于LLaVA

prompt = """您现在是一个面部表情分类器。这里有一些例子。

<img happy.jpg> 描述了一个快乐的表情。
<img http://some_location.com/sad.jpg> 代表了一个悲伤的表情。
<img obama.jpg> 描绘了一个中性的表情。

现在，请识别这个人的面部表情：<img unknown.png>
"""

agent = MultimodalConversableAgent()
user = UserProxyAgent()
user.initiate_chat(agent, message=prompt)

from autogen import UserProxyAgent, config_list_from_json
from autogen.agentchat.contrib.capabilities.teachability import Teachability
from autogen import ConversableAgent  # 作为示例

# 从环境变量或文件中加载 LLM 推理端点
# 参见 https://microsoft.github.io/autogen/docs/FAQ#set-your-api-endpoints
# 和 OAI_CONFIG_LIST_sample
filter_dict = {"model": ["gpt-4"]}  # GPT-3.5 在从用户反馈中学习方面比 GPT-4 不可靠。
config_list = config_list_from_json(env_or_file="OAI_CONFIG_LIST", filter_dict=filter_dict)
llm_config={"config_list": config_list, "timeout": 120}

# 首先实例化任何继承自 ConversableAgent 的代理，这里我们直接使用它来简化操作。
teachable_agent = ConversableAgent(
    name="teachable_agent",  # 名称可以是任意的。
    llm_config=llm_config
)

# 实例化一个 Teachability 对象，其参数都是可选的。
teachability = Teachability(
    reset_db=False,  # 使用 True 强制重置记忆数据库，使用 False 使用现有数据库。
    path_to_db_dir="./tmp/interactive/teachability_db"  # 可以是任何路径，但在群聊中的可教授代理需要唯一的路径。
)

# 现在将 teachability 添加到代理中。
teachability.add_to_agent(teachable_agent)

# 对于这个测试，像往常一样创建一个用户代理。
user = UserProxyAgent("user", human_input_mode="ALWAYS")

# 当用户输入 'exit' 时，该函数将返回。
teachable_agent.initiate_chat(user, message="嗨，我是一个可教授的用户助手！有什么问题吗？")

import autogen
from autogen.agentchat.contrib.retrieve_assistant_agent import RetrieveAssistantAgent
from autogen.agentchat.contrib.retrieve_user_proxy_agent import RetrieveUserProxyAgent

assistant = RetrieveAssistantAgent(
    name="assistant",
    system_message="You are a helpful assistant.",
    llm_config=llm_config,
)

ragproxyagent = RetrieveUserProxyAgent(
    name="ragproxyagent",
    retrieve_config={
        "task": "qa",
        "docs_path": "https://raw.githubusercontent.com/microsoft/autogen/main/README.md",
    },
)

assistant.reset()
ragproxyagent.initiate_chat(assistant, message=ragproxyagent.message_generator, problem="What is autogen?")

assistant.reset()
userproxyagent = autogen.UserProxyAgent(name="userproxyagent")
userproxyagent.initiate_chat(assistant, message="What is autogen?")

from chromadb.utils import embedding_functions

openai_ef = embedding_functions.OpenAIEmbeddingFunction(
                api_key="YOUR_API_KEY",
                model_name="text-embedding-ada-002"
            )

ragproxyagent = RetrieveUserProxyAgent(
    name="ragproxyagent",
    retrieve_config={
        "task": "qa",
        "docs_path": "https://raw.githubusercontent.com/microsoft/autogen/main/README.md",
        "embedding_function": openai_ef,
    },
)

huggingface_ef = embedding_functions.HuggingFaceEmbeddingFunction(
    api_key="YOUR_API_KEY",
    model_name="sentence-transformers/all-MiniLM-L6-v2"
)

from langchain.text_splitter import RecursiveCharacterTextSplitter

recur_spliter = RecursiveCharacterTextSplitter(separators=["\n", "\r", "\t"])

ragproxyagent = RetrieveUserProxyAgent(
    name="ragproxyagent",
    retrieve_config={
        "task": "qa",
        "docs_path": "https://raw.githubusercontent.com/microsoft/autogen/main/README.md",
        "custom_text_split_function": recur_spliter.split_text,
    },
)

# 创建 qdrant 客户端
from qdrant_client import QdrantClient

client = QdrantClient(url="***", api_key="***")

# 包装 RetrieveUserProxyAgent
from litellm import embedding as test_embedding
from autogen.agentchat.contrib.retrieve_user_proxy_agent import RetrieveUserProxyAgent
from qdrant_client.models import SearchRequest, Filter, FieldCondition, MatchText

class QdrantRetrieveUserProxyAgent(RetrieveUserProxyAgent):
    def query_vector_db(
        self,
        query_texts: List[str],
        n_results: int = 10,
        search_string: str = "",
        **kwargs,
    ) -> Dict[str, Union[List[str], List[List[str]]]]:
        # 在这里定义你自己的查询函数
        embed_response = test_embedding('text-embedding-ada-002', input=query_texts)

        all_embeddings: List[List[float]] = []

        for item in embed_response['data']:
            all_embeddings.append(item['embedding'])

        search_queries: List[SearchRequest] = []

        for embedding in all_embeddings:
            search_queries.append(
                SearchRequest(
                    vector=embedding,
                    filter=Filter(
                        must=[
                            FieldCondition(
                                key="page_content",
                                match=MatchText(
                                    text=search_string,
                                )
                            )
                        ]
                    ),
                    limit=n_results,
                    with_payload=True,
                )
            )

        search_response = client.search_batch(
            collection_name="{你的集合名称}",
            requests=search_queries,
        )

        return {
            "ids": [[scored_point.id for scored_point in batch] for batch in search_response],
            "documents": [[scored_point.payload.get('page_content', '') for scored_point in batch] for batch in search_response],
            "metadatas": [[scored_point.payload.get('metadata', {}) for scored_point in batch] for batch in search_response]
        }

    def retrieve_docs(self, problem: str, n_results: int = 20, search_string: str = "", **kwargs):
        results = self.query_vector_db(
            query_texts=[problem],
            n_results=n_results,
            search_string=search_string,
            **kwargs,
        )

        self._results = results


# 使用 QdrantRetrieveUserProxyAgent
qdrantragagent = QdrantRetrieveUserProxyAgent(
    name="ragproxyagent",
    human_input_mode="NEVER",
    max_consecutive_auto_reply=2,
    retrieve_config={
        "task": "qa",
    },
)

qdrantragagent.retrieve_docs("什么是 Autogen?", n_results=10, search_string="autogen")

boss = autogen.UserProxyAgent(
    name="Boss",
    is_termination_msg=termination_msg,
    human_input_mode="TERMINATE",
    system_message="提问问题和分配任务的老板。",
)

boss_aid = RetrieveUserProxyAgent(
    name="Boss_Assistant",
    is_termination_msg=termination_msg,
    system_message="拥有额外内容检索能力，用于解决困难问题的助手。",
    human_input_mode="NEVER",
    max_consecutive_auto_reply=3,
    retrieve_config={
        "task": "qa",
    },
    code_execution_config=False,  # 在这种情况下，我们不想执行代码。
)

coder = autogen.AssistantAgent(
    name="Senior_Python_Engineer",
    is_termination_msg=termination_msg,
    system_message="你是一名高级 Python 工程师。在完成所有工作后回复“TERMINATE”。",
    llm_config={"config_list": config_list, "timeout": 60, "temperature": 0},
)

pm = autogen.AssistantAgent(
    name="Product_Manager",
    is_termination_msg=termination_msg,
    system_message="你是一名产品经理。在完成所有工作后回复“TERMINATE”。",
    llm_config={"config_list": config_list, "timeout": 60, "temperature": 0},
)

reviewer = autogen.AssistantAgent(
    name="Code_Reviewer",
    is_termination_msg=termination_msg,
    system_message="你是一名代码审查员。在完成所有工作后回复“TERMINATE”。",
    llm_config={"config_list": config_list, "timeout": 60, "temperature": 0},
)

def retrieve_content(
    message: Annotated[
        str,
        "经过优化的消息，保留原始含义，可用于检索内容以进行代码生成和问题回答。",
    ],
    n_results: Annotated[int, "结果数量"] = 3,
) -> str:
    boss_aid.n_results = n_results  # 设置要检索的结果数量。
    # 检查是否需要更新上下文。
    update_context_case1, update_context_case2 = boss_aid._check_update_context(message)
    if (update_context_case1 or update_context_case2) and boss_aid.update_context:
        boss_aid.problem = message if not hasattr(boss_aid, "problem") else boss_aid.problem
        _, ret_msg = boss_aid._generate_retrieve_user_reply(message)
    else:
        _context = {"problem": message, "n_results": n_results}
        ret_msg = boss_aid.message_generator(boss_aid, None, _context)
    return ret_msg if ret_msg else message

for caller in [pm, coder, reviewer]:
    d_retrieve_content = caller.register_for_llm(
        description="检索内容以进行代码生成和问题回答。",
        api_style="function"
    )(retrieve_content)

for executor in [boss, pm]:
    executor.register_for_execution()(d_retrieve_content)

groupchat = autogen.GroupChat(
    agents=[boss, pm, coder, reviewer],
    messages=[],
    max_round=12,
    speaker_selection_method="round_robin",
    allow_repeat_speaker=False,
)

llm_config = {"config_list": config_list, "timeout": 60, "temperature": 0}
manager = autogen.GroupChatManager(groupchat=groupchat, llm_config=llm_config)

# 作为用户代理代理开始与老板聊天。
boss.initiate_chat(
    manager,
    message="如何在 FLAML 中使用 Spark 进行并行训练？给我一个示例代码。",
)

# 初始化代理
def initialize_agents(config_list, docs_path=None):
    ...
    return assistant, ragproxyagent

# 初始化聊天
def initiate_chat(config_list, problem, queue, n_results=3):
    ...
    assistant.reset()
    try:
        ragproxyagent.a_initiate_chat(
            assistant, problem=problem, silent=False, n_results=n_results
        )
        messages = ragproxyagent.chat_messages
        messages = [messages[k] for k in messages.keys()][0]
        messages = [m["content"] for m in messages if m["role"] == "user"]
        print("messages: ", messages)
    except Exception as e:
        messages = [str(e)]
    queue.put(messages)

# 将 AutoGen 部分封装成一个函数
def chatbot_reply(input_text):
    """通过终端与代理进行聊天。"""
    queue = mp.Queue()
    process = mp.Process(
        target=initiate_chat,
        args=(config_list, input_text, queue),
    )
    process.start()
    try:
        messages = queue.get(timeout=TIMEOUT)
    except Exception as e:
        messages = [str(e) if len(str(e)) > 0 else "无效的 OpenAI 请求，请检查您的 API 密钥。"]
    finally:
        try:
            process.terminate()
        except:
            pass
    return messages

...

# 使用 Gradio 设置用户界面
with gr.Blocks() as demo:
    ...
    assistant, ragproxyagent = initialize_agents(config_list)

    chatbot = gr.Chatbot(
        [],
        elem_id="chatbot",
        bubble_full_width=False,
        avatar_images=(None, (os.path.join(os.path.dirname(__file__), "autogen.png"))),
        # height=600,
    )

    txt_input = gr.Textbox(
        scale=4,
        show_label=False,
        placeholder="输入文本并按回车键",
        container=False,
    )

    with gr.Row():
        txt_model = gr.Dropdown(
            label="模型",
            choices=[
                "gpt-4",
                "gpt-35-turbo",
                "gpt-3.5-turbo",
            ],
            allow_custom_value=True,
            value="gpt-35-turbo",
            container=True,
        )
        txt_oai_key = gr.Textbox(
            label="OpenAI API 密钥",
            placeholder="输入密钥并按回车键",
            max_lines=1,
            show_label=True,
            value=os.environ.get("OPENAI_API_KEY", ""),
            container=True,
            type="password",
        )
        ...

    clear = gr.ClearButton([txt_input, chatbot])

...

if __name__ == "__main__":
    demo.launch(share=True)

class CompletionResponseChoice(BaseModel):
    index: int
    text: str
    logprobs: Optional[int] = None
    # finish_reason: Optional[Literal["stop", "length"]]

class CompletionResponseStreamChoice(BaseModel):
    index: int
    text: str
    logprobs: Optional[float] = None
    # finish_reason: Optional[Literal["stop", "length"]] = None

from autogen import oai

# 创建一个文本补全请求
response = oai.Completion.create(
    config_list=[
        {
            "model": "chatglm2-6b",
            "base_url": "http://localhost:8000/v1",
            "api_type": "openai",
            "api_key": "NULL", # 只是一个占位符
        }
    ],
    prompt="你好",
)
print(response)

# 创建一个对话补全请求
response = oai.ChatCompletion.create(
    config_list=[
        {
            "model": "chatglm2-6b",
            "base_url": "http://localhost:8000/v1",
            "api_type": "openai",
            "api_key": "NULL",
        }
    ],
    messages=[{"role": "user", "content": "你好"}]
)
print(response)

from autogen import oai

# 创建一个聊天完成请求
response = oai.ChatCompletion.create(
    config_list=[
        {
            "model": "chatglm2-6b",
            "base_url": "http://localhost:8000/v1",
            "api_type": "openai",
            "api_key": "NULL",
        },
        {
            "model": "vicuna-7b-v1.3",
            "base_url": "http://localhost:8000/v1",
            "api_type": "openai",
            "api_key": "NULL",
        }
    ],
    messages=[{"role": "user", "content": "Hi"}]
)
print(response)