diff --git a/python/packages/autogen-agentchat/src/autogen_agentchat/agents/_assistant_agent.py b/python/packages/autogen-agentchat/src/autogen_agentchat/agents/_assistant_agent.py index f60a44dbb90c..4a359eded38c 100644 --- a/python/packages/autogen-agentchat/src/autogen_agentchat/agents/_assistant_agent.py +++ b/python/packages/autogen-agentchat/src/autogen_agentchat/agents/_assistant_agent.py @@ -14,6 +14,7 @@ ) from autogen_core import CancellationToken, FunctionCall +from autogen_core.memory import Memory from autogen_core.model_context import ( ChatCompletionContext, UnboundedChatCompletionContext, @@ -35,6 +36,7 @@ AgentEvent, ChatMessage, HandoffMessage, + MemoryQueryEvent, MultiModalMessage, TextMessage, ToolCallExecutionEvent, @@ -120,6 +122,7 @@ class AssistantAgent(BaseChatAgent): will be returned as the response. Available variables: `{tool_name}`, `{arguments}`, `{result}`. For example, `"{tool_name}: {result}"` will create a summary like `"tool_name: result"`. + memory (Sequence[Memory] | None, optional): The memory store to use for the agent. Defaults to `None`. Raises: ValueError: If tool names are not unique. @@ -240,9 +243,20 @@ def __init__( ) = "You are a helpful AI assistant. Solve tasks using your tools. Reply with TERMINATE when the task has been completed.", reflect_on_tool_use: bool = False, tool_call_summary_format: str = "{result}", + memory: Sequence[Memory] | None = None, ): super().__init__(name=name, description=description) self._model_client = model_client + self._memory = None + if memory is not None: + if isinstance(memory, list): + self._memory = memory + else: + raise TypeError(f"Expected Memory, List[Memory], or None, got {type(memory)}") + + self._system_messages: List[ + SystemMessage | UserMessage | AssistantMessage | FunctionExecutionResultMessage + ] = [] if system_message is None: self._system_messages = [] else: @@ -325,6 +339,17 @@ async def on_messages_stream( # Inner messages. inner_messages: List[AgentEvent | ChatMessage] = [] + # Update the model context with memory content. + if self._memory: + for memory in self._memory: + update_context_result = await memory.update_context(self._model_context) + if update_context_result and len(update_context_result.memories.results) > 0: + memory_query_event_msg = MemoryQueryEvent( + content=update_context_result.memories.results, source=self.name + ) + inner_messages.append(memory_query_event_msg) + yield memory_query_event_msg + # Generate an inference result based on the current model context. llm_messages = self._system_messages + await self._model_context.get_messages() result = await self._model_client.create( diff --git a/python/packages/autogen-agentchat/src/autogen_agentchat/messages.py b/python/packages/autogen-agentchat/src/autogen_agentchat/messages.py index 21fb32d9d584..6069c8ddc8dd 100644 --- a/python/packages/autogen-agentchat/src/autogen_agentchat/messages.py +++ b/python/packages/autogen-agentchat/src/autogen_agentchat/messages.py @@ -8,6 +8,7 @@ class and includes specific fields relevant to the type of message being sent. from typing import List, Literal from autogen_core import FunctionCall, Image +from autogen_core.memory import MemoryContent from autogen_core.models import FunctionExecutionResult, RequestUsage from pydantic import BaseModel, ConfigDict, Field from typing_extensions import Annotated @@ -115,6 +116,15 @@ class UserInputRequestedEvent(BaseAgentEvent): type: Literal["UserInputRequestedEvent"] = "UserInputRequestedEvent" +class MemoryQueryEvent(BaseAgentEvent): + """An event signaling the results of memory queries.""" + + content: List[MemoryContent] + """The memory query results.""" + + type: Literal["MemoryQueryEvent"] = "MemoryQueryEvent" + + ChatMessage = Annotated[ TextMessage | MultiModalMessage | StopMessage | ToolCallSummaryMessage | HandoffMessage, Field(discriminator="type") ] @@ -122,7 +132,8 @@ class UserInputRequestedEvent(BaseAgentEvent): AgentEvent = Annotated[ - ToolCallRequestEvent | ToolCallExecutionEvent | UserInputRequestedEvent, Field(discriminator="type") + ToolCallRequestEvent | ToolCallExecutionEvent | MemoryQueryEvent | UserInputRequestedEvent, + Field(discriminator="type"), ] """Events emitted by agents and teams when they work, not used for agent-to-agent communication.""" @@ -138,5 +149,6 @@ class UserInputRequestedEvent(BaseAgentEvent): "ToolCallExecutionEvent", "ToolCallRequestEvent", "ToolCallSummaryMessage", + "MemoryQueryEvent", "UserInputRequestedEvent", ] diff --git a/python/packages/autogen-agentchat/tests/test_assistant_agent.py b/python/packages/autogen-agentchat/tests/test_assistant_agent.py index ca079ce407b4..930b4f8f7959 100644 --- a/python/packages/autogen-agentchat/tests/test_assistant_agent.py +++ b/python/packages/autogen-agentchat/tests/test_assistant_agent.py @@ -10,6 +10,7 @@ from autogen_agentchat.messages import ( ChatMessage, HandoffMessage, + MemoryQueryEvent, MultiModalMessage, TextMessage, ToolCallExecutionEvent, @@ -17,6 +18,7 @@ ToolCallSummaryMessage, ) from autogen_core import Image +from autogen_core.memory import ListMemory, Memory, MemoryContent, MemoryMimeType, MemoryQueryResult from autogen_core.model_context import BufferedChatCompletionContext from autogen_core.models import LLMMessage from autogen_core.models._model_client import ModelFamily @@ -508,4 +510,85 @@ async def test_model_context(monkeypatch: pytest.MonkeyPatch) -> None: # Check if the mock client is called with only the last two messages. assert len(mock.calls) == 1 - assert len(mock.calls[0]) == 3 # 2 message from the context + 1 system message + # 2 message from the context + 1 system message + assert len(mock.calls[0]) == 3 + + +@pytest.mark.asyncio +async def test_run_with_memory(monkeypatch: pytest.MonkeyPatch) -> None: + model = "gpt-4o-2024-05-13" + chat_completions = [ + ChatCompletion( + id="id1", + choices=[ + Choice( + finish_reason="stop", + index=0, + message=ChatCompletionMessage(content="Hello", role="assistant"), + ) + ], + created=0, + model=model, + object="chat.completion", + usage=CompletionUsage(prompt_tokens=10, completion_tokens=5, total_tokens=0), + ), + ] + b64_image_str = "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAIAAACQd1PeAAAADElEQVR4nGP4//8/AAX+Av4N70a4AAAAAElFTkSuQmCC" + mock = _MockChatCompletion(chat_completions) + monkeypatch.setattr(AsyncCompletions, "create", mock.mock_create) + + # Test basic memory properties and empty context + memory = ListMemory(name="test_memory") + assert memory.name == "test_memory" + + empty_context = BufferedChatCompletionContext(buffer_size=2) + empty_results = await memory.update_context(empty_context) + assert len(empty_results.memories.results) == 0 + + # Test various content types + memory = ListMemory() + await memory.add(MemoryContent(content="text content", mime_type=MemoryMimeType.TEXT)) + await memory.add(MemoryContent(content={"key": "value"}, mime_type=MemoryMimeType.JSON)) + await memory.add(MemoryContent(content=Image.from_base64(b64_image_str), mime_type=MemoryMimeType.IMAGE)) + + # Test query functionality + query_result = await memory.query(MemoryContent(content="", mime_type=MemoryMimeType.TEXT)) + assert isinstance(query_result, MemoryQueryResult) + # Should have all three memories we added + assert len(query_result.results) == 3 + + # Test clear and cleanup + await memory.clear() + empty_query = await memory.query(MemoryContent(content="", mime_type=MemoryMimeType.TEXT)) + assert len(empty_query.results) == 0 + await memory.close() # Should not raise + + # Test invalid memory type + with pytest.raises(TypeError): + AssistantAgent( + "test_agent", + model_client=OpenAIChatCompletionClient(model=model, api_key=""), + memory="invalid", # type: ignore + ) + + # Test with agent + memory2 = ListMemory() + await memory2.add(MemoryContent(content="test instruction", mime_type=MemoryMimeType.TEXT)) + + agent = AssistantAgent( + "test_agent", model_client=OpenAIChatCompletionClient(model=model, api_key=""), memory=[memory2] + ) + + result = await agent.run(task="test task") + assert len(result.messages) > 0 + memory_event = next((msg for msg in result.messages if isinstance(msg, MemoryQueryEvent)), None) + assert memory_event is not None + assert len(memory_event.content) > 0 + assert isinstance(memory_event.content[0], MemoryContent) + + # Test memory protocol + class BadMemory: + pass + + assert not isinstance(BadMemory(), Memory) + assert isinstance(ListMemory(), Memory) diff --git a/python/packages/autogen-core/docs/src/user-guide/agentchat-user-guide/index.md b/python/packages/autogen-core/docs/src/user-guide/agentchat-user-guide/index.md index 7a764c4c7340..5546417eb6d2 100644 --- a/python/packages/autogen-core/docs/src/user-guide/agentchat-user-guide/index.md +++ b/python/packages/autogen-core/docs/src/user-guide/agentchat-user-guide/index.md @@ -92,6 +92,7 @@ tutorial/termination tutorial/custom-agents tutorial/state tutorial/declarative +tutorial/memory ``` ```{toctree} diff --git a/python/packages/autogen-core/docs/src/user-guide/agentchat-user-guide/tutorial/memory.ipynb b/python/packages/autogen-core/docs/src/user-guide/agentchat-user-guide/tutorial/memory.ipynb new file mode 100644 index 000000000000..6a037945b69e --- /dev/null +++ b/python/packages/autogen-core/docs/src/user-guide/agentchat-user-guide/tutorial/memory.ipynb @@ -0,0 +1,188 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Memory \n", + "\n", + "There are several use cases where it is valuable to maintain a _store_ of useful facts that can be intelligently added to the context of the agent just before a specific step. The typically use case here is a RAG pattern where a query is used to retrieve relevant information from a database that is then added to the agent's context.\n", + "\n", + "\n", + "AgentChat provides a {py:class}`~autogen_core.memory.Memory` protocol that can be extended to provide this functionality. The key methods are `query`, `update_context`, `add`, `clear`, and `close`. \n", + "\n", + "- `add`: add new entries to the memory store\n", + "- `query`: retrieve relevant information from the memory store \n", + "- `update_context`: mutate an agent's internal `model_context` by adding the retrieved information (used in the {py:class}`~autogen_agentchat.agents.AssistantAgent` class) \n", + "- `clear`: clear all entries from the memory store\n", + "- `close`: clean up any resources used by the memory store \n", + "\n", + "\n", + "## ListMemory Example\n", + "\n", + "{py:class}~autogen_core.memory.ListMemory is provided as an example implementation of the {py:class}~autogen_core.memory.Memory protocol. It is a simple list-based memory implementation that maintains memories in chronological order, appending the most recent memories to the model's context. The implementation is designed to be straightforward and predictable, making it easy to understand and debug.\n", + "In the following example, we will use ListMemory to maintain a memory bank of user preferences and demonstrate how it can be used to provide consistent context for agent responses over time." + ] + }, + { + "cell_type": "code", + "execution_count": 1, + "metadata": {}, + "outputs": [], + "source": [ + "from autogen_agentchat.agents import AssistantAgent\n", + "from autogen_agentchat.ui import Console\n", + "from autogen_core.memory import ListMemory, MemoryContent, MemoryMimeType\n", + "from autogen_ext.models.openai import OpenAIChatCompletionClient" + ] + }, + { + "cell_type": "code", + "execution_count": 2, + "metadata": {}, + "outputs": [], + "source": [ + "# Initialize user memory\n", + "user_memory = ListMemory()\n", + "\n", + "# Add user preferences to memory\n", + "await user_memory.add(MemoryContent(content=\"The weather should be in metric units\", mime_type=MemoryMimeType.TEXT))\n", + "\n", + "await user_memory.add(MemoryContent(content=\"Meal recipe must be vegan\", mime_type=MemoryMimeType.TEXT))\n", + "\n", + "\n", + "async def get_weather(city: str, units: str = \"imperial\") -> str:\n", + " if units == \"imperial\":\n", + " return f\"The weather in {city} is 73 °F and Sunny.\"\n", + " elif units == \"metric\":\n", + " return f\"The weather in {city} is 23 °C and Sunny.\"\n", + " else:\n", + " return f\"Sorry, I don't know the weather in {city}.\"\n", + "\n", + "\n", + "assistant_agent = AssistantAgent(\n", + " name=\"assistant_agent\",\n", + " model_client=OpenAIChatCompletionClient(\n", + " model=\"gpt-4o-2024-08-06\",\n", + " ),\n", + " tools=[get_weather],\n", + " memory=[user_memory],\n", + ")" + ] + }, + { + "cell_type": "code", + "execution_count": 3, + "metadata": {}, + "outputs": [ + { + "data": { + "text/plain": [ + "TaskResult(messages=[TextMessage(source='user', models_usage=None, content='What is the weather in New York?', type='TextMessage'), MemoryQueryEvent(source='assistant_agent', models_usage=None, content=[MemoryContent(content='The weather should be in metric units', mime_type=, metadata=None, timestamp=None, source=None, score=None), MemoryContent(content='Meal recipe must be vegan', mime_type=, metadata=None, timestamp=None, source=None, score=None)], type='MemoryQueryEvent'), ToolCallRequestEvent(source='assistant_agent', models_usage=RequestUsage(prompt_tokens=123, completion_tokens=20), content=[FunctionCall(id='call_pHq4p89gW6oGjGr3VsVETCYX', arguments='{\"city\":\"New York\",\"units\":\"metric\"}', name='get_weather')], type='ToolCallRequestEvent'), ToolCallExecutionEvent(source='assistant_agent', models_usage=None, content=[FunctionExecutionResult(content='The weather in New York is 23 °C and Sunny.', call_id='call_pHq4p89gW6oGjGr3VsVETCYX')], type='ToolCallExecutionEvent'), ToolCallSummaryMessage(source='assistant_agent', models_usage=None, content='The weather in New York is 23 °C and Sunny.', type='ToolCallSummaryMessage')], stop_reason=None)" + ] + }, + "execution_count": 3, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "# Run the agent with a task.\n", + "stream = assistant_agent.run_stream(task=\"What is the weather in New York?\")\n", + "await Console(stream)" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "We can inspect that the `assistant_agent` model_context is actually updated with the retrieved memory entries. The `transform` method is used to format the retrieved memory entries into a string that can be used by the agent. In this case, we simply concatenate the content of each memory entry into a single string." + ] + }, + { + "cell_type": "code", + "execution_count": 4, + "metadata": {}, + "outputs": [ + { + "data": { + "text/plain": [ + "[UserMessage(content='What is the weather in New York?', source='user', type='UserMessage'),\n", + " SystemMessage(content='\\nRelevant memory content (in chronological order):\\n1. The weather should be in metric units\\n2. Meal recipe must be vegan\\n', type='SystemMessage'),\n", + " AssistantMessage(content=[FunctionCall(id='call_pHq4p89gW6oGjGr3VsVETCYX', arguments='{\"city\":\"New York\",\"units\":\"metric\"}', name='get_weather')], source='assistant_agent', type='AssistantMessage'),\n", + " FunctionExecutionResultMessage(content=[FunctionExecutionResult(content='The weather in New York is 23 °C and Sunny.', call_id='call_pHq4p89gW6oGjGr3VsVETCYX')], type='FunctionExecutionResultMessage')]" + ] + }, + "execution_count": 4, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "await assistant_agent._model_context.get_messages()" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "We see above that the weather is returned in Centigrade as stated in the user preferences. \n", + "\n", + "Similarly, assuming we ask a separate question about generating a meal plan, the agent is able to retrieve relevant information from the memory store and provide a personalized response." + ] + }, + { + "cell_type": "code", + "execution_count": 5, + "metadata": {}, + "outputs": [ + { + "data": { + "text/plain": [ + "TaskResult(messages=[TextMessage(source='user', models_usage=None, content='Write brief meal recipe with broth', type='TextMessage'), MemoryQueryEvent(source='assistant_agent', models_usage=None, content=[MemoryContent(content='The weather should be in metric units', mime_type=, metadata=None, timestamp=None, source=None, score=None), MemoryContent(content='Meal recipe must be vegan', mime_type=, metadata=None, timestamp=None, source=None, score=None)], type='MemoryQueryEvent'), TextMessage(source='assistant_agent', models_usage=RequestUsage(prompt_tokens=208, completion_tokens=253), content=\"Here's a brief vegan meal recipe using broth:\\n\\n**Vegan Mushroom & Herb Broth Soup**\\n\\n**Ingredients:**\\n- 1 tablespoon olive oil\\n- 1 onion, diced\\n- 2 cloves garlic, minced\\n- 250g mushrooms, sliced\\n- 1 carrot, diced\\n- 1 celery stalk, diced\\n- 4 cups vegetable broth\\n- 1 teaspoon thyme\\n- 1 teaspoon rosemary\\n- Salt and pepper to taste\\n- Fresh parsley for garnish\\n\\n**Instructions:**\\n1. Heat the olive oil in a large pot over medium heat. Add the diced onion and garlic, and sauté until the onion becomes translucent.\\n\\n2. Add the sliced mushrooms, carrot, and celery. Continue to sauté until the mushrooms are cooked through and the vegetables begin to soften, about 5 minutes.\\n\\n3. Pour in the vegetable broth. Stir in the thyme and rosemary, and bring the mixture to a boil.\\n\\n4. Reduce the heat to low and let the soup simmer for about 15 minutes, allowing the flavors to meld together.\\n\\n5. Season with salt and pepper to taste.\\n\\n6. Serve hot, garnished with fresh parsley.\\n\\nEnjoy your warm and comforting vegan mushroom & herb broth soup! \\n\\nTERMINATE\", type='TextMessage')], stop_reason=None)" + ] + }, + "execution_count": 5, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "stream = assistant_agent.run_stream(task=\"Write brief meal recipe with broth\")\n", + "await Console(stream)" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Custom Memory Stores (Vector DBs, etc.)\n", + "\n", + "You can build on the `Memory` protocol to implement more complex memory stores. For example, you could implement a custom memory store that uses a vector database to store and retrieve information, or a memory store that uses a machine learning model to generate personalized responses based on the user's preferences etc.\n", + "\n", + "Specifically, you will need to overload the `add`, `query` and `update_context` methods to implement the desired functionality and pass the memory store to your agent.\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.11.9" + } + }, + "nbformat": 4, + "nbformat_minor": 2 +} diff --git a/python/packages/autogen-core/docs/src/user-guide/extensions-user-guide/azure-container-code-executor.ipynb b/python/packages/autogen-core/docs/src/user-guide/extensions-user-guide/azure-container-code-executor.ipynb index c71ee58e9118..692a2afda33e 100644 --- a/python/packages/autogen-core/docs/src/user-guide/extensions-user-guide/azure-container-code-executor.ipynb +++ b/python/packages/autogen-core/docs/src/user-guide/extensions-user-guide/azure-container-code-executor.ipynb @@ -1,280 +1,280 @@ { - "cells": [ - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "# ACA Dynamic Sessions Code Executor\n", - "\n", - "This guide will explain the Azure Container Apps dynamic sessions in Azure Container Apps and show you how to use the Azure Container Code Executor class.\n", - "\n", - "The [Azure Container Apps dynamic sessions](https://learn.microsoft.com/en-us/azure/container-apps/sessions) is a component in the Azure Container Apps service. The environment is hosted on remote Azure instances and will not execute any code locally. The interpreter is capable of executing python code in a jupyter environment with a pre-installed base of commonly used packages. [Custom environments](https://learn.microsoft.com/en-us/azure/container-apps/sessions-custom-container) can be created by users for their applications. Files can additionally be [uploaded to, or downloaded from](https://learn.microsoft.com/en-us/azure/container-apps/sessions-code-interpreter#upload-a-file-to-a-session) each session.\n", - "\n", - "The code interpreter can run multiple sessions of code, each of which are delineated by a session identifier string.\n", - "\n", - "## Create a Container Apps Session Pool\n", - "\n", - "In your Azure portal, create a new `Container App Session Pool` resource with the pool type set to `Python code interpreter` and note the `Pool management endpoint`. The format for the endpoint should be something like `https://{region}.dynamicsessions.io/subscriptions/{subscription_id}/resourceGroups/{resource_group_name}/sessionPools/{session_pool_name}`.\n", - "\n", - "Alternatively, you can use the [Azure CLI to create a session pool.](https://learn.microsoft.com/en-us/azure/container-apps/sessions-code-interpreter#create-a-session-pool-with-azure-cli)\n", - "\n", - "## ACADynamicSessionsCodeExecutor\n", - "\n", - "The {py:class}`~autogen_ext.code_executors.azure.ACADynamicSessionsCodeExecutor` class is a python code executor that creates and executes arbitrary python code on a default Serverless code interpreter session. Its interface is as follows\n", - "\n", - "### Initialization\n", - "\n", - "First, you will need to find or create a credentialing object that implements the {py:class}`~autogen_ext.code_executors.azure.TokenProvider` interface. This is any object that implements the following function\n", - "```python\n", - "def get_token(\n", - " self, *scopes: str, claims: Optional[str] = None, tenant_id: Optional[str] = None, **kwargs: Any\n", - ") -> azure.core.credentials.AccessToken\n", - "```\n", - "An example of such an object is the [azure.identity.DefaultAzureCredential](https://learn.microsoft.com/en-us/python/api/azure-identity/azure.identity.defaultazurecredential?view=azure-python) class.\n", - "\n", - "Lets start by installing that" - ] + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "# ACA Dynamic Sessions Code Executor\n", + "\n", + "This guide will explain the Azure Container Apps dynamic sessions in Azure Container Apps and show you how to use the Azure Container Code Executor class.\n", + "\n", + "The [Azure Container Apps dynamic sessions](https://learn.microsoft.com/en-us/azure/container-apps/sessions) is a component in the Azure Container Apps service. The environment is hosted on remote Azure instances and will not execute any code locally. The interpreter is capable of executing python code in a jupyter environment with a pre-installed base of commonly used packages. [Custom environments](https://learn.microsoft.com/en-us/azure/container-apps/sessions-custom-container) can be created by users for their applications. Files can additionally be [uploaded to, or downloaded from](https://learn.microsoft.com/en-us/azure/container-apps/sessions-code-interpreter#upload-a-file-to-a-session) each session.\n", + "\n", + "The code interpreter can run multiple sessions of code, each of which are delineated by a session identifier string.\n", + "\n", + "## Create a Container Apps Session Pool\n", + "\n", + "In your Azure portal, create a new `Container App Session Pool` resource with the pool type set to `Python code interpreter` and note the `Pool management endpoint`. The format for the endpoint should be something like `https://{region}.dynamicsessions.io/subscriptions/{subscription_id}/resourceGroups/{resource_group_name}/sessionPools/{session_pool_name}`.\n", + "\n", + "Alternatively, you can use the [Azure CLI to create a session pool.](https://learn.microsoft.com/en-us/azure/container-apps/sessions-code-interpreter#create-a-session-pool-with-azure-cli)\n", + "\n", + "## ACADynamicSessionsCodeExecutor\n", + "\n", + "The {py:class}`~autogen_ext.code_executors.azure.ACADynamicSessionsCodeExecutor` class is a python code executor that creates and executes arbitrary python code on a default Serverless code interpreter session. Its interface is as follows\n", + "\n", + "### Initialization\n", + "\n", + "First, you will need to find or create a credentialing object that implements the {py:class}`~autogen_ext.code_executors.azure.TokenProvider` interface. This is any object that implements the following function\n", + "```python\n", + "def get_token(\n", + " self, *scopes: str, claims: Optional[str] = None, tenant_id: Optional[str] = None, **kwargs: Any\n", + ") -> azure.core.credentials.AccessToken\n", + "```\n", + "An example of such an object is the [azure.identity.DefaultAzureCredential](https://learn.microsoft.com/en-us/python/api/azure-identity/azure.identity.defaultazurecredential?view=azure-python) class.\n", + "\n", + "Lets start by installing that" + ] + }, + { + "cell_type": "code", + "execution_count": 1, + "metadata": {}, + "outputs": [], + "source": [ + "# pip install azure.identity" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Next, lets import all the necessary modules and classes for our code" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "import os\n", + "import tempfile\n", + "\n", + "from anyio import open_file\n", + "from autogen_core import CancellationToken\n", + "from autogen_core.code_executor import CodeBlock\n", + "from autogen_ext.code_executors.azure import ACADynamicSessionsCodeExecutor\n", + "from azure.identity import DefaultAzureCredential" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Now, we create our Azure code executor and run some test code along with verification that it ran correctly. We'll create the executor with a temporary working directory to ensure a clean environment as we show how to use each feature" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "cancellation_token = CancellationToken()\n", + "POOL_MANAGEMENT_ENDPOINT = \"...\"\n", + "\n", + "with tempfile.TemporaryDirectory() as temp_dir:\n", + " executor = ACADynamicSessionsCodeExecutor(\n", + " pool_management_endpoint=POOL_MANAGEMENT_ENDPOINT, credential=DefaultAzureCredential(), work_dir=temp_dir\n", + " )\n", + "\n", + " code_blocks = [CodeBlock(code=\"import sys; print('hello world!')\", language=\"python\")]\n", + " code_result = await executor.execute_code_blocks(code_blocks, cancellation_token)\n", + " assert code_result.exit_code == 0 and \"hello world!\" in code_result.output" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Next, lets try uploading some files and verifying their integrity. All files uploaded to the Serverless code interpreter is uploaded into the `/mnt/data` directory. All downloadable files must also be placed in the directory. By default, the current working directory for the code executor is set to `/mnt/data`." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "with tempfile.TemporaryDirectory() as temp_dir:\n", + " test_file_1 = \"test_upload_1.txt\"\n", + " test_file_1_contents = \"test1 contents\"\n", + " test_file_2 = \"test_upload_2.txt\"\n", + " test_file_2_contents = \"test2 contents\"\n", + "\n", + " async with await open_file(os.path.join(temp_dir, test_file_1), \"w\") as f: # type: ignore[syntax]\n", + " await f.write(test_file_1_contents)\n", + " async with await open_file(os.path.join(temp_dir, test_file_2), \"w\") as f: # type: ignore[syntax]\n", + " await f.write(test_file_2_contents)\n", + "\n", + " assert os.path.isfile(os.path.join(temp_dir, test_file_1))\n", + " assert os.path.isfile(os.path.join(temp_dir, test_file_2))\n", + "\n", + " executor = ACADynamicSessionsCodeExecutor(\n", + " pool_management_endpoint=POOL_MANAGEMENT_ENDPOINT, credential=DefaultAzureCredential(), work_dir=temp_dir\n", + " )\n", + " await executor.upload_files([test_file_1, test_file_2], cancellation_token)\n", + "\n", + " file_list = await executor.get_file_list(cancellation_token)\n", + " assert test_file_1 in file_list\n", + " assert test_file_2 in file_list\n", + "\n", + " code_blocks = [\n", + " CodeBlock(\n", + " code=f\"\"\"\n", + "with open(\"{test_file_1}\") as f:\n", + " print(f.read())\n", + "with open(\"{test_file_2}\") as f:\n", + " print(f.read())\n", + "\"\"\",\n", + " language=\"python\",\n", + " )\n", + " ]\n", + " code_result = await executor.execute_code_blocks(code_blocks, cancellation_token)\n", + " assert code_result.exit_code == 0\n", + " assert test_file_1_contents in code_result.output\n", + " assert test_file_2_contents in code_result.output" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Downloading files works in a similar way." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "with tempfile.TemporaryDirectory() as temp_dir:\n", + " test_file_1 = \"test_upload_1.txt\"\n", + " test_file_1_contents = \"test1 contents\"\n", + " test_file_2 = \"test_upload_2.txt\"\n", + " test_file_2_contents = \"test2 contents\"\n", + "\n", + " assert not os.path.isfile(os.path.join(temp_dir, test_file_1))\n", + " assert not os.path.isfile(os.path.join(temp_dir, test_file_2))\n", + "\n", + " executor = ACADynamicSessionsCodeExecutor(\n", + " pool_management_endpoint=POOL_MANAGEMENT_ENDPOINT, credential=DefaultAzureCredential(), work_dir=temp_dir\n", + " )\n", + "\n", + " code_blocks = [\n", + " CodeBlock(\n", + " code=f\"\"\"\n", + "with open(\"{test_file_1}\", \"w\") as f:\n", + " f.write(\"{test_file_1_contents}\")\n", + "with open(\"{test_file_2}\", \"w\") as f:\n", + " f.write(\"{test_file_2_contents}\")\n", + "\"\"\",\n", + " language=\"python\",\n", + " ),\n", + " ]\n", + " code_result = await executor.execute_code_blocks(code_blocks, cancellation_token)\n", + " assert code_result.exit_code == 0\n", + "\n", + " file_list = await executor.get_file_list(cancellation_token)\n", + " assert test_file_1 in file_list\n", + " assert test_file_2 in file_list\n", + "\n", + " await executor.download_files([test_file_1, test_file_2], cancellation_token)\n", + "\n", + " assert os.path.isfile(os.path.join(temp_dir, test_file_1))\n", + " async with await open_file(os.path.join(temp_dir, test_file_1), \"r\") as f: # type: ignore[syntax]\n", + " content = await f.read()\n", + " assert test_file_1_contents in content\n", + " assert os.path.isfile(os.path.join(temp_dir, test_file_2))\n", + " async with await open_file(os.path.join(temp_dir, test_file_2), \"r\") as f: # type: ignore[syntax]\n", + " content = await f.read()\n", + " assert test_file_2_contents in content" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### New Sessions\n", + "\n", + "Every instance of the {py:class}`~autogen_ext.code_executors.azure.ACADynamicSessionsCodeExecutor` class will have a unique session ID. Every call to a particular code executor will be executed on the same session until the {py:meth}`~autogen_ext.code_executors.azure.ACADynamicSessionsCodeExecutor.restart` function is called on it. Previous sessions cannot be reused.\n", + "\n", + "Here we'll run some code on the code session, restart it, then verify that a new session has been opened." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "executor = ACADynamicSessionsCodeExecutor(\n", + " pool_management_endpoint=POOL_MANAGEMENT_ENDPOINT, credential=DefaultAzureCredential()\n", + ")\n", + "\n", + "code_blocks = [CodeBlock(code=\"x = 'abcdefg'\", language=\"python\")]\n", + "code_result = await executor.execute_code_blocks(code_blocks, cancellation_token)\n", + "assert code_result.exit_code == 0\n", + "\n", + "code_blocks = [CodeBlock(code=\"print(x)\", language=\"python\")]\n", + "code_result = await executor.execute_code_blocks(code_blocks, cancellation_token)\n", + "assert code_result.exit_code == 0 and \"abcdefg\" in code_result.output\n", + "\n", + "await executor.restart()\n", + "code_blocks = [CodeBlock(code=\"print(x)\", language=\"python\")]\n", + "code_result = await executor.execute_code_blocks(code_blocks, cancellation_token)\n", + "assert code_result.exit_code != 0 and \"NameError\" in code_result.output" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### Available Packages\n", + "\n", + "Each code execution instance is pre-installed with most of the commonly used packages. However, the list of available packages and versions are not available outside of the execution environment. The packages list on the environment can be retrieved by calling the {py:meth}`~autogen_ext.code_executors.azure.ACADynamicSessionsCodeExecutor.get_available_packages` function on the code executor." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "print(executor.get_available_packages(cancellation_token))" + ] + } + ], + "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.5" + } }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "# pip install azure.identity" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Next, lets import all the necessary modules and classes for our code" - ] - }, - { - "cell_type": "code", - "execution_count": 2, - "metadata": {}, - "outputs": [], - "source": [ - "import os\n", - "import tempfile\n", - "\n", - "from anyio import open_file\n", - "from autogen_core import CancellationToken\n", - "from autogen_core.code_executor import CodeBlock\n", - "from autogen_ext.code_executors.azure import ACADynamicSessionsCodeExecutor\n", - "from azure.identity import DefaultAzureCredential" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Now, we create our Azure code executor and run some test code along with verification that it ran correctly. We'll create the executor with a temporary working directory to ensure a clean environment as we show how to use each feature" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "cancellation_token = CancellationToken()\n", - "POOL_MANAGEMENT_ENDPOINT = \"...\"\n", - "\n", - "with tempfile.TemporaryDirectory() as temp_dir:\n", - " executor = ACADynamicSessionsCodeExecutor(\n", - " pool_management_endpoint=POOL_MANAGEMENT_ENDPOINT, credential=DefaultAzureCredential(), work_dir=temp_dir\n", - " )\n", - "\n", - " code_blocks = [CodeBlock(code=\"import sys; print('hello world!')\", language=\"python\")]\n", - " code_result = await executor.execute_code_blocks(code_blocks, cancellation_token)\n", - " assert code_result.exit_code == 0 and \"hello world!\" in code_result.output" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Next, lets try uploading some files and verifying their integrity. All files uploaded to the Serverless code interpreter is uploaded into the `/mnt/data` directory. All downloadable files must also be placed in the directory. By default, the current working directory for the code executor is set to `/mnt/data`." - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "with tempfile.TemporaryDirectory() as temp_dir:\n", - " test_file_1 = \"test_upload_1.txt\"\n", - " test_file_1_contents = \"test1 contents\"\n", - " test_file_2 = \"test_upload_2.txt\"\n", - " test_file_2_contents = \"test2 contents\"\n", - "\n", - " async with await open_file(os.path.join(temp_dir, test_file_1), \"w\") as f: # type: ignore[syntax]\n", - " await f.write(test_file_1_contents)\n", - " async with await open_file(os.path.join(temp_dir, test_file_2), \"w\") as f: # type: ignore[syntax]\n", - " await f.write(test_file_2_contents)\n", - "\n", - " assert os.path.isfile(os.path.join(temp_dir, test_file_1))\n", - " assert os.path.isfile(os.path.join(temp_dir, test_file_2))\n", - "\n", - " executor = ACADynamicSessionsCodeExecutor(\n", - " pool_management_endpoint=POOL_MANAGEMENT_ENDPOINT, credential=DefaultAzureCredential(), work_dir=temp_dir\n", - " )\n", - " await executor.upload_files([test_file_1, test_file_2], cancellation_token)\n", - "\n", - " file_list = await executor.get_file_list(cancellation_token)\n", - " assert test_file_1 in file_list\n", - " assert test_file_2 in file_list\n", - "\n", - " code_blocks = [\n", - " CodeBlock(\n", - " code=f\"\"\"\n", - "with open(\"{test_file_1}\") as f:\n", - " print(f.read())\n", - "with open(\"{test_file_2}\") as f:\n", - " print(f.read())\n", - "\"\"\",\n", - " language=\"python\",\n", - " )\n", - " ]\n", - " code_result = await executor.execute_code_blocks(code_blocks, cancellation_token)\n", - " assert code_result.exit_code == 0\n", - " assert test_file_1_contents in code_result.output\n", - " assert test_file_2_contents in code_result.output" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Downloading files works in a similar way." - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "with tempfile.TemporaryDirectory() as temp_dir:\n", - " test_file_1 = \"test_upload_1.txt\"\n", - " test_file_1_contents = \"test1 contents\"\n", - " test_file_2 = \"test_upload_2.txt\"\n", - " test_file_2_contents = \"test2 contents\"\n", - "\n", - " assert not os.path.isfile(os.path.join(temp_dir, test_file_1))\n", - " assert not os.path.isfile(os.path.join(temp_dir, test_file_2))\n", - "\n", - " executor = ACADynamicSessionsCodeExecutor(\n", - " pool_management_endpoint=POOL_MANAGEMENT_ENDPOINT, credential=DefaultAzureCredential(), work_dir=temp_dir\n", - " )\n", - "\n", - " code_blocks = [\n", - " CodeBlock(\n", - " code=f\"\"\"\n", - "with open(\"{test_file_1}\", \"w\") as f:\n", - " f.write(\"{test_file_1_contents}\")\n", - "with open(\"{test_file_2}\", \"w\") as f:\n", - " f.write(\"{test_file_2_contents}\")\n", - "\"\"\",\n", - " language=\"python\",\n", - " ),\n", - " ]\n", - " code_result = await executor.execute_code_blocks(code_blocks, cancellation_token)\n", - " assert code_result.exit_code == 0\n", - "\n", - " file_list = await executor.get_file_list(cancellation_token)\n", - " assert test_file_1 in file_list\n", - " assert test_file_2 in file_list\n", - "\n", - " await executor.download_files([test_file_1, test_file_2], cancellation_token)\n", - "\n", - " assert os.path.isfile(os.path.join(temp_dir, test_file_1))\n", - " async with await open_file(os.path.join(temp_dir, test_file_1), \"r\") as f: # type: ignore[syntax]\n", - " content = await f.read()\n", - " assert test_file_1_contents in content\n", - " assert os.path.isfile(os.path.join(temp_dir, test_file_2))\n", - " async with await open_file(os.path.join(temp_dir, test_file_2), \"r\") as f: # type: ignore[syntax]\n", - " content = await f.read()\n", - " assert test_file_2_contents in content" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "### New Sessions\n", - "\n", - "Every instance of the {py:class}`~autogen_ext.code_executors.azure.ACADynamicSessionsCodeExecutor` class will have a unique session ID. Every call to a particular code executor will be executed on the same session until the {py:meth}`~autogen_ext.code_executors.azure.ACADynamicSessionsCodeExecutor.restart` function is called on it. Previous sessions cannot be reused.\n", - "\n", - "Here we'll run some code on the code session, restart it, then verify that a new session has been opened." - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "executor = ACADynamicSessionsCodeExecutor(\n", - " pool_management_endpoint=POOL_MANAGEMENT_ENDPOINT, credential=DefaultAzureCredential()\n", - ")\n", - "\n", - "code_blocks = [CodeBlock(code=\"x = 'abcdefg'\", language=\"python\")]\n", - "code_result = await executor.execute_code_blocks(code_blocks, cancellation_token)\n", - "assert code_result.exit_code == 0\n", - "\n", - "code_blocks = [CodeBlock(code=\"print(x)\", language=\"python\")]\n", - "code_result = await executor.execute_code_blocks(code_blocks, cancellation_token)\n", - "assert code_result.exit_code == 0 and \"abcdefg\" in code_result.output\n", - "\n", - "await executor.restart()\n", - "code_blocks = [CodeBlock(code=\"print(x)\", language=\"python\")]\n", - "code_result = await executor.execute_code_blocks(code_blocks, cancellation_token)\n", - "assert code_result.exit_code != 0 and \"NameError\" in code_result.output" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "### Available Packages\n", - "\n", - "Each code execution instance is pre-installed with most of the commonly used packages. However, the list of available packages and versions are not available outside of the execution environment. The packages list on the environment can be retrieved by calling the {py:meth}`~autogen_ext.code_executors.azure.ACADynamicSessionsCodeExecutor.get_available_packages` function on the code executor." - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "print(executor.get_available_packages(cancellation_token))" - ] - } - ], - "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.5" - } - }, - "nbformat": 4, - "nbformat_minor": 2 + "nbformat": 4, + "nbformat_minor": 2 } diff --git a/python/packages/autogen-core/src/autogen_core/memory/__init__.py b/python/packages/autogen-core/src/autogen_core/memory/__init__.py new file mode 100644 index 000000000000..69a20f24f530 --- /dev/null +++ b/python/packages/autogen-core/src/autogen_core/memory/__init__.py @@ -0,0 +1,11 @@ +from ._base_memory import Memory, MemoryContent, MemoryMimeType, MemoryQueryResult, UpdateContextResult +from ._list_memory import ListMemory + +__all__ = [ + "Memory", + "MemoryContent", + "MemoryQueryResult", + "UpdateContextResult", + "MemoryMimeType", + "ListMemory", +] diff --git a/python/packages/autogen-core/src/autogen_core/memory/_base_memory.py b/python/packages/autogen-core/src/autogen_core/memory/_base_memory.py new file mode 100644 index 000000000000..0ef7ac716ccf --- /dev/null +++ b/python/packages/autogen-core/src/autogen_core/memory/_base_memory.py @@ -0,0 +1,103 @@ +from datetime import datetime +from enum import Enum +from typing import Any, Dict, List, Protocol, Union, runtime_checkable + +from pydantic import BaseModel, ConfigDict + +from .._cancellation_token import CancellationToken +from .._image import Image +from ..model_context import ChatCompletionContext + + +class MemoryMimeType(Enum): + """Supported MIME types for memory content.""" + + TEXT = "text/plain" + JSON = "application/json" + MARKDOWN = "text/markdown" + IMAGE = "image/*" + BINARY = "application/octet-stream" + + +ContentType = Union[str, bytes, Dict[str, Any], Image] + + +class MemoryContent(BaseModel): + content: ContentType + mime_type: MemoryMimeType | str + metadata: Dict[str, Any] | None = None + timestamp: datetime | None = None + source: str | None = None + score: float | None = None + + model_config = ConfigDict(arbitrary_types_allowed=True) + + +class MemoryQueryResult(BaseModel): + results: List[MemoryContent] + + +class UpdateContextResult(BaseModel): + memories: MemoryQueryResult + + +@runtime_checkable +class Memory(Protocol): + """Protocol defining the interface for memory implementations.""" + + @property + def name(self) -> str | None: + """The name of this memory implementation.""" + ... + + async def update_context( + self, + model_context: ChatCompletionContext, + ) -> UpdateContextResult: + """ + Update the provided model context using relevant memory content. + + Args: + model_context: The context to update. + + Returns: + UpdateContextResult containing relevant memories + """ + ... + + async def query( + self, + query: str | MemoryContent, + cancellation_token: CancellationToken | None = None, + **kwargs: Any, + ) -> MemoryQueryResult: + """ + Query the memory store and return relevant entries. + + Args: + query: Query content item + cancellation_token: Optional token to cancel operation + **kwargs: Additional implementation-specific parameters + + Returns: + MemoryQueryResult containing memory entries with relevance scores + """ + ... + + async def add(self, content: MemoryContent, cancellation_token: CancellationToken | None = None) -> None: + """ + Add a new content to memory. + + Args: + content: The memory content to add + cancellation_token: Optional token to cancel operation + """ + ... + + async def clear(self) -> None: + """Clear all entries from memory.""" + ... + + async def close(self) -> None: + """Clean up any resources used by the memory implementation.""" + ... diff --git a/python/packages/autogen-core/src/autogen_core/memory/_list_memory.py b/python/packages/autogen-core/src/autogen_core/memory/_list_memory.py new file mode 100644 index 000000000000..eda206783476 --- /dev/null +++ b/python/packages/autogen-core/src/autogen_core/memory/_list_memory.py @@ -0,0 +1,137 @@ +from typing import Any, List + +from .._cancellation_token import CancellationToken +from ..model_context import ChatCompletionContext +from ..models import SystemMessage +from ._base_memory import Memory, MemoryContent, MemoryQueryResult, UpdateContextResult + + +class ListMemory(Memory): + """Simple chronological list-based memory implementation. + + This memory implementation stores contents in a list and retrieves them in + chronological order. It has an `update_context` method that updates model contexts + by appending all stored memories. + + The memory content can be directly accessed and modified through the content property, + allowing external applications to manage memory contents directly. + + Example: + .. code-block:: python + # Initialize memory + memory = ListMemory(name="chat_history") + + # Add memory content + content = MemoryContent(content="User prefers formal language") + await memory.add(content) + + # Directly modify memory contents + memory.content = [MemoryContent(content="New preference")] + + # Update a model context with memory + memory_contents = await memory.update_context(model_context) + + + Attributes: + name (str): Identifier for this memory instance + content (List[MemoryContent]): Direct access to memory contents + """ + + def __init__(self, name: str | None = None) -> None: + """Initialize ListMemory. + + Args: + name: Optional identifier for this memory instance + """ + self._name = name or "default_list_memory" + self._contents: List[MemoryContent] = [] + + @property + def name(self) -> str: + """Get the memory instance identifier. + + Returns: + str: Memory instance name + """ + return self._name + + @property + def content(self) -> List[MemoryContent]: + """Get the current memory contents. + + Returns: + List[MemoryContent]: List of stored memory contents + """ + return self._contents + + @content.setter + def content(self, value: List[MemoryContent]) -> None: + """Set the memory contents. + + Args: + value: New list of memory contents to store + """ + self._contents = value + + async def update_context( + self, + model_context: ChatCompletionContext, + ) -> UpdateContextResult: + """Update the model context by appending memory content. + + This method mutates the provided model_context by adding all memories as a + SystemMessage. + + Args: + model_context: The context to update. Will be mutated if memories exist. + + Returns: + UpdateContextResult containing the memories that were added to the context + """ + + if not self._contents: + return UpdateContextResult(memories=MemoryQueryResult(results=[])) + + memory_strings = [f"{i}. {str(memory.content)}" for i, memory in enumerate(self._contents, 1)] + + if memory_strings: + memory_context = "\nRelevant memory content (in chronological order):\n" + "\n".join(memory_strings) + "\n" + await model_context.add_message(SystemMessage(content=memory_context)) + + return UpdateContextResult(memories=MemoryQueryResult(results=self._contents)) + + async def query( + self, + query: str | MemoryContent = "", + cancellation_token: CancellationToken | None = None, + **kwargs: Any, + ) -> MemoryQueryResult: + """Return all memories without any filtering. + + Args: + query: Ignored in this implementation + cancellation_token: Optional token to cancel operation + **kwargs: Additional parameters (ignored) + + Returns: + MemoryQueryResult containing all stored memories + """ + _ = query, cancellation_token, kwargs + return MemoryQueryResult(results=self._contents) + + async def add(self, content: MemoryContent, cancellation_token: CancellationToken | None = None) -> None: + """Add new content to memory. + + Args: + content: Memory content to store + cancellation_token: Optional token to cancel operation + """ + self._contents.append(content) + + async def clear(self) -> None: + """Clear all memory content.""" + self._contents = [] + + async def close(self) -> None: + """Cleanup resources if needed.""" + pass diff --git a/python/packages/autogen-core/tests/test_memory.py b/python/packages/autogen-core/tests/test_memory.py new file mode 100644 index 000000000000..41c0d6657fa9 --- /dev/null +++ b/python/packages/autogen-core/tests/test_memory.py @@ -0,0 +1,147 @@ +from datetime import datetime + +import pytest +from autogen_core import CancellationToken +from autogen_core.memory import ( + ListMemory, + Memory, + MemoryContent, + MemoryMimeType, + MemoryQueryResult, + UpdateContextResult, +) +from autogen_core.model_context import BufferedChatCompletionContext, ChatCompletionContext + + +def test_memory_protocol_attributes() -> None: + """Test that Memory protocol has all required attributes.""" + # No changes needed here + assert hasattr(Memory, "name") + assert hasattr(Memory, "update_context") + assert hasattr(Memory, "query") + assert hasattr(Memory, "add") + assert hasattr(Memory, "clear") + assert hasattr(Memory, "close") + + +def test_memory_protocol_runtime_checkable() -> None: + """Test that Memory protocol is properly runtime-checkable.""" + + class ValidMemory: + @property + def name(self) -> str: + return "test" + + async def update_context(self, context: ChatCompletionContext) -> UpdateContextResult: + return UpdateContextResult(memories=MemoryQueryResult(results=[])) + + async def query( + self, query: MemoryContent, cancellation_token: CancellationToken | None = None + ) -> MemoryQueryResult: + return MemoryQueryResult(results=[]) + + async def add(self, content: MemoryContent, cancellation_token: CancellationToken | None = None) -> None: + pass + + async def clear(self) -> None: + pass + + async def close(self) -> None: + pass + + class InvalidMemory: + pass + + assert isinstance(ValidMemory(), Memory) + assert not isinstance(InvalidMemory(), Memory) + + +@pytest.mark.asyncio +async def test_list_memory_empty() -> None: + """Test ListMemory behavior when empty.""" + memory = ListMemory(name="test_memory") + context = BufferedChatCompletionContext(buffer_size=3) + + results = await memory.update_context(context) + context_messages = await context.get_messages() + assert len(results.memories.results) == 0 + assert len(context_messages) == 0 + + query_results = await memory.query(MemoryContent(content="test", mime_type=MemoryMimeType.TEXT)) + assert len(query_results.results) == 0 + + +@pytest.mark.asyncio +async def test_list_memory_add_and_query() -> None: + """Test adding and querying memory contents.""" + memory = ListMemory() + + content1 = MemoryContent(content="test1", mime_type=MemoryMimeType.TEXT, timestamp=datetime.now()) + content2 = MemoryContent(content={"key": "value"}, mime_type=MemoryMimeType.JSON, timestamp=datetime.now()) + + await memory.add(content1) + await memory.add(content2) + + results = await memory.query(MemoryContent(content="query", mime_type=MemoryMimeType.TEXT)) + assert len(results.results) == 2 + assert results.results[0].content == "test1" + assert results.results[1].content == {"key": "value"} + + +@pytest.mark.asyncio +async def test_list_memory_max_memories() -> None: + """Test max_memories limit is enforced.""" + memory = ListMemory() + + for i in range(5): + await memory.add(MemoryContent(content=f"test{i}", mime_type=MemoryMimeType.TEXT)) + + results = await memory.query(MemoryContent(content="query", mime_type=MemoryMimeType.TEXT)) + assert len(results.results) == 5 + + +@pytest.mark.asyncio +async def test_list_memory_update_context() -> None: + """Test context updating with memory contents.""" + memory = ListMemory() + context = BufferedChatCompletionContext(buffer_size=3) + + await memory.add(MemoryContent(content="test1", mime_type=MemoryMimeType.TEXT)) + await memory.add(MemoryContent(content="test2", mime_type=MemoryMimeType.TEXT)) + + results = await memory.update_context(context) + context_messages = await context.get_messages() + assert len(results.memories.results) == 2 + assert len(context_messages) == 1 + assert "test1" in context_messages[0].content + assert "test2" in context_messages[0].content + + +@pytest.mark.asyncio +async def test_list_memory_clear() -> None: + """Test clearing memory contents.""" + memory = ListMemory() + await memory.add(MemoryContent(content="test", mime_type=MemoryMimeType.TEXT)) + await memory.clear() + + results = await memory.query(MemoryContent(content="query", mime_type=MemoryMimeType.TEXT)) + assert len(results.results) == 0 + + +@pytest.mark.asyncio +async def test_list_memory_content_types() -> None: + """Test support for different content types.""" + memory = ListMemory() + text_content = MemoryContent(content="text", mime_type=MemoryMimeType.TEXT) + json_content = MemoryContent(content={"key": "value"}, mime_type=MemoryMimeType.JSON) + binary_content = MemoryContent(content=b"binary", mime_type=MemoryMimeType.BINARY) + + await memory.add(text_content) + await memory.add(json_content) + await memory.add(binary_content) + + results = await memory.query(text_content) + assert len(results.results) == 3 + assert isinstance(results.results[0].content, str) + assert isinstance(results.results[1].content, dict) + assert isinstance(results.results[2].content, bytes) diff --git a/python/packages/autogen-ext/test_filesurfer_agent.html b/python/packages/autogen-ext/test_filesurfer_agent.html new file mode 100644 index 000000000000..8243435009e5 --- /dev/null +++ b/python/packages/autogen-ext/test_filesurfer_agent.html @@ -0,0 +1,9 @@ + + + FileSurfer test file + + +

FileSurfer test H1

+

FileSurfer test body

+ + \ No newline at end of file