test_streaming_simple.py

#!/usr/bin/env python3
"""
Simple test script to verify streaming functionality.
"""

import asyncio
import sys
import logging
from pathlib import Path

# Configure logging to see debug messages
logging.basicConfig(level=logging.INFO, format='%(name)s - %(levelname)s - %(message)s')

# Add project root to path
project_root = Path(__file__).parent
sys.path.insert(0, str(project_root))

from core.brain import AutumnBrain
from core.memory import AutumnMemory
from config.personality import AutumnPersonality

async def simple_streaming_test():
    """Test basic streaming functionality."""
    print("Initializing Autumn...")
    
    # Initialize components
    memory = AutumnMemory()
    personality = AutumnPersonality()
    brain = AutumnBrain(memory, personality)
    
    # Initialize brain
    await brain.initialize()
    
    print("Testing streaming vs non-streaming...")
    
    test_query = "Hello, how are you today?"  # Simple query to trigger local processing
    
    # Test non-streaming first
    print("\n" + "="*50)
    print("NON-STREAMING TEST")
    print("="*50)
    
    import time
    start = time.time()
    response = await brain.process(test_query)
    end = time.time()
    
    print(f"Time: {end-start:.2f}s")
    print(f"Response: {response[:100]}...")
      # Test streaming
    print("\n" + "="*50)
    print("STREAMING TEST")
    print("="*50)
    
    tokens_received = []
    
    async def stream_callback(data):
        print(f"[CALLBACK] Received data: {data}")  # Debug info
        if data["type"] == "token":
            tokens_received.append(data["content"])
            print(data["content"], end="", flush=True)
        elif data["type"] == "sentence":
            print(f"\n[SENTENCE: {len(data['content'])} chars]")
        elif data["type"] == "complete":
            print(f"\n[COMPLETE: {len(data['content'])} total chars]")    
    start = time.time()
    try:
        # Check brain status first
        print(f"Brain ollama_available: {brain.ollama_available}")
        print(f"Brain gemini_available: {brain.gemini_available}")
        
        stream_response = await brain.process_streaming(test_query, stream_callback)
        end = time.time()
        
        print(f"\nTime: {end-start:.2f}s")
        print(f"Tokens received: {len(tokens_received)}")
        print(f"Response length: {len(stream_response)}")
        
    except Exception as e:
        print(f"Streaming error: {e}")
        import traceback
        traceback.print_exc()

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