test_api.py

#!/usr/bin/env python3
"""
Test script for the AI Backend Service API endpoints
"""

import requests
import json
import time

BASE_URL = "http://localhost:8000"

def test_health():
    """Test health endpoint"""
    print("🔍 Testing health endpoint...")
    response = requests.get(f"{BASE_URL}/health")
    print(f"Status: {response.status_code}")
    print(f"Response: {json.dumps(response.json(), indent=2)}")
    print()

def test_root():
    """Test root endpoint"""
    print("🔍 Testing root endpoint...")
    response = requests.get(f"{BASE_URL}/")
    print(f"Status: {response.status_code}")
    print(f"Response: {json.dumps(response.json(), indent=2)}")
    print()

def test_models():
    """Test models endpoint"""
    print("🔍 Testing models endpoint...")
    response = requests.get(f"{BASE_URL}/v1/models")
    print(f"Status: {response.status_code}")
    print(f"Response: {json.dumps(response.json(), indent=2)}")
    print()

def test_chat_completion():
    """Test chat completion endpoint"""
    print("🔍 Testing chat completion endpoint...")
    data = {
        "model": "microsoft/DialoGPT-medium",
        "messages": [
            {"role": "user", "content": "Hello! How are you?"}
        ],
        "max_tokens": 100,
        "temperature": 0.7
    }
    
    response = requests.post(f"{BASE_URL}/v1/chat/completions", json=data)
    print(f"Status: {response.status_code}")
    print(f"Response: {json.dumps(response.json(), indent=2)}")
    print()

def test_completion():
    """Test completion endpoint"""
    print("🔍 Testing completion endpoint...")
    data = {
        "prompt": "The weather today is",
        "max_tokens": 50,
        "temperature": 0.7
    }
    
    response = requests.post(f"{BASE_URL}/v1/completions", json=data)
    print(f"Status: {response.status_code}")
    print(f"Response: {json.dumps(response.json(), indent=2)}")
    print()

def test_streaming_chat():
    """Test streaming chat completion"""
    print("🔍 Testing streaming chat completion...")
    data = {
        "model": "microsoft/DialoGPT-medium",
        "messages": [
            {"role": "user", "content": "Tell me a short joke"}
        ],
        "max_tokens": 100,
        "temperature": 0.7,
        "stream": True
    }
    
    response = requests.post(f"{BASE_URL}/v1/chat/completions", json=data, stream=True)
    print(f"Status: {response.status_code}")
    print("Streaming response:")
    
    for line in response.iter_lines():
        if line:
            line_str = line.decode('utf-8')
            if line_str.startswith('data: '):
                data_part = line_str[6:]  # Remove 'data: ' prefix
                if data_part == '[DONE]':
                    print("Stream completed!")
                    break
                try:
                    chunk_data = json.loads(data_part)
                    if 'choices' in chunk_data and chunk_data['choices']:
                        delta = chunk_data['choices'][0].get('delta', {})
                        if 'content' in delta:
                            print(delta['content'], end='', flush=True)
                except json.JSONDecodeError:
                    pass
    print("\n")

if __name__ == "__main__":
    print("🚀 Testing AI Backend Service API")
    print("=" * 50)
    
    # Wait a moment for service to be ready
    time.sleep(2)
    
    try:
        test_root()
        test_health()
        test_models()
        test_chat_completion()
        test_completion()
        test_streaming_chat()
        
        print("✅ All tests completed!")
        
    except requests.exceptions.ConnectionError:
        print("❌ Could not connect to the service. Make sure it's running on localhost:8000")
    except Exception as e:
        print(f"❌ Test failed with error: {e}")