test_streaming.py

#!/usr/bin/env python3
"""
Test the optimized single-phase streaming functionality
"""

import os
import sys
sys.path.append(os.path.dirname(os.path.abspath(__file__)))

from app import MAIN_SYSTEM_PROMPT, MCP_TOOLS, client, MODEL

def test_streaming_with_tools():
    """
    Test that the single-phase streaming works correctly
    """
    csv_url = "https://mcp-1st-birthday-mlops-agent.hf.space/gradio_api/file=/tmp/gradio/6abcca54f954f2ad99a8f8f330dc6e8082f03ef3090458d97c274efcc76d0170/heart.csv"
    
    user_with_file = f"[Uploaded CSV file URL: {csv_url}]\n\nAnalyze this dataset and show me basic statistics"
    
    print("Testing Single-Phase Streaming...")
    print(f"Input: {user_with_file[:100]}...")
    print("-" * 60)
    
    stream = client.responses.create(
        model=MODEL,
        instructions=MAIN_SYSTEM_PROMPT,
        input=user_with_file,
        tools=MCP_TOOLS,
        reasoning={"effort": "low"},
        stream=True,
    )
    
    print("Streaming Response:")
    print("=" * 60)
    
    final_text = ""
    chunk_count = 0
    
    for event in stream:
        if event.type == "response.output_text.delta":
            chunk_count += 1
            final_text += event.delta
            print(f"[Chunk {chunk_count}] {event.delta[:50]}...")
        elif event.type == "response.completed":
            print("=" * 60)
            print(f"Total chunks: {chunk_count}")
            print(f"Total length: {len(final_text)} characters")
            print("FINAL RESPONSE:")
            print(final_text)
            break
    
    return final_text

def test_streaming_without_tools():
    """
    Test that regular streaming works for non-tool requests
    """
    print("\nTesting Regular Streaming (No Tools)...")
    print("-" * 60)
    
    stream = client.responses.create(
        model=MODEL,
        instructions=MAIN_SYSTEM_PROMPT,
        input="Hello! Can you explain what MLOps is in simple terms?",
        reasoning={"effort": "low"},
        stream=True,
    )
    
    print("Streaming Response:")
    print("=" * 60)
    
    final_text = ""
    chunk_count = 0
    
    for event in stream:
        if event.type == "response.output_text.delta":
            chunk_count += 1
            final_text += event.delta
            print(f"[Chunk {chunk_count}] {event.delta[:50]}...")
        elif event.type == "response.completed":
            print("=" * 60)
            print(f"Total chunks: {chunk_count}")
            print(f"Total length: {len(final_text)} characters")
            print("FINAL RESPONSE:")
            print(final_text)
            break
    
    return final_text

if __name__ == "__main__":
    print("Starting Streaming Tests")
    print("=" * 60)
    
    # Test 1: With MCP tools
    response1 = test_streaming_with_tools()
    
    # Test 2: Without tools  
    response2 = test_streaming_without_tools()
    
    print("\n" + "=" * 60)
    print("✅ Both streaming tests completed successfully!")
    print("✅ Single-phase approach working correctly!")
    print("✅ MCP tools returning natural language responses!")
    print("✅ Response is properly streaming in chunks!")
    print("=" * 60)