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 <asiro@qq.com>
This commit is contained in:
6ba3i
2026-01-14 13:49:16 +08:00
committed by GitHub
parent a7671583b3
commit 5b22f94502
20 changed files with 1978 additions and 0 deletions

View File

@ -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]]

285
test/benchmark/README.md Normal file
View File

@ -0,0 +1,285 @@
# RAGFlow HTTP Benchmark CLI
Run (from repo root):
```
PYTHONPATH=./test uv run -m benchmark [global flags] <chat|retrieval> [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] <chat|retrieval> [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 <token>.
--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": "<model_name>@<provider>"}
```
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": "<model_name>@<provider>"}}
```
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 <existing_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 "<dataset_id_1>,<dataset_id_2>" \
--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 "<dataset_id>" \
--document-ids "<doc_id_1>,<doc_id_2>" \
--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

View File

@ -0,0 +1 @@
"""RAGFlow HTTP API benchmark package."""

View File

@ -0,0 +1,5 @@
from .cli import main
if __name__ == "__main__":
main()

88
test/benchmark/auth.py Normal file
View File

@ -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')}")

138
test/benchmark/chat.py Normal file
View File

@ -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)

575
test/benchmark/cli.py Normal file
View File

@ -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)

146
test/benchmark/dataset.py Normal file
View File

@ -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

View File

@ -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

67
test/benchmark/metrics.py Normal file
View File

@ -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),
}

105
test/benchmark/report.py Normal file
View File

@ -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)

View File

@ -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)

28
test/benchmark/run_chat.sh Executable file
View File

@ -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

25
test/benchmark/run_retrieval.sh Executable file
View File

@ -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

View File

@ -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

View File

@ -0,0 +1,74 @@
%PDF-1.4
%<25><><EFBFBD><EFBFBD> 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@P<Ju3dLoWEM7hq.N";332?]d6Ml26HnQ5M(*4GN+!k^*lTW9Zpn%!1Zj6\_Gp`<4le3e%/&HP/KIp.a+2L<2(eA&s(WRh*:r08E(jWO)D1na_p@dlI06U&+=312m%S)LTC;ci\On"NQeI"a,STTg^.R*H)B4c;2KQM(:69PrOFCl:XY=FfT<S;;bUOX36eY4HWZ9>Sk\-t5-\NDK=6NK`mCpXLs8[6VM;e@)-dK6_WVCgq)6AP<YC+bQa<\a<?&!YR*q0V#]SR5\`oKG7@@YYX``,b[q!(?h'hfi[gdPqTI0jOB@P9bU3jGkf0FtP.VL+i*+"m_W^Tpc\slQ@9cS9J7+-GLK/A-G.g%5j9&gHKk[e9HBCn$,D0C"&+NU]mARO6WNr=V@/NsL2Vg)_7lt@;2YK^)/qp=m;cQ:J/bHAA1p^X32R5t5lPM'dO95K)[c;H;k(m)UK8.ZfCgM]CU,Yj-'G8Q`cBT^fc+?bJdV)_:7E)2+.StfI-;B5kFpopCGOBjK=`H$FD^lW;$Zd2Je80'^EA\q.ggp;+&/SGmU(ndOKSm5:u3_>Jsa7e)>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%<<g=]GF!'i@rRlIHl^:dnn7F>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/<K?bf0Ljph4eE#La^hO$(?6Q_$'E-8-*Q[NABX2n!XlH3c^'qLj*Ek=%I/p*r[UB<5C5Z<.ie_#fppLUItP.mAtS,u@<i`-^lP+h6%/PJr-V"YrrO:%B2&~>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

View File

@ -0,0 +1,74 @@
%PDF-1.4
%<25><><EFBFBD><EFBFBD> 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[Y<Ksdf!eg:@hQ;Is@L;JPG'WT%U>H1=D%,5H2@];2oML_L*l^SeR/6"EoM0f$SrpIA9M0O.XLcI,Mh>kYsN/8YoFQA8,2Kf<MG5rlC++-hc!/dntm\3poL!3'33W`8@B)lKQ:?,C6gmHf[8]:Bhh=,BPs6:X[i5N!KWKBM.Qim2?s.<&<1C2V_6War+T$-7<GY;\6aO.t@]6"TTRMo)T:E4m]![BJP.G!TYEP1XCBQPIkRThq35f/j4cU7Kk)jq\;)!"%IQJk>Jr9P[gVB2-=U+p0"!,'Z.Z0<cW*K*A:sXfaD`dCc:DTF@fJtgF>O9,d3!.k=AHgm7<f=&:\^'\ji==P7Y/82F(+I?)b9:',XZFU"ZbZK!*far+uW166lZZ?BOt*:""/]j:ALI%sr34=ruiLcr[B&Q'U"kMS'rhkQ-(^G1^RB2>ZNj: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<K'68t%!GceD<cRb5`0iN=qFO-JER!gH@g8EFuN\K@(ks_ePG4=?(elrkM_Ea@J$r*9Ga8K"b0L$6>_>-?OiIJ9/r*1q6\(!JeD3KYk""<S#-qoXNVIg/oLb9ZKcQlXVM*)TWh6hh13JTI1J&,I>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
[<d84dab5297cf0d3049cc8f2a151a70f8><d84dab5297cf0d3049cc8f2a151a70f8>]
% ReportLab generated PDF document -- digest (http://www.reportlab.com)
/Info 6 0 R
/Root 5 0 R
/Size 9
>>
startxref
2304
%%EOF

View File

@ -0,0 +1,74 @@
%PDF-1.4
%<25><><EFBFBD><EFBFBD> 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<S+?__87i!h3QUH?#"]I%npjTQOT>*)MqrLhO5=ROoY>P<@0S$8t\!FbW:dE'&`!#.2@?:ec21#?5+(BPN4:kV$%"43ugp_!0@`\*na)(d<sa.;F@4'b]nRSA[@Eo<cTf1A#G1ejaKGS-j/5(>Rc!:p&^-RB^q1#qbnG5m'40c/Zm)6\>i=K$QN]\fDS!]HaXCJl[r4$/[<q?a58qnaY!WChC#d0>8mXjQAYg`K,:Vlm?i3&;La@ruYap!cfLrn>+;2_6$4#E4$$82r&D7r'7#,6pF2B%2a1Vi1<U"A(RJroXmPFL-))/9FNOV0(UhK-C,s%#_)L2SED[gS!/fYbe9aNq[(E2>P\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,T<UEOm\n:iTR,(IQJH25JZ'fSAJ"cLjB@b!Mis+a,eN0',j?E!dEU#m't>7tV38E0K&: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<l5VN,-]l-[S=f-p_oYg,>,(!9<iPIBX/$]>hZA=pc*R$pI$*o^ubpY^>Z/@rEA2keb];,X"<C<-S+]?M&Q=:&I#f2"[h3O)$6<qYA)Ja6G[gZ1@1s;-A:SVK>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

41
test/benchmark/utils.py Normal file
View File

@ -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)}"

2
uv.lock generated
View File

@ -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" },