{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## 模型上下文\n\n模型上下文支持聊天补全消息的存储和检索。\n它总是与模型客户端一起使用，以生成基于LLM的响应。\n\n例如，{py:mod}`~autogen_core.model_context.BufferedChatCompletionContext`\n是一个最近使用(MRU)上下文，存储最近`buffer_size`数量的消息。\n这在避免许多LLM中的上下文溢出时非常有用。\n\n让我们看一个使用\n{py:mod}`~autogen_core.model_context.BufferedChatCompletionContext`的示例。\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 1,
      "metadata": {},
      "outputs": [],
      "source": [
        "from dataclasses import dataclass\n",
        "\n",
        "from autogen_core import AgentId, MessageContext, RoutedAgent, SingleThreadedAgentRuntime, message_handler\n",
        "from autogen_core.model_context import BufferedChatCompletionContext\n",
        "from autogen_core.models import AssistantMessage, ChatCompletionClient, SystemMessage, UserMessage\n",
        "from autogen_ext.models.openai import OpenAIChatCompletionClient"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 2,
      "metadata": {},
      "outputs": [],
      "source": [
        "@dataclass\n",
        "class Message:\n",
        "    content: str"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 3,
      "metadata": {},
      "outputs": [],
      "source": [
        "class SimpleAgentWithContext(RoutedAgent):\n",
        "    def __init__(self, model_client: ChatCompletionClient) -> None:\n",
        "        super().__init__(\"A simple agent\")\n",
        "        self._system_messages = [SystemMessage(content=\"You are a helpful AI assistant.\")]\n",
        "        self._model_client = model_client\n",
        "        self._model_context = BufferedChatCompletionContext(buffer_size=5)\n",
        "\n",
        "    @message_handler\n",
        "    async def handle_user_message(self, message: Message, ctx: MessageContext) -> Message:\n",
        "        # 准备聊天补全模型的输入。\n",
        "        user_message = UserMessage(content=message.content, source=\"user\")\n",
        "        # 向模型上下文添加消息。\n",
        "        await self._model_context.add_message(user_message)\n",
        "        # 生成响应。\n",
        "        response = await self._model_client.create(\n",
        "            self._system_messages + (await self._model_context.get_messages()),\n",
        "            cancellation_token=ctx.cancellation_token,\n",
        "        )\n",
        "        # 返回模型的响应。\n",
        "        assert isinstance(response.content, str)\n",
        "        # 将消息添加到模型上下文中\n",
        "        await self._model_context.add_message(AssistantMessage(content=response.content, source=self.metadata[\"type\"]))\n",
        "        return Message(content=response.content)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "现在让我们尝试在第一个问题之后提出后续问题\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 4,
      "metadata": {},
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Question: Hello, what are some fun things to do in Seattle?\n",
            "Response: Seattle offers a variety of fun activities and attractions. Here are some highlights:\n",
            "\n",
            "1. **Pike Place Market**: Visit this iconic market to explore local vendors, fresh produce, artisanal products, and watch the famous fish throwing.\n",
            "\n",
            "2. **Space Needle**: Take a trip to the observation deck for stunning panoramic views of the city, Puget Sound, and the surrounding mountains.\n",
            "\n",
            "3. **Chihuly Garden and Glass**: Marvel at the stunning glass art installations created by artist Dale Chihuly, located right next to the Space Needle.\n",
            "\n",
            "4. **Seattle Waterfront**: Enjoy a stroll along the waterfront, visit the Seattle Aquarium, and take a ferry ride to nearby islands like Bainbridge Island.\n",
            "\n",
            "5. **Museum of Pop Culture (MoPOP)**: Explore exhibits on music, science fiction, and pop culture in this architecturally striking building.\n",
            "\n",
            "6. **Seattle Art Museum (SAM)**: Discover an extensive collection of art from around the world, including contemporary and Native American art.\n",
            "\n",
            "7. **Gas Works Park**: Relax in this unique park that features remnants of an old gasification plant, offering great views of the Seattle skyline and Lake Union.\n",
            "\n",
            "8. **Discovery Park**: Enjoy nature trails, beaches, and beautiful views of the Puget Sound and the Olympic Mountains in this large urban park.\n",
            "\n",
            "9. **Ballard Locks**: Watch boats navigate the locks and see fish swimming upstream during the salmon migration season.\n",
            "\n",
            "10. **Fremont Troll**: Check out this quirky public art installation under a bridge in the Fremont neighborhood.\n",
            "\n",
            "11. **Underground Tour**: Take an entertaining guided tour through the underground passages of Pioneer Square to learn about Seattle's history.\n",
            "\n",
            "12. **Brewery Tours**: Seattle is known for its craft beer scene. Visit local breweries for tastings and tours.\n",
            "\n",
            "13. **Seattle Center**: Explore the cultural complex that includes the Space Needle, MoPOP, and various festivals and events throughout the year.\n",
            "\n",
            "These are just a few options, and Seattle has something for everyone, whether you're into outdoor activities, culture, history, or food!\n",
            "-----\n",
            "Question: What was the first thing you mentioned?\n",
            "Response: The first thing I mentioned was **Pike Place Market**. It's an iconic market in Seattle known for its local vendors, fresh produce, artisanal products, and the famous fish throwing by the fishmongers. It's a vibrant place full of sights, sounds, and delicious food.\n"
          ]
        }
      ],
      "source": [
        "model_client = OpenAIChatCompletionClient(\n",
        "    model=\"gpt-4o-mini\",\n",
        "    # api_key=\"sk-...\", # 可选，如果你在环境变量中设置了OPENAI_API_KEY\n",
        ")\n",
        "\n",
        "runtime = SingleThreadedAgentRuntime()\n",
        "await SimpleAgentWithContext.register(\n",
        "    runtime,\n",
        "    \"simple_agent_context\",\n",
        "    lambda: SimpleAgentWithContext(model_client=model_client),\n",
        ")\n",
        "# 开始运行时处理消息\n",
        "runtime.start()\n",
        "agent_id = AgentId(\"simple_agent_context\", \"default\")\n",
        "\n",
        "# 第一个问题\n",
        "message = Message(\"Hello, what are some fun things to do in Seattle?\")\n",
        "print(f\"Question: {message.content}\")\n",
        "response = await runtime.send_message(message, agent_id)\n",
        "print(f\"Response: {response.content}\")\n",
        "print(\"-----\")\n",
        "\n",
        "# 第二个问题。\n",
        "message = Message(\"What was the first thing you mentioned?\")\n",
        "print(f\"Question: {message.content}\")\n",
        "response = await runtime.send_message(message, agent_id)\n",
        "print(f\"Response: {response.content}\")\n",
        "\n",
        "# 停止运行时处理消息。\n",
        "await runtime.stop()\n",
        "await model_client.close()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "从第二个响应可以看出，代理现在能够回忆自己之前的响应了。\n"
      ]
    }
  ],
  "metadata": {
    "kernelspec": {
      "display_name": ".venv",
      "language": "python",
      "name": "python3"
    },
    "language_info": {
      "codemirror_mode": {
        "name": "ipython",
        "version": 3
      },
      "file_extension": ".py",
      "mimetype": "text/x-python",
      "name": "python",
      "nbconvert_exporter": "python",
      "pygments_lexer": "ipython3",
      "version": "3.12.7"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 2
}