{ "cells": [ { "cell_type": "markdown", "metadata": {}, "source": [ "# Build an agent with memory\n", "\n", "Create an AI agent that remembers important information across conversations." ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Problem\n", "\n", "You want to build an AI agent that can store and recall important information—user preferences, key facts, or context from previous conversations.\n", "\n", "| Memory type | Example |\n", "|-------------|---------|\n", "| User preferences | \"I prefer Python over JavaScript\" |\n", "| Key facts | \"The project deadline is March 15\" |\n", "| Conversation context | \"We discussed the budget yesterday\" |" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Solution\n", "\n", "**What's in this recipe:**\n", "\n", "- Store memories with embeddings for semantic search\n", "- Retrieve relevant memories based on conversation context\n", "- Use `@pxt.query` for retrieval functions\n", "\n", "This pattern is inspired by [Pixelbot](https://github.com/pixeltable/pixelbot) and [Pixelmemory](https://github.com/pixeltable/pixelmemory)." ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "### Setup" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "%pip install -qU pixeltable openai" ] }, { "cell_type": "code", "execution_count": 30, "metadata": {}, "outputs": [], "source": [ "import getpass\n", "import os\n", "from datetime import datetime\n", "\n", "if 'OPENAI_API_KEY' not in os.environ:\n", " os.environ['OPENAI_API_KEY'] = getpass.getpass('OpenAI API Key: ')" ] }, { "cell_type": "code", "execution_count": 32, "metadata": {}, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "Created directory 'agent_demo'.\n" ] }, { "data": { "text/plain": [ "" ] }, "execution_count": 32, "metadata": {}, "output_type": "execute_result" } ], "source": [ "import pixeltable as pxt\n", "from pixeltable.functions.openai import chat_completions, embeddings\n", "\n", "# Create a fresh directory\n", "pxt.drop_dir('agent_demo', force=True)\n", "pxt.create_dir('agent_demo')" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "### Create memory bank" ] }, { "cell_type": "code", "execution_count": 33, "metadata": {}, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "Created table 'memories'.\n" ] } ], "source": [ "# Create memory bank table\n", "memories = pxt.create_table(\n", " 'agent_demo/memories',\n", " {\n", " 'content': pxt.String, # The memory content\n", " 'category': pxt.String, # Optional category (preference, fact, etc.)\n", " 'created_at': pxt.Timestamp, # When the memory was stored\n", " },\n", ")" ] }, { "cell_type": "code", "execution_count": 34, "metadata": {}, "outputs": [], "source": [ "# Add embedding index for semantic search on content\n", "memories.add_embedding_index(\n", " column='content',\n", " string_embed=embeddings.using(model='text-embedding-3-small'),\n", ")" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "### Define retrieval function" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "# Define a query function to retrieve relevant memories\n", "@pxt.query\n", "def recall_memories(context: str, top_k: int = 3):\n", " \"\"\"Retrieve memories relevant to the current context.\"\"\"\n", " sim = memories.content.similarity(string=context)\n", " return (\n", " memories.where(sim > 0.5)\n", " .order_by(sim, asc=False)\n", " .limit(top_k)\n", " .select(content=memories.content, category=memories.category)\n", " )" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "### Store some memories" ] }, { "cell_type": "code", "execution_count": 36, "metadata": {}, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "Inserting rows into `memories`: 5 rows [00:00, 590.53 rows/s]\n", "Inserted 5 rows with 0 errors.\n" ] }, { "data": { "text/plain": [ "5 rows inserted, 15 values computed." ] }, "execution_count": 36, "metadata": {}, "output_type": "execute_result" } ], "source": [ "# Store some initial memories\n", "initial_memories = [\n", " {\n", " 'content': 'User prefers Python for data analysis',\n", " 'category': 'preference',\n", " 'created_at': datetime.now(),\n", " },\n", " {\n", " 'content': 'The project deadline is March 15, 2024',\n", " 'category': 'fact',\n", " 'created_at': datetime.now(),\n", " },\n", " {\n", " 'content': 'User works at a startup in San Francisco',\n", " 'category': 'fact',\n", " 'created_at': datetime.now(),\n", " },\n", " {\n", " 'content': 'Budget for the ML project is $50,000',\n", " 'category': 'fact',\n", " 'created_at': datetime.now(),\n", " },\n", " {\n", " 'content': 'User prefers concise explanations over detailed ones',\n", " 'category': 'preference',\n", " 'created_at': datetime.now(),\n", " },\n", "]\n", "\n", "memories.insert(initial_memories)" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "### Create conversation table with memory retrieval" ] }, { "cell_type": "code", "execution_count": 37, "metadata": {}, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "Created table 'conversations'.\n" ] } ], "source": [ "# Create conversation table\n", "conversations = pxt.create_table(\n", " 'agent_demo/conversations', {'user_message': pxt.String}\n", ")" ] }, { "cell_type": "code", "execution_count": 38, "metadata": {}, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "Added 0 column values with 0 errors.\n" ] }, { "data": { "text/plain": [ "No rows affected." ] }, "execution_count": 38, "metadata": {}, "output_type": "execute_result" } ], "source": [ "# Add memory retrieval step\n", "conversations.add_computed_column(\n", " relevant_memories=recall_memories(conversations.user_message, top_k=3)\n", ")" ] }, { "cell_type": "code", "execution_count": 39, "metadata": {}, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "Added 0 column values with 0 errors.\n" ] }, { "data": { "text/plain": [ "No rows affected." ] }, "execution_count": 39, "metadata": {}, "output_type": "execute_result" } ], "source": [ "# Build prompt with memories\n", "@pxt.udf\n", "def build_memory_prompt(\n", " user_message: str, relevant_memories: list[dict]\n", ") -> str:\n", " memory_text = '\\n'.join(\n", " [f'- {m[\"content\"]}' for m in relevant_memories]\n", " )\n", " return f\"\"\"You are a helpful assistant with access to the following memories about the user:\n", "\n", "{memory_text}\n", "\n", "Use these memories to personalize your response when relevant.\n", "\n", "User: {user_message}\n", "Assistant:\"\"\"\n", "\n", "\n", "conversations.add_computed_column(\n", " prompt=build_memory_prompt(\n", " conversations.user_message, conversations.relevant_memories\n", " )\n", ")" ] }, { "cell_type": "code", "execution_count": 40, "metadata": {}, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "Added 0 column values with 0 errors.\n", "Added 0 column values with 0 errors.\n" ] }, { "data": { "text/plain": [ "No rows affected." ] }, "execution_count": 40, "metadata": {}, "output_type": "execute_result" } ], "source": [ "# Generate response with memory context\n", "conversations.add_computed_column(\n", " response=chat_completions(\n", " messages=[{'role': 'user', 'content': conversations.prompt}],\n", " model='gpt-4o-mini',\n", " )\n", ")\n", "conversations.add_computed_column(\n", " assistant_reply=conversations.response.choices[0].message.content\n", ")" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "### Chat with memory-aware agent" ] }, { "cell_type": "code", "execution_count": 41, "metadata": {}, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "Inserting rows into `conversations`: 3 rows [00:00, 1047.88 rows/s]\n", "Inserted 3 rows with 0 errors.\n" ] }, { "data": { "text/plain": [ "3 rows inserted, 18 values computed." ] }, "execution_count": 41, "metadata": {}, "output_type": "execute_result" } ], "source": [ "# Test the memory-aware agent\n", "test_messages = [\n", " {\n", " 'user_message': 'What programming language should I use for this project?'\n", " },\n", " {'user_message': 'When do I need to finish this?'},\n", " {'user_message': 'How much can I spend on cloud resources?'},\n", "]\n", "\n", "conversations.insert(test_messages)" ] }, { "cell_type": "code", "execution_count": 42, "metadata": {}, "outputs": [ { "data": { "text/html": [ "\n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", "
user_messagerelevant_memoriesassistant_reply
How much can I spend on cloud resources?[]To help you determine how much you can spend on cloud resources, it would be useful to know your budget and any specific resources you plan to use. If you have a figure in mind or any constraints, please share that, and we can look at options that fit within your financial plan!
When do I need to finish this?[]Could you please provide the specific deadline you're working towards? That way, I can help you plan your tasks more effectively!
What programming language should I use for this project?[]To give the best recommendation, could you share more details about your project? Specifically, what are the goals of the project, any specific platforms you want to target, and your previous experience with programming languages? This will help me suggest the most suitable language for your needs!
" ], "text/plain": [ " user_message relevant_memories \\\n", "0 How much can I spend on cloud resources? [] \n", "1 When do I need to finish this? [] \n", "2 What programming language should I use for thi... [] \n", "\n", " assistant_reply \n", "0 To help you determine how much you can spend o... \n", "1 Could you please provide the specific deadline... \n", "2 To give the best recommendation, could you sha... " ] }, "execution_count": 42, "metadata": {}, "output_type": "execute_result" } ], "source": [ "# View conversations with memory\n", "conversations.select(\n", " conversations.user_message,\n", " conversations.relevant_memories,\n", " conversations.assistant_reply,\n", ").collect()" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Explanation\n", "\n", "**Memory-aware agent architecture:**\n", "\n", "```\n", "User Message → Retrieve Memories → Build Prompt → LLM Response\n", " ↓\n", " Memory Bank (with embeddings)\n", "```\n", "\n", "**Key components:**\n", "\n", "| Component | Purpose |\n", "|-----------|---------|\n", "| Memory table | Store facts, preferences, context |\n", "| Embedding index | Enable semantic memory search |\n", "| `@pxt.query` | Retrieval function for memories |\n", "| Prompt builder | Inject memories into LLM context |\n", "\n", "**Adding new memories:**\n", "\n", "```python\n", "memories.insert([{\n", " 'content': 'New information to remember',\n", " 'category': 'fact',\n", " 'created_at': datetime.now()\n", "}])\n", "```" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## See also\n", "\n", "- [Build a RAG pipeline](https://docs.pixeltable.com/howto/cookbooks/agents/pattern-rag-pipeline) - Document retrieval\n", "- [Use tool calling](https://docs.pixeltable.com/howto/cookbooks/agents/llm-tool-calling) - Function calling with LLMs\n", "- [Pixelbot](https://github.com/pixeltable/pixelbot) - Full agent implementation" ] } ], "metadata": { "kernelspec": { "display_name": "pixeltable", "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.11" } }, "nbformat": 4, "nbformat_minor": 2 }