mirror of
https://github.com/infiniflow/ragflow.git
synced 2026-01-19 03:35:11 +08:00
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:
@ -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
285
test/benchmark/README.md
Normal 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
|
||||
1
test/benchmark/__init__.py
Normal file
1
test/benchmark/__init__.py
Normal file
@ -0,0 +1 @@
|
||||
"""RAGFlow HTTP API benchmark package."""
|
||||
5
test/benchmark/__main__.py
Normal file
5
test/benchmark/__main__.py
Normal file
@ -0,0 +1,5 @@
|
||||
from .cli import main
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
88
test/benchmark/auth.py
Normal file
88
test/benchmark/auth.py
Normal 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
138
test/benchmark/chat.py
Normal 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
575
test/benchmark/cli.py
Normal 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
146
test/benchmark/dataset.py
Normal 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
|
||||
112
test/benchmark/http_client.py
Normal file
112
test/benchmark/http_client.py
Normal 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
67
test/benchmark/metrics.py
Normal 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
105
test/benchmark/report.py
Normal 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)
|
||||
39
test/benchmark/retrieval.py
Normal file
39
test/benchmark/retrieval.py
Normal 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
28
test/benchmark/run_chat.sh
Executable 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
25
test/benchmark/run_retrieval.sh
Executable 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
|
||||
98
test/benchmark/run_retrieval_chat.sh
Executable file
98
test/benchmark/run_retrieval_chat.sh
Executable 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
|
||||
74
test/benchmark/test_docs/Doc1.pdf
Normal file
74
test/benchmark/test_docs/Doc1.pdf
Normal 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
|
||||
74
test/benchmark/test_docs/Doc2.pdf
Normal file
74
test/benchmark/test_docs/Doc2.pdf
Normal 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
|
||||
74
test/benchmark/test_docs/Doc3.pdf
Normal file
74
test/benchmark/test_docs/Doc3.pdf
Normal 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
41
test/benchmark/utils.py
Normal 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
2
uv.lock
generated
@ -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" },
|
||||
|
||||
Reference in New Issue
Block a user