import requests
import json

def test_model(prompt):
    """Test the model with a simple prompt."""
    url = "http://localhost:8000/direct-model"
    
    payload = {
        "prompt": prompt,
        "max_length": 200,
        "num_return_sequences": 1,
        "temperature": 0.7,
        "top_p": 0.9
    }
    
    print(f"Sending prompt: '{prompt}'")
    
    try:
        response = requests.post(url, json=payload)
        
        if response.status_code == 200:
            result = response.json()
            print("\nGenerated text:")
            print(result["generated_texts"][0])
        else:
            print(f"Error: {response.status_code}")
            print(response.text)
    except Exception as e:
        print(f"Error: {str(e)}")
        print("Make sure the FastAPI server is running on http://localhost:8000")

def main():
    print("Simple LLM Test")
    print("===============")
    print("Enter your prompt (or 'quit' to exit):")
    
    while True:
        prompt = input("\n> ")
        
        if prompt.lower() in ["quit", "exit", "q"]:
            break
        
        if not prompt.strip():
            continue
        
        test_model(prompt)

if __name__ == "__main__":
    main()