3_optimization-design-ptn/04_memory-management/01_chat-memory-AutoGen.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "ea566837",
   "metadata": {},
   "source": [
    "# Long-term Chat Memory Agent with AutoGen\n",
    "---\n",
    "\n",
    "This example demonstrates how to create a long-term chat memory agent using the AutoGen framework. The agent is designed to remember past interactions and provide contextually relevant responses over time.\n",
    "\n",
    "- Reference: https://microsoft.github.io/autogen/stable//user-guide/agentchat-user-guide/memory.html"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f2d5666f",
   "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_core.memory import Memory, MemoryContent, MemoryMimeType\n",
    "\n",
    "from autogen_ext.models.openai import AzureOpenAIChatCompletionClient\n",
    "from IPython.display import display, Markdown"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "13563eae",
   "metadata": {},
   "source": [
    "<br>\n",
    "\n",
    "## 🧪 1. Preparation and Define the Agentic Architecture\n",
    "---\n",
    "### Initialize the Memory\n",
    "\n",
    "You can implement your own custom memory store based on the memory protocol. For example, you can implement a custom memory store that uses a vector database to store and retrieve information, or a memory store that uses machine learning models to generate personalized responses based on user preferences, etc. This example implements a memory store based on FAISS, which is not yet supported by AutoGen (as of April 2025). \n",
    "\n",
    "> **Note: There may be errors in the implementation, so please use it only for reference.**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5c21e486",
   "metadata": {},
   "outputs": [],
   "source": [
    "from src.faiss import FAISSVectorMemoryConfig, FAISSVectorMemory\n",
    "\n",
    "MEMORY_DB = \"list\"  # or \"list\"\n",
    "assert MEMORY_DB in [\"faiss\", \"list\"], f\"Unsupported memory db: {MEMORY_DB}\"\n",
    "\n",
    "if MEMORY_DB == \"faiss\":\n",
    "    print(\n",
    "        \"=== Using Custom FAISS vector memory. You can refer to src/faiss.py for more details.\"\n",
    "    )\n",
    "    user_memory = FAISSVectorMemory(\n",
    "        config=FAISSVectorMemoryConfig(\n",
    "            emb_model_name=\"text-embedding-3-large\",\n",
    "            dimension=3072,\n",
    "            use_gpu=False,\n",
    "            top_k=2,\n",
    "            score_threshold=0.3,\n",
    "        )\n",
    "    )\n",
    "elif MEMORY_DB == \"list\":\n",
    "    print(\"=== Using Built-in List memory\")\n",
    "\n",
    "    user_memory = ListMemory()\n",
    "\n",
    "await user_memory.clear()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6bf998bb",
   "metadata": {},
   "source": [
    "### Add some initial memory content"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1f698700",
   "metadata": {},
   "outputs": [],
   "source": [
    "async def add_memory(user_memory, content, metadata, mime_type=MemoryMimeType.TEXT):\n",
    "    await user_memory.add(\n",
    "        MemoryContent(\n",
    "            content=content,\n",
    "            mime_type=mime_type,\n",
    "            metadata=metadata,\n",
    "        )\n",
    "    )"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8adf5114",
   "metadata": {},
   "source": [
    "Baseline: You can add memory with the text format."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c7d84705",
   "metadata": {},
   "outputs": [],
   "source": [
    "await add_memory(\n",
    "    user_memory,\n",
    "    \"Daekeun provides AIML technology support. Hyo is a big fan of Microsoft.\",\n",
    "    {\"user_id\": \"1\"},\n",
    "    MemoryMimeType.TEXT,\n",
    ")\n",
    "await add_memory(\n",
    "    user_memory,\n",
    "    \"Daekeun is interested in AutoGen and Semantic Kernel\",\n",
    "    {\"user_id\": \"1\"},\n",
    "    MemoryMimeType.TEXT,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f0263aec",
   "metadata": {},
   "source": [
    "You can also add memory with JSON format. The memory content can be a list of dictionaries."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c17bc0d4",
   "metadata": {},
   "outputs": [],
   "source": [
    "from azure_genai_utils.tools import BingSearch\n",
    "\n",
    "web_search_tool = BingSearch(\n",
    "    max_results=2,\n",
    "    locale=\"en-US\",\n",
    "    include_news=False,\n",
    "    include_entity=False,\n",
    "    format_output=False,\n",
    ")\n",
    "\n",
    "# Define the tools to be used in the state graph\n",
    "query = \"What is AutoGen's main feature?\"\n",
    "results = web_search_tool.invoke(query)\n",
    "\n",
    "for i in range(1, len(results)):\n",
    "    await add_memory(user_memory, results[i], {\"user_id\": \"1\"}, MemoryMimeType.JSON)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f6ee35d4",
   "metadata": {},
   "source": [
    "### Define the tools"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "47033a79",
   "metadata": {},
   "outputs": [],
   "source": [
    "from autogen_core.tools import FunctionTool\n",
    "\n",
    "\n",
    "def arxiv_search(query: str, max_results: int = 2) -> list:  # type: ignore[type-arg]\n",
    "    \"\"\"\n",
    "    Search Arxiv for papers and return the results including abstracts.\n",
    "    \"\"\"\n",
    "    import arxiv\n",
    "\n",
    "    client = arxiv.Client()\n",
    "    search = arxiv.Search(\n",
    "        query=query, max_results=max_results, sort_by=arxiv.SortCriterion.Relevance\n",
    "    )\n",
    "\n",
    "    results = []\n",
    "    for paper in client.results(search):\n",
    "        results.append(\n",
    "            {\n",
    "                \"title\": paper.title,\n",
    "                \"authors\": [author.name for author in paper.authors],\n",
    "                \"published\": paper.published.strftime(\"%Y-%m-%d\"),\n",
    "                \"abstract\": paper.summary,\n",
    "                \"pdf_url\": paper.pdf_url,\n",
    "            }\n",
    "        )\n",
    "\n",
    "    # # Write results to a file\n",
    "    # with open('arxiv_search_results.json', 'w') as f:\n",
    "    #     json.dump(results, f, indent=2)\n",
    "\n",
    "    return results\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",
    "arxiv_search_tool = FunctionTool(\n",
    "    arxiv_search,\n",
    "    description=\"Search Arxiv for papers related to a given topic, including abstracts\",\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7552c9ec",
   "metadata": {},
   "outputs": [],
   "source": [
    "model_client = AzureOpenAIChatCompletionClient(\n",
    "    model=\"gpt-4o-mini\",\n",
    "    temperature=0.1,\n",
    "    max_tokens=500,\n",
    "    top_p=1,\n",
    ")\n",
    "\n",
    "assistant_agent = AssistantAgent(\n",
    "    name=\"assistant_agent\",\n",
    "    description=\"An assistant agent that can answer questions and provide information.\",\n",
    "    model_client=model_client,\n",
    "    tools=[get_weather, arxiv_search_tool],\n",
    "    memory=[user_memory],\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6cde6076",
   "metadata": {},
   "source": [
    "<br>\n",
    "\n",
    "## 🧪 2. Run the agent with a task.\n",
    "---"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "798ba815",
   "metadata": {},
   "outputs": [],
   "source": [
    "stream = assistant_agent.run_stream(task=\"What is AutoGen?\")\n",
    "await Console(stream)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c6d9a937",
   "metadata": {},
   "outputs": [],
   "source": [
    "message = await assistant_agent._model_context.get_messages()\n",
    "display(Markdown(message[-1].content))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "734617dc",
   "metadata": {},
   "source": [
    "### This below code cell checks if the memory is working."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1f3731e4",
   "metadata": {},
   "outputs": [],
   "source": [
    "stream = assistant_agent.run_stream(task=\"Who is Daekeun?\")\n",
    "await Console(stream)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7795a969",
   "metadata": {},
   "outputs": [],
   "source": [
    "message = await assistant_agent._model_context.get_messages()\n",
    "display(Markdown(message[-1].content))"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "py312-dev",
   "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.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}