92 lines
3.2 KiB
Python
92 lines
3.2 KiB
Python
import concurrent.futures
|
|
import random
|
|
import threading
|
|
|
|
import requests
|
|
|
|
# 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()
|