Build an agent with memory

%pip install -qU pixeltable openai

import getpass
import os
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 chat_completions, embeddings

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

# 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
    },
)

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

# 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(string=context)
    return (
        memories.where(sim > 0.5)
        .order_by(sim, asc=False)
        .limit(top_k)
        .select(content=memories.content, category=memories.category)
    )

# 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)

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

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

# 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
    )
)

# 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
)

# 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)

# View conversations with memory
conversations.select(
    conversations.user_message,
    conversations.relevant_memories,
    conversations.assistant_reply,
).collect()

memories.insert([{
    'content': 'New information to remember',
    'category': 'fact',
    'created_at': datetime.now()
}])