test_api.py

import requests
import json
import sys

def test_api(url, prompt):
    """
    Test the LLM API by sending a prompt and printing the response.
    
    Args:
        url (str): The API URL to test (e.g. http://localhost:8000)
        prompt (str): The prompt to send to the API
    """
    # Add /generate to the URL if it doesn't already end with it
    if not url.endswith('/generate'):
        url = url.rstrip('/') + '/generate'
    
    # Prepare the request payload
    payload = {
        "prompt": prompt
    }
    
    try:
        # Make the request
        print(f"Sending prompt to {url}: {prompt}")
        response = requests.post(url, json=payload)
        
        # Check if the request was successful
        if response.status_code == 200:
            result = response.json()
            print("\nAPI Response:")
            print(f"Response: {result['response']}")
            print("\nAPI test successful!")
        else:
            print(f"\nAPI request failed with status code: {response.status_code}")
            print(f"Error message: {response.text}")
    
    except requests.exceptions.RequestException as e:
        print(f"\nError connecting to API: {e}")
    except json.JSONDecodeError:
        print(f"\nError decoding JSON response: {response.text}")
    except Exception as e:
        print(f"\nUnexpected error: {e}")

if __name__ == "__main__":
    # Default values
    default_url = "http://localhost:8000"
    default_prompt = "Tell me a joke about AI."
    
    # Get URL and prompt from command line arguments
    url = sys.argv[1] if len(sys.argv) > 1 else default_url
    prompt = sys.argv[2] if len(sys.argv) > 2 else default_prompt
    
    # Run the test
    test_api(url, prompt)