locust 2.33.3.dev3__py3-none-any.whl → 2.33.3.dev9__py3-none-any.whl

locust/_version.py

@@ -17,5 +17,5 @@ __version__: str
 __version_tuple__: VERSION_TUPLE
 version_tuple: VERSION_TUPLE
 
-__version__ = version = '2.33.3.dev3'
-__version_tuple__ = version_tuple = (2, 33, 3, 'dev3')
+__version__ = version = '2.33.3.dev9'
+__version_tuple__ = version_tuple = (2, 33, 3, 'dev9')

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)

{locust-2.33.3.dev3.dist-info → locust-2.33.3.dev9.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: locust
-Version: 2.33.3.dev3
+Version: 2.33.3.dev9
 Summary: Developer-friendly load testing framework
 Project-URL: homepage, https://locust.io/
 Project-URL: repository, https://github.com/locustio/locust

{locust-2.33.3.dev3.dist-info → locust-2.33.3.dev9.dist-info}/RECORD

@@ -1,6 +1,6 @@
 locust/__init__.py,sha256=aWeuBPUxONjwNm1xp4v8L4BO14SuYLjscIiwJVX1Ui4,1746
 locust/__main__.py,sha256=vBQ82334kX06ImDbFlPFgiBRiLIinwNk3z8Khs6hd74,31
-locust/_version.py,sha256=SybUtIYL4cgereMxtifIUopX-_x16OzH7xfy0fH5nSk,526
+locust/_version.py,sha256=rtO6T9Z-9Cvc1wxULMvtZ1HAeWQphLTGmdmZ6ypkvu8,526
 locust/argument_parser.py,sha256=ruSJcY4ouERwVopHvHDK-obQD6nlRKDqjLLwje2SH9k,32806
 locust/clients.py,sha256=LHAYjaoRjMXwdOB8KBUYRH9-8Uk4F8NvUsDy5UTxKkA,19352
 locust/debug.py,sha256=7CCm8bIg44uGH2wqBlo1rXBzV2VzwPicLxLewz8r5CQ,5099
@@ -20,6 +20,7 @@ locust/web.py,sha256=oVHqJ_vPQ-gLzMZgc8RwqG44DQ6nzbXpta78ihtVCjo,30122
 locust/contrib/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 locust/contrib/fasthttp.py,sha256=lB1t8syshdpJDrNnrN6j00FF4TszdYY8YUfdg-9mwwE,29275
 locust/contrib/mongodb.py,sha256=1seUYgJOaNKwybYOP9PUEVhgl8hGy-G33f8lFj3R8W8,1246
+locust/contrib/oai.py,sha256=Ot3T8lp31ThckGbNps86oVvq6Vn845Eec0mxhDmONDE,2684
 locust/contrib/postgres.py,sha256=OuMWnGYN10K65Tq2axVESEW25Y0g5gJb0rK90jkcCJg,1230
 locust/rpc/__init__.py,sha256=5YOu-58XSnt-oWWNATgXLTNdYoDkkngwHNXprxkWKSM,99
 locust/rpc/protocol.py,sha256=n-rb3GZQcAlldYDj4E4GuFGylYj_26GSS5U29meft5Y,1282
@@ -50,8 +51,8 @@ locust/webui/dist/assets/graphs-light.png,sha256=7L6pOehXqCojQclzeP91l-LskFQAw_n
 locust/webui/dist/assets/index-DcS9MJSn.js,sha256=CWbthTyP_6GEGZOqPHZXKAA9r7K0CWX32WWSAHKA9O0,1679855
 locust/webui/dist/assets/testruns-dark.png,sha256=np6MvpgJ2gkKQ66SOmukLtjsMtHqTSr5dNfza-2XtCo,267621
 locust/webui/dist/assets/testruns-light.png,sha256=iLAxBZh3kRsfGkcB1-1KSAbFgGji43IqiUrYuJlUoPk,276839
-locust-2.33.3.dev3.dist-info/METADATA,sha256=P9ITTJl4NJO_-oZfUEeIlF0s_VD-N5TtaJsZjuRdQYE,9637
-locust-2.33.3.dev3.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-locust-2.33.3.dev3.dist-info/entry_points.txt,sha256=RAdt8Ku-56m7bFjmdj-MBhbF6h4NX7tVODR9QNnOg0E,44
-locust-2.33.3.dev3.dist-info/licenses/LICENSE,sha256=5hnz-Vpj0Z3kSCQl0LzV2hT1TLc4LHcbpBp3Cy-EuyM,1110
-locust-2.33.3.dev3.dist-info/RECORD,,
+locust-2.33.3.dev9.dist-info/METADATA,sha256=19Te0Z7H0swCOOPhxqVhNrHfitTfPJ2uek21n4xZbj0,9637
+locust-2.33.3.dev9.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+locust-2.33.3.dev9.dist-info/entry_points.txt,sha256=RAdt8Ku-56m7bFjmdj-MBhbF6h4NX7tVODR9QNnOg0E,44
+locust-2.33.3.dev9.dist-info/licenses/LICENSE,sha256=5hnz-Vpj0Z3kSCQl0LzV2hT1TLc4LHcbpBp3Cy-EuyM,1110
+locust-2.33.3.dev9.dist-info/RECORD,,