{
"cells": [
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Model Context\n",
"\n",
"A model context supports storage and retrieval of Chat Completion messages.\n",
"It is always used together with a model client to generate LLM-based responses.\n",
"\n",
"For example, {py:mod}`~autogen_core.model_context.BufferedChatCompletionContext`\n",
"is a most-recent-used (MRU) context that stores the most recent `buffer_size`\n",
"number of messages. This is useful to avoid context overflow in many LLMs.\n",
"\n",
"Let's see an example that uses\n",
"{py:mod}`~autogen_core.model_context.BufferedChatCompletionContext`."
]
},
{
"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",
" # Prepare input to the chat completion model.\n",
" user_message = UserMessage(content=message.content, source=\"user\")\n",
" # Add message to model context.\n",
" await self._model_context.add_message(user_message)\n",
" # Generate a response.\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",
" # Return with the model's response.\n",
" assert isinstance(response.content, str)\n",
" # Add message to model context.\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": [
"Now let's try to ask follow up questions after the first one."
]
},
{
"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": [
"runtime = SingleThreadedAgentRuntime()\n",
"await SimpleAgentWithContext.register(\n",
" runtime,\n",
" \"simple_agent_context\",\n",
" lambda: SimpleAgentWithContext(\n",
" OpenAIChatCompletionClient(\n",
" model=\"gpt-4o-mini\",\n",
" # api_key=\"sk-...\", # Optional if you have an OPENAI_API_KEY set in the environment.\n",
" )\n",
" ),\n",
")\n",
"# Start the runtime processing messages.\n",
"runtime.start()\n",
"agent_id = AgentId(\"simple_agent_context\", \"default\")\n",
"\n",
"# First question.\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",
"# Second question.\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",
"# Stop the runtime processing messages.\n",
"await runtime.stop()"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"From the second response, you can see the agent now can recall its own previous responses."
]
}
],
"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
}