# # Copyright 2026 The InfiniFlow Authors. All Rights Reserved. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # import json import time from typing import Any, List, Optional import multiprocessing as mp from concurrent.futures import ProcessPoolExecutor, as_completed import urllib.parse from pathlib import Path from http_client import HttpClient from lark import Tree from user import encrypt_password, login_user import base64 from Cryptodome.Cipher import PKCS1_v1_5 as Cipher_pkcs1_v1_5 from Cryptodome.PublicKey import RSA try: from requests_toolbelt import MultipartEncoder except Exception as e: # pragma: no cover - fallback without toolbelt print(f"Fallback without belt: {e}") MultipartEncoder = None def encrypt(input_string): pub = "-----BEGIN PUBLIC KEY-----\nMIIBIjANBgkqhkiG9w0BAQEFAAOCAQ8AMIIBCgKCAQEArq9XTUSeYr2+N1h3Afl/z8Dse/2yD0ZGrKwx+EEEcdsBLca9Ynmx3nIB5obmLlSfmskLpBo0UACBmB5rEjBp2Q2f3AG3Hjd4B+gNCG6BDaawuDlgANIhGnaTLrIqWrrcm4EMzJOnAOI1fgzJRsOOUEfaS318Eq9OVO3apEyCCt0lOQK6PuksduOjVxtltDav+guVAA068NrPYmRNabVKRNLJpL8w4D44sfth5RvZ3q9t+6RTArpEtc5sh5ChzvqPOzKGMXW83C95TxmXqpbK6olN4RevSfVjEAgCydH6HN6OhtOQEcnrU97r9H0iZOWwbw3pVrZiUkuRD1R56Wzs2wIDAQAB\n-----END PUBLIC KEY-----" pub_key = RSA.importKey(pub) cipher = Cipher_pkcs1_v1_5.new(pub_key) cipher_text = cipher.encrypt(base64.b64encode(input_string.encode("utf-8"))) return base64.b64encode(cipher_text).decode("utf-8") def _strip_tree_value(value): if isinstance(value, Tree): value = value.children[0] return str(value).strip("'\"") class RAGFlowClient: def __init__(self, http_client: HttpClient, server_type: str): self.http_client = http_client self.server_type = server_type def login_user(self, command): try: response = self.http_client.request("GET", "/system/ping", use_api_base=False, auth_kind="web") if response.status_code == 200 and response.content == b"pong": pass else: print("Server is down") return except Exception as e: print(str(e)) print("Can't access server for login (connection failed)") return email: str = command["email"] 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) self.http_client.login_token = token # Also store as api_key for API endpoint authentication if self.server_type == "user": self.http_client.api_key = token print(f"Login user {email} successfully") except Exception as e: print(str(e)) print("Can't access server for login (connection failed)") def ping_server(self, command): iterations = command.get("iterations", 1) if iterations > 1: response = self.http_client.request("GET", "/system/ping", use_api_base=False, auth_kind="web", iterations=iterations) return response else: response = self.http_client.request("GET", "/system/ping", use_api_base=False, auth_kind="web") if response.status_code == 200 and response.content == b"pong": print("Server is alive") else: print("Server is down") return None def register_user(self, command): if self.server_type != "user": print("This command is only allowed in USER mode") username: str = command["user_name"] nickname: str = command["nickname"] password: str = command["password"] 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="/user/register", json_body=payload, use_api_base=False, auth_kind="web") res_json = response.json() if response.status_code == 200: if res_json["code"] == 0: self._print_table_simple(res_json["data"]) else: print(f"Fail to register user {username}, code: {res_json['code']}, message: {res_json['message']}") else: print(f"Fail to register user {username}, code: {res_json['code']}, message: {res_json['message']}") def list_services(self): if self.server_type != "admin": print("This command is only allowed in ADMIN mode") response = self.http_client.request("GET", "/admin/services", 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 get all services, code: {res_json['code']}, message: {res_json['message']}") pass def show_service(self, command): if self.server_type != "admin": print("This command is only allowed in ADMIN mode") service_id: int = command["number"] 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"] if "status" in res_data and res_data["status"] == "alive": print(f"Service {res_data['service_name']} is alive, ") res_message = res_data["message"] if res_message is None: return elif isinstance(res_message, str): print(res_message) else: data = self._format_service_detail_table(res_message) self._print_table_simple(data) else: print(f"Service {res_data['service_name']} is down, {res_data['message']}") else: print(f"Fail to show service, code: {res_json['code']}, message: {res_json['message']}") def restart_service(self, command): if self.server_type != "admin": print("This command is only allowed in ADMIN mode") # service_id: int = command["number"] print("Restart service isn't implemented") def shutdown_service(self, command): if self.server_type != "admin": print("This command is only allowed in ADMIN mode") # service_id: int = command["number"] print("Shutdown service isn't implemented") def startup_service(self, command): if self.server_type != "admin": print("This command is only allowed in ADMIN mode") # service_id: int = command["number"] print("Startup service isn't implemented") def list_users(self, command): if self.server_type != "admin": print("This command is only allowed in ADMIN mode") response = self.http_client.request("GET", "/admin/users", 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 get all users, code: {res_json['code']}, message: {res_json['message']}") def show_user(self, command): if self.server_type != "admin": print("This command is only allowed in ADMIN mode") username_tree: Tree = command["user_name"] user_name: str = username_tree.children[0].strip("'\"") print(f"Showing user: {user_name}") response = self.http_client.request("GET", f"/admin/users/{user_name}", use_api_base=True, auth_kind="admin") res_json = response.json() if response.status_code == 200: table_data = res_json["data"][0] table_data.pop("avatar") self._print_table_simple(table_data) else: print(f"Fail to get user {user_name}, code: {res_json['code']}, message: {res_json['message']}") def drop_user(self, command): if self.server_type != "admin": print("This command is only allowed in ADMIN mode") username_tree: Tree = command["user_name"] user_name: str = username_tree.children[0].strip("'\"") print(f"Drop user: {user_name}") response = self.http_client.request("DELETE", f"/admin/users/{user_name}", 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 drop user, code: {res_json['code']}, message: {res_json['message']}") def alter_user(self, command): if self.server_type != "admin": print("This command is only allowed in ADMIN mode") user_name_tree: Tree = command["user_name"] user_name: str = user_name_tree.children[0].strip("'\"") 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") res_json = response.json() if response.status_code == 200: print(res_json["message"]) else: print(f"Fail to alter password, code: {res_json['code']}, message: {res_json['message']}") def create_user(self, command): if self.server_type != "admin": print("This command is only allowed in ADMIN mode") user_name_tree: Tree = command["user_name"] user_name: str = user_name_tree.children[0].strip("'\"") password_tree: Tree = command["password"] password: str = password_tree.children[0].strip("'\"") role: str = command["role"] 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") res_json = response.json() if response.status_code == 200: self._print_table_simple(res_json["data"]) else: print(f"Fail to create user {user_name}, code: {res_json['code']}, message: {res_json['message']}") def activate_user(self, command): if self.server_type != "admin": print("This command is only allowed in ADMIN mode") user_name_tree: Tree = command["user_name"] user_name: str = user_name_tree.children[0].strip("'\"") activate_tree: Tree = command["activate_status"] 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") res_json = response.json() if response.status_code == 200: print(res_json["message"]) else: print(f"Fail to alter activate status, code: {res_json['code']}, message: {res_json['message']}") else: print(f"Unknown activate status: {activate_status}.") def grant_admin(self, command): if self.server_type != "admin": print("This command is only allowed in ADMIN mode") 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") 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']}") def revoke_admin(self, command): if self.server_type != "admin": print("This command is only allowed in ADMIN mode") 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") 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']}") def create_role(self, command): if self.server_type != "admin": print("This command is only allowed in ADMIN mode") role_name_tree: Tree = command["role_name"] role_name: str = role_name_tree.children[0].strip("'\"") desc_str: str = "" if "description" in command and command["description"] is not None: desc_tree: Tree = command["description"] 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") res_json = response.json() if response.status_code == 200: self._print_table_simple(res_json["data"]) else: print(f"Fail to create role {role_name}, code: {res_json['code']}, message: {res_json['message']}") def drop_role(self, command): if self.server_type != "admin": print("This command is only allowed in ADMIN mode") 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") res_json = response.json() if response.status_code == 200: self._print_table_simple(res_json["data"]) else: print(f"Fail to drop role {role_name}, code: {res_json['code']}, message: {res_json['message']}") def alter_role(self, command): if self.server_type != "admin": print("This command is only allowed in ADMIN mode") role_name_tree: Tree = command["role_name"] role_name: str = role_name_tree.children[0].strip("'\"") desc_tree: Tree = command["description"] 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") 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']}") 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") res_json = response.json() if response.status_code == 200: self._print_table_simple(res_json["data"]) else: print(f"Fail to list roles, code: {res_json['code']}, message: {res_json['message']}") def show_role(self, command): if self.server_type != "admin": print("This command is only allowed in ADMIN mode") 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") res_json = response.json() if response.status_code == 200: self._print_table_simple(res_json["data"]) else: print(f"Fail to list roles, code: {res_json['code']}, message: {res_json['message']}") def grant_permission(self, command): if self.server_type != "admin": print("This command is only allowed in ADMIN mode") role_name_tree: Tree = command["role_name"] role_name_str: str = role_name_tree.children[0].strip("'\"") resource_tree: Tree = command["resource"] resource_str: str = resource_tree.children[0].strip("'\"") action_tree_list: list = command["actions"] actions: list = [] for action_tree in action_tree_list: 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") 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']}") def revoke_permission(self, command): if self.server_type != "admin": print("This command is only allowed in ADMIN mode") role_name_tree: Tree = command["role_name"] role_name_str: str = role_name_tree.children[0].strip("'\"") resource_tree: Tree = command["resource"] resource_str: str = resource_tree.children[0].strip("'\"") action_tree_list: list = command["actions"] actions: list = [] for action_tree in action_tree_list: 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") 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']}") def alter_user_role(self, command): if self.server_type != "admin": print("This command is only allowed in ADMIN mode") role_name_tree: Tree = command["role_name"] role_name_str: str = role_name_tree.children[0].strip("'\"") 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") 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']}") def show_user_permission(self, command): if self.server_type != "admin": print("This command is only allowed in ADMIN mode") 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") 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']}") def generate_key(self, command: dict[str, Any]) -> None: if self.server_type != "admin": print("This command is only allowed in ADMIN mode") 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") 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']}") def list_keys(self, command: dict[str, Any]) -> None: if self.server_type != "admin": print("This command is only allowed in ADMIN mode") 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") res_json: dict[str, Any] = response.json() if response.status_code == 200: self._print_table_simple(res_json["data"]) else: print(f"Failed to list keys for user {user_name}, code: {res_json['code']}, message: {res_json['message']}") def drop_key(self, command: dict[str, Any]) -> None: if self.server_type != "admin": print("This command is only allowed in ADMIN mode") key_tree: Tree = command["key"] key: str = key_tree.children[0].strip("'\"") username_tree: Tree = command["user_name"] user_name: str = username_tree.children[0].strip("'\"") 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") res_json: dict[str, Any] = response.json() if response.status_code == 200: print(res_json["message"]) else: print(f"Failed to drop key for user {user_name}, code: {res_json['code']}, message: {res_json['message']}") def set_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"]) 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") 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']}") 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") res_json = response.json() if response.status_code == 200: self._print_table_simple(res_json["data"]) else: print(f"Fail to get variable {var_name}, code: {res_json['code']}, message: {res_json['message']}") def list_variables(self, command): if self.server_type != "admin": print("This command is only allowed in ADMIN mode") response = self.http_client.request("GET", "/admin/variables", 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 list variables, code: {res_json['code']}, message: {res_json['message']}") def list_configs(self, command): if self.server_type != "admin": print("This command is only allowed in ADMIN mode") response = self.http_client.request("GET", "/admin/configs", 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 list variables, code: {res_json['code']}, message: {res_json['message']}") def list_environments(self, command): if self.server_type != "admin": print("This command is only allowed in ADMIN mode") response = self.http_client.request("GET", "/admin/environments", 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 list variables, code: {res_json['code']}, message: {res_json['message']}") def show_fingerprint(self, command): if self.server_type != "admin": print("This command is only allowed in ADMIN mode") response = self.http_client.request("GET", "/admin/fingerprint", 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 fingerprint, code: {res_json['code']}, message: {res_json['message']}") def set_license(self, command): 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") res_json = response.json() if response.status_code == 200: print("Set license successfully") else: print(f"Fail to set license, code: {res_json['code']}, message: {res_json['message']}") def set_license_config(self, command): if self.server_type != "admin": 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") res_json = response.json() if response.status_code == 200: print("Set license successfully") else: print(f"Fail to set license, code: {res_json['code']}, message: {res_json['message']}") def show_license(self, command): if self.server_type != "admin": print("This command is only allowed in ADMIN mode") response = self.http_client.request("GET", "/admin/license", 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 license, code: {res_json['code']}, message: {res_json['message']}") def check_license(self, command): if self.server_type != "admin": print("This command is only allowed in ADMIN mode") response = self.http_client.request("GET", "/admin/license?check=true", use_api_base=True, auth_kind="admin") res_json = response.json() if response.status_code == 200: print(res_json["data"]) else: print(f"Fail to show license, code: {res_json['code']}, message: {res_json['message']}") def list_server_configs(self, command): """List server configs by calling /system/configs API and flattening the JSON response.""" response = self.http_client.request("GET", "/system/configs", use_api_base=False, auth_kind="web") res_json = response.json() if res_json.get("code") != 0: print(f"Fail to list server configs, code: {res_json.get('code')}, message: {res_json.get('message')}") return data = res_json.get("data", {}) if not data: print("No server configs found") return # Flatten nested JSON with a.b.c notation def flatten(obj, parent_key=""): items = [] if isinstance(obj, dict): for k, v in obj.items(): new_key = f"{parent_key}.{k}" if parent_key else k if isinstance(v, (dict, list)) and v: items.extend(flatten(v, new_key)) else: items.append({"name": new_key, "value": v}) elif isinstance(obj, list): for i, v in enumerate(obj): new_key = f"{parent_key}[{i}]" if isinstance(v, (dict, list)) and v: items.extend(flatten(v, new_key)) else: items.append({"name": new_key, "value": v}) return items # Reconstruct flattened data and print using _print_table_simple flattened = flatten(data) self._print_table_simple(flattened) def handle_list_datasets(self, command): if self.server_type != "admin": print("This command is only allowed in ADMIN mode") username_tree: Tree = command["user_name"] 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") res_json = response.json() if response.status_code == 200: table_data = res_json["data"] for t in table_data: t.pop("avatar") self._print_table_simple(table_data) else: print(f"Fail to get all datasets of {user_name}, code: {res_json['code']}, message: {res_json['message']}") def handle_list_agents(self, command): if self.server_type != "admin": print("This command is only allowed in ADMIN mode") 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") res_json = response.json() if response.status_code == 200: table_data = res_json["data"] for t in table_data: t.pop("avatar") self._print_table_simple(table_data) else: print(f"Fail to get all agents of {user_name}, code: {res_json['code']}, message: {res_json['message']}") def show_current_user(self, command): if self.server_type != "user": print("This command is only allowed in USER mode") print("show current user") def create_model_provider(self, command): if self.server_type != "user": print("This command is only allowed in USER mode") llm_factory: str = command["provider_name"] api_key: str = command["provider_key"] payload = {"api_key": api_key, "llm_factory": llm_factory} response = self.http_client.request("POST", "/llm/set_api_key", json_body=payload, use_api_base=False, auth_kind="web") res_json = response.json() if response.status_code == 200 and res_json["code"] == 0: print(f"Success to add model provider {llm_factory}") else: print(f"Fail to add model provider {llm_factory}, code: {res_json['code']}, message: {res_json['message']}") def drop_model_provider(self, command): if self.server_type != "user": print("This command is only allowed in USER mode") llm_factory: str = command["provider_name"] payload = {"llm_factory": llm_factory} response = self.http_client.request("POST", "/llm/delete_factory", json_body=payload, use_api_base=False, auth_kind="web") res_json = response.json() if response.status_code == 200 and res_json["code"] == 0: print(f"Success to drop model provider {llm_factory}") else: print( f"Fail to drop model provider {llm_factory}, code: {res_json['code']}, message: {res_json['message']}") def set_default_model(self, command): if self.server_type != "user": print("This command is only allowed in USER mode") model_type: str = command["model_type"] model_id: str = command["model_id"] self._set_default_models(model_type, model_id) def reset_default_model(self, command): if self.server_type != "user": print("This command is only allowed in USER mode") model_type: str = command["model_type"] self._set_default_models(model_type, "") def list_user_datasets(self, command): if self.server_type != "user": print("This command is only allowed in USER mode") iterations = command.get("iterations", 1) if iterations > 1: 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") res_json = response.json() if response.status_code == 200: self._print_table_simple(res_json["data"]) else: print(f"Fail to list datasets, code: {res_json['code']}, message: {res_json['message']}") return None 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"] } 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") res_json = response.json() if response.status_code == 200: self._print_table_simple(res_json["data"]) else: print(f"Fail to create datasets, code: {res_json['code']}, message: {res_json['message']}") def drop_user_dataset(self, command): if self.server_type != "user": print("This command is only allowed in USER mode") dataset_name = command["dataset_name"] dataset_id = self._get_dataset_id(dataset_name) if dataset_id is None: return payload = {"ids": [dataset_id]} response = self.http_client.request("DELETE", "/datasets", json_body=payload, use_api_base=True, auth_kind="web") res_json = response.json() if response.status_code == 200: print(f"Drop dataset {dataset_name} successfully") else: print(f"Fail to drop datasets, code: {res_json['code']}, message: {res_json['message']}") def list_user_dataset_files(self, command_dict): if self.server_type != "user": print("This command is only allowed in USER mode") dataset_name = command_dict["dataset_name"] dataset_id = self._get_dataset_id(dataset_name) if dataset_id is None: return res_json = self._list_documents(dataset_name, dataset_id) if res_json is None: return self._print_table_simple(res_json) def list_user_dataset_documents(self, command_dict): if self.server_type != "user": print("This command is only allowed in USER mode") dataset_name = command_dict["dataset_name"] dataset_id = self._get_dataset_id(dataset_name) if dataset_id is None: return docs = self._list_documents(dataset_name, dataset_id) if docs is None: return if not docs: print(f"No documents found in dataset {dataset_name}") return print(f"Documents in dataset: {dataset_name}") print("-" * 60) # Select key fields for display display_docs = [] for doc in docs: meta_fields = doc.get("meta_fields", {}) # Convert meta_fields dict to string for display meta_fields_str = "" if meta_fields: meta_fields_str = str(meta_fields) display_doc = { "name": doc.get("name", ""), "id": doc.get("id", ""), "size": doc.get("size", 0), "status": doc.get("status", ""), "created_at": doc.get("created_at", ""), } if meta_fields_str: display_doc["meta_fields"] = meta_fields_str display_docs.append(display_doc) self._print_table_simple(display_docs) def list_user_datasets_metadata(self, command_dict): if self.server_type != "user": print("This command is only allowed in USER mode") return dataset_names = command_dict["dataset_names"] valid_datasets = [] for dataset_name in dataset_names: dataset_id = self._get_dataset_id(dataset_name) if dataset_id is None: print(f"Dataset not found: {dataset_name}") continue valid_datasets.append((dataset_name, dataset_id)) if not valid_datasets: print("No valid datasets found") return 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") 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')}") return meta = res_json.get("data", {}) if not meta: print("No metadata found") return 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) }) self._print_table_simple(table_data) def list_user_documents_metadata_summary(self, command_dict): if self.server_type != "user": print("This command is only allowed in USER mode") return dataset_name = command_dict["dataset_name"] doc_ids = command_dict.get("document_ids", []) kb_id = self._get_dataset_id(dataset_name) if kb_id is None: return 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") res_json = response.json() if response.status_code == 200: summary = res_json.get("data", {}).get("summary", {}) if not summary: if doc_ids: print(f"No metadata summary found for documents: {', '.join(doc_ids)}") else: print(f"No metadata summary found in dataset {dataset_name}") return if doc_ids: print(f"Metadata summary for document(s): {', '.join(doc_ids)}") else: print(f"Metadata summary for all documents in dataset: {dataset_name}") print("-" * 60) for field_name, field_info in summary.items(): field_type = field_info.get("type", "unknown") values = field_info.get("values", []) print(f"\nField: {field_name} (type: {field_type})") print(f" Total unique values: {len(values)}") if values: print(" Values:") for value, count in values: print(f" {value}: {count}") else: print(f"Fail to get metadata summary, code: {res_json.get('code')}, message: {res_json.get('message')}") def list_user_agents(self, command): if self.server_type != "user": print("This command is only allowed in USER mode") response = self.http_client.request("GET", "/canvas/list", use_api_base=False, auth_kind="web") res_json = response.json() if response.status_code == 200: self._print_table_simple(res_json["data"]) else: print(f"Fail to list datasets, code: {res_json['code']}, message: {res_json['message']}") def list_user_chats(self, command): if self.server_type != "user": print("This command is only allowed in USER mode") res_json = self._list_chats(command) if res_json is None: return None if "iterations" in command: # for benchmark return res_json self._print_table_simple(res_json) def create_user_chat(self, command): if self.server_type != "user": print("This command is only allowed in USER mode") chat_name = command["chat_name"] default_models = self._get_default_models() or {} payload = { "name": chat_name, "description": "", "icon": "", "dataset_ids": [], "llm_setting": {}, "prompt_config": { "empty_response": "", "prologue": "Hi! I'm your assistant. What can I do for you?", "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.", "refine_multiturn": False, "use_kg": False, "reasoning": False, "parameters": [ { "key": "knowledge", "optional": False } ], "toc_enhance": False, }, "similarity_threshold": 0.2, "top_n": 8, "top_k": 1024, "vector_similarity_weight": 0.3, "rerank_id": default_models.get("rerank_id", ""), } if default_models.get("llm_id"): payload["llm_id"] = default_models["llm_id"] response = self.http_client.request( "POST", "/chats", 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: {chat_name}") else: print(f"Fail to create chat {chat_name}, code: {res_json['code']}, message: {res_json['message']}") def create_dataset_table(self, command): if self.server_type != "user": print("This command is only allowed in USER mode") return dataset_name = command["dataset_name"] vector_size = command.get("vector_size") if not vector_size: print("vector_size is required") return # Get dataset ID by name dataset_id = self._get_dataset_id(dataset_name) if dataset_id is None: return # 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") 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}") else: print(f"Fail to create table for dataset {dataset_name}, code: {res_json.get('code')}, message: {res_json.get('message')}") def drop_dataset_table(self, command): if self.server_type != "user": print("This command is only allowed in USER mode") return dataset_name = command["dataset_name"] # Get dataset ID by name dataset_id = self._get_dataset_id(dataset_name) if dataset_id is None: 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") 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}") else: print(f"Fail to drop table for dataset {dataset_name}, code: {res_json.get('code')}, message: {res_json.get('message')}") def create_metadata_table(self, command): if self.server_type != "user": 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") res_json = response.json() if response.status_code == 200 and res_json.get("code") == 0: print("Success to create metadata table") else: print(f"Fail to create metadata table, code: {res_json.get('code')}, message: {res_json.get('message')}") def drop_metadata_table(self, command): if self.server_type != "user": 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") res_json = response.json() if response.status_code == 200 and res_json.get("code") == 0: print("Success to drop metadata table") else: print(f"Fail to drop metadata table, code: {res_json.get('code')}, message: {res_json.get('message')}") def drop_user_chat(self, command): if self.server_type != "user": print("This command is only allowed in USER mode") chat_name = command["chat_name"] res_json = self._list_chats(command) to_drop_chat_ids = [] for elem in res_json: if elem["name"] == chat_name: to_drop_chat_ids.append(elem["id"]) payload = {"ids": to_drop_chat_ids} response = self.http_client.request( "DELETE", "/chats", 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: {chat_name}") else: print(f"Fail to drop chat {chat_name}, code: {res_json['code']}, message: {res_json['message']}") def _get_chat_id_by_name(self, chat_name): """Get chat (dialog) ID by name.""" res_json = self._list_chats({}) if res_json is None: return None for elem in res_json: if elem["name"] == chat_name: return elem["id"] print(f"Chat '{chat_name}' not found") return None 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") res_json = response.json() if response.status_code == 200 and res_json["code"] == 0: return res_json["data"] else: print(f"Fail to list chat sessions, code: {res_json['code']}, message: {res_json['message']}") return None def create_chat_session(self, command): if self.server_type != "user": print("This command is only allowed in USER mode") chat_name = command["chat_name"] dialog_id = self._get_chat_id_by_name(chat_name) 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") 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']}") def drop_chat_session(self, command): if self.server_type != "user": print("This command is only allowed in USER mode") chat_name = command["chat_name"] session_id = command["session_id"] dialog_id = self._get_chat_id_by_name(chat_name) if dialog_id is None: return sessions = self._list_chat_sessions(dialog_id) if sessions is None: return to_drop_session_ids = [] for session in sessions: if session["id"] == session_id: to_drop_session_ids.append(session["id"]) if not to_drop_session_ids: 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") 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']}") def list_chat_sessions(self, command): if self.server_type != "user": print("This command is only allowed in USER mode") chat_name = command["chat_name"] dialog_id = self._get_chat_id_by_name(chat_name) if dialog_id is None: return sessions = self._list_chat_sessions(dialog_id) if sessions is None: return # Add chat_name to each session for display for session in sessions: session["chat_name"] = chat_name if "iterations" in command: # for benchmark return sessions self._print_table_simple(sessions) def chat_on_session(self, command): if self.server_type != "user": print("This command is only allowed in USER mode") message = command["message"] session_id = command["session_id"] # 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}] } 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}") return print("Assistant: ", end="", flush=True) full_answer = "" for line in response.iter_lines(): if not line: continue line_str = line.decode('utf-8') if not line_str.startswith('data:'): continue data_str = line_str[5:].strip() 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', '')}") return # Check if it's the final message if data_json.get("data") is True: break answer = data_json.get("data", {}).get("answer", "") if answer: print(answer, end="", flush=True) full_answer += answer except json.JSONDecodeError: continue print() # Final newline def list_user_model_providers(self, command): if self.server_type != "user": print("This command is only allowed in USER mode") response = self.http_client.request("GET", "/llm/my_llms", use_api_base=False, auth_kind="web") res_json = response.json() if response.status_code == 200: new_input = [] for key, value in res_json["data"].items(): new_input.append({"model provider": key, "models": value}) self._print_table_simple(new_input) else: print(f"Fail to list model provider, code: {res_json['code']}, message: {res_json['message']}") def list_user_default_models(self, command): if self.server_type != "user": print("This command is only allowed in USER mode") res_json = self._get_default_models() if res_json is None: return else: new_input = [] for key, value in res_json.items(): if key == "asr_id" and value != "": new_input.append({"model_category": "ASR", "model_name": value}) elif key == "embd_id" and value != "": new_input.append({"model_category": "Embedding", "model_name": value}) elif key == "llm_id" and value != "": new_input.append({"model_category": "LLM", "model_name": value}) elif key == "rerank_id" and value != "": new_input.append({"model_category": "Reranker", "model_name": value}) elif key == "tts_id" and value != "": new_input.append({"model_category": "TTS", "model_name": value}) elif key == "img2txt_id" and value != "": new_input.append({"model_category": "VLM", "model_name": value}) else: continue self._print_table_simple(new_input) def parse_dataset_docs(self, command_dict): if self.server_type != "user": print("This command is only allowed in USER mode") dataset_name = command_dict["dataset_name"] dataset_id = self._get_dataset_id(dataset_name) if dataset_id is None: return res_json = self._list_documents(dataset_name, dataset_id) if res_json is None: return document_names = command_dict["document_names"] document_ids = [] to_parse_doc_names = [] for doc in res_json: doc_name = doc["name"] if doc_name in document_names: document_ids.append(doc["id"]) document_names.remove(doc_name) to_parse_doc_names.append(doc_name) if len(document_ids) == 0: print(f"No documents found in {dataset_name}") return if len(document_names) != 0: 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") 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']}") def parse_dataset(self, command_dict): if self.server_type != "user": print("This command is only allowed in USER mode") dataset_name = command_dict["dataset_name"] dataset_id = self._get_dataset_id(dataset_name) if dataset_id is None: return res_json = self._list_documents(dataset_name, dataset_id) if res_json is None: return document_ids = [] for doc in res_json: 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") res_json = response.json() if response.status_code == 200 and res_json["code"] == 0: pass else: print(f"Fail to parse dataset {dataset_name}, code: {res_json['code']}, message: {res_json['message']}") if command_dict["method"] == "async": print(f"Success to start parse dataset {dataset_name}") return else: print(f"Start to parse dataset {dataset_name}, please wait...") if self._wait_parse_done(dataset_name, dataset_id): print(f"Success to parse dataset {dataset_name}") else: print(f"Parse dataset {dataset_name} timeout") def import_docs_into_dataset(self, command_dict): if self.server_type != "user": print("This command is only allowed in USER mode") dataset_name = command_dict["dataset_name"] dataset_id = self._get_dataset_id(dataset_name) if dataset_id is None: return document_paths = command_dict["document_paths"] paths = [Path(p) for p in document_paths] fields = [] file_handles = [] try: for path in paths: fh = path.open("rb") fields.append(("file", (path.name, fh))) file_handles.append(fh) fields.append(("kb_id", dataset_id)) 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 ) res = response.json() if res.get("code") == 0: print(f"Success to import documents into dataset {dataset_name}") else: print(f"Fail to import documents: code: {res['code']}, message: {res['message']}") except Exception as exc: print(f"Fail to import document into dataset: {dataset_name}, error: {exc}") finally: for fh in file_handles: fh.close() def search_on_datasets(self, command_dict): if self.server_type != "user": print("This command is only allowed in USER mode") dataset_names = command_dict["datasets"] dataset_ids = [] for dataset_name in dataset_names: dataset_id = self._get_dataset_id(dataset_name) if dataset_id is None: return dataset_ids.append(dataset_id) payload = { "question": command_dict["question"], "kb_id": dataset_ids, "similarity_threshold": 0.2, "vector_similarity_weight": 0.3, # "top_k": 1024, # "kb_id": command_dict["datasets"][0], } iterations = command_dict.get("iterations", 1) if iterations > 1: response = self.http_client.request("POST", "/chunk/retrieval_test", json_body=payload, use_api_base=False, auth_kind="web", iterations=iterations) return response else: response = self.http_client.request("POST", "/chunk/retrieval_test", json_body=payload, use_api_base=False, 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']}") else: 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": print("This command is only allowed in USER mode") 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") res_json = response.json() if response.status_code == 200: if res_json["code"] == 0: self._print_key_value(res_json["data"]) else: print(f"Fail to get chunk, code: {res_json['code']}, message: {res_json['message']}") else: print(f"Fail to get chunk, code: {res_json['code']}, message: {res_json['message']}") # Internal def insert_dataset_from_file(self, command_dict): if self.server_type != "user": print("This command is only allowed in USER mode") return 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") res_json = response.json() if response.status_code == 200: if res_json["code"] == 0: print(f"Success to insert dataset from file: {file_path}") if res_json.get("data"): self._print_key_value(res_json["data"]) else: print(f"Fail to insert dataset from file, code: {res_json['code']}, message: {res_json['message']}") else: print(f"Fail to insert dataset from file, code: {res_json['code']}, message: {res_json['message']}") # Internal def insert_metadata_from_file(self, command_dict): if self.server_type != "user": print("This command is only allowed in USER mode") return 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") res_json = response.json() if response.status_code == 200: if res_json["code"] == 0: print(f"Success to insert metadata from file: {file_path}") if res_json.get("data"): self._print_key_value(res_json["data"]) else: print(f"Fail to insert metadata from file, code: {res_json['code']}, message: {res_json['message']}") else: print(f"Fail to insert metadata from file, code: {res_json['code']}, message: {res_json['message']}") def update_chunk(self, command_dict): if self.server_type != "user": print("This command is only allowed in USER mode") return chunk_id = command_dict["chunk_id"] dataset_name = command_dict["dataset_name"] json_body_str = command_dict["json_body"] # Get dataset_id from dataset_name dataset_id = self._get_dataset_id(dataset_name) if dataset_id is None: 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") 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')}") return doc_id = None if res_json.get("code") == 0 and res_json.get("data"): doc_id = res_json["data"].get("doc_id") if not doc_id: print(f"Could not find document_id for chunk {chunk_id}") return # Parse json_body try: payload = json.loads(json_body_str) except json.JSONDecodeError as e: print(f"Invalid JSON body: {e}") return # Add IDs to payload payload["dataset_id"] = dataset_id payload["document_id"] = doc_id payload["chunk_id"] = chunk_id # Call POST /v1/chunk/update response = self.http_client.request("POST", "/chunk/update", 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: print(f"Success to update chunk: {chunk_id}") else: print(f"Fail to update chunk, code: {res_json.get('code')}, message: {res_json.get('message')}") 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" ) if response.status_code != 200: return f"Fail to get document info, HTTP {response.status_code}", None res_json = response.json() if res_json.get("code") != 0: return f"Fail to get document info: {res_json.get('message')}", None docs = res_json.get("data", []) if not docs: return f"Document not found: {ids}", None return None, docs def set_metadata(self, command_dict): if self.server_type != "user": print("This command is only allowed in USER mode") return doc_id = command_dict["doc_id"] meta_json_str = command_dict["meta"] # Parse JSON string to dict import json try: meta_fields = json.loads(meta_json_str) except json.JSONDecodeError as e: print(f"Invalid JSON format: {e}") return # Step 1: Get document info to find kb_id (dataset_id) doc_error_msg, docs = self._get_documents_by_ids([doc_id]) if doc_error_msg: print(doc_error_msg) return if len(docs) == 0: print(f"no document found for {doc_id}") return dataset_id = docs[0].get("dataset_id") if not dataset_id: print(f"Dataset ID not found for document: {doc_id}") return # Send meta as JSON string payload = { "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" ) res_json = response.json() if response.status_code == 200: if res_json.get("code") == 0: print(f"Success to set metadata for document: {doc_id}") else: print(f"Fail to set metadata, code: {res_json.get('code')}, message: {res_json.get('message')}") else: print(f"Fail to set metadata, HTTP {response.status_code}: {res_json.get('message', 'no message')}") def remove_tags(self, command_dict): if self.server_type != "user": print("This command is only allowed in USER mode") return dataset_name = command_dict["dataset_name"] dataset_id = self._get_dataset_id(dataset_name) if dataset_id is None: print(f"Dataset not found: {dataset_name}") return tags = command_dict["tags"] payload = { "tags": tags, } 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: print(f"Success to remove tags from dataset: {dataset_name}") else: print(f"Fail to remove tags, code: {res_json.get('code')}, message: {res_json.get('message')}") else: print(f"Fail to remove tags, HTTP {response.status_code}") def remove_chunks(self, command_dict): if self.server_type != "user": print("This command is only allowed in USER mode") return doc_id = command_dict["doc_id"] payload = {"doc_id": doc_id} if command_dict.get("delete_all"): payload["delete_all"] = True 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") res_json = response.json() if response.status_code == 200: if res_json.get("code") == 0: deleted_count = res_json.get("data", 0) print(f"Success to remove chunks from document {doc_id}: {deleted_count} chunks deleted") else: print(f"Fail to remove chunks, code: {res_json.get('code')}, message: {res_json.get('message')}") else: print(f"Fail to remove chunks, HTTP {response.status_code}") def list_chunks(self, command_dict): if self.server_type != "user": print("This command is only allowed in USER mode") return doc_id = command_dict["doc_id"] payload = { "doc_id": doc_id, } # Add optional parameters (only if explicitly provided) if "page" in command_dict: payload["page"] = command_dict["page"] if "size" in command_dict: payload["size"] = command_dict["size"] if "keywords" in command_dict and command_dict["keywords"]: payload["keywords"] = command_dict["keywords"] 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") 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} ---") for key, value in chunk.items(): print(f" {key}: {value}") else: print("No chunks found") else: print(f"Fail to list chunks, code: {res_json['code']}, message: {res_json['message']}") else: print(f"Fail to list chunks, code: {res_json['code']}, message: {res_json['message']}") def show_version(self, command): if self.server_type == "admin": response = self.http_client.request("GET", "/admin/version", use_api_base=True, auth_kind="admin") else: response = self.http_client.request("GET", "/system/version", use_api_base=True, auth_kind="admin") res_json = response.json() if response.status_code == 200: if self.server_type == "admin": self._print_table_simple(res_json["data"]) else: self._print_table_simple({"version": res_json["data"]}) else: print(f"Fail to show version, code: {res_json['code']}, message: {res_json['message']}") def _wait_parse_done(self, dataset_name: str, dataset_id: str): start = time.monotonic() while True: docs = self._list_documents(dataset_name, dataset_id) if docs is None: return False all_done = True for doc in docs: if doc.get("run") != "DONE": print(f"Document {doc["name"]} is not done, status: {doc.get("run")}") all_done = False break if all_done: return True if time.monotonic() - start > 60: return False time.sleep(0.5) def _list_documents(self, dataset_name: str, dataset_id: str): # Use the new RESTful API: GET /api/v1/datasets//documents 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']}") return None return res_json["data"]["docs"] def _get_dataset_id(self, dataset_name: str): response = self.http_client.request("GET", "/datasets", use_api_base=True, auth_kind="web") res_json = response.json() if response.status_code != 200: print(f"Fail to list datasets, code: {res_json['code']}, message: {res_json['message']}") return None dataset_list = res_json["data"] dataset_id: str = "" for dataset in dataset_list: if dataset["name"] == dataset_name: dataset_id = dataset["id"] if dataset_id == "": print(f"Dataset {dataset_name} not found") return None return dataset_id def _list_chats(self, command): iterations = command.get("iterations", 1) if iterations > 1: response = self.http_client.request( "GET", "/chats", use_api_base=True, auth_kind="web", iterations=iterations, ) return response else: response = self.http_client.request( "GET", "/chats", use_api_base=True, auth_kind="web", iterations=iterations, ) res_json = response.json() if response.status_code == 200 and res_json["code"] == 0: return res_json["data"]["chats"] else: print(f"Fail to list chats, code: {res_json['code']}, message: {res_json['message']}") return None def _get_default_models(self): response = self.http_client.request("GET", "/user/tenant_info", use_api_base=False, auth_kind="web") res_json = response.json() if response.status_code == 200: if res_json["code"] == 0: return res_json["data"] else: print(f"Fail to list user default models, code: {res_json['code']}, message: {res_json['message']}") return None else: print(f"Fail to list user default models, HTTP code: {response.status_code}, message: {res_json}") return None def _set_default_models(self, model_type, model_id): current_payload = self._get_default_models() if current_payload is None: return else: current_payload.update({model_type: model_id}) payload = { "tenant_id": current_payload["tenant_id"], "llm_id": current_payload["llm_id"], "embd_id": current_payload["embd_id"], "img2txt_id": current_payload["img2txt_id"], "asr_id": current_payload["asr_id"], "tts_id": current_payload["tts_id"], } response = self.http_client.request("POST", "/user/set_tenant_info", json_body=payload, use_api_base=False, auth_kind="web") res_json = response.json() if response.status_code == 200 and res_json["code"] == 0: print(f"Success to set default llm to {model_type}") else: print(f"Fail to set default llm to {model_type}, code: {res_json['code']}, message: {res_json['message']}") def _format_service_detail_table(self, data): if isinstance(data, list): return data if not all([isinstance(v, list) for v in data.values()]): # normal table return data # handle task_executor heartbeats map, for example {'name': [{'done': 2, 'now': timestamp1}, {'done': 3, 'now': timestamp2}] task_executor_list = [] for k, v in data.items(): # display latest status heartbeats = sorted(v, key=lambda x: x["now"], reverse=True) task_executor_list.append( { "task_executor_name": k, **heartbeats[0], } if heartbeats else {"task_executor_name": k} ) return task_executor_list def _print_table_simple(self, data): if not data: print("No data to print") return if isinstance(data, dict): # handle single row data data = [data] columns = list(set().union(*(d.keys() for d in data))) columns.sort() col_widths = {} def get_string_width(text): half_width_chars = " !\"#$%&'()*+,-./0123456789:;<=>?@ABCDEFGHIJKLMNOPQRSTUVWXYZ[\\]^_`abcdefghijklmnopqrstuvwxyz{|}~\t\n\r" width = 0 for char in text: if char in half_width_chars: width += 1 else: width += 2 return width for col in columns: max_width = get_string_width(str(col)) for item in data: value_len = get_string_width(str(item.get(col, ""))) if value_len > max_width: max_width = value_len col_widths[col] = max(2, max_width) # Generate delimiter separator = "+" + "+".join(["-" * (col_widths[col] + 2) for col in columns]) + "+" # Print header print(separator) header = "|" + "|".join([f" {col:<{col_widths[col]}} " for col in columns]) + "|" print(header) print(separator) # Print data for item in data: row = "|" for col in columns: value = str(item.get(col, "")) if get_string_width(value) > col_widths[col]: value = value[: col_widths[col] - 3] + "..." row += f" {value:<{col_widths[col] - (get_string_width(value) - len(value))}} |" print(row) print(separator) def _print_key_value(self, data: dict): """Print data as key-value pairs (one per line)""" if not data: print("No data to print") return for key, value in data.items(): print(f"{key}: {value}") def run_command(client: RAGFlowClient, command_dict: dict): command_type = command_dict["type"] match command_type: case "benchmark": run_benchmark(client, command_dict) case "login_user": client.login_user(command_dict) case "ping_server": return client.ping_server(command_dict) case "register_user": client.register_user(command_dict) case "list_services": client.list_services() case "show_service": client.show_service(command_dict) case "restart_service": client.restart_service(command_dict) case "shutdown_service": client.shutdown_service(command_dict) case "startup_service": client.startup_service(command_dict) case "list_users": client.list_users(command_dict) case "show_user": client.show_user(command_dict) case "drop_user": client.drop_user(command_dict) case "alter_user": client.alter_user(command_dict) case "create_user": client.create_user(command_dict) case "activate_user": client.activate_user(command_dict) case "list_datasets": client.handle_list_datasets(command_dict) case "list_agents": client.handle_list_agents(command_dict) case "create_role": client.create_role(command_dict) case "drop_role": client.drop_role(command_dict) case "alter_role": client.alter_role(command_dict) case "list_roles": client.list_roles(command_dict) case "show_role": client.show_role(command_dict) case "grant_permission": client.grant_permission(command_dict) case "revoke_permission": client.revoke_permission(command_dict) case "alter_user_role": client.alter_user_role(command_dict) case "show_user_permission": client.show_user_permission(command_dict) case "show_version": client.show_version(command_dict) case "grant_admin": client.grant_admin(command_dict) case "revoke_admin": client.revoke_admin(command_dict) case "generate_key": client.generate_key(command_dict) case "list_keys": client.list_keys(command_dict) case "drop_key": client.drop_key(command_dict) case "set_variable": client.set_variable(command_dict) case "show_variable": client.show_variable(command_dict) case "list_variables": client.list_variables(command_dict) case "list_configs": client.list_configs(command_dict) case "list_environments": client.list_environments(command_dict) case "show_fingerprint": client.show_fingerprint(command_dict) case "set_license": client.set_license(command_dict) case "set_license_config": client.set_license_config(command_dict) case "show_license": client.show_license(command_dict) case "check_license": client.check_license(command_dict) case "list_server_configs": client.list_server_configs(command_dict) case "create_model_provider": client.create_model_provider(command_dict) case "drop_model_provider": client.drop_model_provider(command_dict) case "show_current_user": client.show_current_user(command_dict) case "set_default_model": client.set_default_model(command_dict) case "reset_default_model": client.reset_default_model(command_dict) case "list_user_datasets": return client.list_user_datasets(command_dict) case "create_user_dataset": client.create_user_dataset(command_dict) case "drop_user_dataset": client.drop_user_dataset(command_dict) case "list_user_dataset_files": return client.list_user_dataset_files(command_dict) case "list_user_dataset_documents": return client.list_user_dataset_documents(command_dict) case "list_user_datasets_metadata": return client.list_user_datasets_metadata(command_dict) case "list_user_documents_metadata_summary": return client.list_user_documents_metadata_summary(command_dict) case "list_user_agents": return client.list_user_agents(command_dict) case "list_user_chats": return client.list_user_chats(command_dict) case "create_user_chat": client.create_user_chat(command_dict) case "drop_user_chat": client.drop_user_chat(command_dict) case "create_dataset_table": client.create_dataset_table(command_dict) case "drop_dataset_table": client.drop_dataset_table(command_dict) case "create_metadata_table": client.create_metadata_table(command_dict) case "drop_metadata_table": client.drop_metadata_table(command_dict) case "create_chat_session": client.create_chat_session(command_dict) case "drop_chat_session": client.drop_chat_session(command_dict) case "list_chat_sessions": return client.list_chat_sessions(command_dict) case "chat_on_session": client.chat_on_session(command_dict) case "list_user_model_providers": client.list_user_model_providers(command_dict) case "list_user_default_models": client.list_user_default_models(command_dict) case "parse_dataset_docs": client.parse_dataset_docs(command_dict) case "parse_dataset": client.parse_dataset(command_dict) case "import_docs_into_dataset": client.import_docs_into_dataset(command_dict) case "search_on_datasets": return client.search_on_datasets(command_dict) case "get_chunk": return client.get_chunk(command_dict) case "insert_dataset_from_file": return client.insert_dataset_from_file(command_dict) case "insert_metadata_from_file": return client.insert_metadata_from_file(command_dict) case "update_chunk": return client.update_chunk(command_dict) case "set_metadata": return client.set_metadata(command_dict) case "remove_tags": return client.remove_tags(command_dict) case "remove_chunks": return client.remove_chunks(command_dict) case "list_chunks": return client.list_chunks(command_dict) case "meta": _handle_meta_command(command_dict) case _: print(f"Command '{command_type}' would be executed with API") def _handle_meta_command(command: dict): meta_command = command["command"] args = command.get("args", []) if meta_command in ["?", "h", "help"]: show_help() elif meta_command in ["q", "quit", "exit"]: print("Goodbye!") else: print(f"Meta command '{meta_command}' with args {args}") def show_help(): """Help info""" help_text = """ Commands: LIST SERVICES SHOW SERVICE STARTUP SERVICE SHUTDOWN SERVICE RESTART SERVICE LIST USERS SHOW USER DROP USER CREATE USER ALTER USER PASSWORD ALTER USER ACTIVE LIST DATASETS OF LIST AGENTS OF CREATE ROLE DROP ROLE ALTER ROLE SET DESCRIPTION LIST ROLES SHOW ROLE GRANT ON TO ROLE REVOKE ON TO ROLE ALTER USER SET ROLE SHOW USER PERMISSION SHOW VERSION GRANT ADMIN REVOKE ADMIN GENERATE KEY FOR USER LIST KEYS OF DROP KEY OF User Commands (use -t user): LIST DATASETS LIST DOCUMENTS OF DATASET SEARCH ON DATASETS LIST METADATA OF DATASETS [, ]* LIST METADATA SUMMARY OF DATASET DOCUMENTS [, ]* GET CHUNK LIST CHUNKS OF DOCUMENT [PAGE ] [SIZE ] [KEYWORDS ] [AVAILABLE <0|1>] Meta Commands: \\?, \\h, \\help Show this help \\q, \\quit, \\exit Quit the CLI """ print(help_text) def run_benchmark(client: RAGFlowClient, command_dict: dict): concurrency = command_dict.get("concurrency", 1) iterations = command_dict.get("iterations", 1) command: dict = command_dict["command"] command.update({"iterations": iterations}) command_type = command["type"] if concurrency < 1: print("Concurrency must be greater than 0") return elif concurrency == 1: result = run_command(client, command) success_count: int = 0 response_list = result["response_list"] for response in response_list: match command_type: case "ping_server": if response.status_code == 200: success_count += 1 case _: res_json = response.json() if response.status_code == 200 and res_json["code"] == 0: success_count += 1 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}") 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) } for future in as_completed(future_map): idx = future_map[future] results[idx] = future.result() end_time = time.perf_counter() success_count = 0 for result in results: response_list = result["response_list"] for response in response_list: match command_type: case "ping_server": if response.status_code == 200: success_count += 1 case _: res_json = response.json() if response.status_code == 200 and res_json["code"] == 0: success_count += 1 total_duration = end_time - start_time 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}") pass