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