From 2e2c8f6ca9a672107f28ec75e40393807dbb0ed4 Mon Sep 17 00:00:00 2001 From: Jin Hai Date: Wed, 21 Jan 2026 18:49:52 +0800 Subject: [PATCH] Add more commands to RAGFlow CLI (#12731) ### What problem does this PR solve? This PR is going to make RAGFlow CLI to access RAGFlow as normal user, and work as the a testing tool for RAGFlow server. ### Type of change - [x] New Feature (non-breaking change which adds functionality) --------- Signed-off-by: Jin Hai --- admin/client/http_client.py | 161 ++++ admin/client/parser.py | 609 +++++++++++++ admin/client/pyproject.toml | 3 + admin/client/ragflow_cli.py | 1228 ++------------------------- admin/client/ragflow_client.py | 1453 ++++++++++++++++++++++++++++++++ admin/client/user.py | 65 ++ admin/server/routes.py | 16 +- admin/server/services.py | 25 +- 8 files changed, 2387 insertions(+), 1173 deletions(-) create mode 100644 admin/client/http_client.py create mode 100644 admin/client/parser.py create mode 100644 admin/client/ragflow_client.py create mode 100644 admin/client/user.py diff --git a/admin/client/http_client.py b/admin/client/http_client.py new file mode 100644 index 000000000..97012c297 --- /dev/null +++ b/admin/client/http_client.py @@ -0,0 +1,161 @@ +# +# 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 time +import json +from typing import Any, Dict, Optional, Tuple + +import requests + + +class HttpClient: + def __init__( + self, + host: str = "127.0.0.1", + port: int = 9381, + api_version: str = "v1", + api_key: Optional[str] = None, + connect_timeout: float = 5.0, + read_timeout: float = 60.0, + verify_ssl: bool = False, + ) -> None: + self.host = host + self.port = port + self.api_version = api_version + self.api_key = api_key + self.login_token: str | None = None + self.connect_timeout = connect_timeout + self.read_timeout = read_timeout + self.verify_ssl = verify_ssl + + def api_base(self) -> str: + return f"{self.host}:{self.port}/api/{self.api_version}" + + def non_api_base(self) -> str: + return f"{self.host}:{self.port}/{self.api_version}" + + def build_url(self, path: str, use_api_base: bool = True) -> str: + base = self.api_base() if use_api_base else self.non_api_base() + if self.verify_ssl: + return f"https://{base}/{path.lstrip('/')}" + else: + return f"http://{base}/{path.lstrip('/')}" + + def _headers(self, auth_kind: Optional[str], extra: Optional[Dict[str, str]]) -> Dict[str, str]: + headers = {} + if auth_kind == "api" and self.api_key: + headers["Authorization"] = f"Bearer {self.api_key}" + elif auth_kind == "web" and self.login_token: + headers["Authorization"] = self.login_token + elif auth_kind == "admin" and self.login_token: + headers["Authorization"] = self.login_token + else: + pass + if extra: + headers.update(extra) + return headers + + def request( + self, + method: str, + path: str, + *, + use_api_base: bool = True, + auth_kind: Optional[str] = "api", + headers: Optional[Dict[str, str]] = None, + json_body: Optional[Dict[str, Any]] = None, + data: Any = None, + files: Any = None, + params: Optional[Dict[str, Any]] = None, + stream: bool = False, + iterations: int = 1, + ) -> requests.Response | dict: + url = self.build_url(path, use_api_base=use_api_base) + merged_headers = self._headers(auth_kind, headers) + timeout: Tuple[float, float] = (self.connect_timeout, self.read_timeout) + if iterations > 1: + response_list = [] + total_duration = 0.0 + for _ in range(iterations): + start_time = time.perf_counter() + response = requests.request( + method=method, + url=url, + headers=merged_headers, + json=json_body, + data=data, + files=files, + params=params, + timeout=timeout, + stream=stream, + verify=self.verify_ssl, + ) + end_time = time.perf_counter() + total_duration += end_time - start_time + response_list.append(response) + return {"duration": total_duration, "response_list": response_list} + else: + return requests.request( + method=method, + url=url, + headers=merged_headers, + json=json_body, + data=data, + files=files, + params=params, + timeout=timeout, + stream=stream, + verify=self.verify_ssl, + ) + + + def request_json( + self, + method: str, + path: str, + *, + use_api_base: bool = True, + auth_kind: Optional[str] = "api", + headers: Optional[Dict[str, str]] = None, + json_body: Optional[Dict[str, Any]] = None, + data: Any = None, + files: Any = None, + params: Optional[Dict[str, Any]] = None, + stream: bool = False, + ) -> Dict[str, Any]: + response = self.request( + method, + path, + use_api_base=use_api_base, + auth_kind=auth_kind, + headers=headers, + json_body=json_body, + data=data, + files=files, + params=params, + stream=stream, + ) + try: + return response.json() + except Exception as exc: + raise ValueError(f"Non-JSON response from {path}: {exc}") from exc + + @staticmethod + def parse_json_bytes(raw: bytes) -> Dict[str, Any]: + try: + return json.loads(raw.decode("utf-8")) + except Exception as exc: + raise ValueError(f"Invalid JSON payload: {exc}") from exc diff --git a/admin/client/parser.py b/admin/client/parser.py new file mode 100644 index 000000000..cf6bf31f9 --- /dev/null +++ b/admin/client/parser.py @@ -0,0 +1,609 @@ +# +# Copyright 2025 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. +# + +from lark import Transformer + +GRAMMAR = r""" +start: command + +command: sql_command | meta_command + +sql_command: list_services + | show_service + | startup_service + | shutdown_service + | restart_service + | register_user + | list_users + | show_user + | drop_user + | alter_user + | create_user + | activate_user + | list_datasets + | list_agents + | create_role + | drop_role + | alter_role + | list_roles + | show_role + | grant_permission + | revoke_permission + | alter_user_role + | show_user_permission + | show_version + | grant_admin + | revoke_admin + | set_variable + | show_variable + | list_variables + | list_configs + | list_environments + | generate_key + | list_keys + | drop_key + | show_current_user + | set_default_llm + | set_default_vlm + | set_default_embedding + | set_default_reranker + | set_default_asr + | set_default_tts + | reset_default_llm + | reset_default_vlm + | reset_default_embedding + | reset_default_reranker + | reset_default_asr + | reset_default_tts + | create_model_provider + | drop_model_provider + | create_user_dataset_with_parser + | create_user_dataset_with_pipeline + | drop_user_dataset + | list_user_datasets + | list_user_dataset_files + | list_user_agents + | list_user_chats + | create_user_chat + | drop_user_chat + | list_user_model_providers + | list_user_default_models + | parse_dataset_docs + | parse_dataset_sync + | parse_dataset_async + | import_docs_into_dataset + | search_on_datasets + | benchmark + +// meta command definition +meta_command: "\\" meta_command_name [meta_args] + +meta_command_name: /[a-zA-Z?]+/ +meta_args: (meta_arg)+ + +meta_arg: /[^\\s"']+/ | quoted_string + +// command definition + +REGISTER: "REGISTER"i +LIST: "LIST"i +SERVICES: "SERVICES"i +SHOW: "SHOW"i +CREATE: "CREATE"i +SERVICE: "SERVICE"i +SHUTDOWN: "SHUTDOWN"i +STARTUP: "STARTUP"i +RESTART: "RESTART"i +USERS: "USERS"i +DROP: "DROP"i +USER: "USER"i +ALTER: "ALTER"i +ACTIVE: "ACTIVE"i +ADMIN: "ADMIN"i +PASSWORD: "PASSWORD"i +DATASET: "DATASET"i +DATASETS: "DATASETS"i +OF: "OF"i +AGENTS: "AGENTS"i +ROLE: "ROLE"i +ROLES: "ROLES"i +DESCRIPTION: "DESCRIPTION"i +GRANT: "GRANT"i +REVOKE: "REVOKE"i +ALL: "ALL"i +PERMISSION: "PERMISSION"i +TO: "TO"i +FROM: "FROM"i +FOR: "FOR"i +RESOURCES: "RESOURCES"i +ON: "ON"i +SET: "SET"i +RESET: "RESET"i +VERSION: "VERSION"i +VAR: "VAR"i +VARS: "VARS"i +CONFIGS: "CONFIGS"i +ENVS: "ENVS"i +KEY: "KEY"i +KEYS: "KEYS"i +GENERATE: "GENERATE"i +MODEL: "MODEL"i +MODELS: "MODELS"i +PROVIDER: "PROVIDER"i +PROVIDERS: "PROVIDERS"i +DEFAULT: "DEFAULT"i +CHATS: "CHATS"i +CHAT: "CHAT"i +FILES: "FILES"i +AS: "AS"i +PARSE: "PARSE"i +IMPORT: "IMPORT"i +INTO: "INTO"i +WITH: "WITH"i +PARSER: "PARSER"i +PIPELINE: "PIPELINE"i +SEARCH: "SEARCH"i +CURRENT: "CURRENT"i +LLM: "LLM"i +VLM: "VLM"i +EMBEDDING: "EMBEDDING"i +RERANKER: "RERANKER"i +ASR: "ASR"i +TTS: "TTS"i +ASYNC: "ASYNC"i +SYNC: "SYNC"i +BENCHMARK: "BENCHMARK"i + +list_services: LIST SERVICES ";" +show_service: SHOW SERVICE NUMBER ";" +startup_service: STARTUP SERVICE NUMBER ";" +shutdown_service: SHUTDOWN SERVICE NUMBER ";" +restart_service: RESTART SERVICE NUMBER ";" + +register_user: REGISTER USER quoted_string AS quoted_string PASSWORD quoted_string ";" +list_users: LIST USERS ";" +drop_user: DROP USER quoted_string ";" +alter_user: ALTER USER PASSWORD quoted_string quoted_string ";" +show_user: SHOW USER quoted_string ";" +create_user: CREATE USER quoted_string quoted_string ";" +activate_user: ALTER USER ACTIVE quoted_string status ";" + +list_datasets: LIST DATASETS OF quoted_string ";" +list_agents: LIST AGENTS OF quoted_string ";" + +create_role: CREATE ROLE identifier [DESCRIPTION quoted_string] ";" +drop_role: DROP ROLE identifier ";" +alter_role: ALTER ROLE identifier SET DESCRIPTION quoted_string ";" +list_roles: LIST ROLES ";" +show_role: SHOW ROLE identifier ";" + +grant_permission: GRANT identifier_list ON identifier TO ROLE identifier ";" +revoke_permission: REVOKE identifier_list ON identifier FROM ROLE identifier ";" +alter_user_role: ALTER USER quoted_string SET ROLE identifier ";" +show_user_permission: SHOW USER PERMISSION quoted_string ";" + +show_version: SHOW VERSION ";" + +grant_admin: GRANT ADMIN quoted_string ";" +revoke_admin: REVOKE ADMIN quoted_string ";" + +generate_key: GENERATE KEY FOR USER quoted_string ";" +list_keys: LIST KEYS OF quoted_string ";" +drop_key: DROP KEY quoted_string OF quoted_string ";" + +set_variable: SET VAR identifier identifier ";" +show_variable: SHOW VAR identifier ";" +list_variables: LIST VARS ";" +list_configs: LIST CONFIGS ";" +list_environments: LIST ENVS ";" + +benchmark: BENCHMARK NUMBER NUMBER user_statement + +user_statement: show_current_user + | create_model_provider + | drop_model_provider + | set_default_llm + | set_default_vlm + | set_default_embedding + | set_default_reranker + | set_default_asr + | set_default_tts + | reset_default_llm + | reset_default_vlm + | reset_default_embedding + | reset_default_reranker + | reset_default_asr + | reset_default_tts + | create_user_dataset_with_parser + | create_user_dataset_with_pipeline + | drop_user_dataset + | list_user_datasets + | list_user_dataset_files + | list_user_agents + | list_user_chats + | create_user_chat + | drop_user_chat + | list_user_model_providers + | list_user_default_models + | import_docs_into_dataset + | search_on_datasets + +show_current_user: SHOW CURRENT USER ";" +create_model_provider: CREATE MODEL PROVIDER quoted_string quoted_string ";" +drop_model_provider: DROP MODEL PROVIDER quoted_string ";" +set_default_llm: SET DEFAULT LLM quoted_string ";" +set_default_vlm: SET DEFAULT VLM quoted_string ";" +set_default_embedding: SET DEFAULT EMBEDDING quoted_string ";" +set_default_reranker: SET DEFAULT RERANKER quoted_string ";" +set_default_asr: SET DEFAULT ASR quoted_string ";" +set_default_tts: SET DEFAULT TTS quoted_string ";" + +reset_default_llm: RESET DEFAULT LLM ";" +reset_default_vlm: RESET DEFAULT VLM ";" +reset_default_embedding: RESET DEFAULT EMBEDDING ";" +reset_default_reranker: RESET DEFAULT RERANKER ";" +reset_default_asr: RESET DEFAULT ASR ";" +reset_default_tts: RESET DEFAULT TTS ";" + +list_user_datasets: LIST DATASETS ";" +create_user_dataset_with_parser: CREATE DATASET quoted_string WITH EMBEDDING quoted_string PARSER quoted_string ";" +create_user_dataset_with_pipeline: CREATE DATASET quoted_string WITH EMBEDDING quoted_string PIPELINE quoted_string ";" +drop_user_dataset: DROP DATASET quoted_string ";" +list_user_dataset_files: LIST FILES OF DATASET quoted_string ";" +list_user_agents: LIST AGENTS ";" +list_user_chats: LIST CHATS ";" +create_user_chat: CREATE CHAT quoted_string ";" +drop_user_chat: DROP CHAT quoted_string ";" +list_user_model_providers: LIST MODEL PROVIDERS ";" +list_user_default_models: LIST DEFAULT MODELS ";" +import_docs_into_dataset: IMPORT quoted_string INTO DATASET quoted_string ";" +search_on_datasets: SEARCH quoted_string ON DATASETS quoted_string ";" + +parse_dataset_docs: PARSE quoted_string OF DATASET quoted_string ";" +parse_dataset_sync: PARSE DATASET quoted_string SYNC ";" +parse_dataset_async: PARSE DATASET quoted_string ASYNC ";" + +identifier_list: identifier ("," identifier)* + +identifier: WORD +quoted_string: QUOTED_STRING +status: WORD + +QUOTED_STRING: /'[^']+'/ | /"[^"]+"/ +WORD: /[a-zA-Z0-9_\-\.]+/ +NUMBER: /[0-9]+/ + +%import common.WS +%ignore WS +""" + + +class RAGFlowCLITransformer(Transformer): + def start(self, items): + return items[0] + + def command(self, items): + return items[0] + + def list_services(self, items): + result = {"type": "list_services"} + return result + + def show_service(self, items): + service_id = int(items[2]) + return {"type": "show_service", "number": service_id} + + def startup_service(self, items): + service_id = int(items[2]) + return {"type": "startup_service", "number": service_id} + + def shutdown_service(self, items): + service_id = int(items[2]) + return {"type": "shutdown_service", "number": service_id} + + def restart_service(self, items): + service_id = int(items[2]) + return {"type": "restart_service", "number": service_id} + + def register_user(self, items): + user_name: str = items[2].children[0].strip("'\"") + nickname: str = items[4].children[0].strip("'\"") + password: str = items[6].children[0].strip("'\"") + return {"type": "register_user", "user_name": user_name, "nickname": nickname, "password": password} + + def list_users(self, items): + return {"type": "list_users"} + + def show_user(self, items): + user_name = items[2] + return {"type": "show_user", "user_name": user_name} + + def drop_user(self, items): + user_name = items[2] + return {"type": "drop_user", "user_name": user_name} + + def alter_user(self, items): + user_name = items[3] + new_password = items[4] + return {"type": "alter_user", "user_name": user_name, "password": new_password} + + def create_user(self, items): + user_name = items[2] + password = items[3] + return {"type": "create_user", "user_name": user_name, "password": password, "role": "user"} + + def activate_user(self, items): + user_name = items[3] + activate_status = items[4] + return {"type": "activate_user", "activate_status": activate_status, "user_name": user_name} + + def list_datasets(self, items): + user_name = items[3] + return {"type": "list_datasets", "user_name": user_name} + + def list_agents(self, items): + user_name = items[3] + return {"type": "list_agents", "user_name": user_name} + + def create_role(self, items): + role_name = items[2] + if len(items) > 4: + description = items[4] + return {"type": "create_role", "role_name": role_name, "description": description} + else: + return {"type": "create_role", "role_name": role_name} + + def drop_role(self, items): + role_name = items[2] + return {"type": "drop_role", "role_name": role_name} + + def alter_role(self, items): + role_name = items[2] + description = items[5] + return {"type": "alter_role", "role_name": role_name, "description": description} + + def list_roles(self, items): + return {"type": "list_roles"} + + def show_role(self, items): + role_name = items[2] + return {"type": "show_role", "role_name": role_name} + + def grant_permission(self, items): + action_list = items[1] + resource = items[3] + role_name = items[6] + return {"type": "grant_permission", "role_name": role_name, "resource": resource, "actions": action_list} + + def revoke_permission(self, items): + action_list = items[1] + resource = items[3] + role_name = items[6] + return {"type": "revoke_permission", "role_name": role_name, "resource": resource, "actions": action_list} + + def alter_user_role(self, items): + user_name = items[2] + role_name = items[5] + return {"type": "alter_user_role", "user_name": user_name, "role_name": role_name} + + def show_user_permission(self, items): + user_name = items[3] + return {"type": "show_user_permission", "user_name": user_name} + + def show_version(self, items): + return {"type": "show_version"} + + def grant_admin(self, items): + user_name = items[2] + return {"type": "grant_admin", "user_name": user_name} + + def revoke_admin(self, items): + user_name = items[2] + return {"type": "revoke_admin", "user_name": user_name} + + def generate_key(self, items): + user_name = items[4] + return {"type": "generate_key", "user_name": user_name} + + def list_keys(self, items): + user_name = items[3] + return {"type": "list_keys", "user_name": user_name} + + def drop_key(self, items): + key = items[2] + user_name = items[4] + return {"type": "drop_key", "key": key, "user_name": user_name} + + def set_variable(self, items): + var_name = items[2] + var_value = items[3] + return {"type": "set_variable", "var_name": var_name, "var_value": var_value} + + def show_variable(self, items): + var_name = items[2] + return {"type": "show_variable", "var_name": var_name} + + def list_variables(self, items): + return {"type": "list_variables"} + + def list_configs(self, items): + return {"type": "list_configs"} + + def list_environments(self, items): + return {"type": "list_environments"} + + def create_model_provider(self, items): + provider_name = items[3].children[0].strip("'\"") + provider_key = items[4].children[0].strip("'\"") + return {"type": "create_model_provider", "provider_name": provider_name, "provider_key": provider_key} + + def drop_model_provider(self, items): + provider_name = items[3].children[0].strip("'\"") + return {"type": "drop_model_provider", "provider_name": provider_name} + + def show_current_user(self, items): + return {"type": "show_current_user"} + + def set_default_llm(self, items): + llm_id = items[3].children[0].strip("'\"") + return {"type": "set_default_model", "model_type": "llm_id", "model_id": llm_id} + + def set_default_vlm(self, items): + vlm_id = items[3].children[0].strip("'\"") + return {"type": "set_default_model", "model_type": "img2txt_id", "model_id": vlm_id} + + def set_default_embedding(self, items): + embedding_id = items[3].children[0].strip("'\"") + return {"type": "set_default_model", "model_type": "embd_id", "model_id": embedding_id} + + def set_default_reranker(self, items): + reranker_id = items[3].children[0].strip("'\"") + return {"type": "set_default_model", "model_type": "reranker_id", "model_id": reranker_id} + + def set_default_asr(self, items): + asr_id = items[3].children[0].strip("'\"") + return {"type": "set_default_model", "model_type": "asr_id", "model_id": asr_id} + + def set_default_tts(self, items): + tts_id = items[3].children[0].strip("'\"") + return {"type": "set_default_model", "model_type": "tts_id", "model_id": tts_id} + + def reset_default_llm(self, items): + return {"type": "reset_default_model", "model_type": "llm_id"} + + def reset_default_vlm(self, items): + return {"type": "reset_default_model", "model_type": "img2txt_id"} + + def reset_default_embedding(self, items): + return {"type": "reset_default_model", "model_type": "embd_id"} + + def reset_default_reranker(self, items): + return {"type": "reset_default_model", "model_type": "reranker_id"} + + def reset_default_asr(self, items): + return {"type": "reset_default_model", "model_type": "asr_id"} + + def reset_default_tts(self, items): + return {"type": "reset_default_model", "model_type": "tts_id"} + + def list_user_datasets(self, items): + return {"type": "list_user_datasets"} + + def create_user_dataset_with_parser(self, items): + dataset_name = items[2].children[0].strip("'\"") + embedding = items[5].children[0].strip("'\"") + parser_type = items[7].children[0].strip("'\"") + return {"type": "create_user_dataset", "dataset_name": dataset_name, "embedding": embedding, + "parser_type": parser_type} + + def create_user_dataset_with_pipeline(self, items): + dataset_name = items[2].children[0].strip("'\"") + embedding = items[5].children[0].strip("'\"") + pipeline = items[7].children[0].strip("'\"") + return {"type": "create_user_dataset", "dataset_name": dataset_name, "embedding": embedding, + "pipeline": pipeline} + + def drop_user_dataset(self, items): + dataset_name = items[2].children[0].strip("'\"") + return {"type": "drop_user_dataset", "dataset_name": dataset_name} + + def list_user_dataset_files(self, items): + dataset_name = items[4].children[0].strip("'\"") + return {"type": "list_user_dataset_files", "dataset_name": dataset_name} + + def list_user_agents(self, items): + return {"type": "list_user_agents"} + + def list_user_chats(self, items): + return {"type": "list_user_chats"} + + def create_user_chat(self, items): + chat_name = items[2].children[0].strip("'\"") + return {"type": "create_user_chat", "chat_name": chat_name} + + def drop_user_chat(self, items): + chat_name = items[2].children[0].strip("'\"") + return {"type": "drop_user_chat", "chat_name": chat_name} + + def list_user_model_providers(self, items): + return {"type": "list_user_model_providers"} + + def list_user_default_models(self, items): + return {"type": "list_user_default_models"} + + def parse_dataset_docs(self, items): + document_list_str = items[1].children[0].strip("'\"") + document_names = document_list_str.split(",") + if len(document_names) == 1: + document_names = document_names[0] + document_names = document_names.split(" ") + dataset_name = items[4].children[0].strip("'\"") + return {"type": "parse_dataset_docs", "dataset_name": dataset_name, "document_names": document_names} + + def parse_dataset_sync(self, items): + dataset_name = items[2].children[0].strip("'\"") + return {"type": "parse_dataset", "dataset_name": dataset_name, "method": "sync"} + + def parse_dataset_async(self, items): + dataset_name = items[2].children[0].strip("'\"") + return {"type": "parse_dataset", "dataset_name": dataset_name, "method": "async"} + + def import_docs_into_dataset(self, items): + document_list_str = items[1].children[0].strip("'\"") + document_paths = document_list_str.split(",") + if len(document_paths) == 1: + document_paths = document_paths[0] + document_paths = document_paths.split(" ") + dataset_name = items[4].children[0].strip("'\"") + return {"type": "import_docs_into_dataset", "dataset_name": dataset_name, "document_paths": document_paths} + + def search_on_datasets(self, items): + question = items[1].children[0].strip("'\"") + datasets_str = items[4].children[0].strip("'\"") + datasets = datasets_str.split(",") + if len(datasets) == 1: + datasets = datasets[0] + datasets = datasets.split(" ") + return {"type": "search_on_datasets", "datasets": datasets, "question": question} + + def benchmark(self, items): + concurrency: int = int(items[1]) + iterations: int = int(items[2]) + command = items[3].children[0] + return {"type": "benchmark", "concurrency": concurrency, "iterations": iterations, "command": command} + + def action_list(self, items): + return items + + def meta_command(self, items): + command_name = str(items[0]).lower() + args = items[1:] if len(items) > 1 else [] + + # handle quoted parameter + parsed_args = [] + for arg in args: + if hasattr(arg, "value"): + parsed_args.append(arg.value) + else: + parsed_args.append(str(arg)) + + return {"type": "meta", "command": command_name, "args": parsed_args} + + def meta_command_name(self, items): + return items[0] + + def meta_args(self, items): + return items diff --git a/admin/client/pyproject.toml b/admin/client/pyproject.toml index 3e35d86c1..94413c057 100644 --- a/admin/client/pyproject.toml +++ b/admin/client/pyproject.toml @@ -20,5 +20,8 @@ test = [ "requests-toolbelt>=1.0.0", ] +[tool.setuptools] +py-modules = ["ragflow_cli", "parser"] + [project.scripts] ragflow-cli = "ragflow_cli:main" diff --git a/admin/client/ragflow_cli.py b/admin/client/ragflow_cli.py index a8e50caea..ac3d49b30 100644 --- a/admin/client/ragflow_cli.py +++ b/admin/client/ragflow_cli.py @@ -14,366 +14,24 @@ # limitations under the License. # +import sys import argparse import base64 import getpass -import urllib.parse from cmd import Cmd from typing import Any, Dict, List import requests +import warnings from Cryptodome.Cipher import PKCS1_v1_5 as Cipher_pkcs1_v1_5 from Cryptodome.PublicKey import RSA -from lark import Lark, Transformer, Tree - -GRAMMAR = r""" -start: command - -command: sql_command | meta_command - -sql_command: list_services - | show_service - | startup_service - | shutdown_service - | restart_service - | list_users - | show_user - | drop_user - | alter_user - | create_user - | activate_user - | list_datasets - | list_agents - | create_role - | drop_role - | alter_role - | list_roles - | show_role - | grant_permission - | revoke_permission - | alter_user_role - | show_user_permission - | show_version - | grant_admin - | revoke_admin - | set_variable - | show_variable - | list_variables - | list_configs - | list_environments - | generate_key - | list_keys - | drop_key - | list_user_datasets - | list_user_agents - | list_user_chats - | list_user_model_providers - | list_user_default_models - -// meta command definition -meta_command: "\\" meta_command_name [meta_args] - -meta_command_name: /[a-zA-Z?]+/ -meta_args: (meta_arg)+ - -meta_arg: /[^\\s"']+/ | quoted_string - -// command definition - -LIST: "LIST"i -SERVICES: "SERVICES"i -SHOW: "SHOW"i -CREATE: "CREATE"i -SERVICE: "SERVICE"i -SHUTDOWN: "SHUTDOWN"i -STARTUP: "STARTUP"i -RESTART: "RESTART"i -USERS: "USERS"i -DROP: "DROP"i -USER: "USER"i -ALTER: "ALTER"i -ACTIVE: "ACTIVE"i -ADMIN: "ADMIN"i -PASSWORD: "PASSWORD"i -DATASETS: "DATASETS"i -OF: "OF"i -AGENTS: "AGENTS"i -ROLE: "ROLE"i -ROLES: "ROLES"i -DESCRIPTION: "DESCRIPTION"i -GRANT: "GRANT"i -REVOKE: "REVOKE"i -ALL: "ALL"i -PERMISSION: "PERMISSION"i -TO: "TO"i -FROM: "FROM"i -FOR: "FOR"i -RESOURCES: "RESOURCES"i -ON: "ON"i -SET: "SET"i -VERSION: "VERSION"i -VAR: "VAR"i -VARS: "VARS"i -CONFIGS: "CONFIGS"i -ENVS: "ENVS"i -KEY: "KEY"i -KEYS: "KEYS"i -GENERATE: "GENERATE"i -MODEL: "MODEL"i -MODELS: "MODELS"i -PROVIDERS: "PROVIDERS"i -DEFAULT: "DEFAULT"i -CHATS: "CHATS"i - -list_services: LIST SERVICES ";" -show_service: SHOW SERVICE NUMBER ";" -startup_service: STARTUP SERVICE NUMBER ";" -shutdown_service: SHUTDOWN SERVICE NUMBER ";" -restart_service: RESTART SERVICE NUMBER ";" - -list_users: LIST USERS ";" -drop_user: DROP USER quoted_string ";" -alter_user: ALTER USER PASSWORD quoted_string quoted_string ";" -show_user: SHOW USER quoted_string ";" -create_user: CREATE USER quoted_string quoted_string ";" -activate_user: ALTER USER ACTIVE quoted_string status ";" - -list_datasets: LIST DATASETS OF quoted_string ";" -list_agents: LIST AGENTS OF quoted_string ";" - -create_role: CREATE ROLE identifier [DESCRIPTION quoted_string] ";" -drop_role: DROP ROLE identifier ";" -alter_role: ALTER ROLE identifier SET DESCRIPTION quoted_string ";" -list_roles: LIST ROLES ";" -show_role: SHOW ROLE identifier ";" - -grant_permission: GRANT action_list ON identifier TO ROLE identifier ";" -revoke_permission: REVOKE action_list ON identifier FROM ROLE identifier ";" -alter_user_role: ALTER USER quoted_string SET ROLE identifier ";" -show_user_permission: SHOW USER PERMISSION quoted_string ";" - -show_version: SHOW VERSION ";" - -grant_admin: GRANT ADMIN quoted_string ";" -revoke_admin: REVOKE ADMIN quoted_string ";" - -generate_key: GENERATE KEY FOR USER quoted_string ";" -list_keys: LIST KEYS OF quoted_string ";" -drop_key: DROP KEY quoted_string OF quoted_string ";" - -set_variable: SET VAR identifier identifier ";" -show_variable: SHOW VAR identifier ";" -list_variables: LIST VARS ";" -list_configs: LIST CONFIGS ";" -list_environments: LIST ENVS ";" - -list_user_datasets: LIST DATASETS ";" -list_user_agents: LIST AGENTS ";" -list_user_chats: LIST CHATS ";" -list_user_model_providers: LIST MODEL PROVIDERS ";" -list_user_default_models: LIST DEFAULT MODELS ";" - -action_list: identifier ("," identifier)* - -identifier: WORD -quoted_string: QUOTED_STRING -status: WORD - -QUOTED_STRING: /'[^']+'/ | /"[^"]+"/ -WORD: /[a-zA-Z0-9_\-\.]+/ -NUMBER: /[0-9]+/ - -%import common.WS -%ignore WS -""" - - -class RAGFlowCLITransformer(Transformer): - def start(self, items): - return items[0] - - def command(self, items): - return items[0] - - def list_services(self, items): - result = {"type": "list_services"} - return result - - def show_service(self, items): - service_id = int(items[2]) - return {"type": "show_service", "number": service_id} - - def startup_service(self, items): - service_id = int(items[2]) - return {"type": "startup_service", "number": service_id} - - def shutdown_service(self, items): - service_id = int(items[2]) - return {"type": "shutdown_service", "number": service_id} - - def restart_service(self, items): - service_id = int(items[2]) - return {"type": "restart_service", "number": service_id} - - def list_users(self, items): - return {"type": "list_users"} - - def show_user(self, items): - user_name = items[2] - return {"type": "show_user", "user_name": user_name} - - def drop_user(self, items): - user_name = items[2] - return {"type": "drop_user", "user_name": user_name} - - def alter_user(self, items): - user_name = items[3] - new_password = items[4] - return {"type": "alter_user", "user_name": user_name, "password": new_password} - - def create_user(self, items): - user_name = items[2] - password = items[3] - return {"type": "create_user", "user_name": user_name, "password": password, "role": "user"} - - def activate_user(self, items): - user_name = items[3] - activate_status = items[4] - return {"type": "activate_user", "activate_status": activate_status, "user_name": user_name} - - def list_datasets(self, items): - user_name = items[3] - return {"type": "list_datasets", "user_name": user_name} - - def list_agents(self, items): - user_name = items[3] - return {"type": "list_agents", "user_name": user_name} - - def create_role(self, items): - role_name = items[2] - if len(items) > 4: - description = items[4] - return {"type": "create_role", "role_name": role_name, "description": description} - else: - return {"type": "create_role", "role_name": role_name} - - def drop_role(self, items): - role_name = items[2] - return {"type": "drop_role", "role_name": role_name} - - def alter_role(self, items): - role_name = items[2] - description = items[5] - return {"type": "alter_role", "role_name": role_name, "description": description} - - def list_roles(self, items): - return {"type": "list_roles"} - - def show_role(self, items): - role_name = items[2] - return {"type": "show_role", "role_name": role_name} - - def grant_permission(self, items): - action_list = items[1] - resource = items[3] - role_name = items[6] - return {"type": "grant_permission", "role_name": role_name, "resource": resource, "actions": action_list} - - def revoke_permission(self, items): - action_list = items[1] - resource = items[3] - role_name = items[6] - return {"type": "revoke_permission", "role_name": role_name, "resource": resource, "actions": action_list} - - def alter_user_role(self, items): - user_name = items[2] - role_name = items[5] - return {"type": "alter_user_role", "user_name": user_name, "role_name": role_name} - - def show_user_permission(self, items): - user_name = items[3] - return {"type": "show_user_permission", "user_name": user_name} - - def show_version(self, items): - return {"type": "show_version"} - - def grant_admin(self, items): - user_name = items[2] - return {"type": "grant_admin", "user_name": user_name} - - def revoke_admin(self, items): - user_name = items[2] - return {"type": "revoke_admin", "user_name": user_name} - - def generate_key(self, items): - user_name = items[4] - return {"type": "generate_key", "user_name": user_name} - - def list_keys(self, items): - user_name = items[3] - return {"type": "list_keys", "user_name": user_name} - - def drop_key(self, items): - key = items[2] - user_name = items[4] - return {"type": "drop_key", "key": key, "user_name": user_name} - - def set_variable(self, items): - var_name = items[2] - var_value = items[3] - return {"type": "set_variable", "var_name": var_name, "var_value": var_value} - - def show_variable(self, items): - var_name = items[2] - return {"type": "show_variable", "var_name": var_name} - - def list_variables(self, items): - return {"type": "list_variables"} - - def list_configs(self, items): - return {"type": "list_configs"} - - def list_environments(self, items): - return {"type": "list_environments"} - - def list_user_datasets(self, items): - return {"type": "list_user_datasets"} - - def list_user_agents(self, items): - return {"type": "list_user_agents"} - - def list_user_chats(self, items): - return {"type": "list_user_chats"} - - def list_user_model_providers(self, items): - return {"type": "list_user_model_providers"} - - def list_user_default_models(self, items): - return {"type": "list_user_default_models"} - - def action_list(self, items): - return items - - def meta_command(self, items): - command_name = str(items[0]).lower() - args = items[1:] if len(items) > 1 else [] - - # handle quoted parameter - parsed_args = [] - for arg in args: - if hasattr(arg, "value"): - parsed_args.append(arg.value) - else: - parsed_args.append(str(arg)) - - return {"type": "meta", "command": command_name, "args": parsed_args} - - def meta_command_name(self, items): - return items[0] - - def meta_args(self, items): - return items +from lark import Lark, Tree +from parser import GRAMMAR, RAGFlowCLITransformer +from http_client import HttpClient +from ragflow_client import RAGFlowClient, run_command +from user import login_user +warnings.filterwarnings("ignore", category=getpass.GetPassWarning) def encrypt(input_string): pub = "-----BEGIN PUBLIC KEY-----\nMIIBIjANBgkqhkiG9w0BAQEFAAOCAQ8AMIIBCgKCAQEArq9XTUSeYr2+N1h3Afl/z8Dse/2yD0ZGrKwx+EEEcdsBLca9Ynmx3nIB5obmLlSfmskLpBo0UACBmB5rEjBp2Q2f3AG3Hjd4B+gNCG6BDaawuDlgANIhGnaTLrIqWrrcm4EMzJOnAOI1fgzJRsOOUEfaS318Eq9OVO3apEyCCt0lOQK6PuksduOjVxtltDav+guVAA068NrPYmRNabVKRNLJpL8w4D44sfth5RvZ3q9t+6RTArpEtc5sh5ChzvqPOzKGMXW83C95TxmXqpbK6olN4RevSfVjEAgCydH6HN6OhtOQEcnrU97r9H0iZOWwbw3pVrZiUkuRD1R56Wzs2wIDAQAB\n-----END PUBLIC KEY-----" @@ -388,44 +46,7 @@ def encode_to_base64(input_string): return base64_encoded.decode("utf-8") -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 -Meta Commands: -\\?, \\h, \\help Show this help -\\q, \\quit, \\exit Quit the CLI - """ - print(help_text) class RAGFlowCLI(Cmd): @@ -437,10 +58,10 @@ class RAGFlowCLI(Cmd): self.account = "admin@ragflow.io" self.account_password: str = "admin" self.session = requests.Session() - self.access_token: str = "" self.host: str = "" self.port: int = 0 self.mode: str = "admin" + self.ragflow_client = None intro = r"""Type "\h" for help.""" prompt = "ragflow> " @@ -486,22 +107,14 @@ class RAGFlowCLI(Cmd): except Exception as e: return {"type": "error", "message": f"Parse error: {str(e)}"} - def verify_auth(self, arguments: dict, single_command: bool): - self.host = arguments["host"] - self.port = arguments["port"] - # Determine mode and username - self.mode = arguments.get("type", "admin") - username = arguments.get("username", "admin@ragflow.io") - self.account = username - - # Set login endpoint based on mode - if self.mode == "admin": - url = f"http://{self.host}:{self.port}/api/v1/admin/login" - print("Attempt to access server for admin login") - else: # user mode - url = f"http://{self.host}:{self.port}/v1/user/login" - print("Attempt to access server for user login") + def verify_auth(self, arguments: dict, single_command: bool, auth: bool): + server_type = arguments.get("type", "admin") + http_client = HttpClient(arguments["host"], arguments["port"]) + if not auth: + self.ragflow_client = RAGFlowClient(http_client, server_type) + return True + user_name = arguments["username"] attempt_count = 3 if single_command: attempt_count = 1 @@ -513,27 +126,15 @@ class RAGFlowCLI(Cmd): return False if single_command: - account_passwd = arguments["password"] + user_password = arguments["password"] else: - account_passwd = getpass.getpass(f"password for {self.account}: ").strip() + user_password = getpass.getpass(f"password for {user_name}: ").strip() + try: - self.account_password = encrypt(account_passwd) - response = self.session.post(url, json={"email": self.account, "password": self.account_password}) - if response.status_code == 200: - res_json = response.json() - error_code = res_json.get("code", -1) - if error_code == 0: - self.session.headers.update( - {"Content-Type": "application/json", "Authorization": response.headers["Authorization"], - "User-Agent": "RAGFlow-CLI/0.23.1"}) - print("Authentication successful.") - return True - else: - error_message = res_json.get("message", "Unknown error") - print(f"Authentication failed: {error_message}, try again") - continue - else: - print(f"Bad response，status: {response.status_code}, password is wrong") + token = login_user(http_client, server_type, user_name, user_password) + http_client.login_token = token + self.ragflow_client = RAGFlowClient(http_client, server_type) + return True except Exception as e: print(str(e)) print("Can't access server for login (connection failed)") @@ -610,39 +211,31 @@ class RAGFlowCLI(Cmd): print(separator) - def run_interactive(self): + def run_interactive(self, args): self.is_interactive = True + if self.verify_auth(args, single_command=False, auth=args["auth"]): + print(r""" + ____ ___ ______________ ________ ____ + / __ \/ | / ____/ ____/ /___ _ __ / ____/ / / _/ + / /_/ / /| |/ / __/ /_ / / __ \ | /| / / / / / / / / + / _, _/ ___ / /_/ / __/ / / /_/ / |/ |/ / / /___/ /____/ / + /_/ |_/_/ |_\____/_/ /_/\____/|__/|__/ \____/_____/___/ + """) + self.cmdloop() + print("RAGFlow command line interface - Type '\\?' for help, '\\q' to quit") - while True: - try: - command = input("ragflow> ").strip() - if not command: - continue + def run_single_command(self, args): + self.is_interactive = False + if self.verify_auth(args, single_command=True, auth=args["auth"]): + command = args["command"] + result = self.parse_command(command) + self.execute_command(result) - print(f"command: {command}") - result = self.parse_command(command) - self.execute_command(result) - - if isinstance(result, Tree): - continue - - if result.get("type") == "meta" and result.get("command") in ["q", "quit", "exit"]: - break - - except KeyboardInterrupt: - print("\nUse '\\q' to quit") - except EOFError: - print("\nGoodbye!") - break - - def run_single_command(self, command: str): - result = self.parse_command(command) - self.execute_command(result) def parse_connection_args(self, args: List[str]) -> Dict[str, Any]: parser = argparse.ArgumentParser(description="RAGFlow CLI Client", add_help=False) - parser.add_argument("-h", "--host", default="localhost", help="Admin or RAGFlow service host") + parser.add_argument("-h", "--host", default="127.0.0.1", help="Admin or RAGFlow service host") parser.add_argument("-p", "--port", type=int, default=9381, help="Admin or RAGFlow service port") parser.add_argument("-w", "--password", default="admin", type=str, help="Superuser password") parser.add_argument("-t", "--type", default="admin", type=str, help="CLI mode, admin or user") @@ -656,27 +249,38 @@ class RAGFlowCLI(Cmd): if parsed_args.type == "admin": if username is None: username = "admin@ragflow.io" - else: # user mode + + if remaining_args: + if remaining_args[0] == "command": + command_str = ' '.join(remaining_args[1:]) + ';' + auth = True + if remaining_args[1] == "register": + auth = False + else: + if username is None: + print("Error: username (-u) is required in user mode") + return {"error": "Username required"} + return { + "host": parsed_args.host, + "port": parsed_args.port, + "password": parsed_args.password, + "type": parsed_args.type, + "username": username, + "command": command_str, + "auth": auth + } + else: + return {"error": "Invalid command"} + else: if username is None: print("Error: username (-u) is required in user mode") return {"error": "Username required"} - - if remaining_args: - command = remaining_args[0] - return { - "host": parsed_args.host, - "port": parsed_args.port, - "password": parsed_args.password, - "type": parsed_args.type, - "username": username, - "command": command - } - else: return { "host": parsed_args.host, "port": parsed_args.port, "type": parsed_args.type, "username": username, + "auth": True } except SystemExit: return {"error": "Invalid connection arguments"} @@ -693,690 +297,9 @@ class RAGFlowCLI(Cmd): command_dict = parsed_command # print(f"Parsed command: {command_dict}") - - command_type = command_dict["type"] - - match command_type: - case "list_services": - self._handle_list_services(command_dict) - case "show_service": - self._handle_show_service(command_dict) - case "restart_service": - self._handle_restart_service(command_dict) - case "shutdown_service": - self._handle_shutdown_service(command_dict) - case "startup_service": - self._handle_startup_service(command_dict) - case "list_users": - self._handle_list_users(command_dict) - case "show_user": - self._handle_show_user(command_dict) - case "drop_user": - self._handle_drop_user(command_dict) - case "alter_user": - self._handle_alter_user(command_dict) - case "create_user": - self._handle_create_user(command_dict) - case "activate_user": - self._handle_activate_user(command_dict) - case "list_datasets": - self._handle_list_datasets(command_dict) - case "list_agents": - self._handle_list_agents(command_dict) - case "create_role": - self._create_role(command_dict) - case "drop_role": - self._drop_role(command_dict) - case "alter_role": - self._alter_role(command_dict) - case "list_roles": - self._list_roles(command_dict) - case "show_role": - self._show_role(command_dict) - case "grant_permission": - self._grant_permission(command_dict) - case "revoke_permission": - self._revoke_permission(command_dict) - case "alter_user_role": - self._alter_user_role(command_dict) - case "show_user_permission": - self._show_user_permission(command_dict) - case "show_version": - self._show_version(command_dict) - case "grant_admin": - self._grant_admin(command_dict) - case "revoke_admin": - self._revoke_admin(command_dict) - case "generate_key": - self._generate_key(command_dict) - case "list_keys": - self._list_keys(command_dict) - case "drop_key": - self._drop_key(command_dict) - case "set_variable": - self._set_variable(command_dict) - case "show_variable": - self._show_variable(command_dict) - case "list_variables": - self._list_variables(command_dict) - case "list_configs": - self._list_configs(command_dict) - case "list_environments": - self._list_environments(command_dict) - case "list_user_datasets": - self._list_user_datasets(command_dict) - case "list_user_agents": - self._list_user_agents(command_dict) - case "list_user_chats": - self._list_user_chats(command_dict) - case "list_user_model_providers": - self._list_user_model_providers(command_dict) - case "list_user_default_models": - self._list_user_default_models(command_dict) - case "meta": - self._handle_meta_command(command_dict) - case _: - print(f"Command '{command_type}' would be executed with API") - - def _handle_list_services(self, command): - if self.mode != "admin": - print("This command is only allowed in ADMIN mode") - - url = f"http://{self.host}:{self.port}/api/v1/admin/services" - response = self.session.get(url) - 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']}") - - def _handle_show_service(self, command): - if self.mode != "admin": - print("This command is only allowed in ADMIN mode") - - service_id: int = command["number"] - - url = f"http://{self.host}:{self.port}/api/v1/admin/services/{service_id}" - response = self.session.get(url) - 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, ") - if isinstance(res_data["message"], str): - print(res_data["message"]) - else: - data = self._format_service_detail_table(res_data["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 _handle_restart_service(self, command): - if self.mode != "admin": - print("This command is only allowed in ADMIN mode") - - service_id: int = command["number"] - print(f"Restart service {service_id}") - - def _handle_shutdown_service(self, command): - if self.mode != "admin": - print("This command is only allowed in ADMIN mode") - - service_id: int = command["number"] - print(f"Shutdown service {service_id}") - - def _handle_startup_service(self, command): - if self.mode != "admin": - print("This command is only allowed in ADMIN mode") - - service_id: int = command["number"] - print(f"Startup service {service_id}") - - def _handle_list_users(self, command): - if self.mode != "admin": - print("This command is only allowed in ADMIN mode") - - url = f"http://{self.host}:{self.port}/api/v1/admin/users" - response = self.session.get(url) - 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 _handle_show_user(self, command): - if self.mode != "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}") - url = f"http://{self.host}:{self.port}/api/v1/admin/users/{user_name}" - response = self.session.get(url) - 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 _handle_drop_user(self, command): - if self.mode != "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}") - url = f"http://{self.host}:{self.port}/api/v1/admin/users/{user_name}" - response = self.session.delete(url) - 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 _handle_alter_user(self, command): - if self.mode != "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: ******") - url = f"http://{self.host}:{self.port}/api/v1/admin/users/{user_name}/password" - response = self.session.put(url, json={"new_password": encrypt(password)}) - 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 _handle_create_user(self, command): - if self.mode != "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}") - url = f"http://{self.host}:{self.port}/api/v1/admin/users" - response = self.session.post(url, json={"user_name": user_name, "password": encrypt(password), "role": role}) - 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 _handle_activate_user(self, command): - if self.mode != "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()}.") - url = f"http://{self.host}:{self.port}/api/v1/admin/users/{user_name}/activate" - response = self.session.put(url, json={"activate_status": activate_status}) - 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.mode != "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("'\"") - url = f"http://{self.host}:{self.port}/api/v1/admin/users/{user_name}/admin" - # print(f"Grant admin: {url}") - # return - response = self.session.put(url) - 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.mode != "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("'\"") - url = f"http://{self.host}:{self.port}/api/v1/admin/users/{user_name}/admin" - # print(f"Revoke admin: {url}") - # return - response = self.session.delete(url) - 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 _generate_key(self, command: dict[str, Any]) -> None: - username_tree: Tree = command["user_name"] - user_name: str = username_tree.children[0].strip("'\"") - print(f"Generating API key for user: {user_name}") - url: str = f"http://{self.host}:{self.port}/api/v1/admin/users/{user_name}/new_token" - response: requests.Response = self.session.post(url) - 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: - username_tree: Tree = command["user_name"] - user_name: str = username_tree.children[0].strip("'\"") - print(f"Listing API keys for user: {user_name}") - url: str = f"http://{self.host}:{self.port}/api/v1/admin/users/{user_name}/token_list" - response: requests.Response = self.session.get(url) - 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: - 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="") - url: str = f"http://{self.host}:{self.port}/api/v1/admin/users/{user_name}/token/{encoded_key}" - response: requests.Response = self.session.delete(url) - 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.mode != "admin": - print("This command is only allowed in ADMIN mode") - - var_name_tree: Tree = command["var_name"] - var_name = var_name_tree.children[0].strip("'\"") - var_value_tree: Tree = command["var_value"] - var_value = var_value_tree.children[0].strip("'\"") - url = f"http://{self.host}:{self.port}/api/v1/admin/variables" - response = self.session.put(url, json={"var_name": var_name, "var_value": var_value}) - 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.mode != "admin": - print("This command is only allowed in ADMIN mode") - - var_name_tree: Tree = command["var_name"] - var_name = var_name_tree.children[0].strip("'\"") - url = f"http://{self.host}:{self.port}/api/v1/admin/variables" - response = self.session.get(url, json={"var_name": var_name}) - 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.mode != "admin": - print("This command is only allowed in ADMIN mode") - - url = f"http://{self.host}:{self.port}/api/v1/admin/variables" - response = self.session.get(url) - 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.mode != "admin": - print("This command is only allowed in ADMIN mode") - - url = f"http://{self.host}:{self.port}/api/v1/admin/configs" - response = self.session.get(url) - 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.mode != "admin": - print("This command is only allowed in ADMIN mode") - - url = f"http://{self.host}:{self.port}/api/v1/admin/environments" - response = self.session.get(url) - 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 _handle_list_datasets(self, command): - if self.mode != "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}") - url = f"http://{self.host}:{self.port}/api/v1/admin/users/{user_name}/datasets" - response = self.session.get(url) - 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.mode != "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}") - url = f"http://{self.host}:{self.port}/api/v1/admin/users/{user_name}/agents" - response = self.session.get(url) - 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 _create_role(self, command): - if self.mode != "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: - desc_tree: Tree = command["description"] - desc_str = desc_tree.children[0].strip("'\"") - - print(f"create role name: {role_name}, description: {desc_str}") - url = f"http://{self.host}:{self.port}/api/v1/admin/roles" - response = self.session.post(url, json={"role_name": role_name, "description": desc_str}) - 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.mode != "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}") - url = f"http://{self.host}:{self.port}/api/v1/admin/roles/{role_name}" - response = self.session.delete(url) - 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.mode != "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}") - url = f"http://{self.host}:{self.port}/api/v1/admin/roles/{role_name}" - response = self.session.put(url, json={"description": desc_str}) - 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.mode != "admin": - print("This command is only allowed in ADMIN mode") - - url = f"http://{self.host}:{self.port}/api/v1/admin/roles" - response = self.session.get(url) - 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.mode != "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}") - url = f"http://{self.host}:{self.port}/api/v1/admin/roles/{role_name}/permission" - response = self.session.get(url) - 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.mode != "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}") - url = f"http://{self.host}:{self.port}/api/v1/admin/roles/{role_name_str}/permission" - response = self.session.post(url, json={"actions": actions, "resource": resource_str}) - 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.mode != "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}") - url = f"http://{self.host}:{self.port}/api/v1/admin/roles/{role_name_str}/permission" - response = self.session.delete(url, json={"actions": actions, "resource": resource_str}) - 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.mode != "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}") - url = f"http://{self.host}:{self.port}/api/v1/admin/users/{user_name_str}/role" - response = self.session.put(url, json={"role_name": role_name_str}) - 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.mode != "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}") - url = f"http://{self.host}:{self.port}/api/v1/admin/users/{user_name_str}/permission" - response = self.session.get(url) - 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 _show_version(self, command): - if self.mode == "admin": - url = f"http://{self.host}:{self.port}/api/v1/admin/version" - else: - url = f"http://{self.host}:{self.port}/v1/system/version" - - response = self.session.get(url) - res_json = response.json() - if response.status_code == 200: - if self.mode == "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 _list_user_datasets(self, command): - if self.mode != "user": - print("This command is only allowed in USER mode") - - url = f"http://{self.host}:{self.port}/v1/kb/list" - response = self.session.post(url) - 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_agents(self, command): - if self.mode != "user": - print("This command is only allowed in USER mode") - - url = f"http://{self.host}:{self.port}/v1/canvas/list" - response = self.session.get(url) - 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.mode != "user": - print("This command is only allowed in USER mode") - - url = f"http://{self.host}:{self.port}/v1/dialog/next" - response = self.session.get(url) - 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_model_providers(self, command): - if self.mode != "user": - print("This command is only allowed in USER mode") - - url = f"http://{self.host}:{self.port}/v1/llm/my_llms" - response = self.session.get(url) - 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) - - def _list_user_default_models(self, command): - if self.mode != "user": - print("This command is only allowed in USER mode") - - url = f"http://{self.host}:{self.port}/v1/user/tenant_info" - response = self.session.get(url) - res_json = response.json() - if response.status_code == 200: - new_input = [] - for key, value in res_json["data"].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 _handle_meta_command(self, command): - 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}") - + run_command(self.ragflow_client, command_dict, self.is_interactive) def main(): - import sys cli = RAGFlowCLI() @@ -1386,23 +309,16 @@ def main(): return if "command" in args: + # single command mode + # for user mode, api key or password is ok + # for admin mode, only password if "password" not in args: print("Error: password is missing") return - if cli.verify_auth(args, single_command=True): - command: str = args["command"] - # print(f"Run single command: {command}") - cli.run_single_command(command) + + cli.run_single_command(args) else: - if cli.verify_auth(args, single_command=False): - print(r""" - ____ ___ ______________ ________ ____ - / __ \/ | / ____/ ____/ /___ _ __ / ____/ / / _/ - / /_/ / /| |/ / __/ /_ / / __ \ | /| / / / / / / / / - / _, _/ ___ / /_/ / __/ / / /_/ / |/ |/ / / /___/ /____/ / - /_/ |_/_/ |_\____/_/ /_/\____/|__/|__/ \____/_____/___/ - """) - cli.cmdloop() + cli.run_interactive(args) if __name__ == "__main__": diff --git a/admin/client/ragflow_client.py b/admin/client/ragflow_client.py new file mode 100644 index 000000000..6830b279a --- /dev/null +++ b/admin/client/ragflow_client.py @@ -0,0 +1,1453 @@ +# +# 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 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 + +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") + + +class RAGFlowClient: + def __init__(self, http_client: HttpClient, server_type: str): + self.http_client = http_client + self.server_type = server_type + + 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_tree: Tree = command["var_name"] + var_name = var_name_tree.children[0].strip("'\"") + var_value_tree: Tree = command["var_value"] + var_value = var_value_tree.children[0].strip("'\"") + 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_tree: Tree = command["var_name"] + var_name = var_name_tree.children[0].strip("'\"") + 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 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("POST", "/dialog/next", use_api_base=False, auth_kind="web", + iterations=iterations) + return response + else: + response = self.http_client.request("POST", "/kb/list", use_api_base=False, auth_kind="web") + res_json = response.json() + if response.status_code == 200: + self._print_table_simple(res_json["data"]["kbs"]) + 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"], + "embd_id": command["embedding"] + } + if "parser_id" in command: + payload["parser_id"] = command["parser"] + if "pipeline" in command: + payload["pipeline_id"] = command["pipeline"] + response = self.http_client.request("POST", "/kb/create", json_body=payload, 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 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 = {"kb_id": dataset_id} + response = self.http_client.request("POST", "/kb/rm", json_body=payload, use_api_base=False, 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_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") + ''' + description + : + "" + icon + : + "" + language + : + "English" + llm_id + : + "glm-4-flash@ZHIPU-AI" + llm_setting + : + {} + name + : + "xx" + prompt_config + : + {empty_response: "", prologue: "Hi! I'm your assistant. What can I do for you?", quote: true,…} + empty_response + : + "" + keyword + : + false + parameters + : + [{key: "knowledge", optional: false}] + prologue + : + "Hi! I'm your assistant. What can I do for you?" + quote + : + true + reasoning + : + false + refine_multiturn + : + 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." + toc_enhance + : + false + tts + : + false + use_kg + : + false + similarity_threshold + : + 0.2 + top_n + : + 8 + vector_similarity_weight + : + 0.3 + ''' + chat_name = command["chat_name"] + payload = { + "description": "", + "icon": "", + "language": "English", + "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, + "vector_similarity_weight": 0.3 + } + + payload.update({"name": chat_name}) + response = self.http_client.request("POST", "/dialog/set", 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 create chat: {chat_name}") + else: + print(f"Fail to create chat {chat_name}, code: {res_json['code']}, message: {res_json['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 = {"dialog_ids": to_drop_chat_ids} + response = self.http_client.request("POST", "/dialog/rm", 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 chat: {chat_name}") + else: + print(f"Fail to drop chat {chat_name}, code: {res_json['code']}, message: {res_json['message']}") + + 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", "/document/run", 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 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", "/document/run", json_body=payload, use_api_base=False, + 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", + "/document/upload", + headers=headers, + data=encoder, + json_body=None, + params=None, + stream=False, + auth_kind="web", + use_api_base=False + ) + 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 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=False, 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") != "3": + 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): + response = self.http_client.request("POST", f"/document/list?kb_id={dataset_id}", use_api_base=False, + 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("POST", "/kb/list", use_api_base=False, 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"]["kbs"] + 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("POST", "/dialog/next", use_api_base=False, auth_kind="web", + iterations=iterations) + return response + else: + response = self.http_client.request("POST", "/dialog/next", use_api_base=False, auth_kind="web", + iterations=iterations) + res_json = response.json() + if response.status_code == 200 and res_json["code"] == 0: + return res_json["data"]["dialogs"] + else: + print(f"Fail to list datasets, 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 and 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 + + 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 run_command(client: RAGFlowClient, command_dict: dict, is_interactive: bool): + command_type = command_dict["type"] + + match command_type: + case "benchmark": + run_benchmark(client, command_dict, is_interactive) + case "register_user": + if is_interactive: + print("Register user command is not supported in interactive mode") + return + 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 "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_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 "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 "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 + +Meta Commands: +\\?, \\h, \\help Show this help +\\q, \\quit, \\exit Quit the CLI + """ + print(help_text) + + +def run_benchmark(client: RAGFlowClient, command_dict: dict, is_interactive: bool): + concurrency = command_dict.get("concurrency", 1) + iterations = command_dict.get("iterations", 1) + command: dict = command_dict["command"] + command.update({"iterations": iterations}) + if concurrency < 1: + print("Concurrency must be greater than 0") + return + elif concurrency == 1: + result = run_command(client, command, is_interactive) + response_list = result["response_list"] + total_duration = result["duration"] + success_count = 0 + for response in response_list: + res_json = response.json() + if response.status_code == 200 and res_json["code"] == 0: + success_count += 1 + 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, + is_interactive + ): 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: + 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 diff --git a/admin/client/user.py b/admin/client/user.py new file mode 100644 index 000000000..823e2a130 --- /dev/null +++ b/admin/client/user.py @@ -0,0 +1,65 @@ +# +# Copyright 2025 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. +# + +from http_client import HttpClient + + +class AuthException(Exception): + def __init__(self, message, code=401): + super().__init__(message) + self.code = code + self.message = message + + +def encrypt_password(password_plain: str) -> str: + try: + from api.utils.crypt import crypt + except Exception as exc: + raise AuthException( + "Password encryption unavailable; install pycryptodomex (uv sync --python 3.12 --group test)." + ) from exc + return crypt(password_plain) + + +def register_user(client: HttpClient, email: str, nickname: str, password: str) -> None: + password_enc = encrypt_password(password) + payload = {"email": email, "nickname": nickname, "password": password_enc} + res = client.request_json("POST", "/user/register", use_api_base=False, auth_kind=None, json_body=payload) + if res.get("code") == 0: + return + msg = res.get("message", "") + if "has already registered" in msg: + return + raise AuthException(f"Register failed: {msg}") + + +def login_user(client: HttpClient, server_type: str, email: str, password: str) -> str: + password_enc = encrypt_password(password) + payload = {"email": email, "password": password_enc} + if server_type == "admin": + response = client.request("POST", "/admin/login", use_api_base=True, auth_kind=None, json_body=payload) + else: + response = client.request("POST", "/user/login", use_api_base=False, auth_kind=None, json_body=payload) + try: + res = response.json() + except Exception as exc: + raise AuthException(f"Login failed: invalid JSON response ({exc})") from exc + if res.get("code") != 0: + raise AuthException(f"Login failed: {res.get('message')}") + token = response.headers.get("Authorization") + if not token: + raise AuthException("Login failed: missing Authorization header") + return token diff --git a/admin/server/routes.py b/admin/server/routes.py index d7d3e53f8..c3eacc009 100644 --- a/admin/server/routes.py +++ b/admin/server/routes.py @@ -484,7 +484,7 @@ def get_environments(): return error_response(str(e), 500) -@admin_bp.route("/users//new_token", methods=["POST"]) +@admin_bp.route("/users//keys", methods=["POST"]) @login_required @check_admin_auth def generate_user_api_key(username: str) -> tuple[Response, int]: @@ -496,10 +496,10 @@ def generate_user_api_key(username: str) -> tuple[Response, int]: if not tenants: return error_response("Tenant not found!", 404) tenant_id: str = tenants[0]["tenant_id"] - token: str = generate_confirmation_token() + key: str = generate_confirmation_token() obj: dict[str, Any] = { "tenant_id": tenant_id, - "token": token, + "token": key, "beta": generate_confirmation_token().replace("ragflow-", "")[:32], "create_time": current_timestamp(), "create_date": datetime_format(datetime.now()), @@ -507,7 +507,7 @@ def generate_user_api_key(username: str) -> tuple[Response, int]: "update_date": None, } - if not UserMgr.save_api_token(obj): + if not UserMgr.save_api_key(obj): return error_response("Failed to generate API key!", 500) return success_response(obj, "API key generated successfully") except AdminException as e: @@ -516,7 +516,7 @@ def generate_user_api_key(username: str) -> tuple[Response, int]: return error_response(str(e), 500) -@admin_bp.route("/users//token_list", methods=["GET"]) +@admin_bp.route("/users//keys", methods=["GET"]) @login_required @check_admin_auth def get_user_api_keys(username: str) -> tuple[Response, int]: @@ -529,12 +529,12 @@ def get_user_api_keys(username: str) -> tuple[Response, int]: return error_response(str(e), 500) -@admin_bp.route("/users//token/", methods=["DELETE"]) +@admin_bp.route("/users//keys/", methods=["DELETE"]) @login_required @check_admin_auth -def delete_user_api_key(username: str, token: str) -> tuple[Response, int]: +def delete_user_api_key(username: str, key: str) -> tuple[Response, int]: try: - deleted = UserMgr.delete_api_token(username, token) + deleted = UserMgr.delete_api_key(username, key) if deleted: return success_response(None, "API key deleted successfully") else: diff --git a/admin/server/services.py b/admin/server/services.py index 8b4e23476..d44361eeb 100644 --- a/admin/server/services.py +++ b/admin/server/services.py @@ -159,21 +159,21 @@ class UserMgr: # tenant_id is typically the same as user_id for the owner tenant tenant_id: str = usr.id - # Query all API tokens for this tenant - api_tokens: Any = APITokenService.query(tenant_id=tenant_id) + # Query all API keys for this tenant + api_keys: Any = APITokenService.query(tenant_id=tenant_id) result: list[dict[str, Any]] = [] - for token_obj in api_tokens: - result.append(token_obj.to_dict()) + for key in api_keys: + result.append(key.to_dict()) return result @staticmethod - def save_api_token(api_token: dict[str, Any]) -> bool: - return APITokenService.save(**api_token) + def save_api_key(api_key: dict[str, Any]) -> bool: + return APITokenService.save(**api_key) @staticmethod - def delete_api_token(username: str, token: str) -> bool: + def delete_api_key(username: str, key: str) -> bool: # use email to find user. check exist and unique. user_list: list[Any] = UserService.query_user_by_email(username) if not user_list: @@ -185,8 +185,8 @@ class UserMgr: # tenant_id is typically the same as user_id for the owner tenant tenant_id: str = usr.id - # Delete the API token - deleted_count: int = APITokenService.filter_delete([APIToken.tenant_id == tenant_id, APIToken.token == token]) + # Delete the API key + deleted_count: int = APITokenService.filter_delete([APIToken.tenant_id == tenant_id, APIToken.token == key]) return deleted_count > 0 @staticmethod @@ -305,6 +305,13 @@ class ServiceMgr: raise AdminException(f"invalid service_index: {service_idx}") service_config = configs[service_idx] + + # exclude retrieval service if retrieval_type is not matched + doc_engine = os.getenv("DOC_ENGINE", "elasticsearch") + if service_config.service_type == "retrieval": + if service_config.retrieval_type != doc_engine: + raise AdminException(f"invalid service_index: {service_idx}") + service_info = {"name": service_config.name, "detail_func_name": service_config.detail_func_name} detail_func = getattr(health_utils, service_info.get("detail_func_name"))