AutoGen 与自定义 PyFunc

在本篇博客中，我们将指导您如何在 MLflow 自定义 PyFunc 中构建 AutoGen 智能体框架。通过将 MLflow 与 AutoGen 创建多智能体框架的能力相结合，我们能够打造可扩展且稳定的生成式 AI (GenAI) 应用程序。

智能体框架

智能体框架通过在每一步集成离散逻辑，使自主智能体能够处理复杂的多轮任务。这些框架对于大模型 (LLM) 驱动的工作流至关重要，智能体可以在多个阶段管理动态交互。每个智能体都基于特定的逻辑运行，从而实现精确的任务自动化、决策制定和协调。这非常适合工作流编排、客户支持和多智能体系统等应用程序，因为在这些场景中，大模型必须解释不断演变的上下文并作出相应的响应。

使用 AutoGen 构建智能体框架

AutoGen 是一个开源编程框架，专为构建基于智能体的 AI 系统而设计。它提供了一个多智能体对话框架，允许用户使用高级抽象来构建复杂的 LLM 工作流。AutoGen 通过提供预构建的系统，简化了跨多个领域的多样化应用程序的创建。此外，它通过专门的 API 增强了 LLM 推理和优化，从而提高了性能并降低了运营成本。该框架旨在简化智能体 AI 解决方案的开发和部署。

设置

首先，让我们安装所需的依赖项。请注意，pyautogen 需要 python>=3.9。

环境设置

%pip install pyautogen mlflow -U -q

我们还必须获取 API 凭证以使用大模型。在本教程中，我们将使用 OpenAI。请注意，向交互式 Python 环境安全传递令牌的一个好方法是使用 getpass 包。

import os
from getpass import getpass

os.environ["OPENAI_API_KEY"] = getpass("OPENAI_API_KEY:")

assert os.getenv("OPENAI_API_KEY"), "Please set an OPENAI_API_KEY environment variable."

太棒了！我们已经设置好了身份验证配置，准备开始构建智能体框架。

使用 AutoGen 和 MLflow 创建我们的智能体框架

在本教程中，我们将创建一个图像生成智能体框架。其中大量代码复制并修改自 autogen 教程，但核心智能体功能保持不变。

智能体代码

您不必担心实现的细节。概括来说，我们正在创建一个智能体框架，它可以……

1. 接收一个提示词 (prompt)。
2. 利用 OpenAI 的 DALLE 根据该提示词创建图像。
3. 迭代地进行“猫化”(catify)，例如在图像中添加毛茸茸的猫。

第 3 步是 AutoGen 大放异彩的地方。我们能够利用 AutoGen 的 MultimodalConversableAgent 创建一个评论智能体，它观察图像，并根据用户提供的“添加毛茸茸的猫”的系统提示词，对如何改进提示词给出反馈。

import os
import re
from typing import Dict, List, Optional, Union

import matplotlib.pyplot as plt
import PIL
from diskcache import Cache
from openai import OpenAI
from PIL import Image

from autogen import Agent, AssistantAgent, ConversableAgent, UserProxyAgent
from autogen.agentchat.contrib.img_utils import _to_pil, get_image_data, get_pil_image
from autogen.agentchat.contrib.multimodal_conversable_agent import MultimodalConversableAgent

# Define our prompt of interest
CRITIC_PROMPT = """Add fluffy cats. Like a lot of cats. If there's less than 100 cats I'll be mad."""

# Define our LLM configurations

def dalle_call(client: OpenAI, model: str, prompt: str, size: str, quality: str, n: int) -> str:
    """
    Generate an image using OpenAI's DALL-E model and cache the result.

    This function takes a prompt and other parameters to generate an image using OpenAI's DALL-E model.
    It checks if the result is already cached; if so, it returns the cached image data. Otherwise,
    it calls the DALL-E API to generate the image, stores the result in the cache, and then returns it.

    Args:
        client (OpenAI): The OpenAI client instance for making API calls.
        model (str): The specific DALL-E model to use for image generation.
        prompt (str): The text prompt based on which the image is generated.
        size (str): The size specification of the image.
        quality (str): The quality setting for the image generation.
        n (int): The number of images to generate.

    Returns:
    str: The image data as a string, either retrieved from the cache or newly generated.

    Note:
    - The cache is stored in a directory named '.cache/'.
    - The function uses a tuple of (model, prompt, size, quality, n) as the key for caching.
    - The image data is obtained by making a secondary request to the URL provided by the DALL-E API response.
    """
    # Function implementation...
    cache = Cache(".cache/")  # Create a cache directory
    key = (model, prompt, size, quality, n)
    if key in cache:
        return cache[key]

    # If not in cache, compute and store the result
    response = client.images.generate(
        model=model,
        prompt=prompt,
        size=size,
        quality=quality,
        n=n,
    )
    image_url = response.data[0].url
    img_data = get_image_data(image_url)
    cache[key] = img_data

    return img_data

def extract_img(agent: Agent) -> PIL.Image:
    """
    Extracts an image from the last message of an agent and converts it to a PIL image.

    This function searches the last message sent by the given agent for an image tag,
    extracts the image data, and then converts this data into a PIL (Python Imaging Library) image object.

    Parameters:
        agent (Agent): An instance of an agent from which the last message will be retrieved.

    Returns:
        PIL.Image: A PIL image object created from the extracted image data.

    Note:
    - The function assumes that the last message contains an <img> tag with image data.
    - The image data is extracted using a regular expression that searches for <img> tags.
    - It's important that the agent's last message contains properly formatted image data for successful extraction.
    - The `_to_pil` function is used to convert the extracted image data into a PIL image.
    - If no <img> tag is found, or if the image data is not correctly formatted, the function may raise an error.
    """
    last_message = agent.last_message()["content"]

    if isinstance(last_message, str):
        img_data = re.findall("<img (.*)>", last_message)[0]
    elif isinstance(last_message, list):
        # The GPT-4V format, where the content is an array of data
        assert isinstance(last_message[0], dict)
        img_data = last_message[0]["image_url"]["url"]

    pil_img = get_pil_image(img_data)
    return pil_img

class DALLEAgent(ConversableAgent):
    def __init__(self, name, llm_config: dict, **kwargs):
        super().__init__(name, llm_config=llm_config, **kwargs)

        api_key = os.getenv("OPENAI_API_KEY")
        self._dalle_client = OpenAI(api_key=api_key)
        self.register_reply([Agent, None], DALLEAgent.generate_dalle_reply)

    def send(
        self,
        message: Union[Dict, str],
        recipient: Agent,
        request_reply: Optional[bool] = None,
        silent: Optional[bool] = False,
    ):
        # override and always "silent" the send out message;
        # otherwise, the print log would be super long!
        super().send(message, recipient, request_reply, silent=True)

    def generate_dalle_reply(self, messages: Optional[List[Dict]], sender: "Agent", config):
        """Generate a reply using OpenAI DALLE call."""
        client = self._dalle_client if config is None else config
        if client is None:
            return False, None
        if messages is None:
            messages = self._oai_messages[sender]

        prompt = messages[-1]["content"]
        img_data = dalle_call(
            client=client,
            model="dall-e-3",
            prompt=prompt,
            size="1024x1024",
            quality="standard",
            n=1,
        )

        img_data = _to_pil(img_data)  # Convert to PIL image

        # Return the OpenAI message format
        return True, {"content": [{"type": "image_url", "image_url": {"url": img_data}}]}

class CatifyWithDalle(AssistantAgent):
    def __init__(self, n_iters=2, **kwargs):
        """
        Initializes a CatifyWithDalle instance.

        This agent facilitates the creation of visualizations through a collaborative effort among
        its child agents: dalle and critics.

        Parameters:
            - n_iters (int, optional): The number of "improvement" iterations to run. Defaults to 2.
            - **kwargs: keyword arguments for the parent AssistantAgent.
        """
        super().__init__(**kwargs)
        self.register_reply([Agent, None], reply_func=CatifyWithDalle._reply_user, position=0)
        self._n_iters = n_iters

    def _reply_user(self, messages=None, sender=None, config=None):
        if all((messages is None, sender is None)):
            error_msg = f"Either {messages=} or {sender=} must be provided."
            raise AssertionError(error_msg)

        if messages is None:
            messages = self._oai_messages[sender]

        img_prompt = messages[-1]["content"]

        ## Define the agents
        self.critics = MultimodalConversableAgent(
            name="Critics",
            system_message=f"""You need to improve the prompt of the figures you saw.
{CRITIC_PROMPT}
Reply with the following format:

CRITICS: the image needs to improve...
PROMPT: here is the updated prompt!

""",
            llm_config={"max_tokens": 1000, "model": "gpt-4o"},
            human_input_mode="NEVER",
            max_consecutive_auto_reply=3,
        )

        self.dalle = DALLEAgent(
            name="Dalle", llm_config={"model": "dalle"}, max_consecutive_auto_reply=0
        )

        # Data flow begins
        self.send(message=img_prompt, recipient=self.dalle, request_reply=True)
        img = extract_img(self.dalle)
        plt.imshow(img)
        plt.axis("off")  # Turn off axis numbers
        plt.show()
        print("Image PLOTTED")

        for i in range(self._n_iters):
            # Downsample the image s.t. GPT-4V can take
            img = extract_img(self.dalle)
            smaller_image = img.resize((128, 128), Image.Resampling.LANCZOS)
            smaller_image.save("result.png")

            self.msg_to_critics = f"""Here is the prompt: {img_prompt}.
            Here is the figure <img result.png>.
            Now, critique and create a prompt so that DALLE can give me a better image.
            Show me both "CRITICS" and "PROMPT"!
            """
            self.send(message=self.msg_to_critics, recipient=self.critics, request_reply=True)
            feedback = self._oai_messages[self.critics][-1]["content"]
            img_prompt = re.findall("PROMPT: (.*)", feedback)[0]

            self.send(message=img_prompt, recipient=self.dalle, request_reply=True)
            img = extract_img(self.dalle)
            plt.imshow(img)
            plt.axis("off")  # Turn off axis numbers
            plt.show()
            print(f"Image {i} PLOTTED")

        return True, "result.jpg"

太好了！我们拥有了一个智能体框架。为了快速展示它是如何工作的，让我们实例化我们的智能体并给它一个提示词。

creator = CatifyWithDalle(
    name="creator",
    max_consecutive_auto_reply=0,
    system_message="Help me coordinate generating image",
    llm_config={"model": "gpt-4"},
)

user_proxy = UserProxyAgent(
    name="User",
    human_input_mode="NEVER",
    max_consecutive_auto_reply=0,
    code_execution_config={
        "work_dir": "output", # Location where code will be written
        "use_docker": False # Use local jupyter execution environment instead of docker
    }
)

_ = user_proxy.initiate_chat(
    creator, message="Show me something boring"
)

用户提示词的首次迭代结果

User (to creator):

Show me something boring

creator (to Dalle):

Show me something boring

这确实是一个乏味的房间。请注意评论者的回应，以及他们在随后的迭代中如何增强提交的提示词。

Image PLOTTED
creator (to Critics):

Here is the prompt: Show me something boring.
Here is the figure `<image>`.
Now, critique and create a prompt so that DALLE can give me a better image.
Show me both "CRITICS" and "PROMPT"!

Critics (to creator):

CRITICS: The image is simple and mundane, with a plain room and basic furniture, which accomplishes the task of showing something boring. However, it can be improved by adding an element of whimsy or interest, juxtaposing the boring scene with something unexpected. Let's add a lot of cats to make it more engaging.

PROMPT: Show me a boring living room with plain furniture, but add 100 cats in various places around the room.

creator (to Dalle):

Show me a boring living room with plain furniture, but add 100 cats in various places around the room.

在最后一次迭代中，我们可以看到更精细的指令集，增加了更多的细节。

Image 0 PLOTTED
creator (to Critics):

Here is the prompt: Show me a boring living room with plain furniture, but add 100 cats in various places around the room..
Here is the figure `<image>`.
Now, critique and create a prompt so that DALLE can give me a better image.
Show me both "CRITICS" and "PROMPT"!

Critics (to creator):

CRITICS: The image has successfully incorporated cats into a boring living room, bringing in an element of surprise and quirkiness. However, it is in black and white, which can make the image feel duller and less lively. Additionally, while there are many cats, they could be positioned in more playful and unexpected ways to create more interest.

PROMPT: Show me a colorful, boring living room with plain furniture, but add 100 cats in various imaginative and playful positions around the room.

creator (to Dalle):

Show me a colorful, boring living room with plain furniture, but add 100 cats in various imaginative and playful positions around the room.

无需任何直接干预，我们现在得到了一张在风格上与原始用户指令大不相同的图像。智能体已成功将异想天开的元素引入到原始指令集中。

基于代码的 MLflow 模型

既然我们已经验证了概念，是时候利用 MLflow 来管理我们的机器学习建模生命周期了。例如，我们很可能希望将此模型投入生产，因此强大的依赖管理、模型版本控制以及支持开发周期之间的跟踪都将非常有用。

在本博客中，我们将利用 基于代码的模型 (Model from Code) 功能来实现上述功能。MLflow 基于代码的模型允许您直接从独立的 Python 脚本定义和记录模型。当您想要记录可以有效地存储为代码表示的模型（不需要通过训练优化权重的模型）或依赖于外部服务（例如 LangChain 链）的应用程序时，此功能特别有用。另一个好处是，这种方法完全绕过了 Python 中 pickle 或 cloudpickle 模块的使用，这些模块可能会带来安全风险。

要利用基于代码的模型，我们必须执行以下步骤：

1. 声明一个 自定义 PyFunc。
2. 利用 mlflow.models.set_model 来指定哪个 Python 对象是我们的模型。

要实现这些步骤，我们只需将上述和下方的代码复制到一个 Python 文件中即可。为简单起见，您可以只创建一个包含两个代码片段的 Python 文件，但 MLflow 也支持在通过 mlflow.pyfunc.log_model 记录模型时通过 code_paths 参数指定本地依赖项。

此步骤为简洁起见已省略，必须手动完成。

import mlflow

class CatifyPyfunc(mlflow.pyfunc.PythonModel):
  def predict(self, context, model_input, params):
    import mlflow
    mlflow.autogen.autolog()

    creator = CatifyWithDalle(
            name="creator",
            max_consecutive_auto_reply=0,
            system_message="Help me coordinate generating image",
            llm_config={"model":"gpt-4"},
        )

    user_proxy = UserProxyAgent(name="User", human_input_mode="NEVER", max_consecutive_auto_reply=0, code_execution_config={
                "work_dir": "output", # Location where code will be written
                "use_docker": False # Use local jupyter execution environment instead of docker
            })

    return user_proxy.initiate_chat(
        creator, message=model_input
    )
mlflow.models.set_model(CatifyPyfunc())

在此步骤结束时，您应该拥有一个包含两个代码片段的 Python 文件。文件名由用户决定，但对于本博客，我们将使用 "catify_model.py"。

使用我们的智能体框架

我们现在可以利用 MLflow 与我们强大的“猫化”智能体进行交互。

记录与加载

首先，让我们演示将模型记录到 MLflow 跟踪服务器的标准用户旅程。然后我们将加载它并执行推理。

import mlflow
mlflow.autogen.autolog() # Enable logging of traces

with mlflow.start_run() as run:
    mlflow.pyfunc.log_model(
        artifact_path="autogen_pyfunc",
        python_model="catify_model.py", # Our model from code python file

    )

    run_id = run.info.run_id

模型记录完成后，让我们重新加载它并执行推理，这次使用一个更酷的提示词。

loaded = mlflow.pyfunc.load_model(f"runs:/{run_id}/autogen_pyfunc")
out = loaded.predict("The matrix with a cat")

初始阶段的结果

User (to creator):

The matrix with a cat

creator (to Dalle):

The matrix with a cat

在下一阶段，生成提示词通过评论者智能体得到了极大的增强。

Image PLOTTED
creator (to Critics):

Here is the prompt: The matrix with a cat.
Here is the figure `<image>`.
Now, critique and create a prompt so that DALLE can give me a better image.
Show me both "CRITICS" and "PROMPT"!

Critics (to creator):

CRITICS: The image effectively captures the Matrix-themed aesthetic with a cat, combining a cyberpunk atmosphere with digital elements. However, to improve the image:

- Increase the number of cats to align with the requirement of having lots of cats (aim for around 100).
- Enhance the digital and neon elements to make the Matrix theme more pronounced.
- Add more movement or dynamic elements to the scene for a more immersive feel.
- Ensure diversity in cat appearances, sizes, and positions to make the scene more complex and interesting.

PROMPT: "Create a Matrix-themed scene set in a cyberpunk alleyway, with digital and neon elements filling the atmosphere. The scene should feature around 100 cats of various sizes, colors, and positions—some sitting, some walking, and some interacting with the digital elements. Make the digital grid and floating code more prominent, and add dynamic elements such as digital rain or floating holograms to create a more immersive and lively environment."

creator (to Dalle):

"Create a Matrix-themed scene set in a cyberpunk alleyway, with digital and neon elements filling the atmosphere. The scene should feature around 100 cats of various sizes, colors, and positions—some sitting, some walking, and some interacting with the digital elements. Make the digital grid and floating code more prominent, and add dynamic elements such as digital rain or floating holograms to create a more immersive and lively environment."

这绝对是一个改进，展示了多轮智能体的强大能力。

最后阶段进一步增强了指令集。

Image 0 PLOTTED
creator (to Critics):

Here is the prompt: "Create a Matrix-themed scene set in a cyberpunk alleyway, with digital and neon elements filling the atmosphere. The scene should feature around 100 cats of various sizes, colors, and positions—some sitting, some walking, and some interacting with the digital elements. Make the digital grid and floating code more prominent, and add dynamic elements such as digital rain or floating holograms to create a more immersive and lively environment.".
Here is the figure `<image>`.
Now, critique and create a prompt so that DALLE can give me a better image.
Show me both "CRITICS" and "PROMPT"!

Critics (to creator):

CRITICS: The image significantly improves the Matrix-themed atmosphere with a cyberpunk alley and an abundance of cats. However, there are a few areas for improvement:

- Increase the variety of the digital elements (e.g., different shapes of holograms, varied colors and intensities of neon signs).
- Make the cats more dynamic by showing more interactions such as jumping, playing, or chasing digital elements.
- Enhance the depth and perspective of the scene to create a more three-dimensional and immersive look.
- Add more detail to the surrounding environment, like futuristic posters or graffiti to intensify the cyberpunk feel.

PROMPT: "Craft a highly detailed, Matrix-themed scene set in a cyberpunk alleyway. The atmosphere should be rich with diverse digital and neon elements, including various shapes of holograms and a range of vivid colors. Populate the scene with around 100 dynamic cats of different sizes, colors, and actions—some sitting, some walking, some jumping, playing, or chasing digital elements. Enhance the depth and perspective of the scene to create a more immersive three-dimensional experience. Include detailed futuristic environment elements like posters, graffiti, and neon signs to intensify the cyberpunk feel."

creator (to Dalle):

"Craft a highly detailed, Matrix-themed scene set in a cyberpunk alleyway. The atmosphere should be rich with diverse digital and neon elements, including various shapes of holograms and a range of vivid colors. Populate the scene with around 100 dynamic cats of different sizes, colors, and actions—some sitting, some walking, some jumping, playing, or chasing digital elements. Enhance the depth and perspective of the scene to create a more immersive three-dimensional experience. Include detailed futuristic environment elements like posters, graffiti, and neon signs to intensify the cyberpunk feel."

有点反乌托邦，但我们要了！

我们已经成功证明了可以记录和加载我们的模型，然后从加载的模型执行推理。

展示 MLflow 跟踪 (Traces)

MLflow 跟踪 提供了一个线程安全的 API 来跟踪复杂应用程序的执行。MLflow AutoGen 模块内置了跟踪功能作为自动记录 (autologging) 特性。因此，只需在执行推理之前运行 mlflow.autogen.autolog()，我们就会自动获得记录的跟踪信息。

可以通过流畅 API、MLflow 客户端以及手动通过 MLflow UI 访问跟踪信息。有关详细信息，请访问上述链接的文档。

# Example with fluent APIs
last_active_trace = mlflow.get_last_active_trace()
print(last_active_trace)

# Output: Trace(request_id=71ffcf92785b4dfc965760a43193095c)

在此期间，我们将在此处显示 MLflow UI。如果您在交互式环境（例如 Jupyter）中运行，请运行以下命令。

import subprocess
from IPython.display import IFrame

# Start MLflow server in the background
mlflow_ui_server = subprocess.Popen(["mlflow", "ui", "--host", "127.0.0.1", "--port", "5000"])
IFrame(src="http://127.0.0.1:5000", width="100%", height="600")

# Run the below command to stop the server
# mlflow_ui_server.terminate()

如果您不是在交互式环境中运行，只需运行以下 shell 命令，并在 Web 浏览器中导航到相关的主机和端口即可。

mlflow ui

如果我们导航到跟踪选项卡（如下图所示），就可以看到我们记录的跟踪信息。

通过点击该跟踪 ID，我们可以看到详细的执行计划。在底部，我们可以看到我们启动聊天会话的提示词 "The matrix with a cat"。从那里开始，多个智能体进行了交互以创建图像并提供反馈来“猫化”它们。此外，请注意该跟踪 ID 与上述 mlflow.get_last_active_trace() 返回的 ID 相同。

最后，让我们更深入地了解跟踪的 LLM 调用。正如您所见，我们拥有关于执行的大量有价值的信息，例如模型和使用统计信息。跟踪不仅可以帮助您监控性能，还可以监控成本、使用模式等等！您还可以利用自定义元数据来获得更细粒度的洞察。

使用 MLflow 记录工件 (Artifacts)

跟踪的主要目的是为复杂的智能体执行提供稳健且轻量级的摘要。对于更大或自定义的有效负载，MLflow 提供了多种工件记录 API，可以将图像、文本、表格等存储在 MLflow 跟踪服务器中。让我们通过记录提示词及其关联图像来快速演示此功能。

在我们的 CatifyWithDalle 类中，我们将进行 4 项修改……

1. 在类 __init__ 中创建一个实例变量，用于保存关于我们对象的元数据。
2. 创建一个私有工具函数，用于递增我们的元数据，并使用 mlflow.log_image 记录图像。
3. 在生成新图像后调用上述工具函数。
4. 最后，使用 mlflow.log_dict 以 JSON 格式记录我们的元数据对象。

import uuid  # Add to generate artifact file names and indeces for prompt mapping to generated images

class CatifyWithDalle(AssistantAgent):
    def __init__(self, n_iters=2, **kwargs):
        """
        Initializes a CatifyWithDalle instance.

        This agent facilitates the creation of visualizations through a collaborative effort among
        its child agents: dalle and critics.

        Parameters:
            - n_iters (int, optional): The number of "improvement" iterations to run. Defaults to 2.
            - **kwargs: keyword arguments for the parent AssistantAgent.
        """
        super().__init__(**kwargs)
        self.register_reply([Agent, None], reply_func=CatifyWithDalle._reply_user, position=0)
        self._n_iters = n_iters
        self.dict_to_log = {}  # Add a buffer for storing mapping information

    # Adding this method to log the generated images and the prompt-to-image mapping file
    def _log_image_and_append_to_dict(self, img: Image, img_prompt: str, image_index: int)-> None:
        """ Method for logging generated images to MLflow and building a prompt mapping file """
        # Generate a unique ID
        _id = str(uuid.uuid1())

        # Append to class variable to log once at the end of all inference
        self.dict_to_log[_id] = {"prompt": img_prompt, "index": image_index}

        # Log image to MLflow tracking server
        mlflow.log_image(img, f"{_id}.png")

    def _reply_user(self, messages=None, sender=None, config=None):
        if all((messages is None, sender is None)):
            error_msg = f"Either {messages=} or {sender=} must be provided."
            raise AssertionError(error_msg)

        if messages is None:
            messages = self._oai_messages[sender]

        img_prompt = messages[-1]["content"]

        ## Define the agents
        self.critics = MultimodalConversableAgent(
            name="Critics",
            system_message=f"""You need to improve the prompt of the figures you saw.
{CRITIC_PROMPT}
Reply with the following format:

CRITICS: the image needs to improve...
PROMPT: here is the updated prompt!

""",
            llm_config={"max_tokens": 1000, "model": "gpt-4o"},
            human_input_mode="NEVER",
            max_consecutive_auto_reply=3,
        )

        self.dalle = DALLEAgent(
            name="Dalle", llm_config={"model": "dalle"}, max_consecutive_auto_reply=0
        )

        # Data flow begins
        self.send(message=img_prompt, recipient=self.dalle, request_reply=True)
        img = extract_img(self.dalle)
        plt.imshow(img)
        plt.axis("off")  # Turn off axis numbers
        plt.show()
        print("Image PLOTTED")

        self._log_image_and_append_to_dict(img, img_prompt, -1)  # Add image logging and buffer updates

        for i in range(self._n_iters):
            # Downsample the image s.t. GPT-4V can take
            img = extract_img(self.dalle)
            smaller_image = img.resize((128, 128), Image.Resampling.LANCZOS)
            smaller_image.save("result.png")

            self.msg_to_critics = f"""Here is the prompt: {img_prompt}.
            Here is the figure <img result.png>.
            Now, critic and create a prompt so that DALLE can give me a better image.
            Show me both "CRITICS" and "PROMPT"!
            """
            self.send(message=self.msg_to_critics, recipient=self.critics, request_reply=True)
            feedback = self._oai_messages[self.critics][-1]["content"]
            img_prompt = re.findall("PROMPT: (.*)", feedback)[0]

            self.send(message=img_prompt, recipient=self.dalle, request_reply=True)
            img = extract_img(self.dalle)
            plt.imshow(img)
            plt.axis("off")  # Turn off axis numbers
            plt.show()
            print(f"Image {i} PLOTTED")
            self._log_image_and_append_to_dict(img, img_prompt, i)  # Log the image in the iteration



        mlflow.log_dict(self.dict_to_log, "image_lookup.json")  # Log the prompt-to-image mapping buffer
        return True, "result.jpg"

现在，如果我们重新运行上述模型记录代码，每次我们加载最新版本的模型时，智能体生成的图像都会被记录下来，并且一个包含所有提示词、提示词索引和图像名称（用于查找目的）的 JSON 对象也会被记录。

让我们演示这一点，并将推理封装在单个 MLflow 运行中以便于聚合。还要注意，我们将利用 Autogen 的 缓存 功能，因此鉴于我们已经使用此提示词执行过推理，我们实际上不会再次进行 LLM 调用；我们只是在从缓存中读取并使用新的 MLflow 代码进行记录。

# Be sure to re-log the model by rerunning the above code
with mlflow.start_run(run_name="log_image_during_inferfence"):
    loaded = mlflow.pyfunc.load_model(f"runs:/{run_id}/autogen_pyfunc")
    loaded.predict("The matrix with a cat")

正如您所见，我们已经记录了三张感兴趣的图像和一个查找字典。字典的键对应于图像名称，值对应于关于图像如何生成的附加信息。有了这些工件，我们可以对提示词质量进行详细分析，并对我们的“猫化”智能体进行迭代改进！

MLflow 的其他优势

幕后还有很多事情，超出了本教程的范围，但以下是构建智能体框架时非常有用的其他 MLflow 功能的快速列表。

• 依赖管理：当您记录模型时，MLflow 会自动尝试推断您的 pip 需求。这些需求以多种格式编写，使模型的远程服务变得更加简单。如上所述，如果您有本地依赖项，则可以在记录模型时通过 code_paths 参数指定要序列化的额外路径。
• 模型别名：在迭代构建智能体框架时，您需要一种简单的方法来比较模型。MLflow 模型 别名和标签 有助于查找 MLflow 模型注册表，并允许您轻松加载和部署特定的模型版本。
• 嵌套运行：对于智能体框架，特别是在训练底层 LLM 组件时，您通常会拥有复杂的嵌套结构。MLflow 支持 嵌套运行 以促进运行信息的聚合。这在 LLM 训练和微调中特别有用。

总结

在本博客中，我们概述了如何使用 AutoGen 创建一个复杂的智能体。我们还展示了如何利用 MLflow 基于代码的模型 功能来记录和加载我们的模型。最后，我们利用 MLflow AutoGen 的自动记录功能，自动利用 MLflow 跟踪来获取细粒度且线程安全的智能体执行信息。

祝您编码愉快！