Add OpenAIUser and example.

Author: cyberw

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,69 @@
+# 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
+
+
+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,
+            )
+
+        # dont forget to set OPENAI_API_KEY in your environment
+        super().__init__(
+            *args,
+            **kwargs,
+            http_client=httpx.Client(event_hooks={"request": [request_start], "response": [request_end]}),
+            # convenience for when running in Locust Cloud, where only LOCUST_* env vars are forwarded:
+            api_key=os.getenv("LOCUST_OPENAI_API_KEY", None),
+        )
+
+    @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)