6.Integrate a Memory

# Upsonic Docs: Using Memory in your Agent
# https://docs.upsonic.ai/guides/using_memory_in_your_agent

# Imports
from upsonic import Agent
from upsonic import Task
from upsonic.storage import InMemoryStorage, Memory

# Create a storage provider
storage = InMemoryStorage(max_sessions=100, max_profiles=50)

# Create a Memory instance with all memory types enabled
memory = Memory(
    storage=storage,
    session_id="user_session_123",
    user_id="user_456",
    full_session_memory=True,      # Remember complete conversations
    summary_memory=True,           # Maintain conversation summaries
    user_analysis_memory=True,     # Build user profiles
    dynamic_user_profile=True,     # Automatically adapt profile schema
    num_last_messages=10,          # Limit context to last 10 messages
    debug=True                     # Enable debug logging
)

# Create an Agent with Memory
personal_assistant = Agent(
    name="Personal Assistant",
    role="Personal AI Assistant with Memory",
    goal="Provide personalized assistance by remembering user preferences and conversation history",
    instructions="""
    Always reference previous conversations when relevant.
    Adapt your responses based on the user's profile and preferences.
    Use the conversation summary to provide context-aware assistance.
    """,
    memory=memory
)

# Example tasks that demonstrate memory capabilities
task1 = Task(
    description="I prefer to be called 'Alex' and I'm interested in machine learning",
)

task2 = Task(
    description="What did we discuss in our previous conversation?",
)

task3 = Task(
    description="Based on what you know about me, suggest a learning path for AI",
)

# Execute tasks to build memory
print("=== Task 1: Building User Profile ===")
result1 = personal_assistant.do(task1)
print(f"Response: {result1}")

print("\n=== Task 2: Testing Memory Recall ===")
result2 = personal_assistant.do(task2)
print(f"Response: {result2}")

print("\n=== Task 3: Using Memory for Personalization ===")
result3 = personal_assistant.do(task3)
print(f"Response: {result3}")

# Check memory statistics
print(f"\n=== Memory Statistics ===")
print(f"Session ID: {memory.session_id}")
print(f"User ID: {memory.user_id}")
print(f"Full Session Memory: {memory.full_session_memory_enabled}")
print(f"Summary Memory: {memory.summary_memory_enabled}")
print(f"User Analysis Memory: {memory.user_analysis_memory_enabled}")

# PostgreSQL Storage for Production
from upsonic.storage import PostgresStorage

postgres_storage = PostgresStorage(
    sessions_table_name="agent_sessions",
    profiles_table_name="user_profiles",
    db_url="postgresql://user:pass@localhost:5432/upsonic_db",
    schema="public"
)

# Redis Storage for High Performance
from upsonic.storage import RedisStorage

redis_storage = RedisStorage(
    prefix="upsonic",
    host="localhost",
    port=6379,
    expire=86400  # 24 hours TTL
)

# JSON Storage for Simple Applications
from upsonic.storage import JSONStorage

json_storage = JSONStorage(
    directory_path="./data/memory",
    pretty_print=True
)

from pydantic import BaseModel, Field
from upsonic.storage import Memory

class CustomUserProfile(BaseModel):
    expertise_level: str = Field(description="User's expertise level in the domain")
    preferred_language: str = Field(description="User's preferred programming language")
    learning_goals: list[str] = Field(description="User's learning objectives")
    communication_style: str = Field(description="User's preferred communication style")

# Create Memory with custom profile schema
memory_with_custom_schema = Memory(
    storage=storage,
    session_id="session_789",
    user_id="user_101",
    user_analysis_memory=True,
    user_profile_schema=CustomUserProfile,
    dynamic_user_profile=False  # Use our custom schema
)

from upsonic.models.factory import ModelFactory

# Get a model provider for memory operations
model_provider = ModelFactory.create("gpt-4o")

memory_with_model = Memory(
    storage=storage,
    session_id="session_456",
    user_id="user_789",
    summary_memory=True,
    user_analysis_memory=True,
    model_provider=model_provider,  # For generating summaries and analyzing traits
    debug=True
)

# Create unique session IDs for different conversation contexts
import uuid

session_id = f"chat_session_{uuid.uuid4()}"
user_id = f"user_{uuid.uuid4()}"

memory = Memory(
    storage=storage,
    session_id=session_id,
    user_id=user_id,
    full_session_memory=True
)

# Check memory statistics
print(f"Memory Configuration:")
print(f"- Full Session Memory: {memory.full_session_memory_enabled}")
print(f"- Summary Memory: {memory.summary_memory_enabled}")
print(f"- User Analysis Memory: {memory.user_analysis_memory_enabled}")
print(f"- Max Messages: {memory.num_last_messages}")
print(f"- Feed Tool Results: {memory.feed_tool_call_results}")

from upsonic import Team

# Create a shared memory instance
shared_memory = Memory(
    storage=storage,
    session_id="team_session_123",
    user_id="user_456",
    full_session_memory=True,
    summary_memory=True,
    model_provider=model_provider  # Required for summary generation
)

# Create agents with shared memory
research_agent = Agent(
    name="Research Agent",
    role="Research Specialist",
    memory=shared_memory
)

writing_agent = Agent(
    name="Writing Agent", 
    role="Content Writer",
    memory=shared_memory
)

# Create a team with shared memory
research_team = Team(
    agents=[research_agent, writing_agent],
    memory=shared_memory,
    tasks=[task1, task2]
)

memory = Memory(
    storage=storage,
    session_id="debug_session",
    user_id="debug_user",
    full_session_memory=True,
    debug=True  # Enable detailed logging
)