test_memory_integration.py

#!/usr/bin/env python3
"""
Test script to verify MemoryOS integration with Echo Mind
"""

import asyncio
import json
from memory_service import EchoMindMemoryService

async def test_memory_service():
    """Test the memory service functionality"""
    print("🧠 Testing Echo Mind Memory Service")
    print("=" * 50)
    
    # Initialize memory service
    memory_service = EchoMindMemoryService()
    
    # Test 1: Basic initialization
    print("\n1. Testing initialization...")
    stats = memory_service.get_memory_stats()
    print(f"   Status: {stats['status']}")
    print(f"   User ID: {stats['user_id']}")
    
    # Test 2: Store conversation turns
    print("\n2. Testing conversation storage...")
    
    conversations = [
        ("Hello, I'm John", "Hi John! Nice to meet you. How can I help you today?"),
        ("I love playing football", "That's awesome! Football is a great sport. Do you play in a team or just for fun?"),
        ("I play for my local team", "That's fantastic! How long have you been playing for your local team?"),
        ("About 3 years now", "Three years is a good amount of time! You must have developed some great skills and friendships.")
    ]
    
    for user_input, assistant_response in conversations:
        success = await memory_service.store_conversation_turn(
            user_input=user_input,
            assistant_response=assistant_response,
            metadata={"test": True}
        )
        print(f"   Stored: '{user_input[:30]}...' -> {success}")
    
    # Test 3: Get memory stats
    print("\n3. Testing memory statistics...")
    stats = memory_service.get_memory_stats()
    print(f"   Conversation turns: {stats['conversation_turns']}")
    print(f"   Entities tracked: {stats['entities_tracked']}")
    print(f"   Top entities: {stats['top_entities']}")
    
    # Test 4: Search memories
    print("\n4. Testing memory search...")
    
    search_queries = ["football", "team", "John"]
    
    for query in search_queries:
        memories = await memory_service.search_memories(query, limit=2)
        print(f"   Search '{query}': Found {len(memories)} memories")
        for memory in memories:
            print(f"     - User: {memory['user_input']}")
            print(f"       Assistant: {memory['assistant_response'][:50]}...")
    
    # Test 5: Context building
    print("\n5. Testing context building...")
    context = await memory_service.build_context_prompt("Tell me about football")
    print(f"   Context length: {len(context)} characters")
    if context:
        print(f"   Context preview: {context[:100]}...")
    
    # Test 6: Entity extraction
    print("\n6. Testing entity extraction...")
    test_text = "I met Sarah at the Manchester United game yesterday"
    entities = await memory_service.extract_and_store_entities(test_text)
    print(f"   Extracted entities: {entities}")
    
    print("\n✅ Memory service test completed!")
    return True

async def test_memory_integration_simulation():
    """Simulate a conversation with memory integration"""
    print("\n🎭 Simulating conversation with memory...")
    print("=" * 50)
    
    memory_service = EchoMindMemoryService()
    
    # Simulate a multi-turn conversation
    conversation_turns = [
        "Hi, I'm Alex and I work as a software engineer",
        "I'm working on a Python project using FastAPI",
        "The project is about building an AI avatar system",
        "Can you remind me what I told you about my work?"
    ]
    
    for i, user_input in enumerate(conversation_turns, 1):
        print(f"\nTurn {i}:")
        print(f"User: {user_input}")
        
        # Get context for this input
        context = await memory_service.build_context_prompt(user_input)
        
        # Simulate assistant response
        if "remind me" in user_input.lower():
            # Search memories for work-related information
            memories = await memory_service.search_memories("work software engineer Python", limit=3)
            if memories:
                assistant_response = "Based on our conversation, you told me that you're Alex, a software engineer working on a Python project using FastAPI to build an AI avatar system."
            else:
                assistant_response = "I don't have any previous information about your work."
        else:
            assistant_response = f"Thank you for sharing that information, Alex. I'll remember that you {user_input.lower()}."
        
        print(f"Assistant: {assistant_response}")
        
        # Store the conversation turn
        await memory_service.store_conversation_turn(
            user_input=user_input,
            assistant_response=assistant_response
        )
        
        # Extract entities
        entities = await memory_service.extract_and_store_entities(user_input)
        if entities:
            print(f"Entities extracted: {entities}")
    
    # Final stats
    print(f"\nFinal memory stats:")
    stats = memory_service.get_memory_stats()
    print(f"  Conversations: {stats['conversation_turns']}")
    print(f"  Entities: {stats['entities_tracked']}")
    print(f"  Top entities: {stats['top_entities'][:3]}")
    
    print("\n✅ Conversation simulation completed!")

if __name__ == "__main__":
    async def main():
        await test_memory_service()
        await test_memory_integration_simulation()
        print("\n🎉 All memory integration tests passed!")
    
    asyncio.run(main())