From 5b22f9450214bf847515281adff458218d5c95ad Mon Sep 17 00:00:00 2001 From: 6ba3i <112825897+6ba3i@users.noreply.github.com> Date: Wed, 14 Jan 2026 13:49:16 +0800 Subject: [PATCH] Feat: Benchmark CLI additions and documentation (#12536) ### What problem does this PR solve? This PR adds a dedicated HTTP benchmark CLI for RAGFlow chat and retrieval endpoints so we can measure latency/QPS. ### Type of change - [x] Documentation Update - [x] Other (please describe): Adds a CLI benchmarking tool for chat/retrieval latency/QPS --------- Co-authored-by: Liu An --- pyproject.toml | 1 + test/benchmark/README.md | 285 +++++++++++++ test/benchmark/__init__.py | 1 + test/benchmark/__main__.py | 5 + test/benchmark/auth.py | 88 ++++ test/benchmark/chat.py | 138 +++++++ test/benchmark/cli.py | 575 +++++++++++++++++++++++++++ test/benchmark/dataset.py | 146 +++++++ test/benchmark/http_client.py | 112 ++++++ test/benchmark/metrics.py | 67 ++++ test/benchmark/report.py | 105 +++++ test/benchmark/retrieval.py | 39 ++ test/benchmark/run_chat.sh | 28 ++ test/benchmark/run_retrieval.sh | 25 ++ test/benchmark/run_retrieval_chat.sh | 98 +++++ test/benchmark/test_docs/Doc1.pdf | 74 ++++ test/benchmark/test_docs/Doc2.pdf | 74 ++++ test/benchmark/test_docs/Doc3.pdf | 74 ++++ test/benchmark/utils.py | 41 ++ uv.lock | 2 + 20 files changed, 1978 insertions(+) create mode 100644 test/benchmark/README.md create mode 100644 test/benchmark/__init__.py create mode 100644 test/benchmark/__main__.py create mode 100644 test/benchmark/auth.py create mode 100644 test/benchmark/chat.py create mode 100644 test/benchmark/cli.py create mode 100644 test/benchmark/dataset.py create mode 100644 test/benchmark/http_client.py create mode 100644 test/benchmark/metrics.py create mode 100644 test/benchmark/report.py create mode 100644 test/benchmark/retrieval.py create mode 100755 test/benchmark/run_chat.sh create mode 100755 test/benchmark/run_retrieval.sh create mode 100755 test/benchmark/run_retrieval_chat.sh create mode 100644 test/benchmark/test_docs/Doc1.pdf create mode 100644 test/benchmark/test_docs/Doc2.pdf create mode 100644 test/benchmark/test_docs/Doc3.pdf create mode 100644 test/benchmark/utils.py diff --git a/pyproject.toml b/pyproject.toml index f8e5338f6..4ba8a8b78 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -169,6 +169,7 @@ test = [ "reportlab>=4.4.1", "requests>=2.32.2", "requests-toolbelt>=1.0.0", + "pycryptodomex==3.20.0", ] [[tool.uv.index]] diff --git a/test/benchmark/README.md b/test/benchmark/README.md new file mode 100644 index 000000000..847ae457b --- /dev/null +++ b/test/benchmark/README.md @@ -0,0 +1,285 @@ +# RAGFlow HTTP Benchmark CLI + +Run (from repo root): +``` + PYTHONPATH=./test uv run -m benchmark [global flags] [command flags] + Global flags can be placed before or after the command. +``` + +If you run from another directory: +``` + PYTHONPATH=/Directory_name/ragflow/test uv run -m benchmark [global flags] [command flags] +``` + +JSON args: + For --dataset-payload, --chat-payload, --messages-json, --extra-body, --payload + - Pass inline JSON: '{"key": "value"}' + - Or use a file: '@/path/to/file.json' + +Global flags +``` + --base-url + Base server URL. + Env: RAGFLOW_BASE_URL or HOST_ADDRESS + --api-version + API version string (default: v1). + Env: RAGFLOW_API_VERSION + --api-key + API key for Authorization: Bearer . + --connect-timeout + Connect timeout seconds (default: 5.0). + --read-timeout + Read timeout seconds (default: 60.0). + --no-verify-ssl + Disable SSL verification. + --iterations + Iterations per benchmark (default: 1). + --concurrency + Number of concurrent requests (default: 1). Uses multiprocessing. + --json + Output JSON report (plain stdout). + --print-response + Print response content per iteration (stdout). With --json, responses are included in the JSON output. + --response-max-chars + Truncate printed responses to N chars (0 = no limit). +``` + +Auth and bootstrap flags (used when --api-key is not provided) +``` + --login-email + Login email. + Env: RAGFLOW_EMAIL + --login-nickname + Nickname for registration. If omitted, defaults to email prefix when registering. + Env: RAGFLOW_NICKNAME + --login-password + Login password (encrypted client-side). Requires pycryptodomex in the test group. + --allow-register + Attempt /user/register before login (best effort). + --token-name + Optional API token name for /system/new_token. + --bootstrap-llm + Ensure LLM factory API key is configured via /llm/set_api_key. + --llm-factory + LLM factory name for bootstrap. + Env: RAGFLOW_LLM_FACTORY + --llm-api-key + LLM API key for bootstrap. + Env: ZHIPU_AI_API_KEY + --llm-api-base + Optional LLM API base URL. + Env: RAGFLOW_LLM_API_BASE + --set-tenant-info + Set tenant defaults via /user/set_tenant_info. + --tenant-llm-id + Tenant chat model ID. + Env: RAGFLOW_TENANT_LLM_ID + --tenant-embd-id + Tenant embedding model ID. + Env: RAGFLOW_TENANT_EMBD_ID + --tenant-img2txt-id + Tenant image2text model ID. + Env: RAGFLOW_TENANT_IMG2TXT_ID + --tenant-asr-id + Tenant ASR model ID (default empty). + Env: RAGFLOW_TENANT_ASR_ID + --tenant-tts-id + Tenant TTS model ID. + Env: RAGFLOW_TENANT_TTS_ID +``` + +Dataset/document flags (shared by chat and retrieval) +``` + --dataset-id + Existing dataset ID. + --dataset-ids + Comma-separated dataset IDs. + --dataset-name + Dataset name when creating a new dataset. + Env: RAGFLOW_DATASET_NAME + --dataset-payload + JSON body for dataset creation (see API docs). + --document-path + Document path to upload (repeatable). + --document-paths-file + File containing document paths, one per line. + --parse-timeout + Document parse timeout seconds (default: 120.0). + --parse-interval + Document parse poll interval seconds (default: 1.0). + --teardown + Delete created resources after run. +``` + +Chat command flags +``` + --chat-id + Existing chat ID. If omitted, a chat is created. + --chat-name + Chat name when creating a new chat. + Env: RAGFLOW_CHAT_NAME + --chat-payload + JSON body for chat creation (see API docs). + --model + Model field for OpenAI-compatible completion request. + Env: RAGFLOW_CHAT_MODEL + --message + Single user message (required unless --messages-json is provided). + --messages-json + JSON list of OpenAI-format messages (required unless --message is provided). + --extra-body + JSON extra_body for OpenAI-compatible request. +``` + +Retrieval command flags +``` + --question + Retrieval question (required unless provided in --payload). + --payload + JSON body for /api/v1/retrieval (see API docs). + --document-ids + Comma-separated document IDs for retrieval. +``` + +Model selection guidance + - Embedding model is tied to the dataset. + Set during dataset creation using --dataset-payload: +``` + {"name": "...", "embedding_model": "@"} +``` + Or set tenant defaults via --set-tenant-info with --tenant-embd-id. + - Chat model is tied to the chat assistant. + Set during chat creation using --chat-payload: +``` + {"name": "...", "llm": {"model_name": "@"}} +``` + Or set tenant defaults via --set-tenant-info with --tenant-llm-id. + - --model is required by the OpenAI-compatible endpoint but does not override + the chat assistant's configured model on the server. + +What this CLI can do + - This is a benchmark CLI. It always runs either a chat or retrieval benchmark + and prints a report. + - It can create datasets, upload documents, trigger parsing, and create chats + as part of a benchmark run (setup for the benchmark). + - It is not a general admin CLI; there are no standalone "create-only" or + "manage" commands. Use the reports to capture created IDs for reuse. + +Do I need the dataset ID? + - If the CLI creates a dataset, it uses the returned dataset ID internally. + You do not need to supply it for that same run. + - The report prints "Created Dataset ID" so you can reuse it later with + --dataset-id or --dataset-ids. + - Dataset name is only used at creation time. Selection is always by ID. + +Examples + +Example: chat benchmark creating dataset + upload + parse + chat (login + register) +``` + PYTHONPATH=./test uv run -m benchmark chat \ + --base-url http://127.0.0.1:9380 \ + --allow-register \ + --login-email "qa@infiniflow.org" \ + --login-password "123" \ + --bootstrap-llm \ + --llm-factory ZHIPU-AI \ + --llm-api-key $ZHIPU_AI_API_KEY \ + --dataset-name "bench_dataset" \ + --dataset-payload '{"name":"bench_dataset","embedding_model":"BAAI/bge-small-en-v1.5@Builtin"}' \ + --document-path test/benchmark/test_docs/Doc1.pdf \ + --document-path test/benchmark/test_docs/Doc2.pdf \ + --document-path test/benchmark/test_docs/Doc3.pdf \ + --chat-name "bench_chat" \ + --chat-payload '{"name":"bench_chat","llm":{"model_name":"glm-4-flash@ZHIPU-AI"}}' \ + --message "What is the purpose of RAGFlow?" \ + --model "glm-4-flash@ZHIPU-AI" +``` + +Example: chat benchmark with existing dataset + chat id (no creation) +``` + PYTHONPATH=./test uv run -m benchmark chat \ + --base-url http://127.0.0.1:9380 \ + --chat-id \ + --login-email "qa@infiniflow.org" \ + --login-password "123" \ + --message "What is the purpose of RAGFlow?" \ + --model "glm-4-flash@ZHIPU-AI" +``` + +Example: retrieval benchmark creating dataset + upload + parse +``` + PYTHONPATH=./test uv run -m benchmark retrieval \ + --base-url http://127.0.0.1:9380 \ + --allow-register \ + --login-email "qa@infiniflow.org" \ + --login-password "123" \ + --bootstrap-llm \ + --llm-factory ZHIPU-AI \ + --llm-api-key $ZHIPU_AI_API_KEY \ + --dataset-name "bench_dataset" \ + --dataset-payload '{"name":"bench_dataset","embedding_model":"BAAI/bge-small-en-v1.5@Builtin"}' \ + --document-path test/benchmark/test_docs/Doc1.pdf \ + --document-path test/benchmark/test_docs/Doc2.pdf \ + --document-path test/benchmark/test_docs/Doc3.pdf \ + --question "What does RAG mean?" +``` + +Example: retrieval benchmark with existing dataset IDs +``` + PYTHONPATH=./test uv run -m benchmark retrieval \ + --base-url http://127.0.0.1:9380 \ + --login-email "qa@infiniflow.org" \ + --login-password "123" \ + --dataset-ids "," \ + --question "What does RAG mean?" +``` + +Example: retrieval benchmark with existing dataset IDs and document IDs +``` + PYTHONPATH=./test uv run -m benchmark retrieval \ + --base-url http://127.0.0.1:9380 \ + --login-email "qa@infiniflow.org" \ + --login-password "123" \ + --dataset-id "" \ + --document-ids "," \ + --question "What does RAG mean?" +``` + +Quick scripts + +These scripts create a dataset, +upload/parse docs from test/benchmark/test_docs, run the benchmark, and clean up. +The both script runs retrieval then chat on the same dataset, then deletes it. + +- Make sure to run ```uv sync --python 3.12 --group test ``` before running the commands. +- It is also necessary to run these commands prior to initializing your containers if you plan on using the built-in embedded model: ```echo -e "TEI_MODEL=BAAI/bge-small-en-v1.5" >> docker/.env``` + and ```echo -e "COMPOSE_PROFILES=\${COMPOSE_PROFILES},tei-cpu" >> docker/.env``` + +Chat only: +``` + ./test/benchmark/run_chat.sh +``` + +Retrieval only: +``` + ./test/benchmark/run_retrieval.sh +``` + +Both (retrieval then chat on the same dataset): +``` + ./test/benchmark/run_retrieval_chat.sh +``` + +Requires: + - ZHIPU_AI_API_KEY exported in your shell. + +Defaults used: + - Base URL: http://127.0.0.1:9380 + - Login: qa@infiniflow.org / 123 (with allow-register) + - LLM bootstrap: ZHIPU-AI with $ZHIPU_AI_API_KEY + - Dataset: bench_dataset (BAAI/bge-small-en-v1.5@Builtin) + - Chat: bench_chat (glm-4-flash@ZHIPU-AI) + - Chat message: "What is the purpose of RAGFlow?" + - Retrieval question: "What does RAG mean?" + - Iterations: 1 + - concurrency:f 4 diff --git a/test/benchmark/__init__.py b/test/benchmark/__init__.py new file mode 100644 index 000000000..06603ac05 --- /dev/null +++ b/test/benchmark/__init__.py @@ -0,0 +1 @@ +"""RAGFlow HTTP API benchmark package.""" diff --git a/test/benchmark/__main__.py b/test/benchmark/__main__.py new file mode 100644 index 000000000..2f05ddc22 --- /dev/null +++ b/test/benchmark/__main__.py @@ -0,0 +1,5 @@ +from .cli import main + + +if __name__ == "__main__": + main() diff --git a/test/benchmark/auth.py b/test/benchmark/auth.py new file mode 100644 index 000000000..307dd4ed8 --- /dev/null +++ b/test/benchmark/auth.py @@ -0,0 +1,88 @@ +from typing import Any, Dict, Optional + +from .http_client import HttpClient + + +class AuthError(RuntimeError): + pass + + +def encrypt_password(password_plain: str) -> str: + try: + from api.utils.crypt import crypt + except Exception as exc: + raise AuthError( + "Password encryption unavailable; install pycryptodomex (uv sync --python 3.12 --group test)." + ) from exc + return crypt(password_plain) + +def register_user(client: HttpClient, email: str, nickname: str, password_enc: str) -> None: + payload = {"email": email, "nickname": nickname, "password": password_enc} + res = client.request_json("POST", "/user/register", use_api_base=False, auth_kind=None, json_body=payload) + if res.get("code") == 0: + return + msg = res.get("message", "") + if "has already registered" in msg: + return + raise AuthError(f"Register failed: {msg}") + + +def login_user(client: HttpClient, email: str, password_enc: str) -> str: + payload = {"email": email, "password": password_enc} + response = client.request("POST", "/user/login", use_api_base=False, auth_kind=None, json_body=payload) + try: + res = response.json() + except Exception as exc: + raise AuthError(f"Login failed: invalid JSON response ({exc})") from exc + if res.get("code") != 0: + raise AuthError(f"Login failed: {res.get('message')}") + token = response.headers.get("Authorization") + if not token: + raise AuthError("Login failed: missing Authorization header") + return token + + +def create_api_token(client: HttpClient, login_token: str, token_name: Optional[str] = None) -> str: + client.login_token = login_token + params = {"name": token_name} if token_name else None + res = client.request_json("POST", "/system/new_token", use_api_base=False, auth_kind="login", params=params) + if res.get("code") != 0: + raise AuthError(f"API token creation failed: {res.get('message')}") + token = res.get("data", {}).get("token") + if not token: + raise AuthError("API token creation failed: missing token in response") + return token + + +def get_my_llms(client: HttpClient) -> Dict[str, Any]: + res = client.request_json("GET", "/llm/my_llms", use_api_base=False, auth_kind="login") + if res.get("code") != 0: + raise AuthError(f"Failed to list LLMs: {res.get('message')}") + return res.get("data", {}) + + +def set_llm_api_key( + client: HttpClient, + llm_factory: str, + api_key: str, + base_url: Optional[str] = None, +) -> None: + payload = {"llm_factory": llm_factory, "api_key": api_key} + if base_url: + payload["base_url"] = base_url + res = client.request_json("POST", "/llm/set_api_key", use_api_base=False, auth_kind="login", json_body=payload) + if res.get("code") != 0: + raise AuthError(f"Failed to set LLM API key: {res.get('message')}") + + +def get_tenant_info(client: HttpClient) -> Dict[str, Any]: + res = client.request_json("GET", "/user/tenant_info", use_api_base=False, auth_kind="login") + if res.get("code") != 0: + raise AuthError(f"Failed to get tenant info: {res.get('message')}") + return res.get("data", {}) + + +def set_tenant_info(client: HttpClient, payload: Dict[str, Any]) -> None: + res = client.request_json("POST", "/user/set_tenant_info", use_api_base=False, auth_kind="login", json_body=payload) + if res.get("code") != 0: + raise AuthError(f"Failed to set tenant info: {res.get('message')}") diff --git a/test/benchmark/chat.py b/test/benchmark/chat.py new file mode 100644 index 000000000..52146314c --- /dev/null +++ b/test/benchmark/chat.py @@ -0,0 +1,138 @@ +import json +import time +from typing import Any, Dict, List, Optional + +from .http_client import HttpClient +from .metrics import ChatSample + + +class ChatError(RuntimeError): + pass + + +def delete_chat(client: HttpClient, chat_id: str) -> None: + payload = {"ids": [chat_id]} + res = client.request_json("DELETE", "/chats", json_body=payload) + if res.get("code") != 0: + raise ChatError(f"Delete chat failed: {res.get('message')}") + + +def create_chat( + client: HttpClient, + name: str, + dataset_ids: Optional[List[str]] = None, + payload: Optional[Dict[str, Any]] = None, +) -> Dict[str, Any]: + body = dict(payload or {}) + if "name" not in body: + body["name"] = name + if dataset_ids is not None and "dataset_ids" not in body: + body["dataset_ids"] = dataset_ids + res = client.request_json("POST", "/chats", json_body=body) + if res.get("code") != 0: + raise ChatError(f"Create chat failed: {res.get('message')}") + return res.get("data", {}) + + +def get_chat(client: HttpClient, chat_id: str) -> Dict[str, Any]: + res = client.request_json("GET", "/chats", params={"id": chat_id}) + if res.get("code") != 0: + raise ChatError(f"Get chat failed: {res.get('message')}") + data = res.get("data", []) + if not data: + raise ChatError("Chat not found") + return data[0] + + +def resolve_model(model: Optional[str], chat_data: Optional[Dict[str, Any]]) -> str: + if model: + return model + if chat_data: + llm = chat_data.get("llm") or {} + llm_name = llm.get("model_name") + if llm_name: + return llm_name + raise ChatError("Model name is required; provide --model or use a chat with llm.model_name.") + + +def _parse_stream_error(response) -> Optional[str]: + content_type = response.headers.get("Content-Type", "") + if "text/event-stream" in content_type: + return None + try: + payload = response.json() + except Exception: + return f"Unexpected non-stream response (status {response.status_code})" + if payload.get("code") not in (0, None): + return payload.get("message", "Unknown error") + return f"Unexpected non-stream response (status {response.status_code})" + + +def stream_chat_completion( + client: HttpClient, + chat_id: str, + model: str, + messages: List[Dict[str, Any]], + extra_body: Optional[Dict[str, Any]] = None, +) -> ChatSample: + payload: Dict[str, Any] = {"model": model, "messages": messages, "stream": True} + if extra_body: + payload["extra_body"] = extra_body + t0 = time.perf_counter() + response = client.request( + "POST", + f"/chats_openai/{chat_id}/chat/completions", + json_body=payload, + stream=True, + ) + error = _parse_stream_error(response) + if error: + response.close() + return ChatSample(t0=t0, t1=None, t2=None, error=error) + + t1: Optional[float] = None + t2: Optional[float] = None + stream_error: Optional[str] = None + content_parts: List[str] = [] + try: + for raw_line in response.iter_lines(decode_unicode=True): + if raw_line is None: + continue + line = raw_line.strip() + if not line or not line.startswith("data:"): + continue + data = line[5:].strip() + if not data: + continue + if data == "[DONE]": + t2 = time.perf_counter() + break + try: + chunk = json.loads(data) + except Exception as exc: + stream_error = f"Invalid JSON chunk: {exc}" + t2 = time.perf_counter() + break + choices = chunk.get("choices") or [] + choice = choices[0] if choices else {} + delta = choice.get("delta") or {} + content = delta.get("content") + if t1 is None and isinstance(content, str) and content != "": + t1 = time.perf_counter() + if isinstance(content, str) and content: + content_parts.append(content) + finish_reason = choice.get("finish_reason") + if finish_reason: + t2 = time.perf_counter() + break + finally: + response.close() + + if t2 is None: + t2 = time.perf_counter() + response_text = "".join(content_parts) if content_parts else None + if stream_error: + return ChatSample(t0=t0, t1=t1, t2=t2, error=stream_error, response_text=response_text) + if t1 is None: + return ChatSample(t0=t0, t1=None, t2=t2, error="No assistant content received", response_text=response_text) + return ChatSample(t0=t0, t1=t1, t2=t2, error=None, response_text=response_text) diff --git a/test/benchmark/cli.py b/test/benchmark/cli.py new file mode 100644 index 000000000..53a04321b --- /dev/null +++ b/test/benchmark/cli.py @@ -0,0 +1,575 @@ +import argparse +import json +import os +import multiprocessing as mp +import time +from concurrent.futures import ProcessPoolExecutor, as_completed +from pathlib import Path +from typing import Any, Dict, List, Optional + +from . import auth +from .auth import AuthError +from .chat import ChatError, create_chat, delete_chat, get_chat, resolve_model, stream_chat_completion +from .dataset import ( + DatasetError, + create_dataset, + dataset_has_chunks, + delete_dataset, + extract_document_ids, + list_datasets, + parse_documents, + upload_documents, + wait_for_parse_done, +) +from .http_client import HttpClient +from .metrics import ChatSample, RetrievalSample, summarize +from .report import chat_report, retrieval_report +from .retrieval import RetrievalError, build_payload, run_retrieval as run_retrieval_request +from .utils import eprint, load_json_arg, split_csv + + +def _parse_args() -> argparse.Namespace: + base_parser = argparse.ArgumentParser(add_help=False) + base_parser.add_argument( + "--base-url", + default=os.getenv("RAGFLOW_BASE_URL") or os.getenv("HOST_ADDRESS"), + help="Base URL (env: RAGFLOW_BASE_URL or HOST_ADDRESS)", + ) + base_parser.add_argument( + "--api-version", + default=os.getenv("RAGFLOW_API_VERSION", "v1"), + help="API version (default: v1)", + ) + base_parser.add_argument("--api-key", help="API key (Bearer token)") + base_parser.add_argument("--connect-timeout", type=float, default=5.0, help="Connect timeout seconds") + base_parser.add_argument("--read-timeout", type=float, default=60.0, help="Read timeout seconds") + base_parser.add_argument("--no-verify-ssl", action="store_false", dest="verify_ssl", help="Disable SSL verification") + base_parser.add_argument("--iterations", type=int, default=1, help="Number of iterations") + base_parser.add_argument("--concurrency", type=int, default=1, help="Concurrency") + base_parser.add_argument("--json", action="store_true", help="Print JSON report (optional)") + base_parser.add_argument("--print-response", action="store_true", help="Print response content per iteration") + base_parser.add_argument( + "--response-max-chars", + type=int, + default=0, + help="Truncate printed response to N chars (0 = no limit)", + ) + + # Auth/login options + base_parser.add_argument("--login-email", default=os.getenv("RAGFLOW_EMAIL"), help="Login email") + base_parser.add_argument("--login-nickname", default=os.getenv("RAGFLOW_NICKNAME"), help="Nickname for registration") + base_parser.add_argument("--login-password", help="Login password (encrypted client-side)") + base_parser.add_argument("--allow-register", action="store_true", help="Attempt /user/register before login") + base_parser.add_argument("--token-name", help="Optional API token name") + base_parser.add_argument("--bootstrap-llm", action="store_true", help="Ensure LLM factory API key is configured") + base_parser.add_argument("--llm-factory", default=os.getenv("RAGFLOW_LLM_FACTORY"), help="LLM factory name") + base_parser.add_argument("--llm-api-key", default=os.getenv("ZHIPU_AI_API_KEY"), help="LLM API key") + base_parser.add_argument("--llm-api-base", default=os.getenv("RAGFLOW_LLM_API_BASE"), help="LLM API base URL") + base_parser.add_argument("--set-tenant-info", action="store_true", help="Set tenant default model IDs") + base_parser.add_argument("--tenant-llm-id", default=os.getenv("RAGFLOW_TENANT_LLM_ID"), help="Tenant chat model ID") + base_parser.add_argument("--tenant-embd-id", default=os.getenv("RAGFLOW_TENANT_EMBD_ID"), help="Tenant embedding model ID") + base_parser.add_argument("--tenant-img2txt-id", default=os.getenv("RAGFLOW_TENANT_IMG2TXT_ID"), help="Tenant image2text model ID") + base_parser.add_argument("--tenant-asr-id", default=os.getenv("RAGFLOW_TENANT_ASR_ID", ""), help="Tenant ASR model ID") + base_parser.add_argument("--tenant-tts-id", default=os.getenv("RAGFLOW_TENANT_TTS_ID"), help="Tenant TTS model ID") + + # Dataset/doc options + base_parser.add_argument("--dataset-id", help="Existing dataset ID") + base_parser.add_argument("--dataset-ids", help="Comma-separated dataset IDs") + base_parser.add_argument("--dataset-name", default=os.getenv("RAGFLOW_DATASET_NAME"), help="Dataset name when creating") + base_parser.add_argument("--dataset-payload", help="Dataset payload JSON or @file") + base_parser.add_argument("--document-path", action="append", help="Document path (repeatable)") + base_parser.add_argument("--document-paths-file", help="File with document paths, one per line") + base_parser.add_argument("--parse-timeout", type=float, default=120.0, help="Parse timeout seconds") + base_parser.add_argument("--parse-interval", type=float, default=1.0, help="Parse poll interval seconds") + base_parser.add_argument("--teardown", action="store_true", help="Delete created resources after run") + + parser = argparse.ArgumentParser(description="RAGFlow HTTP API benchmark", parents=[base_parser]) + subparsers = parser.add_subparsers(dest="command", required=True) + + chat_parser = subparsers.add_parser( + "chat", + help="Chat streaming latency benchmark", + parents=[base_parser], + add_help=False, + ) + chat_parser.add_argument("--chat-id", help="Existing chat ID") + chat_parser.add_argument("--chat-name", default=os.getenv("RAGFLOW_CHAT_NAME"), help="Chat name when creating") + chat_parser.add_argument("--chat-payload", help="Chat payload JSON or @file") + chat_parser.add_argument("--model", default=os.getenv("RAGFLOW_CHAT_MODEL"), help="Model name for OpenAI endpoint") + chat_parser.add_argument("--message", help="User message") + chat_parser.add_argument("--messages-json", help="Messages JSON or @file") + chat_parser.add_argument("--extra-body", help="extra_body JSON or @file") + + retrieval_parser = subparsers.add_parser( + "retrieval", + help="Retrieval latency benchmark", + parents=[base_parser], + add_help=False, + ) + retrieval_parser.add_argument("--question", help="Retrieval question") + retrieval_parser.add_argument("--payload", help="Retrieval payload JSON or @file") + retrieval_parser.add_argument("--document-ids", help="Comma-separated document IDs") + + return parser.parse_args() + + +def _load_paths(args: argparse.Namespace) -> List[str]: + paths = [] + if args.document_path: + paths.extend(args.document_path) + if args.document_paths_file: + file_path = Path(args.document_paths_file) + for line in file_path.read_text(encoding="utf-8").splitlines(): + line = line.strip() + if line: + paths.append(line) + return paths + + +def _truncate_text(text: str, max_chars: int) -> str: + if max_chars and len(text) > max_chars: + return f"{text[:max_chars]}...[truncated]" + return text + + +def _format_chat_response(sample: ChatSample, max_chars: int) -> str: + if sample.error: + text = f"[error] {sample.error}" + if sample.response_text: + text = f"{text} | {sample.response_text}" + else: + text = sample.response_text or "" + if not text: + text = "(empty)" + return _truncate_text(text, max_chars) + + +def _format_retrieval_response(sample: RetrievalSample, max_chars: int) -> str: + if sample.response is not None: + text = json.dumps(sample.response, ensure_ascii=False, sort_keys=True) + if sample.error: + text = f"[error] {sample.error} | {text}" + elif sample.error: + text = f"[error] {sample.error}" + else: + text = "(empty)" + return _truncate_text(text, max_chars) + + +def _chat_worker( + base_url: str, + api_version: str, + api_key: str, + connect_timeout: float, + read_timeout: float, + verify_ssl: bool, + chat_id: str, + model: str, + messages: List[Dict[str, Any]], + extra_body: Optional[Dict[str, Any]], +) -> ChatSample: + client = HttpClient( + base_url=base_url, + api_version=api_version, + api_key=api_key, + connect_timeout=connect_timeout, + read_timeout=read_timeout, + verify_ssl=verify_ssl, + ) + return stream_chat_completion(client, chat_id, model, messages, extra_body) + + +def _retrieval_worker( + base_url: str, + api_version: str, + api_key: str, + connect_timeout: float, + read_timeout: float, + verify_ssl: bool, + payload: Dict[str, Any], +) -> RetrievalSample: + client = HttpClient( + base_url=base_url, + api_version=api_version, + api_key=api_key, + connect_timeout=connect_timeout, + read_timeout=read_timeout, + verify_ssl=verify_ssl, + ) + return run_retrieval_request(client, payload) + + +def _ensure_auth(client: HttpClient, args: argparse.Namespace) -> None: + if args.api_key: + client.api_key = args.api_key + return + if not args.login_email: + raise AuthError("Missing API key and login email") + if not args.login_password: + raise AuthError("Missing login password") + + password_enc = auth.encrypt_password(args.login_password) + + if args.allow_register: + nickname = args.login_nickname or args.login_email.split("@")[0] + try: + auth.register_user(client, args.login_email, nickname, password_enc) + except AuthError as exc: + eprint(f"Register warning: {exc}") + + login_token = auth.login_user(client, args.login_email, password_enc) + client.login_token = login_token + + if args.bootstrap_llm: + if not args.llm_factory: + raise AuthError("Missing --llm-factory for bootstrap") + if not args.llm_api_key: + raise AuthError("Missing --llm-api-key for bootstrap") + existing = auth.get_my_llms(client) + if args.llm_factory not in existing: + auth.set_llm_api_key(client, args.llm_factory, args.llm_api_key, args.llm_api_base) + + if args.set_tenant_info: + if not args.tenant_llm_id or not args.tenant_embd_id: + raise AuthError("Missing --tenant-llm-id or --tenant-embd-id for tenant setup") + tenant = auth.get_tenant_info(client) + tenant_id = tenant.get("tenant_id") + if not tenant_id: + raise AuthError("Tenant info missing tenant_id") + payload = { + "tenant_id": tenant_id, + "llm_id": args.tenant_llm_id, + "embd_id": args.tenant_embd_id, + "img2txt_id": args.tenant_img2txt_id or "", + "asr_id": args.tenant_asr_id or "", + "tts_id": args.tenant_tts_id, + } + auth.set_tenant_info(client, payload) + + api_key = auth.create_api_token(client, login_token, args.token_name) + client.api_key = api_key + + +def _prepare_dataset( + client: HttpClient, + args: argparse.Namespace, + needs_dataset: bool, + document_paths: List[str], +) -> Dict[str, Any]: + created = {} + dataset_ids = split_csv(args.dataset_ids) or [] + dataset_id = args.dataset_id + dataset_payload = load_json_arg(args.dataset_payload, "dataset-payload") if args.dataset_payload else None + + if dataset_id: + dataset_ids = [dataset_id] + elif dataset_ids: + dataset_id = dataset_ids[0] + elif needs_dataset or document_paths: + if not args.dataset_name and not (dataset_payload and dataset_payload.get("name")): + raise DatasetError("Missing --dataset-name or dataset payload name") + name = args.dataset_name or dataset_payload.get("name") + data = create_dataset(client, name, dataset_payload) + dataset_id = data.get("id") + if not dataset_id: + raise DatasetError("Dataset creation did not return id") + dataset_ids = [dataset_id] + created["Created Dataset ID"] = dataset_id + return { + "dataset_id": dataset_id, + "dataset_ids": dataset_ids, + "dataset_payload": dataset_payload, + "created": created, + } + + +def _maybe_upload_and_parse( + client: HttpClient, + dataset_id: str, + document_paths: List[str], + parse_timeout: float, + parse_interval: float, +) -> List[str]: + if not document_paths: + return [] + docs = upload_documents(client, dataset_id, document_paths) + doc_ids = extract_document_ids(docs) + if not doc_ids: + raise DatasetError("No document IDs returned after upload") + parse_documents(client, dataset_id, doc_ids) + wait_for_parse_done(client, dataset_id, doc_ids, parse_timeout, parse_interval) + return doc_ids + + +def _ensure_dataset_has_chunks(client: HttpClient, dataset_id: str) -> None: + datasets = list_datasets(client, dataset_id=dataset_id) + if not datasets: + raise DatasetError("Dataset not found") + if not dataset_has_chunks(datasets[0]): + raise DatasetError("Dataset has no parsed chunks; upload and parse documents first.") + + +def _cleanup(client: HttpClient, created: Dict[str, str], teardown: bool) -> None: + if not teardown: + return + chat_id = created.get("Created Chat ID") + if chat_id: + try: + delete_chat(client, chat_id) + except Exception as exc: + eprint(f"Cleanup warning: failed to delete chat {chat_id}: {exc}") + dataset_id = created.get("Created Dataset ID") + if dataset_id: + try: + delete_dataset(client, dataset_id) + except Exception as exc: + eprint(f"Cleanup warning: failed to delete dataset {dataset_id}: {exc}") + + +def run_chat(client: HttpClient, args: argparse.Namespace) -> int: + document_paths = _load_paths(args) + needs_dataset = bool(document_paths) + dataset_info = _prepare_dataset(client, args, needs_dataset, document_paths) + created = dict(dataset_info["created"]) + dataset_id = dataset_info["dataset_id"] + dataset_ids = dataset_info["dataset_ids"] + doc_ids = [] + if dataset_id and document_paths: + doc_ids = _maybe_upload_and_parse(client, dataset_id, document_paths, args.parse_timeout, args.parse_interval) + created["Created Document IDs"] = ",".join(doc_ids) + if dataset_id and not document_paths: + _ensure_dataset_has_chunks(client, dataset_id) + if dataset_id and not document_paths and dataset_ids: + _ensure_dataset_has_chunks(client, dataset_id) + + chat_payload = load_json_arg(args.chat_payload, "chat-payload") if args.chat_payload else None + chat_id = args.chat_id + if not chat_id: + if not args.chat_name and not (chat_payload and chat_payload.get("name")): + raise ChatError("Missing --chat-name or chat payload name") + chat_name = args.chat_name or chat_payload.get("name") + chat_data = create_chat(client, chat_name, dataset_ids or [], chat_payload) + chat_id = chat_data.get("id") + if not chat_id: + raise ChatError("Chat creation did not return id") + created["Created Chat ID"] = chat_id + chat_data = get_chat(client, chat_id) + model = resolve_model(args.model, chat_data) + + messages = None + if args.messages_json: + messages = load_json_arg(args.messages_json, "messages-json") + if not messages: + if not args.message: + raise ChatError("Missing --message or --messages-json") + messages = [{"role": "user", "content": args.message}] + extra_body = load_json_arg(args.extra_body, "extra-body") if args.extra_body else None + + samples: List[ChatSample] = [] + responses: List[str] = [] + start_time = time.perf_counter() + if args.concurrency <= 1: + for _ in range(args.iterations): + samples.append(stream_chat_completion(client, chat_id, model, messages, extra_body)) + else: + results: List[Optional[ChatSample]] = [None] * args.iterations + mp_context = mp.get_context("spawn") + with ProcessPoolExecutor(max_workers=args.concurrency, mp_context=mp_context) as executor: + future_map = { + executor.submit( + _chat_worker, + client.base_url, + client.api_version, + client.api_key or "", + client.connect_timeout, + client.read_timeout, + client.verify_ssl, + chat_id, + model, + messages, + extra_body, + ): idx + for idx in range(args.iterations) + } + for future in as_completed(future_map): + idx = future_map[future] + results[idx] = future.result() + samples = [sample for sample in results if sample is not None] + total_duration = time.perf_counter() - start_time + if args.print_response: + for idx, sample in enumerate(samples, start=1): + rendered = _format_chat_response(sample, args.response_max_chars) + if args.json: + responses.append(rendered) + else: + print(f"Response[{idx}]: {rendered}") + + total_latencies = [s.total_latency for s in samples if s.total_latency is not None and s.error is None] + first_latencies = [s.first_token_latency for s in samples if s.first_token_latency is not None and s.error is None] + success = len(total_latencies) + failure = len(samples) - success + errors = [s.error for s in samples if s.error] + + total_stats = summarize(total_latencies) + first_stats = summarize(first_latencies) + if args.json: + payload = { + "interface": "chat", + "concurrency": args.concurrency, + "iterations": args.iterations, + "success": success, + "failure": failure, + "model": model, + "total_latency": total_stats, + "first_token_latency": first_stats, + "errors": [e for e in errors if e], + "created": created, + "total_duration_s": total_duration, + "qps": (args.iterations / total_duration) if total_duration > 0 else None, + } + if args.print_response: + payload["responses"] = responses + print(json.dumps(payload, sort_keys=True)) + else: + report = chat_report( + interface="chat", + concurrency=args.concurrency, + total_duration_s=total_duration, + iterations=args.iterations, + success=success, + failure=failure, + model=model, + total_stats=total_stats, + first_token_stats=first_stats, + errors=[e for e in errors if e], + created=created, + ) + print(report, end="") + _cleanup(client, created, args.teardown) + return 0 if failure == 0 else 1 + + +def run_retrieval(client: HttpClient, args: argparse.Namespace) -> int: + document_paths = _load_paths(args) + needs_dataset = True + dataset_info = _prepare_dataset(client, args, needs_dataset, document_paths) + created = dict(dataset_info["created"]) + dataset_id = dataset_info["dataset_id"] + dataset_ids = dataset_info["dataset_ids"] + if not dataset_ids: + raise RetrievalError("dataset_ids required for retrieval") + + doc_ids = [] + if dataset_id and document_paths: + doc_ids = _maybe_upload_and_parse(client, dataset_id, document_paths, args.parse_timeout, args.parse_interval) + created["Created Document IDs"] = ",".join(doc_ids) + + payload_override = load_json_arg(args.payload, "payload") if args.payload else None + question = args.question + if not question and (payload_override is None or "question" not in payload_override): + raise RetrievalError("Missing --question or retrieval payload question") + document_ids = split_csv(args.document_ids) if args.document_ids else None + + payload = build_payload(question, dataset_ids, document_ids, payload_override) + + samples: List[RetrievalSample] = [] + responses: List[str] = [] + start_time = time.perf_counter() + if args.concurrency <= 1: + for _ in range(args.iterations): + samples.append(run_retrieval_request(client, payload)) + else: + results: List[Optional[RetrievalSample]] = [None] * args.iterations + mp_context = mp.get_context("spawn") + with ProcessPoolExecutor(max_workers=args.concurrency, mp_context=mp_context) as executor: + future_map = { + executor.submit( + _retrieval_worker, + client.base_url, + client.api_version, + client.api_key or "", + client.connect_timeout, + client.read_timeout, + client.verify_ssl, + payload, + ): idx + for idx in range(args.iterations) + } + for future in as_completed(future_map): + idx = future_map[future] + results[idx] = future.result() + samples = [sample for sample in results if sample is not None] + total_duration = time.perf_counter() - start_time + if args.print_response: + for idx, sample in enumerate(samples, start=1): + rendered = _format_retrieval_response(sample, args.response_max_chars) + if args.json: + responses.append(rendered) + else: + print(f"Response[{idx}]: {rendered}") + + latencies = [s.latency for s in samples if s.latency is not None and s.error is None] + success = len(latencies) + failure = len(samples) - success + errors = [s.error for s in samples if s.error] + + stats = summarize(latencies) + if args.json: + payload = { + "interface": "retrieval", + "concurrency": args.concurrency, + "iterations": args.iterations, + "success": success, + "failure": failure, + "latency": stats, + "errors": [e for e in errors if e], + "created": created, + "total_duration_s": total_duration, + "qps": (args.iterations / total_duration) if total_duration > 0 else None, + } + if args.print_response: + payload["responses"] = responses + print(json.dumps(payload, sort_keys=True)) + else: + report = retrieval_report( + interface="retrieval", + concurrency=args.concurrency, + total_duration_s=total_duration, + iterations=args.iterations, + success=success, + failure=failure, + stats=stats, + errors=[e for e in errors if e], + created=created, + ) + print(report, end="") + _cleanup(client, created, args.teardown) + return 0 if failure == 0 else 1 + + +def main() -> None: + args = _parse_args() + if not args.base_url: + raise SystemExit("Missing --base-url or HOST_ADDRESS") + if args.iterations < 1: + raise SystemExit("--iterations must be >= 1") + if args.concurrency < 1: + raise SystemExit("--concurrency must be >= 1") + client = HttpClient( + base_url=args.base_url, + api_version=args.api_version, + api_key=args.api_key, + connect_timeout=args.connect_timeout, + read_timeout=args.read_timeout, + verify_ssl=args.verify_ssl, + ) + try: + _ensure_auth(client, args) + if args.command == "chat": + raise SystemExit(run_chat(client, args)) + if args.command == "retrieval": + raise SystemExit(run_retrieval(client, args)) + raise SystemExit("Unknown command") + except (AuthError, DatasetError, ChatError, RetrievalError) as exc: + eprint(f"Error: {exc}") + raise SystemExit(2) diff --git a/test/benchmark/dataset.py b/test/benchmark/dataset.py new file mode 100644 index 000000000..e349bddfb --- /dev/null +++ b/test/benchmark/dataset.py @@ -0,0 +1,146 @@ +from pathlib import Path +from typing import Any, Dict, Iterable, List, Optional + +from .http_client import HttpClient + +try: + from requests_toolbelt import MultipartEncoder +except Exception: # pragma: no cover - fallback without toolbelt + MultipartEncoder = None + + +class DatasetError(RuntimeError): + pass + + +def create_dataset(client: HttpClient, name: str, payload: Optional[Dict[str, Any]] = None) -> Dict[str, Any]: + body = dict(payload or {}) + if "name" not in body: + body["name"] = name + res = client.request_json("POST", "/datasets", json_body=body) + if res.get("code") != 0: + raise DatasetError(f"Create dataset failed: {res.get('message')}") + return res.get("data", {}) + + +def list_datasets(client: HttpClient, dataset_id: Optional[str] = None, name: Optional[str] = None) -> List[Dict[str, Any]]: + params = {} + if dataset_id is not None: + params["id"] = dataset_id + if name is not None: + params["name"] = name + res = client.request_json("GET", "/datasets", params=params or None) + if res.get("code") != 0: + raise DatasetError(f"List datasets failed: {res.get('message')}") + return res.get("data", []) + + +def delete_dataset(client: HttpClient, dataset_id: str) -> None: + payload = {"ids": [dataset_id]} + res = client.request_json("DELETE", "/datasets", json_body=payload) + if res.get("code") != 0: + raise DatasetError(f"Delete dataset failed: {res.get('message')}") + + +def upload_documents(client: HttpClient, dataset_id: str, file_paths: Iterable[str]) -> List[Dict[str, Any]]: + paths = [Path(p) for p in file_paths] + if MultipartEncoder is None: + files = [("file", (p.name, p.open("rb"))) for p in paths] + try: + response = client.request( + "POST", + f"/datasets/{dataset_id}/documents", + headers=None, + data=None, + json_body=None, + files=files, + params=None, + stream=False, + auth_kind="api", + ) + finally: + for _, (_, fh) in files: + fh.close() + res = response.json() + else: + fields = [] + file_handles = [] + try: + for path in paths: + fh = path.open("rb") + fields.append(("file", (path.name, fh))) + file_handles.append(fh) + encoder = MultipartEncoder(fields=fields) + headers = {"Content-Type": encoder.content_type} + response = client.request( + "POST", + f"/datasets/{dataset_id}/documents", + headers=headers, + data=encoder, + json_body=None, + params=None, + stream=False, + auth_kind="api", + ) + res = response.json() + finally: + for fh in file_handles: + fh.close() + if res.get("code") != 0: + raise DatasetError(f"Upload documents failed: {res.get('message')}") + return res.get("data", []) + + +def parse_documents(client: HttpClient, dataset_id: str, document_ids: List[str]) -> Dict[str, Any]: + payload = {"document_ids": document_ids} + res = client.request_json("POST", f"/datasets/{dataset_id}/chunks", json_body=payload) + if res.get("code") != 0: + raise DatasetError(f"Parse documents failed: {res.get('message')}") + return res + + +def list_documents(client: HttpClient, dataset_id: str, params: Optional[Dict[str, Any]] = None) -> Dict[str, Any]: + res = client.request_json("GET", f"/datasets/{dataset_id}/documents", params=params) + if res.get("code") != 0: + raise DatasetError(f"List documents failed: {res.get('message')}") + return res.get("data", {}) + + +def wait_for_parse_done( + client: HttpClient, + dataset_id: str, + document_ids: Optional[List[str]], + timeout: float, + interval: float, +) -> None: + import time + + start = time.monotonic() + while True: + data = list_documents(client, dataset_id) + docs = data.get("docs", []) + target_ids = set(document_ids or []) + all_done = True + for doc in docs: + if target_ids and doc.get("id") not in target_ids: + continue + if doc.get("run") != "DONE": + all_done = False + break + if all_done: + return + if time.monotonic() - start > timeout: + raise DatasetError("Document parsing timeout") + time.sleep(max(interval, 0.1)) + + +def extract_document_ids(documents: Iterable[Dict[str, Any]]) -> List[str]: + return [doc["id"] for doc in documents if "id" in doc] + + +def dataset_has_chunks(dataset_info: Dict[str, Any]) -> bool: + for key in ("chunk_count", "chunk_num"): + value = dataset_info.get(key) + if isinstance(value, int) and value > 0: + return True + return False diff --git a/test/benchmark/http_client.py b/test/benchmark/http_client.py new file mode 100644 index 000000000..c8b1a91a7 --- /dev/null +++ b/test/benchmark/http_client.py @@ -0,0 +1,112 @@ +import json +from typing import Any, Dict, Optional, Tuple + +import requests + + +class HttpClient: + def __init__( + self, + base_url: str, + api_version: str = "v1", + api_key: Optional[str] = None, + login_token: Optional[str] = None, + connect_timeout: float = 5.0, + read_timeout: float = 60.0, + verify_ssl: bool = True, + ) -> None: + self.base_url = base_url.rstrip("/") + self.api_version = api_version + self.api_key = api_key + self.login_token = login_token + self.connect_timeout = connect_timeout + self.read_timeout = read_timeout + self.verify_ssl = verify_ssl + + def api_base(self) -> str: + return f"{self.base_url}/api/{self.api_version}" + + def non_api_base(self) -> str: + return f"{self.base_url}/{self.api_version}" + + def build_url(self, path: str, use_api_base: bool = True) -> str: + base = self.api_base() if use_api_base else self.non_api_base() + return f"{base}/{path.lstrip('/')}" + + def _headers(self, auth_kind: Optional[str], extra: Optional[Dict[str, str]]) -> Dict[str, str]: + headers = {} + if auth_kind == "api" and self.api_key: + headers["Authorization"] = f"Bearer {self.api_key}" + elif auth_kind == "login" and self.login_token: + headers["Authorization"] = self.login_token + if extra: + headers.update(extra) + return headers + + def request( + self, + method: str, + path: str, + *, + use_api_base: bool = True, + auth_kind: Optional[str] = "api", + headers: Optional[Dict[str, str]] = None, + json_body: Optional[Dict[str, Any]] = None, + data: Any = None, + files: Any = None, + params: Optional[Dict[str, Any]] = None, + stream: bool = False, + ) -> requests.Response: + url = self.build_url(path, use_api_base=use_api_base) + merged_headers = self._headers(auth_kind, headers) + timeout: Tuple[float, float] = (self.connect_timeout, self.read_timeout) + return requests.request( + method=method, + url=url, + headers=merged_headers, + json=json_body, + data=data, + files=files, + params=params, + timeout=timeout, + stream=stream, + verify=self.verify_ssl, + ) + + def request_json( + self, + method: str, + path: str, + *, + use_api_base: bool = True, + auth_kind: Optional[str] = "api", + headers: Optional[Dict[str, str]] = None, + json_body: Optional[Dict[str, Any]] = None, + data: Any = None, + files: Any = None, + params: Optional[Dict[str, Any]] = None, + stream: bool = False, + ) -> Dict[str, Any]: + response = self.request( + method, + path, + use_api_base=use_api_base, + auth_kind=auth_kind, + headers=headers, + json_body=json_body, + data=data, + files=files, + params=params, + stream=stream, + ) + try: + return response.json() + except Exception as exc: + raise ValueError(f"Non-JSON response from {path}: {exc}") from exc + + @staticmethod + def parse_json_bytes(raw: bytes) -> Dict[str, Any]: + try: + return json.loads(raw.decode("utf-8")) + except Exception as exc: + raise ValueError(f"Invalid JSON payload: {exc}") from exc diff --git a/test/benchmark/metrics.py b/test/benchmark/metrics.py new file mode 100644 index 000000000..02183ec49 --- /dev/null +++ b/test/benchmark/metrics.py @@ -0,0 +1,67 @@ +import math +from dataclasses import dataclass +from typing import Any, List, Optional + + +@dataclass +class ChatSample: + t0: float + t1: Optional[float] + t2: Optional[float] + error: Optional[str] = None + response_text: Optional[str] = None + + @property + def first_token_latency(self) -> Optional[float]: + if self.t1 is None: + return None + return self.t1 - self.t0 + + @property + def total_latency(self) -> Optional[float]: + if self.t2 is None: + return None + return self.t2 - self.t0 + + +@dataclass +class RetrievalSample: + t0: float + t1: Optional[float] + error: Optional[str] = None + response: Optional[Any] = None + + @property + def latency(self) -> Optional[float]: + if self.t1 is None: + return None + return self.t1 - self.t0 + + +def _percentile(sorted_values: List[float], p: float) -> Optional[float]: + if not sorted_values: + return None + n = len(sorted_values) + k = max(0, math.ceil((p / 100.0) * n) - 1) + return sorted_values[k] + + +def summarize(values: List[float]) -> dict: + if not values: + return { + "count": 0, + "avg": None, + "min": None, + "p50": None, + "p90": None, + "p95": None, + } + sorted_vals = sorted(values) + return { + "count": len(values), + "avg": sum(values) / len(values), + "min": sorted_vals[0], + "p50": _percentile(sorted_vals, 50), + "p90": _percentile(sorted_vals, 90), + "p95": _percentile(sorted_vals, 95), + } diff --git a/test/benchmark/report.py b/test/benchmark/report.py new file mode 100644 index 000000000..64008deb2 --- /dev/null +++ b/test/benchmark/report.py @@ -0,0 +1,105 @@ +from typing import Dict, List, Optional + + +def _fmt_seconds(value: Optional[float]) -> str: + if value is None: + return "n/a" + return f"{value:.4f}s" + + +def _fmt_ms(value: Optional[float]) -> str: + if value is None: + return "n/a" + return f"{value * 1000.0:.2f}ms" + + +def _fmt_qps(qps: Optional[float]) -> str: + if qps is None or qps <= 0: + return "n/a" + return f"{qps:.2f}" + + +def _calc_qps(total_duration_s: Optional[float], total_requests: int) -> Optional[float]: + if total_duration_s is None or total_duration_s <= 0: + return None + return total_requests / total_duration_s + + +def render_report(lines: List[str]) -> str: + return "\n".join(lines).strip() + "\n" + + +def chat_report( + *, + interface: str, + concurrency: int, + total_duration_s: Optional[float], + iterations: int, + success: int, + failure: int, + model: str, + total_stats: Dict[str, Optional[float]], + first_token_stats: Dict[str, Optional[float]], + errors: List[str], + created: Dict[str, str], +) -> str: + lines = [ + f"Interface: {interface}", + f"Concurrency: {concurrency}", + f"Iterations: {iterations}", + f"Success: {success}", + f"Failure: {failure}", + f"Model: {model}", + ] + for key, value in created.items(): + lines.append(f"{key}: {value}") + lines.extend( + [ + "Latency (total): " + f"avg={_fmt_ms(total_stats['avg'])}, min={_fmt_ms(total_stats['min'])}, " + f"p50={_fmt_ms(total_stats['p50'])}, p90={_fmt_ms(total_stats['p90'])}, p95={_fmt_ms(total_stats['p95'])}", + "Latency (first token): " + f"avg={_fmt_ms(first_token_stats['avg'])}, min={_fmt_ms(first_token_stats['min'])}, " + f"p50={_fmt_ms(first_token_stats['p50'])}, p90={_fmt_ms(first_token_stats['p90'])}, p95={_fmt_ms(first_token_stats['p95'])}", + f"Total Duration: {_fmt_seconds(total_duration_s)}", + f"QPS (requests / total duration): {_fmt_qps(_calc_qps(total_duration_s, iterations))}", + ] + ) + if errors: + lines.append("Errors: " + "; ".join(errors[:5])) + return render_report(lines) + + +def retrieval_report( + *, + interface: str, + concurrency: int, + total_duration_s: Optional[float], + iterations: int, + success: int, + failure: int, + stats: Dict[str, Optional[float]], + errors: List[str], + created: Dict[str, str], +) -> str: + lines = [ + f"Interface: {interface}", + f"Concurrency: {concurrency}", + f"Iterations: {iterations}", + f"Success: {success}", + f"Failure: {failure}", + ] + for key, value in created.items(): + lines.append(f"{key}: {value}") + lines.extend( + [ + "Latency: " + f"avg={_fmt_ms(stats['avg'])}, min={_fmt_ms(stats['min'])}, " + f"p50={_fmt_ms(stats['p50'])}, p90={_fmt_ms(stats['p90'])}, p95={_fmt_ms(stats['p95'])}", + f"Total Duration: {_fmt_seconds(total_duration_s)}", + f"QPS (requests / total duration): {_fmt_qps(_calc_qps(total_duration_s, iterations))}", + ] + ) + if errors: + lines.append("Errors: " + "; ".join(errors[:5])) + return render_report(lines) diff --git a/test/benchmark/retrieval.py b/test/benchmark/retrieval.py new file mode 100644 index 000000000..c2a488004 --- /dev/null +++ b/test/benchmark/retrieval.py @@ -0,0 +1,39 @@ +import time +from typing import Any, Dict, List, Optional + +from .http_client import HttpClient +from .metrics import RetrievalSample + + +class RetrievalError(RuntimeError): + pass + + +def build_payload( + question: str, + dataset_ids: List[str], + document_ids: Optional[List[str]] = None, + payload: Optional[Dict[str, Any]] = None, +) -> Dict[str, Any]: + body = dict(payload or {}) + if "question" not in body: + body["question"] = question + if "dataset_ids" not in body: + body["dataset_ids"] = dataset_ids + if document_ids is not None and "document_ids" not in body: + body["document_ids"] = document_ids + return body + + +def run_retrieval(client: HttpClient, payload: Dict[str, Any]) -> RetrievalSample: + t0 = time.perf_counter() + response = client.request("POST", "/retrieval", json_body=payload, stream=False) + raw = response.content + t1 = time.perf_counter() + try: + res = client.parse_json_bytes(raw) + except Exception as exc: + return RetrievalSample(t0=t0, t1=t1, error=f"Invalid JSON response: {exc}") + if res.get("code") != 0: + return RetrievalSample(t0=t0, t1=t1, error=res.get("message"), response=res) + return RetrievalSample(t0=t0, t1=t1, error=None, response=res) diff --git a/test/benchmark/run_chat.sh b/test/benchmark/run_chat.sh new file mode 100755 index 000000000..54c232748 --- /dev/null +++ b/test/benchmark/run_chat.sh @@ -0,0 +1,28 @@ +#!/usr/bin/env bash +set -euo pipefail + +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +REPO_ROOT="$(cd "${SCRIPT_DIR}/../.." && pwd)" + +: "${ZHIPU_AI_API_KEY:?ZHIPU_AI_API_KEY is required}" + +PYTHONPATH="${REPO_ROOT}/test" uv run -m benchmark chat \ + --base-url http://127.0.0.1:9380 \ + --allow-register \ + --login-email "qa@infiniflow.org" \ + --login-password "123" \ + --bootstrap-llm \ + --llm-factory ZHIPU-AI \ + --llm-api-key "$ZHIPU_AI_API_KEY" \ + --dataset-name "bench_dataset" \ + --dataset-payload '{"name":"bench_dataset","embedding_model":"BAAI/bge-small-en-v1.5@Builtin"}' \ + --document-path "${SCRIPT_DIR}/test_docs/Doc1.pdf" \ + --document-path "${SCRIPT_DIR}/test_docs/Doc2.pdf" \ + --document-path "${SCRIPT_DIR}/test_docs/Doc3.pdf" \ + --chat-name "bench_chat" \ + --chat-payload '{"name":"bench_chat","llm":{"model_name":"glm-4-flash@ZHIPU-AI"}}' \ + --message "What is the purpose of RAGFlow?" \ + --model "glm-4-flash@ZHIPU-AI" \ + --iterations 10 \ + --concurrency 8 \ + --teardown diff --git a/test/benchmark/run_retrieval.sh b/test/benchmark/run_retrieval.sh new file mode 100755 index 000000000..238cd039c --- /dev/null +++ b/test/benchmark/run_retrieval.sh @@ -0,0 +1,25 @@ +#!/usr/bin/env bash +set -euo pipefail + +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +REPO_ROOT="$(cd "${SCRIPT_DIR}/../.." && pwd)" + +: "${ZHIPU_AI_API_KEY:?ZHIPU_AI_API_KEY is required}" + +PYTHONPATH="${REPO_ROOT}/test" uv run -m benchmark retrieval \ + --base-url http://127.0.0.1:9380 \ + --allow-register \ + --login-email "qa@infiniflow.org" \ + --login-password "123" \ + --bootstrap-llm \ + --llm-factory ZHIPU-AI \ + --llm-api-key "$ZHIPU_AI_API_KEY" \ + --dataset-name "bench_dataset" \ + --dataset-payload '{"name":"bench_dataset","embedding_model":"BAAI/bge-small-en-v1.5@Builtin"}' \ + --document-path "${SCRIPT_DIR}/test_docs/Doc1.pdf" \ + --document-path "${SCRIPT_DIR}/test_docs/Doc2.pdf" \ + --document-path "${SCRIPT_DIR}/test_docs/Doc3.pdf" \ + --question "What does RAG mean?" \ + --iterations 10 \ + --concurrency 8 \ + --teardown diff --git a/test/benchmark/run_retrieval_chat.sh b/test/benchmark/run_retrieval_chat.sh new file mode 100755 index 000000000..9cd531803 --- /dev/null +++ b/test/benchmark/run_retrieval_chat.sh @@ -0,0 +1,98 @@ +#!/usr/bin/env bash +set -euo pipefail + +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +REPO_ROOT="$(cd "${SCRIPT_DIR}/../.." && pwd)" + +: "${ZHIPU_AI_API_KEY:?ZHIPU_AI_API_KEY is required}" + +BASE_URL="http://127.0.0.1:9380" +LOGIN_EMAIL="qa@infiniflow.org" +LOGIN_PASSWORD="123" +DATASET_PAYLOAD='{"name":"bench_dataset","embedding_model":"BAAI/bge-small-en-v1.5@Builtin"}' +CHAT_PAYLOAD='{"name":"bench_chat","llm":{"model_name":"glm-4-flash@ZHIPU-AI"}}' +DATASET_ID="" + +cleanup_dataset() { + if [[ -z "${DATASET_ID}" ]]; then + return + fi + set +e + BENCH_BASE_URL="${BASE_URL}" \ + BENCH_LOGIN_EMAIL="${LOGIN_EMAIL}" \ + BENCH_LOGIN_PASSWORD="${LOGIN_PASSWORD}" \ + BENCH_DATASET_ID="${DATASET_ID}" \ + PYTHONPATH="${REPO_ROOT}/test" uv run python - <<'PY' +import os +import sys + +from benchmark import auth +from benchmark.auth import AuthError +from benchmark.dataset import delete_dataset +from benchmark.http_client import HttpClient + +base_url = os.environ["BENCH_BASE_URL"] +email = os.environ["BENCH_LOGIN_EMAIL"] +password = os.environ["BENCH_LOGIN_PASSWORD"] +dataset_id = os.environ["BENCH_DATASET_ID"] + +client = HttpClient(base_url=base_url, api_version="v1") + +try: + password_enc = auth.encrypt_password(password) + nickname = email.split("@")[0] + try: + auth.register_user(client, email, nickname, password_enc) + except AuthError as exc: + print(f"Register warning: {exc}", file=sys.stderr) + login_token = auth.login_user(client, email, password_enc) + client.login_token = login_token + client.api_key = auth.create_api_token(client, login_token, None) + delete_dataset(client, dataset_id) +except Exception as exc: + print(f"Cleanup warning: failed to delete dataset {dataset_id}: {exc}", file=sys.stderr) +PY +} + +trap cleanup_dataset EXIT + +retrieval_output="$(PYTHONPATH="${REPO_ROOT}/test" uv run -m benchmark retrieval \ + --base-url "${BASE_URL}" \ + --allow-register \ + --login-email "${LOGIN_EMAIL}" \ + --login-password "${LOGIN_PASSWORD}" \ + --bootstrap-llm \ + --llm-factory ZHIPU-AI \ + --llm-api-key "${ZHIPU_AI_API_KEY}" \ + --dataset-name "bench_dataset" \ + --dataset-payload "${DATASET_PAYLOAD}" \ + --document-path "${SCRIPT_DIR}/test_docs/Doc1.pdf" \ + --document-path "${SCRIPT_DIR}/test_docs/Doc2.pdf" \ + --document-path "${SCRIPT_DIR}/test_docs/Doc3.pdf" \ + --iterations 10 \ + --concurrency 8 \ + --question "What does RAG mean?")" +printf '%s\n' "${retrieval_output}" + +DATASET_ID="$(printf '%s\n' "${retrieval_output}" | sed -n 's/^Created Dataset ID: //p' | head -n 1)" +if [[ -z "${DATASET_ID}" ]]; then + echo "Failed to parse Created Dataset ID from retrieval output." >&2 + exit 1 +fi + +PYTHONPATH="${REPO_ROOT}/test" uv run -m benchmark chat \ + --base-url "${BASE_URL}" \ + --allow-register \ + --login-email "${LOGIN_EMAIL}" \ + --login-password "${LOGIN_PASSWORD}" \ + --bootstrap-llm \ + --llm-factory ZHIPU-AI \ + --llm-api-key "${ZHIPU_AI_API_KEY}" \ + --dataset-id "${DATASET_ID}" \ + --chat-name "bench_chat" \ + --chat-payload "${CHAT_PAYLOAD}" \ + --message "What is the purpose of RAGFlow?" \ + --model "glm-4-flash@ZHIPU-AI" \ + --iterations 10 \ + --concurrency 8 \ + --teardown diff --git a/test/benchmark/test_docs/Doc1.pdf b/test/benchmark/test_docs/Doc1.pdf new file mode 100644 index 000000000..2bc4232fb --- /dev/null +++ b/test/benchmark/test_docs/Doc1.pdf @@ -0,0 +1,74 @@ +%PDF-1.4 +%東京 ReportLab Generated PDF document http://www.reportlab.com +1 0 obj +<< +/F1 2 0 R /F2 3 0 R +>> +endobj +2 0 obj +<< +/BaseFont /Helvetica /Encoding /WinAnsiEncoding /Name /F1 /Subtype /Type1 /Type /Font +>> +endobj +3 0 obj +<< +/BaseFont /Helvetica-Bold /Encoding /WinAnsiEncoding /Name /F2 /Subtype /Type1 /Type /Font +>> +endobj +4 0 obj +<< +/Contents 8 0 R /MediaBox [ 0 0 612 792 ] /Parent 7 0 R /Resources << +/Font 1 0 R /ProcSet [ /PDF /Text /ImageB /ImageC /ImageI ] +>> /Rotate 0 /Trans << + +>> + /Type /Page +>> +endobj +5 0 obj +<< +/PageMode /UseNone /Pages 7 0 R /Type /Catalog +>> +endobj +6 0 obj +<< +/Author (\(anonymous\)) /CreationDate (D:20260113015512+00'00') /Creator (\(unspecified\)) /Keywords () /ModDate (D:20260113015512+00'00') /Producer (ReportLab PDF Library - www.reportlab.com) + /Subject (\(unspecified\)) /Title (PDF 1: Purpose of RAGFlow) /Trapped /False +>> +endobj +7 0 obj +<< +/Count 1 /Kids [ 4 0 R ] /Type /Pages +>> +endobj +8 0 obj +<< +/Filter [ /ASCII85Decode /FlateDecode ] /Length 1210 +>> +stream +Gat=)bDt=8']%):guY@$.&J7SP+@MSJcWHM"C>\-;qDe#M:@X:f^@AcOL,7tgE]$l^^l=sFj@PSk\-t5-\NDK=6NK`mCpXLs8[6VM;e@)-dK6_WVCgq)6APJsa7e)>X/[J!/Q,.G$Hbfo!A$T$\Zf_qjDL/RgU!iOKT#=,>.5'.j(;6rC1Z"V8fO2KJ/m2-l@*\1[0,ZuguqLg'5^jn&aISTnjut'`[K3S7]LkfmKZd=Wtk/%Yo<`YA[apV%Q^tmFN>bXU6SkErJjqDg%<V4ZRn>)$hSg4ETUuO4*k*gugYsVuF$FboN9!>'*[%!>%>6F`BmC<5cq+\h_#oj[Wf=s4-[c1F7#(b;\rVM>&u3-(!n9t3_]a:-77)*[#_t8p&&CI!:!K]:(Ljhs:,M/8k%"0rde\5$!<(t&2VCeI?5[cLPfDWFS1pOg>&cS"@4Elk@)0hH:!l&Th97pLMX%j5=>4o&WQrps4[8`J?sqZ#`?5WV,2rGTC:3M*Ba^+pkE<."s)?ggfu'2mQ+o2L@SqQPcDCt_+R@hL,PFj/endstream +endobj +xref +0 9 +0000000000 65535 f +0000000073 00000 n +0000000114 00000 n +0000000221 00000 n +0000000333 00000 n +0000000526 00000 n +0000000594 00000 n +0000000889 00000 n +0000000948 00000 n +trailer +<< +/ID +[<2a9ef7a3c6d757889a2da50ec37e89c2><2a9ef7a3c6d757889a2da50ec37e89c2>] +% ReportLab generated PDF document -- digest (http://www.reportlab.com) + +/Info 6 0 R +/Root 5 0 R +/Size 9 +>> +startxref +2249 +%%EOF diff --git a/test/benchmark/test_docs/Doc2.pdf b/test/benchmark/test_docs/Doc2.pdf new file mode 100644 index 000000000..1bf0c3e4b --- /dev/null +++ b/test/benchmark/test_docs/Doc2.pdf @@ -0,0 +1,74 @@ +%PDF-1.4 +%東京 ReportLab Generated PDF document http://www.reportlab.com +1 0 obj +<< +/F1 2 0 R /F2 3 0 R +>> +endobj +2 0 obj +<< +/BaseFont /Helvetica /Encoding /WinAnsiEncoding /Name /F1 /Subtype /Type1 /Type /Font +>> +endobj +3 0 obj +<< +/BaseFont /Helvetica-Bold /Encoding /WinAnsiEncoding /Name /F2 /Subtype /Type1 /Type /Font +>> +endobj +4 0 obj +<< +/Contents 8 0 R /MediaBox [ 0 0 612 792 ] /Parent 7 0 R /Resources << +/Font 1 0 R /ProcSet [ /PDF /Text /ImageB /ImageC /ImageI ] +>> /Rotate 0 /Trans << + +>> + /Type /Page +>> +endobj +5 0 obj +<< +/PageMode /UseNone /Pages 7 0 R /Type /Catalog +>> +endobj +6 0 obj +<< +/Author (\(anonymous\)) /CreationDate (D:20260113015512+00'00') /Creator (\(unspecified\)) /Keywords () /ModDate (D:20260113015512+00'00') /Producer (ReportLab PDF Library - www.reportlab.com) + /Subject (\(unspecified\)) /Title (PDF 2: Definition of RAG) /Trapped /False +>> +endobj +7 0 obj +<< +/Count 1 /Kids [ 4 0 R ] /Type /Pages +>> +endobj +8 0 obj +<< +/Filter [ /ASCII85Decode /FlateDecode ] /Length 1266 +>> +stream +Gat=)gQ&$u&:L1SW5hKj-%I^pni"U%7P!RDg%>l_Ta[KCDPk$KYQ$ea)XPB`m6lg_=Z4W,FY7gX.)dQlht98,bDq$tasbGIGG)$.$JPGR^[:Gqn$;&sZ:]/;dQH4B!3W6i\ngbR1qQ^;BfD\_:0B$2nm(Hj7?MDd'"<&oh>d$8M=R;',9O;?6!H_M5bB-jJ"YV)[/5tL>Wn)brl!u3NcJF<5=_uH#6CW:L(jka[YH1=D%,5H2@];2oML_L*l^SeR/6"EoM0f$SrpIA9M0O.XLcI,Mh>kYsN/8YoFQA8,2KfJr9P[gVB2-=U+p0"!,'Z.Z0O9,d3!.k=AHgm7ZNj:Ij^hnkcpc')KH^V2>*XrJm+lnM`X'2Yj5oAtPT1RF&+cDE56%iZC\3`P(2_H?CT9)O:%(d"MlQCGVTQD0OS([dh%@5@OduphF7N!!g^_3fAo_Rd[;l.FR[`?`a8pI8TdAqTpBY[nB7WF"BJQ4%qdkf*X`6@i'mI=lD]tp#K']^b6*T/<7]3;(UNkl$_?O\.;P_8j'BlDd:JI=6m5`6UGPRL6o:3Uf+=IRKGoNV'b0E[)`B[%Gk!r[r[G2Ch_-L(jd@,kU^=1SHV-r`_Aq!%pA^ia7_5]/s_XHl[*u32OMB_>-?OiIJ9/r*1q6\(!JeD3KYk""VcAK:!YX]cI0"MIEL7*bbWZ4^K/!g]B+.XIPL9g"PTZekf(.mA+!>;!pLB$.*+NZg9+<3bcVh>V@lof[=:'UVn[@AIOgTrd)fnmJS$r&kDL*5`W7Spd!\EYfDERI:mE(eGJQ4jHk"a]>"B)M2l`~>endstream +endobj +xref +0 9 +0000000000 65535 f +0000000073 00000 n +0000000114 00000 n +0000000221 00000 n +0000000333 00000 n +0000000526 00000 n +0000000594 00000 n +0000000888 00000 n +0000000947 00000 n +trailer +<< +/ID +[] +% ReportLab generated PDF document -- digest (http://www.reportlab.com) + +/Info 6 0 R +/Root 5 0 R +/Size 9 +>> +startxref +2304 +%%EOF diff --git a/test/benchmark/test_docs/Doc3.pdf b/test/benchmark/test_docs/Doc3.pdf new file mode 100644 index 000000000..d36b581c2 --- /dev/null +++ b/test/benchmark/test_docs/Doc3.pdf @@ -0,0 +1,74 @@ +%PDF-1.4 +%東京 ReportLab Generated PDF document http://www.reportlab.com +1 0 obj +<< +/F1 2 0 R /F2 3 0 R +>> +endobj +2 0 obj +<< +/BaseFont /Helvetica /Encoding /WinAnsiEncoding /Name /F1 /Subtype /Type1 /Type /Font +>> +endobj +3 0 obj +<< +/BaseFont /Helvetica-Bold /Encoding /WinAnsiEncoding /Name /F2 /Subtype /Type1 /Type /Font +>> +endobj +4 0 obj +<< +/Contents 8 0 R /MediaBox [ 0 0 612 792 ] /Parent 7 0 R /Resources << +/Font 1 0 R /ProcSet [ /PDF /Text /ImageB /ImageC /ImageI ] +>> /Rotate 0 /Trans << + +>> + /Type /Page +>> +endobj +5 0 obj +<< +/PageMode /UseNone /Pages 7 0 R /Type /Catalog +>> +endobj +6 0 obj +<< +/Author (\(anonymous\)) /CreationDate (D:20260113015512+00'00') /Creator (\(unspecified\)) /Keywords () /ModDate (D:20260113015512+00'00') /Producer (ReportLab PDF Library - www.reportlab.com) + /Subject (\(unspecified\)) /Title (PDF 3: InfiniFlow Company History) /Trapped /False +>> +endobj +7 0 obj +<< +/Count 1 /Kids [ 4 0 R ] /Type /Pages +>> +endobj +8 0 obj +<< +/Filter [ /ASCII85Decode /FlateDecode ] /Length 1281 +>> +stream +Gat=)9lo&I&A@C2m*Y4I/.%I`h4'nbM.\luUsZF`Lg3i66%B/)fugME!i-N-NoQOf`ZIPK]5QSRqj9]:T/?JKH[n==nCq4U.>(>&BccHAq&B=eX<(I-Hu`N4%*F7B'j[Kq8U,-a'jXO0e^gthW=ji*m40pRl_I)akBTL'D]6.^pu`Z.^Vda"KbAT-5J]G;%pU.5K*)MqrLhO5=ROoY>P<@0S$8t\!FbW:dE'&`!#.2@?:ec21#?5+(BPN4:kV$%"43ugp_!0@`\*na)(dRc!:p&^-RB^q1#qbnG5m'40c/Zm)6\>i=K$QN]\fDS!]HaXCJl[r4$/[8mXjQAYg`K,:Vlm?i3&;La@ruYap!cfLrn>+;2_6$4#E4$$82r&D7r'7#,6pF2B%2a1Vi1P\T,kS,%Z)XKElg,Q8o=Efk:BM?,i93-_cIQFHR@j@/2Q=D.=#;'WsQH$+P_)0fQU)#QsSS@S]Nl:Ei',S;-Kqj>@nX9@Z7Z?c%g7FsDfL@9@9LGLgY+<;Haelk/JK,T7tV38E0K&:j#*5%,PUfH7g%D1!>`e@\Ft=_i@sk'eN"8"#OOSD^t77@u(]X[3@P[dG*&pD*#%TX[*@DWDaM\qef'9isoh,0&09#aHVl!O2g_P2)HqB7D>Z](+j5q(;UeQ@o+@t+.WT`f[(N7UqlIc1pfVrfeSIaQ6,4:0AM>&)j!t%opjf&5>uCn,(!9hZA=pc*R$pI$*o^ubpY^>Z/@rEA2keb];,X"5;c>UiJ`ogQ/6)=n\3Z:bN9j,7nj7k)PmH^#eXJUmlZ,-<5rl@pArAfnAp^8JYd3dKl@c*?L(nFljuBR.!\W[otSW6/!$d!ff10OZro2`dd/F?YO'XS)jqb2IM7Ao@r:,mYsr*VXU>S"#$*GO-=>a/ninKc4[Ck^)`dh16N%8\q6_NqNIL*SLgNkl"-a^!5oN:QEqOR,jo6>?Yq&Zlcl-YqfEsHa""W2iO(D?YH'4YkeR$P.a'~>endstream +endobj +xref +0 9 +0000000000 65535 f +0000000073 00000 n +0000000114 00000 n +0000000221 00000 n +0000000333 00000 n +0000000526 00000 n +0000000594 00000 n +0000000897 00000 n +0000000956 00000 n +trailer +<< +/ID +[<3b8236972136a5b3b2d0ba0cc5dfa592><3b8236972136a5b3b2d0ba0cc5dfa592>] +% ReportLab generated PDF document -- digest (http://www.reportlab.com) + +/Info 6 0 R +/Root 5 0 R +/Size 9 +>> +startxref +2328 +%%EOF diff --git a/test/benchmark/utils.py b/test/benchmark/utils.py new file mode 100644 index 000000000..d46641344 --- /dev/null +++ b/test/benchmark/utils.py @@ -0,0 +1,41 @@ +import json +import sys +import time +from pathlib import Path + + +def eprint(*args, **kwargs): + print(*args, file=sys.stderr, **kwargs) + + +def load_json_arg(value, name): + if value is None: + return None + if isinstance(value, dict): + return value + if isinstance(value, str) and value.startswith("@"): + path = Path(value[1:]) + try: + return json.loads(path.read_text(encoding="utf-8")) + except Exception as exc: + raise ValueError(f"Failed to read {name} from {path}: {exc}") from exc + try: + return json.loads(value) + except Exception as exc: + raise ValueError(f"Invalid JSON for {name}: {exc}") from exc + + +def split_csv(value): + if value is None: + return None + if isinstance(value, list): + return value + if isinstance(value, str): + items = [item.strip() for item in value.split(",")] + return [item for item in items if item] + return [value] + + +def unique_name(prefix): + return f"{prefix}_{int(time.time() * 1000)}" + diff --git a/uv.lock b/uv.lock index 1e88de682..734489174 100644 --- a/uv.lock +++ b/uv.lock @@ -6174,6 +6174,7 @@ test = [ { name = "hypothesis" }, { name = "openpyxl" }, { name = "pillow" }, + { name = "pycryptodomex" }, { name = "pytest" }, { name = "pytest-asyncio" }, { name = "pytest-cov" }, @@ -6307,6 +6308,7 @@ test = [ { name = "hypothesis", specifier = ">=6.132.0" }, { name = "openpyxl", specifier = ">=3.1.5" }, { name = "pillow", specifier = ">=10.4.0,<13.0.0" }, + { name = "pycryptodomex", specifier = "==3.20.0" }, { name = "pytest", specifier = ">=8.3.5" }, { name = "pytest-asyncio", specifier = ">=1.3.0" }, { name = "pytest-cov", specifier = ">=7.0.0" },