tests/test_complete_integration.py

#!/usr/bin/env python3
"""
Complete integration test for Agent Bayko LlamaIndex workflow
Tests ReActAgent, FunctionTools, Memory, and LLM integration
"""

import os
import asyncio
import json
from pathlib import Path

# Load environment variables
try:
    from dotenv import load_dotenv

    load_dotenv()
except ImportError:
    pass

from agents.bayko_workflow import create_agent_bayko


async def test_complete_integration():
    """Test complete Bayko workflow integration"""

    print("🏆 Testing Complete Agent Bayko LlamaIndex Integration")
    print("=" * 70)
    print("🎯 Demonstrating: ReActAgent + FunctionTools + Memory + LLM")
    print("=" * 70)

    # Test 1: Create workflow with LLM
    if os.getenv("OPENAI_API_KEY"):
        print("\n1️⃣ Creating Bayko Workflow with LlamaIndex ReActAgent")
        print("-" * 50)

        try:
            # Create workflow
            bayko_workflow = create_agent_bayko(
                openai_api_key=os.getenv("OPENAI_API_KEY")
            )

            # Initialize session
            session_id = "hackathon_demo_001"
            bayko_workflow.initialize_session(session_id)

            print(
                f"✅ Workflow created with LLM: {bayko_workflow.llm is not None}"
            )
            print(f"✅ ReActAgent created: {bayko_workflow.agent is not None}")
            print(f"✅ Tools available: {len(bayko_workflow.tools)}")
            print(f"✅ Session initialized: {session_id}")

            # List available tools
            print("\n🛠️ Available LlamaIndex FunctionTools:")
            for i, tool in enumerate(bayko_workflow.tools, 1):
                print(
                    f"   {i}. {tool.metadata.name}: {tool.metadata.description[:60]}..."
                )

            # Test 2: Test individual tools
            print("\n2️⃣ Testing Individual LlamaIndex FunctionTools")
            print("-" * 50)

            # Test enhanced prompt generation
            print("🤖 Testing generate_enhanced_prompt tool...")
            prompt_result = (
                bayko_workflow.bayko_tools.generate_enhanced_prompt_tool(
                    description="A melancholic K-pop idol walking in rain",
                    style_tags='["whimsical", "soft_lighting", "watercolor"]',
                    mood="melancholy",
                )
            )
            prompt_data = json.loads(prompt_result)
            print(
                f"✨ Enhanced prompt: {prompt_data['enhanced_prompt'][:100]}..."
            )
            print(f"🎨 LLM used: {prompt_data['llm_used']}")

            # Test session info
            print("\n📊 Testing get_session_info tool...")
            session_result = bayko_workflow.bayko_tools.get_session_info_tool()
            session_data = json.loads(session_result)
            print(f"📋 Session ID: {session_data['session_id']}")
            print(f"🧠 Memory size: {session_data['memory_size']}")
            print(f"🤖 LLM available: {session_data['llm_available']}")

            # Test 3: Test ReActAgent workflow
            print("\n3️⃣ Testing LlamaIndex ReActAgent Workflow")
            print("-" * 50)

            # Create test request
            test_request = {
                "prompt": prompt_data["enhanced_prompt"],
                "original_prompt": "A melancholic K-pop idol walking in rain",
                "style_tags": ["whimsical", "soft_lighting", "watercolor"],
                "panels": 2,
                "language": "english",
                "extras": ["narration"],
                "session_id": session_id,
            }

            print("🎯 Processing request through ReActAgent...")
            workflow_result = bayko_workflow.process_generation_request(
                test_request
            )
            print(
                f"🎉 Workflow completed: {len(workflow_result)} chars response"
            )

            # Test 4: Verify session data persistence
            print("\n4️⃣ Testing Session Data Persistence")
            print("-" * 50)

            # Check if session directory was created
            session_dir = Path(f"storyboard/{session_id}")
            if session_dir.exists():
                print(f"✅ Session directory created: {session_dir}")

                # Check for LLM data
                llm_dir = session_dir / "llm_data"
                if llm_dir.exists():
                    llm_files = list(llm_dir.glob("*.json"))
                    print(f"💾 LLM data files: {len(llm_files)}")
                    for file in llm_files:
                        print(f"   📄 {file.name}")
                else:
                    print("⚠️ No LLM data directory found")

                # Check for agent data
                agents_dir = session_dir / "agents"
                if agents_dir.exists():
                    agent_files = list(agents_dir.glob("*.json"))
                    print(f"🤖 Agent data files: {len(agent_files)}")
                    for file in agent_files:
                        print(f"   📄 {file.name}")
                else:
                    print("⚠️ No agents data directory found")
            else:
                print("⚠️ Session directory not found")

            # Test 5: Memory integration
            print("\n5️⃣ Testing LlamaIndex Memory Integration")
            print("-" * 50)

            if bayko_workflow.bayko_agent.memory:
                memory_history = (
                    bayko_workflow.bayko_agent.memory.get_history()
                )
                print(f"🧠 Memory entries: {len(memory_history)}")

                # Show recent memory entries
                for i, entry in enumerate(memory_history[-3:], 1):
                    print(
                        f"   {i}. {entry['role']}: {entry['content'][:50]}..."
                    )
            else:
                print("⚠️ No memory system found")

            print("\n🏆 HACKATHON DEMO SUMMARY")
            print("=" * 50)
            print("✅ LlamaIndex ReActAgent: WORKING")
            print("✅ LlamaIndex FunctionTools: WORKING")
            print("✅ LlamaIndex Memory: WORKING")
            print("✅ OpenAI LLM Integration: WORKING")
            print("✅ Session Management: WORKING")
            print("✅ Multi-Agent Workflow: WORKING")
            print("\n🎯 Ready for LlamaIndex Prize Submission!")

        except Exception as e:
            print(f"❌ Integration test failed: {e}")
            import traceback

            traceback.print_exc()

    else:
        print("❌ OPENAI_API_KEY not found - cannot test LLM integration")
        print("🔄 Testing fallback mode...")

        # Test fallback mode
        bayko_workflow = create_agent_bayko(openai_api_key=None)
        bayko_workflow.initialize_session("fallback_session")

        print(f"✅ Fallback workflow created")
        print(f"⚠️ LLM available: {bayko_workflow.llm is not None}")
        print(f"⚠️ ReActAgent available: {bayko_workflow.agent is not None}")

        # Test fallback generation
        test_request = {
            "prompt": "A simple test prompt",
            "panels": 2,
            "session_id": "fallback_session",
        }

        result = bayko_workflow.process_generation_request(test_request)
        print(f"🔄 Fallback result: {result[:100]}...")


if __name__ == "__main__":
    asyncio.run(test_complete_integration())