{"repo": "aws-samples/sample-kolya-br-proxy", "files": {"benchmark/locustfile.py": {"899ec855828974416e292987c32e569e60297fec": {"t": 1785674422.0715644, "n": 7282, "h": "346c239567efe3f2", "c": "\"\"\"Locust benchmark for kolya-br-proxy — three API endpoints.\n\nUsage:\n    # Web UI\n    BENCHMARK_API_TOKEN=sk-ant-api03_xxx locust -f benchmark/locustfile.py --host https://api.kbp.kolya.fun\n\n    # Headless\n    BENCHMARK_API_TOKEN=sk-ant-api03_xxx locust -f benchmark/locustfile.py \\\n        --host https://api.kbp.kolya.fun --headless -u 10 -r 2 -t 5m --csv=results/run\n\n    # Single endpoint\n    BENCHMARK_API_TOKEN=sk-ant-api03_xxx locust -f benchmark/locustfile.py \\\n        --host https://api.kbp.kolya.fun OpenAIUser\n\"\"\"\n\nimport json\n\nfrom locust import HttpUser, between, tag, task\n\nfrom benchmark.config import (\n    ANTHROPIC_MODEL,\n    API_TOKEN,\n    GEMINI_MODEL,\n    MAX_TOKENS,\n    OPENAI_MODEL,\n    PROMPT_SIZE,\n    TEMPERATURE,\n    THINKING_BUDGET,\n)\nfrom benchmark.prompts import anthropic_payload, gemini_payload, openai_messages\nfrom benchmark.sse_client import fire_ttft_event, stream_and_measure\n\n\nclass OpenAIUser(HttpUser):\n    \"\"\"Load test /v1/chat/completions (OpenAI format).\"\"\"\n\n    wait_time = between(1, 3)\n    weight = 1\n\n    def on_start(self):\n        self.client.headers.update(\n            {\n                \"Authorization\": f\"Bearer {API_TOKEN}\",\n                \"Content-Type\": \"application/json\",\n            }\n        )\n        self._messages = openai_messages(PROMPT_SIZE)\n\n    def _body(self, stream: bool) -> dict:\n        return {\n            \"model\": OPENAI_MODEL,\n            \"messages\": self._messages,\n            \"max_tokens\": MAX_TOKENS,\n            \"temperature\": TEMPERATURE,\n            \"stream\": stream,\n        }\n\n    @tag(\"openai\")\n    @task(3)\n    def stream_chat(self):\n        name = \"/v1/chat/completions [stream]\"\n        with self.client.post(\n            \"/v1/chat/completions\",\n            json=self._body(stream=True),\n            stream=True,\n            catch_response=True,\n            name=name,\n        ) as resp:\n            if resp.status_code != 200:\n                resp.failure(f\"HTTP {resp.status_code}\")\n                return\n            result = stream_and_measure(resp, \"openai\")\n            if result[\"error\"]:\n                resp.failure(result[\"error\"])\n            else:\n                resp.success()\n            if result[\"ttft_s\"] is not None:\n                fire_ttft_event(name, result[\"ttft_s\"] * 1000, self.environment)\n\n    @tag(\"openai\")\n    @task(1)\n    def nonstream_chat(self):\n        with self.client.post(\n            \"/v1/chat/completions\",\n            json=self._body(stream=False),\n            catch_response=True,\n            name=\"/v1/chat/completions [non-stream]\",\n        ) as resp:\n            if resp.status_code != 200:\n                resp.failure(f\"HTTP {resp.status_code}\")\n                return\n            try:\n                data = resp.json()\n                if \"choices\" not in data:\n                    resp.failure(\"No choices in response\")\n                else:\n                    resp.success()\n            except json.JSONDecodeError:\n                resp.failure(\"Invalid JSON response\")\n\n\nclass AnthropicUser(HttpUser):\n    \"\"\"Load test /v1/messages (Anthropic format).\"\"\"\n\n    wait_time = between(1, 3)\n    weight = 1\n\n    def on_start(self):\n        self.client.headers.update(\n            {\n                \"x-api-key\": API_TOKEN,\n                \"Content-Type\": \"application/json\",\n                \"anthropic-version\": \"2023-06-01\",\n            }\n        )\n        self._payload = anthropic_payload(\n            PROMPT_SIZE, ANTHROPIC_MODEL, MAX_TOKENS, THINKING_BUDGET\n        )\n\n    def _body(self, stream: bool) -> dict:\n        return {**self._payload, \"stream\": stream}\n\n    @tag(\"anthropic\")\n    @task(3)\n    def stream_messages(self):\n        name = \"/v1/messages [stream]\"\n        with self.client.post(\n            \"/v1/messages\",\n            json=self._body(stream=True),\n            stream=True,\n            catch_response=True,\n            name=name,\n        ) as resp:\n            if resp.status_code != 200:\n                resp.failure(f\"HTTP {resp.status_code}\")\n                return\n            result = stream_and_measure(resp, \"anthropic\")\n            if result[\"error\"]:\n                resp.failure(result[\"error\"])\n            else:\n                resp.success()\n            if result[\"ttft_thinking_s\"] is not None:\n                fire_ttft_event(\n                    f\"{name} thinking\",\n                    result[\"ttft_thinking_s\"] * 1000,\n                    self.environment,\n                )\n            if result[\"ttft_s\"] is not None:\n                fire_ttft_event(name, result[\"ttft_s\"] * 1000, self.environment)\n\n    @tag(\"anthropic\")\n    @task(1)\n    def nonstream_messages(self):\n        with self.client.post(\n            \"/v1/messages\",\n            json=self._body(stream=False),\n            catch_response=True,\n            name=\"/v1/messages [non-stream]\",\n        ) as resp:\n            if resp.status_code != 200:\n                resp.failure(f\"HTTP {resp.status_code}\")\n                return\n            try:\n                data = resp.json()\n                if \"content\" not in data:\n                    resp.failure(\"No content in response\")\n                else:\n                    resp.success()\n            except json.JSONDecodeError:\n                resp.failure(\"Invalid JSON response\")\n\n\nclass GeminiUser(HttpUser):\n    \"\"\"Load test /v1beta/models/{model}:generateContent (Gemini format).\"\"\"\n\n    wait_time = between(1, 3)\n    weight = 1\n\n    def on_start(self):\n        self.client.headers.update(\n            {\n                \"x-goog-api-key\": API_TOKEN,\n                \"Content-Type\": \"application/json\",\n            }\n        )\n        self._payload = gemini_payload(PROMPT_SIZE, MAX_TOKENS, TEMPERATURE)\n\n    @tag(\"gemini\")\n    @task(3)\n    def stream_generate(self):\n        name = f\"/v1beta/models/{GEMINI_MODEL}:streamGenerateContent\"\n        url = f\"/v1beta/models/{GEMINI_MODEL}:streamGenerateContent?alt=sse\"\n        with self.client.post(\n            url,\n            json=self._payload,\n            stream=True,\n            catch_response=True,\n            name=name,\n        ) as resp:\n            if resp.status_code != 200:\n                resp.failure(f\"HTTP {resp.status_code}\")\n                return\n            result = stream_and_measure(resp, \"gemini\")\n            if result[\"error\"]:\n                resp.failure(result[\"error\"])\n            else:\n                resp.success()\n            if result[\"ttft_s\"] is not None:\n                fire_ttft_event(name, result[\"ttft_s\"] * 1000, self.environment)\n\n    @tag(\"gemini\")\n    @task(1)\n    def nonstream_generate(self):\n        url = f\"/v1beta/models/{GEMINI_MODEL}:generateContent\"\n        with self.client.post(\n            url,\n            json=self._payload,\n            catch_response=True,\n            name=f\"/v1beta/models/{GEMINI_MODEL}:generateContent\",\n        ) as resp:\n            if resp.status_code != 200:\n                resp.failure(f\"HTTP {resp.status_code}\")\n                return\n            try:\n                data = resp.json()\n                if \"candidates\" not in data:\n                    resp.failure(\"No candidates in response\")\n                else:\n                    resp.success()\n            except json.JSONDecodeError:\n                resp.failure(\"Invalid JSON response\")\n"}}}}