test_memory_fix.py

#!/usr/bin/env python3
"""
Test script to verify memory system functionality.
Tests both direct database saving and the agent orchestration flow.
"""

import asyncio
import os
import sys
import uuid
from datetime import datetime

# Add the apps/api directory to the path
sys.path.insert(0, os.path.join(os.path.dirname(__file__), 'apps', 'api'))

from memory.direct_memory_core import DirectMemoryCore, DirectMemoryConfig
from services.unified_memory_knowledge_service import UnifiedMemoryKnowledgeService

async def test_direct_memory_save():
    """Test direct memory saving to database."""
    print("🧠 Testing direct memory save...")

    # Test data
    project_id = "5a05f05d-3668-44e8-a733-2b19346c3061"
    agent_id = "e1a79bdb-74b1-48c8-9d86-6bd8b39de84c"  # Research Specialist Agent

    config = DirectMemoryConfig(
        database_url=os.getenv("DATABASE_URL"),
        project_id=project_id
    )

    memory_core = DirectMemoryCore(config)

    # Test memory save
    result = await memory_core.add_memory_with_graph(
        agent_id=agent_id,
        content="Test memory for debugging memory system issue",
        entities=[
            ("Agent", "Test Agent"),
            ("Task", "Memory System Debug"),
            ("System", "Memory Core")
        ],
        relationships=[
            ("Test Agent", "debugged", "Memory System Debug"),
            ("Memory System Debug", "affects", "System"),
            ("Agent", "uses", "Memory Core")
        ],
        metadata={
            "test": True,
            "created_at": datetime.utcnow().isoformat(),
            "debug_session": True
        },
        importance=0.8,
        project_id=project_id
    )

    print(f"💾 Direct memory save result: {result}")
    return result.get("success", False)

async def test_unified_memory_service():
    """Test the unified memory service."""
    print("🔄 Testing unified memory service...")

    service = UnifiedMemoryKnowledgeService()

    # Test data
    project_id = "5a05f05d-3668-44e8-a733-2b19346c3061"
    agent_id = "e1a79bdb-74b1-48c8-9d86-6bd8b39de84c"  # Research Specialist Agent

    try:
        result = await service.memory_core.add_memory_with_graph(
            agent_id=agent_id,
            content="Test memory via unified service for memory debugging",
            entities=[
                ("Service", "Unified Memory Service"),
                ("Debug", "Memory System Test"),
                ("Agent", "Test Agent")
            ],
            relationships=[
                ("Test Agent", "uses", "Unified Memory Service"),
                ("Unified Memory Service", "performs", "Memory System Test"),
                ("Memory System Test", "validates", "Debug")
            ],
            metadata={
                "project_id": project_id,
                "service_test": True,
                "timestamp": datetime.utcnow().isoformat()
            },
            importance=0.9
        )

        print(f"🔄 Unified service result: {result}")
        return result.get("success", False)

    except Exception as e:
        print(f"❌ Unified service failed: {e}")
        return False

async def check_database_memories():
    """Check if memories were actually saved to database."""
    print("🔍 Checking database for saved memories...")

    try:
        from models.base import get_db_session
        from sqlalchemy import text

        async with get_db_session() as session:
            result = await session.execute(text("""
                SELECT COUNT(*) as total_memories,
                       COUNT(DISTINCT agent_id) as unique_agents,
                       MAX(created_at) as latest_memory
                FROM memories
                WHERE project_id = '5a05f05d-3668-44e8-a733-2b19346c3061'
            """))

            row = result.fetchone()
            if row:
                print(f"📊 Database stats:")
                print(f"   - Total memories: {row[0]}")
                print(f"   - Unique agents: {row[1]}")
                print(f"   - Latest memory: {row[2]}")
                return row[0] > 0
            else:
                print("❌ No memory data found")
                return False

    except Exception as e:
        print(f"❌ Database check failed: {e}")
        return False

async def main():
    """Run all memory system tests."""
    print("🚀 Starting memory system diagnostic tests...")
    print("=" * 50)

    # Test 1: Direct memory save
    direct_success = await test_direct_memory_save()
    print()

    # Test 2: Unified service
    service_success = await test_unified_memory_service()
    print()

    # Test 3: Database verification
    db_success = await check_database_memories()
    print()

    # Summary
    print("=" * 50)
    print("🎯 Test Results Summary:")
    print(f"   ✅ Direct Memory Save: {'PASS' if direct_success else 'FAIL'}")
    print(f"   ✅ Unified Service: {'PASS' if service_success else 'FAIL'}")
    print(f"   ✅ Database Check: {'PASS' if db_success else 'FAIL'}")

    if direct_success and service_success and db_success:
        print("\n🎉 All tests passed! Memory system is working.")
        return True
    else:
        print("\n❌ Some tests failed. Memory system needs attention.")
        return False

if __name__ == "__main__":
    # Load environment variables
    from dotenv import load_dotenv
    load_dotenv()

    # Run tests
    success = asyncio.run(main())
    sys.exit(0 if success else 1)