First commit

scripts/stress_test.py

@@ -0,0 +1,91 @@
+import requests
+import time
+import random
+import concurrent.futures
+import threading
+
+# URL for the endpoint
+url = "http://localhost:8000/sigma-rag"
+
+# List of Spanish banking questions
+spanish_questions = [
+    "¿Cuáles son los beneficios de una tarjeta de crédito?",
+    "¿Cómo puedo abrir una cuenta de ahorros?",
+    "¿Qué es una hipoteca y cómo funciona?",
+    "¿Cuáles son las tasas de interés para un préstamo personal?",
+    "¿Cómo puedo solicitar un préstamo para un coche?",
+    "¿Qué es la banca en línea y cómo me registro?",
+    "¿Cómo puedo reportar una tarjeta de crédito perdida o robada?",
+    "¿Qué es el phishing y cómo puedo protegerme?",
+    "¿Cuáles son los diferentes tipos de cuentas corrientes que ofrecen?",
+    "¿Cómo puedo transferir dinero a una cuenta internacional?",
+]
+
+# A threading Event to signal all threads to stop
+stop_event = threading.Event()
+
+def send_request(question, request_id):
+    """Sends a single request and handles the response."""
+    if stop_event.is_set():
+        return
+
+    data = {"sessionInfo": {"parameters": {"query": question}}}
+    try:
+        response = requests.post(url, json=data)
+
+        if stop_event.is_set():
+            return
+
+        if response.status_code == 500:
+            print(f"Request {request_id}: Received 500 error with question: '{question}'.")
+            print("Stopping stress test.")
+            stop_event.set()
+        else:
+            print(f"Request {request_id}: Successful with status code {response.status_code}.")
+
+    except requests.exceptions.RequestException as e:
+        if not stop_event.is_set():
+            print(f"Request {request_id}: An error occurred: {e}")
+            stop_event.set()
+
+def main():
+    """Runs the stress test with parallel requests."""
+    num_workers = 30  # Number of parallel requests
+    print(f"Starting stress test with {num_workers} parallel workers. Press Ctrl+C to stop.")
+
+    with concurrent.futures.ThreadPoolExecutor(max_workers=num_workers) as executor:
+        futures = {
+            executor.submit(send_request, random.choice(spanish_questions), i)
+            for i in range(1, num_workers + 1)
+        }
+        request_id_counter = num_workers + 1
+
+        try:
+            while not stop_event.is_set():
+                # Wait for any future to complete
+                done, _ = concurrent.futures.wait(
+                    futures, return_when=concurrent.futures.FIRST_COMPLETED
+                )
+
+                for future in done:
+                    # Remove the completed future
+                    futures.remove(future)
+
+                    # If we are not stopping, submit a new one
+                    if not stop_event.is_set():
+                        futures.add(
+                            executor.submit(
+                                send_request,
+                                random.choice(spanish_questions),
+                                request_id_counter,
+                            )
+                        )
+                        request_id_counter += 1
+        except KeyboardInterrupt:
+            print("\nKeyboard interrupt received. Stopping threads.")
+            stop_event.set()
+
+    print("Stress test finished.")
+
+if __name__ == "__main__":
+    main()