import requests
import time

# Create chat
try:
    chat = requests.post("http://localhost:4000/chats", json={"title": "RAG Tool Test"}).json()
    chat_id = chat["id"]
except Exception as e:
    print(f"Failed to create chat: {e}")
    exit(1)

def ask(question):
    print(f"\nUser: {question}")
    t0 = time.time()
    try:
        res = requests.post(f"http://localhost:4000/chats/{chat_id}/messages", json={"content": question})
        duration = time.time() - t0
        if res.status_code == 200:
            data = res.json()
            print(f"Assistant ({duration:.1f}s): {data['assistant_message']['content']}")
        else:
            print(f"Error {res.status_code}: {res.text}")
    except Exception as e:
        print(f"Request failed: {e}")

# Test 1: Math Tool (Should verify tools still work)
ask("Calculate 55 * 4")

# Test 2: RAG Tool (Should verify retrieve_documents is called)
ask("What specific data is in the uploaded document?")

# Test 3: Rate Limit Resilience (Hammer the API)
print("\n--- Stress Test (Rate Limit) ---")
for i in range(5):
    ask(f"Quick question {i}: What is 1+{i}?")