Skip to main content
Open in Kaggle  Open in Colab  Download Notebook
This documentation page is also available as an interactive notebook. You can launch the notebook in Kaggle or Colab, or download it for use with an IDE or local Jupyter installation, by clicking one of the above links.
Create an AI agent that remembers important information across conversations.

​
Problem

You want to build an AI agent that can store and recall important information—user preferences, key facts, or context from previous conversations.

​
Solution

What’s in this recipe:
  • Store memories with embeddings for semantic search
  • Retrieve relevant memories based on conversation context
  • Use @pxt.query for retrieval functions
This pattern is inspired by Pixelbot and Pixelmemory.

​
Setup

%pip install -qU pixeltable openai

import os
import getpass
from datetime import datetime

if 'OPENAI_API_KEY' not in os.environ:
    os.environ['OPENAI_API_KEY'] = getpass.getpass('OpenAI API Key: ')

import pixeltable as pxt
from pixeltable.functions.openai import embeddings, chat_completions

# Create a fresh directory
pxt.drop_dir('agent_demo', force=True)
pxt.create_dir('agent_demo')

Created directory ‘agent_demo’.
<pixeltable.catalog.dir.Dir at 0x32a0c6390>

​
Create memory bank

# Create memory bank table
memories = pxt.create_table(
    'agent_demo.memories',
    {
        'content': pxt.String,           # The memory content
        'category': pxt.String,          # Optional category (preference, fact, etc.)
        'created_at': pxt.Timestamp,     # When the memory was stored
    }
)

Created table ‘memories’.
# Add embedding index for semantic search on content
memories.add_embedding_index(
    column="content",
    string_embed=embeddings.using(model="text-embedding-3-small")
)

​
Define retrieval function

# Define a query function to retrieve relevant memories
@pxt.query
def recall_memories(context: str, top_k: int = 3):
    """Retrieve memories relevant to the current context."""
    sim = memories.content.similarity(context)
    return (
        memories.where(sim > 0.5)
        .order_by(sim, asc=False)
        .limit(top_k)
        .select(
            content=memories.content,
            category=memories.category
        )
    )

​
Store some memories

# Store some initial memories
initial_memories = [
    {'content': 'User prefers Python for data analysis', 'category': 'preference', 'created_at': datetime.now()},
    {'content': 'The project deadline is March 15, 2024', 'category': 'fact', 'created_at': datetime.now()},
    {'content': 'User works at a startup in San Francisco', 'category': 'fact', 'created_at': datetime.now()},
    {'content': 'Budget for the ML project is $50,000', 'category': 'fact', 'created_at': datetime.now()},
    {'content': 'User prefers concise explanations over detailed ones', 'category': 'preference', 'created_at': datetime.now()},
]

memories.insert(initial_memories)

Inserting rows into `memories`: 5 rows [00:00, 590.53 rows/s]
Inserted 5 rows with 0 errors.
5 rows inserted, 15 values computed.

​
Create conversation table with memory retrieval

# Create conversation table
conversations = pxt.create_table(
    'agent_demo.conversations',
    {'user_message': pxt.String}
)

Created table ‘conversations’.
# Add memory retrieval step
conversations.add_computed_column(
    relevant_memories=recall_memories(conversations.user_message, top_k=3)
)

Added 0 column values with 0 errors.
No rows affected.
# Build prompt with memories
@pxt.udf
def build_memory_prompt(user_message: str, relevant_memories: list[dict]) -> str:
    memory_text = '\n'.join([f"- {m['content']}" for m in relevant_memories])
    return f"""You are a helpful assistant with access to the following memories about the user:

{memory_text}

Use these memories to personalize your response when relevant.

User: {user_message}
Assistant:"""

conversations.add_computed_column(
    prompt=build_memory_prompt(conversations.user_message, conversations.relevant_memories)
)

Added 0 column values with 0 errors.
No rows affected.
# Generate response with memory context
conversations.add_computed_column(
    response=chat_completions(
        messages=[{'role': 'user', 'content': conversations.prompt}],
        model='gpt-4o-mini'
    )
)
conversations.add_computed_column(
    assistant_reply=conversations.response.choices[0].message.content
)

Added 0 column values with 0 errors.
Added 0 column values with 0 errors.
No rows affected.

​
Chat with memory-aware agent

# Test the memory-aware agent
test_messages = [
    {'user_message': 'What programming language should I use for this project?'},
    {'user_message': 'When do I need to finish this?'},
    {'user_message': 'How much can I spend on cloud resources?'},
]

conversations.insert(test_messages)

Inserting rows into `conversations`: 3 rows [00:00, 1047.88 rows/s]
Inserted 3 rows with 0 errors.
3 rows inserted, 18 values computed.
# View conversations with memory
conversations.select(
    conversations.user_message,
    conversations.relevant_memories,
    conversations.assistant_reply
).collect()

​
Explanation

Memory-aware agent architecture: 
User Message → Retrieve Memories → Build Prompt → LLM Response
                    ↓
            Memory Bank (with embeddings)
Key components:
Adding new memories: 
memories.insert([{
    'content': 'New information to remember',
    'category': 'fact',
    'created_at': datetime.now()
}])

​
See also