zlei6/ragflow
1
0
Fork 0
mirror of https://github.com/infiniflow/ragflow.git synced 2026-08-08 16:38:01 +08:00
Code Issues Packages Projects Releases Wiki Activity

Refactor: reformat all code for lefthook using ruff and gofmt (#16585)

Browse Source
This commit is contained in:
Wang Qi
2026-07-03 12:53:39 +08:00
committed by GitHub
parent 19fcb4a981
commit 6a4b9be426
588 changed files with 11123 additions and 15412 deletions
Download Patch File Download Diff File Expand all files Collapse all files

320
admin/client/ragflow_client.py
View File

@@ -71,6 +71,7 @@ class RAGFlowClient:
user_password: str = command.get("password")
if not user_password:
import getpass
user_password = getpass.getpass("Password: ")
try:
token = login_user(self.http_client, self.server_type, email, user_password)
@@ -86,8 +87,7 @@ class RAGFlowClient:
def ping_server(self, command):
iterations = command.get("iterations", 1)
if iterations > 1:
response = self.http_client.request("GET", "/system/ping", use_api_base=True, auth_kind="web",
iterations=iterations)
response = self.http_client.request("GET", "/system/ping", use_api_base=True, auth_kind="web", iterations=iterations)
return response
else:
response = self.http_client.request("GET", "/system/ping", use_api_base=True, auth_kind="web")
@@ -106,8 +106,7 @@ class RAGFlowClient:
enc_password = encrypt_password(password)
print(f"Register user: {nickname}, email: {username}, password: ******")
payload = {"email": username, "nickname": nickname, "password": enc_password}
response = self.http_client.request(method="POST", path="/users",
json_body=payload, use_api_base=True, auth_kind="web")
response = self.http_client.request(method="POST", path="/users", json_body=payload, use_api_base=True, auth_kind="web")
res_json = response.json()
if response.status_code == 200:
if res_json["code"] == 0:
@@ -135,8 +134,7 @@ class RAGFlowClient:
service_id: int = command["number"]
response = self.http_client.request("GET", f"/admin/services/{service_id}", use_api_base=True,
auth_kind="admin")
response = self.http_client.request("GET", f"/admin/services/{service_id}", use_api_base=True, auth_kind="admin")
res_json = response.json()
if response.status_code == 200:
res_data = res_json["data"]
@@ -226,9 +224,7 @@ class RAGFlowClient:
password_tree: Tree = command["password"]
password: str = password_tree.children[0].strip("'\"")
print(f"Alter user: {user_name}, password: ******")
response = self.http_client.request("PUT", f"/admin/users/{user_name}/password",
json_body={"new_password": encrypt_password(password)}, use_api_base=True,
auth_kind="admin")
response = self.http_client.request("PUT", f"/admin/users/{user_name}/password", json_body={"new_password": encrypt_password(password)}, use_api_base=True, auth_kind="admin")
res_json = response.json()
if response.status_code == 200:
print(res_json["message"])
@@ -247,9 +243,7 @@ class RAGFlowClient:
print(f"Create user: {user_name}, password: ******, role: {role}")
# enpass1 = encrypt(password)
enc_password = encrypt_password(password)
response = self.http_client.request(method="POST", path="/admin/users",
json_body={"username": user_name, "password": enc_password, "role": role},
use_api_base=True, auth_kind="admin")
response = self.http_client.request(method="POST", path="/admin/users", json_body={"username": user_name, "password": enc_password, "role": role}, use_api_base=True, auth_kind="admin")
res_json = response.json()
if response.status_code == 200:
self._print_table_simple(res_json["data"])
@@ -266,9 +260,7 @@ class RAGFlowClient:
activate_status: str = activate_tree.children[0].strip("'\"")
if activate_status.lower() in ["on", "off"]:
print(f"Alter user {user_name} activate status, turn {activate_status.lower()}.")
response = self.http_client.request("PUT", f"/admin/users/{user_name}/activate",
json_body={"activate_status": activate_status}, use_api_base=True,
auth_kind="admin")
response = self.http_client.request("PUT", f"/admin/users/{user_name}/activate", json_body={"activate_status": activate_status}, use_api_base=True, auth_kind="admin")
res_json = response.json()
if response.status_code == 200:
print(res_json["message"])
@@ -283,14 +275,12 @@ class RAGFlowClient:
user_name_tree: Tree = command["user_name"]
user_name: str = user_name_tree.children[0].strip("'\"")
response = self.http_client.request("PUT", f"/admin/users/{user_name}/admin", use_api_base=True,
auth_kind="admin")
response = self.http_client.request("PUT", f"/admin/users/{user_name}/admin", use_api_base=True, auth_kind="admin")
res_json = response.json()
if response.status_code == 200:
print(res_json["message"])
else:
print(
f"Fail to grant {user_name} admin authorization, code: {res_json['code']}, message: {res_json['message']}")
print(f"Fail to grant {user_name} admin authorization, code: {res_json['code']}, message: {res_json['message']}")
def revoke_admin(self, command):
if self.server_type != "admin":
@@ -298,14 +288,12 @@ class RAGFlowClient:
user_name_tree: Tree = command["user_name"]
user_name: str = user_name_tree.children[0].strip("'\"")
response = self.http_client.request("DELETE", f"/admin/users/{user_name}/admin", use_api_base=True,
auth_kind="admin")
response = self.http_client.request("DELETE", f"/admin/users/{user_name}/admin", use_api_base=True, auth_kind="admin")
res_json = response.json()
if response.status_code == 200:
print(res_json["message"])
else:
print(
f"Fail to revoke {user_name} admin authorization, code: {res_json['code']}, message: {res_json['message']}")
print(f"Fail to revoke {user_name} admin authorization, code: {res_json['code']}, message: {res_json['message']}")
def create_role(self, command):
if self.server_type != "admin":
@@ -319,10 +307,7 @@ class RAGFlowClient:
desc_str = desc_tree.children[0].strip("'\"")
print(f"create role name: {role_name}, description: {desc_str}")
response = self.http_client.request("POST", "/admin/roles",
json_body={"role_name": role_name, "description": desc_str},
use_api_base=True,
auth_kind="admin")
response = self.http_client.request("POST", "/admin/roles", json_body={"role_name": role_name, "description": desc_str}, use_api_base=True, auth_kind="admin")
res_json = response.json()
if response.status_code == 200:
self._print_table_simple(res_json["data"])
@@ -336,9 +321,7 @@ class RAGFlowClient:
role_name_tree: Tree = command["role_name"]
role_name: str = role_name_tree.children[0].strip("'\"")
print(f"drop role name: {role_name}")
response = self.http_client.request("DELETE", f"/admin/roles/{role_name}",
use_api_base=True,
auth_kind="admin")
response = self.http_client.request("DELETE", f"/admin/roles/{role_name}", use_api_base=True, auth_kind="admin")
res_json = response.json()
if response.status_code == 200:
self._print_table_simple(res_json["data"])
@@ -355,24 +338,18 @@ class RAGFlowClient:
desc_str: str = desc_tree.children[0].strip("'\"")
print(f"alter role name: {role_name}, description: {desc_str}")
response = self.http_client.request("PUT", f"/admin/roles/{role_name}",
json_body={"description": desc_str},
use_api_base=True,
auth_kind="admin")
response = self.http_client.request("PUT", f"/admin/roles/{role_name}", json_body={"description": desc_str}, use_api_base=True, auth_kind="admin")
res_json = response.json()
if response.status_code == 200:
self._print_table_simple(res_json["data"])
else:
print(
f"Fail to update role {role_name} with description: {desc_str}, code: {res_json['code']}, message: {res_json['message']}")
print(f"Fail to update role {role_name} with description: {desc_str}, code: {res_json['code']}, message: {res_json['message']}")
def list_roles(self, command):
if self.server_type != "admin":
print("This command is only allowed in ADMIN mode")
response = self.http_client.request("GET", "/admin/roles",
use_api_base=True,
auth_kind="admin")
response = self.http_client.request("GET", "/admin/roles", use_api_base=True, auth_kind="admin")
res_json = response.json()
if response.status_code == 200:
self._print_table_simple(res_json["data"])
@@ -386,9 +363,7 @@ class RAGFlowClient:
role_name_tree: Tree = command["role_name"]
role_name: str = role_name_tree.children[0].strip("'\"")
print(f"show role: {role_name}")
response = self.http_client.request("GET", f"/admin/roles/{role_name}/permission",
use_api_base=True,
auth_kind="admin")
response = self.http_client.request("GET", f"/admin/roles/{role_name}/permission", use_api_base=True, auth_kind="admin")
res_json = response.json()
if response.status_code == 200:
self._print_table_simple(res_json["data"])
@@ -409,15 +384,12 @@ class RAGFlowClient:
action_str: str = action_tree.children[0].strip("'\"")
actions.append(action_str)
print(f"grant role_name: {role_name_str}, resource: {resource_str}, actions: {actions}")
response = self.http_client.request("POST", f"/admin/roles/{role_name_str}/permission",
json_body={"actions": actions, "resource": resource_str}, use_api_base=True,
auth_kind="admin")
response = self.http_client.request("POST", f"/admin/roles/{role_name_str}/permission", json_body={"actions": actions, "resource": resource_str}, use_api_base=True, auth_kind="admin")
res_json = response.json()
if response.status_code == 200:
self._print_table_simple(res_json["data"])
else:
print(
f"Fail to grant role {role_name_str} with {actions} on {resource_str}, code: {res_json['code']}, message: {res_json['message']}")
print(f"Fail to grant role {role_name_str} with {actions} on {resource_str}, code: {res_json['code']}, message: {res_json['message']}")
def revoke_permission(self, command):
if self.server_type != "admin":
@@ -433,15 +405,12 @@ class RAGFlowClient:
action_str: str = action_tree.children[0].strip("'\"")
actions.append(action_str)
print(f"revoke role_name: {role_name_str}, resource: {resource_str}, actions: {actions}")
response = self.http_client.request("DELETE", f"/admin/roles/{role_name_str}/permission",
json_body={"actions": actions, "resource": resource_str}, use_api_base=True,
auth_kind="admin")
response = self.http_client.request("DELETE", f"/admin/roles/{role_name_str}/permission", json_body={"actions": actions, "resource": resource_str}, use_api_base=True, auth_kind="admin")
res_json = response.json()
if response.status_code == 200:
self._print_table_simple(res_json["data"])
else:
print(
f"Fail to revoke role {role_name_str} with {actions} on {resource_str}, code: {res_json['code']}, message: {res_json['message']}")
print(f"Fail to revoke role {role_name_str} with {actions} on {resource_str}, code: {res_json['code']}, message: {res_json['message']}")
def alter_user_role(self, command):
if self.server_type != "admin":
@@ -452,15 +421,12 @@ class RAGFlowClient:
user_name_tree: Tree = command["user_name"]
user_name_str: str = user_name_tree.children[0].strip("'\"")
print(f"alter_user_role user_name: {user_name_str}, role_name: {role_name_str}")
response = self.http_client.request("PUT", f"/admin/users/{user_name_str}/role",
json_body={"role_name": role_name_str}, use_api_base=True,
auth_kind="admin")
response = self.http_client.request("PUT", f"/admin/users/{user_name_str}/role", json_body={"role_name": role_name_str}, use_api_base=True, auth_kind="admin")
res_json = response.json()
if response.status_code == 200:
self._print_table_simple(res_json["data"])
else:
print(
f"Fail to alter user: {user_name_str} to role {role_name_str}, code: {res_json['code']}, message: {res_json['message']}")
print(f"Fail to alter user: {user_name_str} to role {role_name_str}, code: {res_json['code']}, message: {res_json['message']}")
def show_user_permission(self, command):
if self.server_type != "admin":
@@ -469,14 +435,12 @@ class RAGFlowClient:
user_name_tree: Tree = command["user_name"]
user_name_str: str = user_name_tree.children[0].strip("'\"")
print(f"show_user_permission user_name: {user_name_str}")
response = self.http_client.request("GET", f"/admin/users/{user_name_str}/permission", use_api_base=True,
auth_kind="admin")
response = self.http_client.request("GET", f"/admin/users/{user_name_str}/permission", use_api_base=True, auth_kind="admin")
res_json = response.json()
if response.status_code == 200:
self._print_table_simple(res_json["data"])
else:
print(
f"Fail to show user: {user_name_str} permission, code: {res_json['code']}, message: {res_json['message']}")
print(f"Fail to show user: {user_name_str} permission, code: {res_json['code']}, message: {res_json['message']}")
def generate_key(self, command: dict[str, Any]) -> None:
if self.server_type != "admin":
@@ -485,14 +449,12 @@ class RAGFlowClient:
username_tree: Tree = command["user_name"]
user_name: str = username_tree.children[0].strip("'\"")
print(f"Generating API key for user: {user_name}")
response = self.http_client.request("POST", f"/admin/users/{user_name}/keys", use_api_base=True,
auth_kind="admin")
response = self.http_client.request("POST", f"/admin/users/{user_name}/keys", use_api_base=True, auth_kind="admin")
res_json: dict[str, Any] = response.json()
if response.status_code == 200:
self._print_table_simple(res_json["data"])
else:
print(
f"Failed to generate key for user {user_name}, code: {res_json['code']}, message: {res_json['message']}")
print(f"Failed to generate key for user {user_name}, code: {res_json['code']}, message: {res_json['message']}")
def list_keys(self, command: dict[str, Any]) -> None:
if self.server_type != "admin":
@@ -501,8 +463,7 @@ class RAGFlowClient:
username_tree: Tree = command["user_name"]
user_name: str = username_tree.children[0].strip("'\"")
print(f"Listing API keys for user: {user_name}")
response = self.http_client.request("GET", f"/admin/users/{user_name}/keys", use_api_base=True,
auth_kind="admin")
response = self.http_client.request("GET", f"/admin/users/{user_name}/keys", use_api_base=True, auth_kind="admin")
res_json: dict[str, Any] = response.json()
if response.status_code == 200:
self._print_table_simple(res_json["data"])
@@ -520,8 +481,7 @@ class RAGFlowClient:
print(f"Dropping API key for user: {user_name}")
# URL encode the key to handle special characters
encoded_key: str = urllib.parse.quote(key, safe="")
response = self.http_client.request("DELETE", f"/admin/users/{user_name}/keys/{encoded_key}", use_api_base=True,
auth_kind="admin")
response = self.http_client.request("DELETE", f"/admin/users/{user_name}/keys/{encoded_key}", use_api_base=True, auth_kind="admin")
res_json: dict[str, Any] = response.json()
if response.status_code == 200:
print(res_json["message"])
@@ -534,23 +494,19 @@ class RAGFlowClient:
var_name = _strip_tree_value(command["var_name"])
var_value = _strip_tree_value(command["var_value"])
response = self.http_client.request("PUT", "/admin/variables",
json_body={"var_name": var_name, "var_value": var_value}, use_api_base=True,
auth_kind="admin")
response = self.http_client.request("PUT", "/admin/variables", json_body={"var_name": var_name, "var_value": var_value}, use_api_base=True, auth_kind="admin")
res_json = response.json()
if response.status_code == 200:
print(res_json["message"])
else:
print(
f"Fail to set variable {var_name} to {var_value}, code: {res_json['code']}, message: {res_json['message']}")
print(f"Fail to set variable {var_name} to {var_value}, code: {res_json['code']}, message: {res_json['message']}")
def show_variable(self, command):
if self.server_type != "admin":
print("This command is only allowed in ADMIN mode")
var_name = _strip_tree_value(command["var_name"])
response = self.http_client.request(method="GET", path="/admin/variables", json_body={"var_name": var_name},
use_api_base=True, auth_kind="admin")
response = self.http_client.request(method="GET", path="/admin/variables", json_body={"var_name": var_name}, use_api_base=True, auth_kind="admin")
res_json = response.json()
if response.status_code == 200:
self._print_table_simple(res_json["data"])
@@ -604,8 +560,7 @@ class RAGFlowClient:
if self.server_type != "admin":
print("This command is only allowed in ADMIN mode")
license = command["license"]
response = self.http_client.request("POST", "/admin/license", json_body={"license": license}, use_api_base=True,
auth_kind="admin")
response = self.http_client.request("POST", "/admin/license", json_body={"license": license}, use_api_base=True, auth_kind="admin")
res_json = response.json()
if response.status_code == 200:
print("Set license successfully")
@@ -617,9 +572,7 @@ class RAGFlowClient:
print("This command is only allowed in ADMIN mode")
value1 = command["value1"]
value2 = command["value2"]
response = self.http_client.request("POST", "/admin/license/config",
json_body={"value1": value1, "value2": value2}, use_api_base=True,
auth_kind="admin")
response = self.http_client.request("POST", "/admin/license/config", json_body={"value1": value1, "value2": value2}, use_api_base=True, auth_kind="admin")
res_json = response.json()
if response.status_code == 200:
print("Set license successfully")
@@ -690,8 +643,7 @@ class RAGFlowClient:
user_name: str = username_tree.children[0].strip("'\"")
print(f"Listing all datasets of user: {user_name}")
response = self.http_client.request("GET", f"/admin/users/{user_name}/datasets", use_api_base=True,
auth_kind="admin")
response = self.http_client.request("GET", f"/admin/users/{user_name}/datasets", use_api_base=True, auth_kind="admin")
res_json = response.json()
if response.status_code == 200:
table_data = res_json["data"]
@@ -708,8 +660,7 @@ class RAGFlowClient:
username_tree: Tree = command["user_name"]
user_name: str = username_tree.children[0].strip("'\"")
print(f"Listing all agents of user: {user_name}")
response = self.http_client.request("GET", f"/admin/users/{user_name}/agents", use_api_base=True,
auth_kind="admin")
response = self.http_client.request("GET", f"/admin/users/{user_name}/agents", use_api_base=True, auth_kind="admin")
res_json = response.json()
if response.status_code == 200:
table_data = res_json["data"]
@@ -733,8 +684,7 @@ class RAGFlowClient:
# Step 1: Add provider
provider_payload = {"provider_name": provider_name}
provider_response = self.http_client.request("PUT", "/providers", json_body=provider_payload,
use_api_base=True, auth_kind="web")
provider_response = self.http_client.request("PUT", "/providers", json_body=provider_payload, use_api_base=True, auth_kind="web")
provider_res = provider_response.json()
if provider_response.status_code == 200 and provider_res.get("code") == 0:
print(f"Success to add provider {provider_name}")
@@ -747,15 +697,8 @@ class RAGFlowClient:
return
# Step 2: Add instance
instance_payload = {
"instance_name": "default",
"api_key": api_key,
"region": "default",
"base_url": ""
}
instance_response = self.http_client.request("POST", f"/providers/{provider_name}/instances",
json_body=instance_payload, use_api_base=True,
auth_kind="web")
instance_payload = {"instance_name": "default", "api_key": api_key, "region": "default", "base_url": ""}
instance_response = self.http_client.request("POST", f"/providers/{provider_name}/instances", json_body=instance_payload, use_api_base=True, auth_kind="web")
instance_res = instance_response.json()
if instance_response.status_code == 200 and instance_res.get("code") == 0:
print(f"Success to add instance for provider {provider_name}")
@@ -771,8 +714,7 @@ class RAGFlowClient:
print("This command is only allowed in USER mode")
return
provider_name: str = command["provider_name"]
response = self.http_client.request("DELETE", f"/providers/{provider_name}", use_api_base=True,
auth_kind="web")
response = self.http_client.request("DELETE", f"/providers/{provider_name}", use_api_base=True, auth_kind="web")
res_json = response.json()
if response.status_code == 200 and res_json.get("code") == 0:
print(f"Success to drop model provider {provider_name}")
@@ -810,8 +752,7 @@ class RAGFlowClient:
"model_type": model_type,
"model_name": model_name,
}
response = self.http_client.request("PATCH", "/models/default", json_body=payload, use_api_base=True,
auth_kind="web")
response = self.http_client.request("PATCH", "/models/default", json_body=payload, use_api_base=True, auth_kind="web")
res_json = response.json()
if response.status_code == 200 and res_json.get("code") == 0:
print(f"Success to set default {model_type} to {model_id}")
@@ -830,8 +771,7 @@ class RAGFlowClient:
return
payload = {"model_type": model_type}
response = self.http_client.request("PATCH", "/models/default", json_body=payload, use_api_base=True,
auth_kind="web")
response = self.http_client.request("PATCH", "/models/default", json_body=payload, use_api_base=True, auth_kind="web")
res_json = response.json()
if response.status_code == 200 and res_json.get("code") == 0:
print(f"Success to reset default {model_type}")
@@ -861,8 +801,7 @@ class RAGFlowClient:
iterations = command.get("iterations", 1)
if iterations > 1:
response = self.http_client.request("GET", "/datasets", use_api_base=True, auth_kind="web",
iterations=iterations)
response = self.http_client.request("GET", "/datasets", use_api_base=True, auth_kind="web", iterations=iterations)
return response
else:
response = self.http_client.request("GET", "/datasets", use_api_base=True, auth_kind="web")
@@ -876,16 +815,12 @@ class RAGFlowClient:
def create_user_dataset(self, command):
if self.server_type != "user":
print("This command is only allowed in USER mode")
payload = {
"name": command["dataset_name"],
"embedding_model": command["embedding"]
}
payload = {"name": command["dataset_name"], "embedding_model": command["embedding"]}
if "parser_id" in command:
payload["chunk_method"] = command["parser"]
if "pipeline" in command:
payload["pipeline_id"] = command["pipeline"]
response = self.http_client.request("POST", "/datasets", json_body=payload, use_api_base=True,
auth_kind="web")
response = self.http_client.request("POST", "/datasets", json_body=payload, use_api_base=True, auth_kind="web")
res_json = response.json()
if response.status_code == 200:
self._print_table_simple(res_json["data"])
@@ -981,8 +916,7 @@ class RAGFlowClient:
dataset_ids = [dataset_id for _, dataset_id in valid_datasets]
kb_ids_param = ",".join(dataset_ids)
response = self.http_client.request("GET", f"/kb/get_meta?kb_ids={kb_ids_param}",
use_api_base=False, auth_kind="web")
response = self.http_client.request("GET", f"/kb/get_meta?kb_ids={kb_ids_param}", use_api_base=False, auth_kind="web")
res_json = response.json()
if response.status_code != 200:
print(f"Fail to get metadata, code: {res_json.get('code')}, message: {res_json.get('message')}")
@@ -996,11 +930,7 @@ class RAGFlowClient:
table_data = []
for field_name, values_dict in meta.items():
for value, docs in values_dict.items():
table_data.append({
"field": field_name,
"value": value,
"doc_ids": ", ".join(docs)
})
table_data.append({"field": field_name, "value": value, "doc_ids": ", ".join(docs)})
self._print_table_simple(table_data)
def list_user_documents_metadata_summary(self, command_dict):
@@ -1018,8 +948,7 @@ class RAGFlowClient:
payload = {"kb_id": kb_id}
if doc_ids:
payload["doc_ids"] = doc_ids
response = self.http_client.request("POST", "/document/metadata/summary", json_body=payload,
use_api_base=False, auth_kind="web")
response = self.http_client.request("POST", "/document/metadata/summary", json_body=payload, use_api_base=False, auth_kind="web")
res_json = response.json()
if response.status_code == 200:
summary = res_json.get("data", {}).get("summary", {})
@@ -1086,16 +1015,11 @@ class RAGFlowClient:
"quote": True,
"keyword": False,
"tts": False,
"system": "You are an intelligent assistant. Your primary function is to answer questions based strictly on the provided knowledge base.\n\n **Essential Rules:**\n - Your answer must be derived **solely** from this knowledge base: `{knowledge}`.\n - **When information is available**: Summarize the content to give a detailed answer.\n - **When information is unavailable**: Your response must contain this exact sentence: \"The answer you are looking for is not found in the knowledge base!\"\n - **Always consider** the entire conversation history.",
"system": 'You are an intelligent assistant. Your primary function is to answer questions based strictly on the provided knowledge base.\n\n **Essential Rules:**\n - Your answer must be derived **solely** from this knowledge base: `{knowledge}`.\n - **When information is available**: Summarize the content to give a detailed answer.\n - **When information is unavailable**: Your response must contain this exact sentence: "The answer you are looking for is not found in the knowledge base!"\n - **Always consider** the entire conversation history.',
"refine_multiturn": False,
"use_kg": False,
"reasoning": False,
"parameters": [
{
"key": "knowledge",
"optional": False
}
],
"parameters": [{"key": "knowledge", "optional": False}],
"toc_enhance": False,
},
"similarity_threshold": 0.2,
@@ -1136,8 +1060,7 @@ class RAGFlowClient:
# Build payload
payload = {"kb_id": dataset_id, "vector_size": vector_size}
# Call API
response = self.http_client.request("POST", "/kb/doc_engine_table", json_body=payload,
use_api_base=False, auth_kind="web")
response = self.http_client.request("POST", "/kb/doc_engine_table", json_body=payload, use_api_base=False, auth_kind="web")
res_json = response.json()
if response.status_code == 200 and res_json.get("code") == 0:
print(f"Success to create table for dataset: {dataset_name}")
@@ -1155,8 +1078,7 @@ class RAGFlowClient:
return
# Call API to delete table
payload = {"kb_id": dataset_id}
response = self.http_client.request("DELETE", "/kb/doc_engine_table", json_body=payload,
use_api_base=False, auth_kind="web")
response = self.http_client.request("DELETE", "/kb/doc_engine_table", json_body=payload, use_api_base=False, auth_kind="web")
res_json = response.json()
if response.status_code == 200 and res_json.get("code") == 0:
print(f"Success to drop table for dataset: {dataset_name}")
@@ -1168,8 +1090,7 @@ class RAGFlowClient:
print("This command is only allowed in USER mode")
return
# Call API to create metadata table
response = self.http_client.request("POST", "/tenant/doc_engine_metadata_table",
use_api_base=False, auth_kind="web")
response = self.http_client.request("POST", "/tenant/doc_engine_metadata_table", use_api_base=False, auth_kind="web")
res_json = response.json()
if response.status_code == 200 and res_json.get("code") == 0:
print("Success to create metadata table")
@@ -1181,8 +1102,7 @@ class RAGFlowClient:
print("This command is only allowed in USER mode")
return
# Call API to delete metadata table
response = self.http_client.request("DELETE", "/tenant/doc_engine_metadata_table",
use_api_base=False, auth_kind="web")
response = self.http_client.request("DELETE", "/tenant/doc_engine_metadata_table", use_api_base=False, auth_kind="web")
res_json = response.json()
if response.status_code == 200 and res_json.get("code") == 0:
print("Success to drop metadata table")
@@ -1225,8 +1145,7 @@ class RAGFlowClient:
def _list_chat_sessions(self, dialog_id):
"""List all sessions (conversations) for a given dialog."""
response = self.http_client.request("GET", f"/chats/{dialog_id}/conversations", use_api_base=True,
auth_kind="web")
response = self.http_client.request("GET", f"/chats/{dialog_id}/conversations", use_api_base=True, auth_kind="web")
res_json = response.json()
if response.status_code == 200 and res_json["code"] == 0:
return res_json["data"]
@@ -1242,14 +1161,12 @@ class RAGFlowClient:
if dialog_id is None:
return
payload = {"name": "New conversation"}
response = self.http_client.request("POST", f"/chats/{dialog_id}/conversations", json_body=payload,
use_api_base=True, auth_kind="web")
response = self.http_client.request("POST", f"/chats/{dialog_id}/conversations", json_body=payload, use_api_base=True, auth_kind="web")
res_json = response.json()
if response.status_code == 200 and res_json["code"] == 0:
print(f"Success to create chat session for chat: {chat_name}")
else:
print(
f"Fail to create chat session for chat {chat_name}, code: {res_json['code']}, message: {res_json['message']}")
print(f"Fail to create chat session for chat {chat_name}, code: {res_json['code']}, message: {res_json['message']}")
def drop_chat_session(self, command):
if self.server_type != "user":
@@ -1270,14 +1187,12 @@ class RAGFlowClient:
print(f"Chat session '{session_id}' not found in chat '{chat_name}'")
return
payload = {"ids": to_drop_session_ids}
response = self.http_client.request("DELETE", f"/chats/{dialog_id}/conversations", json_body=payload,
use_api_base=True, auth_kind="web")
response = self.http_client.request("DELETE", f"/chats/{dialog_id}/conversations", json_body=payload, use_api_base=True, auth_kind="web")
res_json = response.json()
if response.status_code == 200 and res_json["code"] == 0:
print(f"Success to drop chat session '{session_id}' from chat: {chat_name}")
else:
print(
f"Fail to drop chat session '{session_id}' from chat {chat_name}, code: {res_json['code']}, message: {res_json['message']}")
print(f"Fail to drop chat session '{session_id}' from chat {chat_name}, code: {res_json['code']}, message: {res_json['message']}")
def list_chat_sessions(self, command):
if self.server_type != "user":
@@ -1305,13 +1220,9 @@ class RAGFlowClient:
# Prepare payload for completion API
# Note: stream parameter is not sent, server defaults to stream=True
payload = {
"session_id": session_id,
"messages": [{"role": "user", "content": message}]
}
payload = {"session_id": session_id, "messages": [{"role": "user", "content": message}]}
response = self.http_client.request("POST", "/chat/completions", json_body=payload,
use_api_base=True, auth_kind="web", stream=True)
response = self.http_client.request("POST", "/chat/completions", json_body=payload, use_api_base=True, auth_kind="web", stream=True)
if response.status_code != 200:
print(f"Fail to chat on session, status code: {response.status_code}")
@@ -1322,17 +1233,16 @@ class RAGFlowClient:
for line in response.iter_lines():
if not line:
continue
line_str = line.decode('utf-8')
if not line_str.startswith('data:'):
line_str = line.decode("utf-8")
if not line_str.startswith("data:"):
continue
data_str = line_str[5:].strip()
if data_str == '[DONE]':
if data_str == "[DONE]":
break
try:
data_json = json.loads(data_str)
if data_json.get("code") != 0:
print(
f"\nFail to chat on session, code: {data_json.get('code')}, message: {data_json.get('message', '')}")
print(f"\nFail to chat on session, code: {data_json.get('code')}, message: {data_json.get('message', '')}")
return
# Check if it's the final message
if data_json.get("data") is True:
@@ -1416,14 +1326,12 @@ class RAGFlowClient:
print(f"Documents {document_names} not found in {dataset_name}")
payload = {"doc_ids": document_ids, "run": 1}
response = self.http_client.request("POST", "/documents/ingest", json_body=payload, use_api_base=True,
auth_kind="web")
response = self.http_client.request("POST", "/documents/ingest", json_body=payload, use_api_base=True, auth_kind="web")
res_json = response.json()
if response.status_code == 200 and res_json["code"] == 0:
print(f"Success to parse {to_parse_doc_names} of {dataset_name}")
else:
print(
f"Fail to parse documents {res_json["data"]["docs"]}, code: {res_json['code']}, message: {res_json['message']}")
print(f"Fail to parse documents {res_json['data']['docs']}, code: {res_json['code']}, message: {res_json['message']}")
def parse_dataset(self, command_dict):
if self.server_type != "user":
@@ -1442,8 +1350,7 @@ class RAGFlowClient:
document_ids.append(doc["id"])
payload = {"doc_ids": document_ids, "run": 1}
response = self.http_client.request("POST", "/documents/ingest", json_body=payload, use_api_base=True,
auth_kind="web")
response = self.http_client.request("POST", "/documents/ingest", json_body=payload, use_api_base=True, auth_kind="web")
res_json = response.json()
if response.status_code == 200 and res_json["code"] == 0:
pass
@@ -1483,15 +1390,7 @@ class RAGFlowClient:
encoder = MultipartEncoder(fields=fields)
headers = {"Content-Type": encoder.content_type}
response = self.http_client.request(
"POST",
f"/datasets/{dataset_id}/documents?return_raw_files=true",
headers=headers,
data=encoder,
json_body=None,
params=None,
stream=False,
auth_kind="web",
use_api_base=True
"POST", f"/datasets/{dataset_id}/documents?return_raw_files=true", headers=headers, data=encoder, json_body=None, params=None, stream=False, auth_kind="web", use_api_base=True
)
res = response.json()
if res.get("code") == 0:
@@ -1526,22 +1425,18 @@ class RAGFlowClient:
}
iterations = command_dict.get("iterations", 1)
if iterations > 1:
response = self.http_client.request("POST", "/retrieval", json_body=payload, use_api_base=True,
auth_kind="web", iterations=iterations)
response = self.http_client.request("POST", "/retrieval", json_body=payload, use_api_base=True, auth_kind="web", iterations=iterations)
return response
else:
response = self.http_client.request("POST", "/retrieval", json_body=payload, use_api_base=True,
auth_kind="web")
response = self.http_client.request("POST", "/retrieval", json_body=payload, use_api_base=True, auth_kind="web")
res_json = response.json()
if response.status_code == 200:
if res_json["code"] == 0:
self._print_table_simple(res_json["data"]["chunks"])
else:
print(
f"Fail to search datasets: {dataset_names}, code: {res_json['code']}, message: {res_json['message']}")
print(f"Fail to search datasets: {dataset_names}, code: {res_json['code']}, message: {res_json['message']}")
else:
print(
f"Fail to search datasets: {dataset_names}, code: {res_json['code']}, message: {res_json['message']}")
print(f"Fail to search datasets: {dataset_names}, code: {res_json['code']}, message: {res_json['message']}")
def get_chunk(self, command_dict):
if self.server_type != "user":
@@ -1549,8 +1444,7 @@ class RAGFlowClient:
return
chunk_id = command_dict["chunk_id"]
response = self.http_client.request("GET", f"/chunk/get?chunk_id={chunk_id}", use_api_base=False,
auth_kind="web")
response = self.http_client.request("GET", f"/chunk/get?chunk_id={chunk_id}", use_api_base=False, auth_kind="web")
res_json = response.json()
if response.status_code == 200:
if res_json["code"] == 0:
@@ -1568,8 +1462,7 @@ class RAGFlowClient:
file_path = command_dict["file_path"]
payload = {"file_path": file_path}
response = self.http_client.request("POST", "/kb/insert_from_file", json_body=payload,
use_api_base=False, auth_kind="web")
response = self.http_client.request("POST", "/kb/insert_from_file", json_body=payload, use_api_base=False, auth_kind="web")
res_json = response.json()
if response.status_code == 200:
if res_json["code"] == 0:
@@ -1589,8 +1482,7 @@ class RAGFlowClient:
file_path = command_dict["file_path"]
payload = {"file_path": file_path}
response = self.http_client.request("POST", "/tenant/insert_metadata_from_file", json_body=payload,
use_api_base=False, auth_kind="web")
response = self.http_client.request("POST", "/tenant/insert_metadata_from_file", json_body=payload, use_api_base=False, auth_kind="web")
res_json = response.json()
if response.status_code == 200:
if res_json["code"] == 0:
@@ -1617,8 +1509,7 @@ class RAGFlowClient:
return
# Get doc_id from chunk_id via GET /chunk/get
response = self.http_client.request("GET", f"/chunk/get?chunk_id={chunk_id}", use_api_base=False,
auth_kind="web")
response = self.http_client.request("GET", f"/chunk/get?chunk_id={chunk_id}", use_api_base=False, auth_kind="web")
res_json = response.json()
if response.status_code != 200:
print(f"Fail to get chunk info, code: {res_json.get('code')}, message: {res_json.get('message')}")
@@ -1655,14 +1546,8 @@ class RAGFlowClient:
else:
print(f"Fail to update chunk, HTTP {response.status_code}")
def _get_documents_by_ids(self, ids:list[str]):
response = self.http_client.request(
"POST",
"/document/infos",
json_body={"doc_ids": ids},
use_api_base=False,
auth_kind="web"
)
def _get_documents_by_ids(self, ids: list[str]):
response = self.http_client.request("POST", "/document/infos", json_body={"doc_ids": ids}, use_api_base=False, auth_kind="web")
if response.status_code != 200:
return f"Fail to get document info, HTTP {response.status_code}", None
@@ -1687,6 +1572,7 @@ class RAGFlowClient:
# Parse JSON string to dict
import json
try:
meta_fields = json.loads(meta_json_str)
except json.JSONDecodeError as e:
@@ -1713,13 +1599,7 @@ class RAGFlowClient:
"meta_fields": meta_fields,
}
response = self.http_client.request(
"PATCH",
f"/datasets/{dataset_id}/documents/{doc_id}",
json_body=payload,
use_api_base=True,
auth_kind="web"
)
response = self.http_client.request("PATCH", f"/datasets/{dataset_id}/documents/{doc_id}", json_body=payload, use_api_base=True, auth_kind="web")
res_json = response.json()
if response.status_code == 200:
@@ -1747,8 +1627,7 @@ class RAGFlowClient:
"tags": tags,
}
response = self.http_client.request("POST", f"/kb/{dataset_id}/rm_tags", json_body=payload,
use_api_base=False, auth_kind="web")
response = self.http_client.request("POST", f"/kb/{dataset_id}/rm_tags", json_body=payload, use_api_base=False, auth_kind="web")
res_json = response.json()
if response.status_code == 200:
if res_json.get("code") == 0:
@@ -1771,8 +1650,7 @@ class RAGFlowClient:
elif command_dict.get("chunk_ids"):
payload["chunk_ids"] = command_dict["chunk_ids"]
response = self.http_client.request("POST", "/chunk/rm", json_body=payload,
use_api_base=False, auth_kind="web")
response = self.http_client.request("POST", "/chunk/rm", json_body=payload, use_api_base=False, auth_kind="web")
res_json = response.json()
if response.status_code == 200:
if res_json.get("code") == 0:
@@ -1803,15 +1681,14 @@ class RAGFlowClient:
if "available_int" in command_dict:
payload["available_int"] = command_dict["available_int"]
response = self.http_client.request("POST", "/chunk/list", json_body=payload, use_api_base=False,
auth_kind="web")
response = self.http_client.request("POST", "/chunk/list", json_body=payload, use_api_base=False, auth_kind="web")
res_json = response.json()
if response.status_code == 200:
if res_json["code"] == 0:
chunks = res_json["data"]["chunks"]
if chunks:
for i, chunk in enumerate(chunks):
print(f"\n--- Chunk {i+1} ---")
print(f"\n--- Chunk {i + 1} ---")
for key, value in chunk.items():
print(f" {key}: {value}")
else:
@@ -1845,7 +1722,7 @@ class RAGFlowClient:
all_done = True
for doc in docs:
if doc.get("run") != "DONE":
print(f"Document {doc["name"]} is not done, status: {doc.get("run")}")
print(f"Document {doc['name']} is not done, status: {doc.get('run')}")
all_done = False
break
if all_done:
@@ -1856,16 +1733,10 @@ class RAGFlowClient:
def _list_documents(self, dataset_name: str, dataset_id: str):
# Use the new RESTful API: GET /api/v1/datasets/<dataset_id>/documents
response = self.http_client.request(
"GET",
f"/datasets/{dataset_id}/documents",
use_api_base=True,
auth_kind="web"
)
response = self.http_client.request("GET", f"/datasets/{dataset_id}/documents", use_api_base=True, auth_kind="web")
res_json = response.json()
if response.status_code != 200:
print(
f"Fail to list files from dataset {dataset_name}, code: {res_json['code']}, message: {res_json['message']}")
print(f"Fail to list files from dataset {dataset_name}, code: {res_json['code']}, message: {res_json['message']}")
return None
return res_json["data"]["docs"]
@@ -2254,22 +2125,14 @@ def run_benchmark(client: RAGFlowClient, command_dict: dict):
total_duration = result["duration"]
qps = iterations / total_duration if total_duration > 0 else None
print(f"command: {command}, Concurrency: {concurrency}, iterations: {iterations}")
print(
f"total duration: {total_duration:.4f}s, QPS: {qps}, COMMAND_COUNT: {iterations}, SUCCESS: {success_count}, FAILURE: {iterations - success_count}")
print(f"total duration: {total_duration:.4f}s, QPS: {qps}, COMMAND_COUNT: {iterations}, SUCCESS: {success_count}, FAILURE: {iterations - success_count}")
pass
else:
results: List[Optional[dict]] = [None] * concurrency
mp_context = mp.get_context("spawn")
start_time = time.perf_counter()
with ProcessPoolExecutor(max_workers=concurrency, mp_context=mp_context) as executor:
future_map = {
executor.submit(
run_command,
client,
command
): idx
for idx in range(concurrency)
}
future_map = {executor.submit(run_command, client, command): idx for idx in range(concurrency)}
for future in as_completed(future_map):
idx = future_map[future]
results[idx] = future.result()
@@ -2291,7 +2154,6 @@ def run_benchmark(client: RAGFlowClient, command_dict: dict):
total_command_count = iterations * concurrency
qps = total_command_count / total_duration if total_duration > 0 else None
print(f"command: {command}, Concurrency: {concurrency} , iterations: {iterations}")
print(
f"total duration: {total_duration:.4f}s, QPS: {qps}, COMMAND_COUNT: {total_command_count}, SUCCESS: {success_count}, FAILURE: {total_command_count - success_count}")
print(f"total duration: {total_duration:.4f}s, QPS: {qps}, COMMAND_COUNT: {total_command_count}, SUCCESS: {success_count}, FAILURE: {total_command_count - success_count}")
pass