Files
agent/scripts/stress_test.py
Anibal Angulo a53f8fcf62 First commit
2026-02-18 19:57:43 +00:00

92 lines
3.2 KiB
Python

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()