smart-models/Placebo_AI
0
1import asyncio2import time3import requests4import aiohttp5 6BASE_URL = "http://localhost:8000"7# NOTE: You MUST provide a valid JWT token here if your API is secured.8# If testing without Auth, ensure the backend route allows it.9TEST_JWT_TOKEN = "your_test_jwt_token_here" 10 11async def fetch_chat(session, prompt):12 headers = {"Authorization": f"Bearer {TEST_JWT_TOKEN}"}13 payload = {"message": prompt, "role": "MBBS Student/Doctor", "history": []}14 15 start_time = time.time()16 try:17 # Note: Streaming responses make latency calculation complex. 18 # We are just measuring Time-to-First-Token/Connection.19 async with session.post(f"{BASE_URL}/chat", json=payload, headers=headers) as response:20 latency = time.time() - start_time21 return response.status, latency22 except Exception as e:23 return 500, time.time() - start_time24 25async def test_load_stress():26 print("--- Phase 4: Load & Stress Testing ---")27 prompts = ["What is Aspirin?", "Define Appendicitis", "Mechanism of Metformin", "Explain Hypertension"]28 29 print(f"[TEST] Firing {len(prompts)} concurrent requests to stress Test VRAM...")30 async with aiohttp.ClientSession() as session:31 tasks = [fetch_chat(session, p) for p in prompts]32 results = await asyncio.gather(*tasks)33 34 for i, (status, latency) in enumerate(results):35 print(f"Request {i+1}: Status {status} | Latency: {latency:.2f}s")36 37 success_rate = sum(1 for status, _ in results if status == 200) / len(results) * 10038 print(f"--- Load Test Completed ---")39 print(f"Success Rate: {success_rate}%")40 41if __name__ == "__main__":42 asyncio.run(test_load_stress())43 