Merge pull request #3081 from locustio/add-openai-User-and-example

Add OpenAI User and example

Author: cyberw

Diff for: docs/testing-other-systems.rst

@@ -58,6 +58,17 @@ REST
 
 See :ref:`FastHttpUser <rest>`
 
+OpenAI
+======
+
+Performance/load testing AI services is a little different. While you could call the OpenAI API using HttpUser or FastHttpUser, it is often convenient to use `the SDK <https://github.com/openai/openai-python>`_.
+
+.. literalinclude:: ../examples/openai_ex.py
+
+.. note::
+
+    OpenAIUser is experimental and may change without notice.
+
 Other examples
 ==============
 

Diff for: examples/openai_ex.py

@@ -0,0 +1,28 @@
+# You need to install the openai package and set OPENAI_API_KEY env var to run this
+
+# OpenAIUser tracks the number of output tokens in the response_length field,
+# because it is more useful than the actual payload size. This field is available to event handlers,
+# but only graphed in Locust Cloud.
+
+from locust import run_single_user, task
+from locust.contrib.oai import OpenAIUser
+
+
+class MyUser(OpenAIUser):
+    @task
+    def t(self):
+        self.client.responses.create(
+            model="gpt-4o",
+            instructions="You are a coding assistant that speaks like it were a Monty Python skit.",
+            input="How do I check if a Python object is an instance of a class?",
+        )
+        with self.client.rename_request("mini"):  # here's how to rename requests
+            self.client.responses.create(
+                model="gpt-4o-mini",
+                instructions="You are a coding assistant that speaks like it were a Monty Python skit.",
+                input="How do I check if a Python object is an instance of a class?",
+            )
+
+
+if __name__ == "__main__":
+    run_single_user(MyUser)

Diff for: locust/contrib/oai.py

@@ -0,0 +1,73 @@
+# Note: this User is experimental and may change without notice.
+# The filename is oai.py so it doesnt clash with the openai package.
+from locust.user import User
+
+import os
+import time
+from collections.abc import Generator
+from contextlib import contextmanager
+
+import httpx
+from openai import OpenAI  # dont forget to install openai
+
+# convenience for when running in Locust Cloud, where only LOCUST_* env vars are forwarded
+if "LOCUST_OPENAI_API_KEY" in os.environ:
+    os.environ["OPENAI_API_KEY"] = os.environ["LOCUST_OPENAI_API_KEY"]
+
+if not "OPENAI_API_KEY" in os.environ:
+    raise Exception("You need to set OPENAI_API_KEY or LOCUST_OPENAI_API_KEY env var to use OpenAIUser")
+
+
+class OpenAIClient(OpenAI):
+    def __init__(self, request_event, user, *args, **kwargs):
+        self.request_name = None  # used to override url-based request names
+        self.user = user  # currently unused, but could be useful later
+
+        def request_start(request):
+            request.start_time = time.time()
+            request.start_perf_counter = time.perf_counter()
+
+        def request_end(response):
+            exception = None
+            response.read()
+            response_time = (time.perf_counter() - response.request.start_perf_counter) * 1000
+            try:
+                response.raise_for_status()
+            except httpx.HTTPStatusError as e:
+                exception = e
+            request_event.fire(
+                request_type=response.request.method,
+                name=self.request_name if self.request_name else response.url.path,
+                context={},
+                response=response,
+                exception=exception,
+                start_time=response.request.start_time,
+                response_time=response_time,
+                # Store the number of output tokens as response_length instead of the actual payload size because it is more useful
+                response_length=response.json().get("usage", {}).get("output_tokens", 0),
+                url=response.url,
+            )
+
+        super().__init__(
+            *args,
+            **kwargs,
+            http_client=httpx.Client(event_hooks={"request": [request_start], "response": [request_end]}),
+        )
+
+    @contextmanager
+    def rename_request(self, name: str) -> Generator[None]:
+        """Group requests using the "with" keyword"""
+
+        self.request_name = name
+        try:
+            yield
+        finally:
+            self.request_name = None
+
+
+class OpenAIUser(User):
+    abstract = True
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.client = OpenAIClient(self.environment.events.request, user=self)