infiniflow/ragflow
1
0
Fork 0
mirror of https://github.com/infiniflow/ragflow.git synced 2026-10-11 21:39:54 +00:00
Code Issues Packages Projects Releases Wiki Activity

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

Browse Source
This commit is contained in:
Wang Qi
2026-07-03 12:53:39 +08:00
committed by GitHub
parent 19fcb4a981
commit 6a4b9be426
588 changed files with 11123 additions and 15412 deletions
Download Patch File Download Diff File Expand all files Collapse all files
admin/client/http_client.py
+33 -33
View File
@@ -25,14 +25,14 @@ 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,
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
@@ -71,19 +71,19 @@ class HttpClient:
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,
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)
@@ -144,18 +144,18 @@ class HttpClient:
# )
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,
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,
admin/client/parser.py
+19 -21
View File
@@ -336,8 +336,8 @@ 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 ";"
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_dataset_documents: LIST DOCUMENTS OF DATASET quoted_string ";"
@@ -640,15 +640,13 @@ class RAGFlowCLITransformer(Transformer):
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}
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}
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("'\"")
@@ -666,7 +664,7 @@ class RAGFlowCLITransformer(Transformer):
dataset_names = []
dataset_names.append(items[4].children[0].strip("'\""))
for i in range(5, len(items)):
if items[i] and hasattr(items[i], 'children') and items[i].children:
if items[i] and hasattr(items[i], "children") and items[i].children:
dataset_names.append(items[i].children[0].strip("'\""))
return {"type": "list_user_datasets_metadata", "dataset_names": dataset_names}
@@ -675,7 +673,7 @@ class RAGFlowCLITransformer(Transformer):
doc_ids = []
if len(items) > 6 and items[6] == "DOCUMENTS":
for i in range(7, len(items)):
if items[i] and hasattr(items[i], 'children') and items[i].children:
if items[i] and hasattr(items[i], "children") and items[i].children:
doc_id = items[i].children[0].strip("'\"")
doc_ids.append(doc_id)
return {"type": "list_user_documents_metadata_summary", "dataset_name": dataset_name, "document_ids": doc_ids}
@@ -698,17 +696,17 @@ class RAGFlowCLITransformer(Transformer):
dataset_name = None
vector_size = None
for i, item in enumerate(items):
if hasattr(item, 'data') and item.data == 'quoted_string':
if hasattr(item, "data") and item.data == "quoted_string":
dataset_name = item.children[0].strip("'\"")
if hasattr(item, 'type') and item.type == 'NUMBER':
if i > 0 and items[i-1].type == 'SIZE' and items[i-2].type == 'VECTOR':
if hasattr(item, "type") and item.type == "NUMBER":
if i > 0 and items[i - 1].type == "SIZE" and items[i - 2].type == "VECTOR":
vector_size = int(item)
return {"type": "create_dataset_table", "dataset_name": dataset_name, "vector_size": vector_size}
def drop_dataset_table(self, items):
dataset_name = None
for item in items:
if hasattr(item, 'data') and item.data == 'quoted_string':
if hasattr(item, "data") and item.data == "quoted_string":
dataset_name = item.children[0].strip("'\"")
return {"type": "drop_dataset_table", "dataset_name": dataset_name}
@@ -792,7 +790,7 @@ class RAGFlowCLITransformer(Transformer):
def update_chunk(self, items):
def get_quoted_value(item):
if hasattr(item, 'children') and item.children:
if hasattr(item, "children") and item.children:
return item.children[0].strip("'\"")
return str(item).strip("'\"")
@@ -813,16 +811,16 @@ class RAGFlowCLITransformer(Transformer):
for i in range(2, len(items)):
item = items[i]
# Check for FROM token to stop
if hasattr(item, 'type') and item.type == 'FROM':
if hasattr(item, "type") and item.type == "FROM":
break
if hasattr(item, 'children') and item.children:
if hasattr(item, "children") and item.children:
tag = item.children[0].strip("'\"")
tags.append(tag)
# Find dataset_name: quoted_string after DATASET
dataset_name = None
for i, item in enumerate(items):
# Check if item is a DATASET token
if hasattr(item, 'type') and item.type == 'DATASET':
if hasattr(item, "type") and item.type == "DATASET":
# Next item should be quoted_string
dataset_name = items[i + 1].children[0].strip("'\"")
break
@@ -835,10 +833,10 @@ class RAGFlowCLITransformer(Transformer):
# Check if it's "REMOVE ALL CHUNKS"
for item in items:
if hasattr(item, 'type') and item.type == 'ALL':
if hasattr(item, "type") and item.type == "ALL":
# Find doc_id
for j, inner_item in enumerate(items):
if hasattr(inner_item, 'type') and inner_item.type == 'DOCUMENT':
if hasattr(inner_item, "type") and inner_item.type == "DOCUMENT":
doc_id = items[j + 1].children[0].strip("'\"")
return {"type": "remove_chunks", "doc_id": doc_id, "delete_all": True}
@@ -846,12 +844,12 @@ class RAGFlowCLITransformer(Transformer):
chunk_ids = []
doc_id = None
for i, item in enumerate(items):
if hasattr(item, 'type') and item.type == 'DOCUMENT':
if hasattr(item, "type") and item.type == "DOCUMENT":
doc_id = items[i + 1].children[0].strip("'\"")
elif hasattr(item, 'children') and item.children:
elif hasattr(item, "children") and item.children:
val = item.children[0].strip("'\"")
# Skip if it's "FROM" or "DOCUMENT"
if val.upper() in ['FROM', 'DOCUMENT']:
if val.upper() in ["FROM", "DOCUMENT"]:
continue
chunk_ids.append(val)
admin/client/ragflow_cli.py
+9 -26
View File
@@ -36,6 +36,7 @@ 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-----"
pub_key = RSA.importKey(pub)
@@ -49,9 +50,6 @@ def encode_to_base64(input_string):
return base64_encoded.decode("utf-8")
class RAGFlowCLI(Cmd):
def __init__(self):
super().__init__()
@@ -240,9 +238,9 @@ class RAGFlowCLI(Cmd):
print(r"""
____ ___ ______________ ________ ____
/ __ \/ | / ____/ ____/ /___ _ __ / ____/ / / _/
/ /_/ / /| |/ / __/ /_ / / __ \ | /| / / / / / / / /
/ _, _/ ___ / /_/ / __/ / / /_/ / |/ |/ / / /___/ /____/ /
/_/ |_/_/ |_\____/_/ /_/\____/|__/|__/ \____/_____/___/
/ /_/ / /| |/ / __/ /_ / / __ \ | /| / / / / / / / /
/ _, _/ ___ / /_/ / __/ / / /_/ / |/ |/ / / /___/ /____/ /
/_/ |_/_/ |_\____/_/ /_/\____/|__/|__/ \____/_____/___/
""")
self.cmdloop()
@@ -254,15 +252,13 @@ class RAGFlowCLI(Cmd):
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="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")
parser.add_argument("-u", "--username", default=None,
help="Username (email). In admin mode defaults to [email protected], in user mode required.")
parser.add_argument("-u", "--username", default=None, help="Username (email). In admin mode defaults to [email protected], in user mode required.")
parser.add_argument("command", nargs="?", help="Single command")
try:
parsed_args, remaining_args = parser.parse_known_args(args)
@@ -274,7 +270,7 @@ class RAGFlowCLI(Cmd):
if remaining_args:
if remaining_args[0] == "command":
command_str = ' '.join(remaining_args[1:]) + ';'
command_str = " ".join(remaining_args[1:]) + ";"
auth = True
if remaining_args[1] == "register":
auth = False
@@ -282,28 +278,14 @@ class RAGFlowCLI(Cmd):
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
}
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:
auth = True
if username is None:
auth = False
return {
"host": parsed_args.host,
"port": parsed_args.port,
"type": parsed_args.type,
"username": username,
"auth": auth
}
return {"host": parsed_args.host, "port": parsed_args.port, "type": parsed_args.type, "username": username, "auth": auth}
except SystemExit:
return {"error": "Invalid connection arguments"}
@@ -321,6 +303,7 @@ class RAGFlowCLI(Cmd):
# print(f"Parsed command: {command_dict}")
run_command(self.ragflow_client, command_dict)
def main():
cli = RAGFlowCLI()
admin/client/ragflow_client.py
+91 -229
View File
@@ -71,6 +71,7 @@ class RAGFlowClient:
user_password: str = command.get("password")
if not user_password:
import getpass
user_password = getpass.getpass("Password: ")
try:
token = login_user(self.http_client, self.server_type, email, user_password)
@@ -86,8 +87,7 @@ class RAGFlowClient:
def ping_server(self, command):
iterations = command.get("iterations", 1)
if iterations > 1:
response = self.http_client.request("GET", "/system/ping", use_api_base=True, auth_kind="web",
iterations=iterations)
response = self.http_client.request("GET", "/system/ping", use_api_base=True, auth_kind="web", iterations=iterations)
return response
else:
response = self.http_client.request("GET", "/system/ping", use_api_base=True, auth_kind="web")
@@ -106,8 +106,7 @@ class RAGFlowClient:
enc_password = encrypt_password(password)
print(f"Register user: {nickname}, email: {username}, password: ******")
payload = {"email": username, "nickname": nickname, "password": enc_password}
response = self.http_client.request(method="POST", path="/users",
json_body=payload, use_api_base=True, auth_kind="web")
response = self.http_client.request(method="POST", path="/users", json_body=payload, use_api_base=True, auth_kind="web")
res_json = response.json()
if response.status_code == 200:
if res_json["code"] == 0:
@@ -135,8 +134,7 @@ class RAGFlowClient:
service_id: int = command["number"]
response = self.http_client.request("GET", f"/admin/services/{service_id}", use_api_base=True,
auth_kind="admin")
response = self.http_client.request("GET", f"/admin/services/{service_id}", use_api_base=True, auth_kind="admin")
res_json = response.json()
if response.status_code == 200:
res_data = res_json["data"]
@@ -226,9 +224,7 @@ class RAGFlowClient:
password_tree: Tree = command["password"]
password: str = password_tree.children[0].strip("'\"")
print(f"Alter user: {user_name}, password: ******")
response = self.http_client.request("PUT", f"/admin/users/{user_name}/password",
json_body={"new_password": encrypt_password(password)}, use_api_base=True,
auth_kind="admin")
response = self.http_client.request("PUT", f"/admin/users/{user_name}/password", json_body={"new_password": encrypt_password(password)}, use_api_base=True, auth_kind="admin")
res_json = response.json()
if response.status_code == 200:
print(res_json["message"])
@@ -247,9 +243,7 @@ class RAGFlowClient:
print(f"Create user: {user_name}, password: ******, role: {role}")
# enpass1 = encrypt(password)
enc_password = encrypt_password(password)
response = self.http_client.request(method="POST", path="/admin/users",
json_body={"username": user_name, "password": enc_password, "role": role},
use_api_base=True, auth_kind="admin")
response = self.http_client.request(method="POST", path="/admin/users", json_body={"username": user_name, "password": enc_password, "role": role}, use_api_base=True, auth_kind="admin")
res_json = response.json()
if response.status_code == 200:
self._print_table_simple(res_json["data"])
@@ -266,9 +260,7 @@ class RAGFlowClient:
activate_status: str = activate_tree.children[0].strip("'\"")
if activate_status.lower() in ["on", "off"]:
print(f"Alter user {user_name} activate status, turn {activate_status.lower()}.")
response = self.http_client.request("PUT", f"/admin/users/{user_name}/activate",
json_body={"activate_status": activate_status}, use_api_base=True,
auth_kind="admin")
response = self.http_client.request("PUT", f"/admin/users/{user_name}/activate", json_body={"activate_status": activate_status}, use_api_base=True, auth_kind="admin")
res_json = response.json()
if response.status_code == 200:
print(res_json["message"])
@@ -283,14 +275,12 @@ class RAGFlowClient:
user_name_tree: Tree = command["user_name"]
user_name: str = user_name_tree.children[0].strip("'\"")
response = self.http_client.request("PUT", f"/admin/users/{user_name}/admin", use_api_base=True,
auth_kind="admin")
response = self.http_client.request("PUT", f"/admin/users/{user_name}/admin", use_api_base=True, auth_kind="admin")
res_json = response.json()
if response.status_code == 200:
print(res_json["message"])
else:
print(
f"Fail to grant {user_name} admin authorization, code: {res_json['code']}, message: {res_json['message']}")
print(f"Fail to grant {user_name} admin authorization, code: {res_json['code']}, message: {res_json['message']}")
def revoke_admin(self, command):
if self.server_type != "admin":
@@ -298,14 +288,12 @@ class RAGFlowClient:
user_name_tree: Tree = command["user_name"]
user_name: str = user_name_tree.children[0].strip("'\"")
response = self.http_client.request("DELETE", f"/admin/users/{user_name}/admin", use_api_base=True,
auth_kind="admin")
response = self.http_client.request("DELETE", f"/admin/users/{user_name}/admin", use_api_base=True, auth_kind="admin")
res_json = response.json()
if response.status_code == 200:
print(res_json["message"])
else:
print(
f"Fail to revoke {user_name} admin authorization, code: {res_json['code']}, message: {res_json['message']}")
print(f"Fail to revoke {user_name} admin authorization, code: {res_json['code']}, message: {res_json['message']}")
def create_role(self, command):
if self.server_type != "admin":
@@ -319,10 +307,7 @@ class RAGFlowClient:
desc_str = desc_tree.children[0].strip("'\"")
print(f"create role name: {role_name}, description: {desc_str}")
response = self.http_client.request("POST", "/admin/roles",
json_body={"role_name": role_name, "description": desc_str},
use_api_base=True,
auth_kind="admin")
response = self.http_client.request("POST", "/admin/roles", json_body={"role_name": role_name, "description": desc_str}, use_api_base=True, auth_kind="admin")
res_json = response.json()
if response.status_code == 200:
self._print_table_simple(res_json["data"])
@@ -336,9 +321,7 @@ class RAGFlowClient:
role_name_tree: Tree = command["role_name"]
role_name: str = role_name_tree.children[0].strip("'\"")
print(f"drop role name: {role_name}")
response = self.http_client.request("DELETE", f"/admin/roles/{role_name}",
use_api_base=True,
auth_kind="admin")
response = self.http_client.request("DELETE", f"/admin/roles/{role_name}", use_api_base=True, auth_kind="admin")
res_json = response.json()
if response.status_code == 200:
self._print_table_simple(res_json["data"])
@@ -355,24 +338,18 @@ class RAGFlowClient:
desc_str: str = desc_tree.children[0].strip("'\"")
print(f"alter role name: {role_name}, description: {desc_str}")
response = self.http_client.request("PUT", f"/admin/roles/{role_name}",
json_body={"description": desc_str},
use_api_base=True,
auth_kind="admin")
response = self.http_client.request("PUT", f"/admin/roles/{role_name}", json_body={"description": desc_str}, use_api_base=True, auth_kind="admin")
res_json = response.json()
if response.status_code == 200:
self._print_table_simple(res_json["data"])
else:
print(
f"Fail to update role {role_name} with description: {desc_str}, code: {res_json['code']}, message: {res_json['message']}")
print(f"Fail to update role {role_name} with description: {desc_str}, code: {res_json['code']}, message: {res_json['message']}")
def list_roles(self, command):
if self.server_type != "admin":
print("This command is only allowed in ADMIN mode")
response = self.http_client.request("GET", "/admin/roles",
use_api_base=True,
auth_kind="admin")
response = self.http_client.request("GET", "/admin/roles", use_api_base=True, auth_kind="admin")
res_json = response.json()
if response.status_code == 200:
self._print_table_simple(res_json["data"])
@@ -386,9 +363,7 @@ class RAGFlowClient:
role_name_tree: Tree = command["role_name"]
role_name: str = role_name_tree.children[0].strip("'\"")
print(f"show role: {role_name}")
response = self.http_client.request("GET", f"/admin/roles/{role_name}/permission",
use_api_base=True,
auth_kind="admin")
response = self.http_client.request("GET", f"/admin/roles/{role_name}/permission", use_api_base=True, auth_kind="admin")
res_json = response.json()
if response.status_code == 200:
self._print_table_simple(res_json["data"])
@@ -409,15 +384,12 @@ class RAGFlowClient:
action_str: str = action_tree.children[0].strip("'\"")
actions.append(action_str)
print(f"grant role_name: {role_name_str}, resource: {resource_str}, actions: {actions}")
response = self.http_client.request("POST", f"/admin/roles/{role_name_str}/permission",
json_body={"actions": actions, "resource": resource_str}, use_api_base=True,
auth_kind="admin")
response = self.http_client.request("POST", f"/admin/roles/{role_name_str}/permission", json_body={"actions": actions, "resource": resource_str}, use_api_base=True, auth_kind="admin")
res_json = response.json()
if response.status_code == 200:
self._print_table_simple(res_json["data"])
else:
print(
f"Fail to grant role {role_name_str} with {actions} on {resource_str}, code: {res_json['code']}, message: {res_json['message']}")
print(f"Fail to grant role {role_name_str} with {actions} on {resource_str}, code: {res_json['code']}, message: {res_json['message']}")
def revoke_permission(self, command):
if self.server_type != "admin":
@@ -433,15 +405,12 @@ class RAGFlowClient:
action_str: str = action_tree.children[0].strip("'\"")
actions.append(action_str)
print(f"revoke role_name: {role_name_str}, resource: {resource_str}, actions: {actions}")
response = self.http_client.request("DELETE", f"/admin/roles/{role_name_str}/permission",
json_body={"actions": actions, "resource": resource_str}, use_api_base=True,
auth_kind="admin")
response = self.http_client.request("DELETE", f"/admin/roles/{role_name_str}/permission", json_body={"actions": actions, "resource": resource_str}, use_api_base=True, auth_kind="admin")
res_json = response.json()
if response.status_code == 200:
self._print_table_simple(res_json["data"])
else:
print(
f"Fail to revoke role {role_name_str} with {actions} on {resource_str}, code: {res_json['code']}, message: {res_json['message']}")
print(f"Fail to revoke role {role_name_str} with {actions} on {resource_str}, code: {res_json['code']}, message: {res_json['message']}")
def alter_user_role(self, command):
if self.server_type != "admin":
@@ -452,15 +421,12 @@ class RAGFlowClient:
user_name_tree: Tree = command["user_name"]
user_name_str: str = user_name_tree.children[0].strip("'\"")
print(f"alter_user_role user_name: {user_name_str}, role_name: {role_name_str}")
response = self.http_client.request("PUT", f"/admin/users/{user_name_str}/role",
json_body={"role_name": role_name_str}, use_api_base=True,
auth_kind="admin")
response = self.http_client.request("PUT", f"/admin/users/{user_name_str}/role", json_body={"role_name": role_name_str}, use_api_base=True, auth_kind="admin")
res_json = response.json()
if response.status_code == 200:
self._print_table_simple(res_json["data"])
else:
print(
f"Fail to alter user: {user_name_str} to role {role_name_str}, code: {res_json['code']}, message: {res_json['message']}")
print(f"Fail to alter user: {user_name_str} to role {role_name_str}, code: {res_json['code']}, message: {res_json['message']}")
def show_user_permission(self, command):
if self.server_type != "admin":
@@ -469,14 +435,12 @@ class RAGFlowClient:
user_name_tree: Tree = command["user_name"]
user_name_str: str = user_name_tree.children[0].strip("'\"")
print(f"show_user_permission user_name: {user_name_str}")
response = self.http_client.request("GET", f"/admin/users/{user_name_str}/permission", use_api_base=True,
auth_kind="admin")
response = self.http_client.request("GET", f"/admin/users/{user_name_str}/permission", use_api_base=True, auth_kind="admin")
res_json = response.json()
if response.status_code == 200:
self._print_table_simple(res_json["data"])
else:
print(
f"Fail to show user: {user_name_str} permission, code: {res_json['code']}, message: {res_json['message']}")
print(f"Fail to show user: {user_name_str} permission, code: {res_json['code']}, message: {res_json['message']}")
def generate_key(self, command: dict[str, Any]) -> None:
if self.server_type != "admin":
@@ -485,14 +449,12 @@ class RAGFlowClient:
username_tree: Tree = command["user_name"]
user_name: str = username_tree.children[0].strip("'\"")
print(f"Generating API key for user: {user_name}")
response = self.http_client.request("POST", f"/admin/users/{user_name}/keys", use_api_base=True,
auth_kind="admin")
response = self.http_client.request("POST", f"/admin/users/{user_name}/keys", use_api_base=True, auth_kind="admin")
res_json: dict[str, Any] = response.json()
if response.status_code == 200:
self._print_table_simple(res_json["data"])
else:
print(
f"Failed to generate key for user {user_name}, code: {res_json['code']}, message: {res_json['message']}")
print(f"Failed to generate key for user {user_name}, code: {res_json['code']}, message: {res_json['message']}")
def list_keys(self, command: dict[str, Any]) -> None:
if self.server_type != "admin":
@@ -501,8 +463,7 @@ class RAGFlowClient:
username_tree: Tree = command["user_name"]
user_name: str = username_tree.children[0].strip("'\"")
print(f"Listing API keys for user: {user_name}")
response = self.http_client.request("GET", f"/admin/users/{user_name}/keys", use_api_base=True,
auth_kind="admin")
response = self.http_client.request("GET", f"/admin/users/{user_name}/keys", use_api_base=True, auth_kind="admin")
res_json: dict[str, Any] = response.json()
if response.status_code == 200:
self._print_table_simple(res_json["data"])
@@ -520,8 +481,7 @@ class RAGFlowClient:
print(f"Dropping API key for user: {user_name}")
# URL encode the key to handle special characters
encoded_key: str = urllib.parse.quote(key, safe="")
response = self.http_client.request("DELETE", f"/admin/users/{user_name}/keys/{encoded_key}", use_api_base=True,
auth_kind="admin")
response = self.http_client.request("DELETE", f"/admin/users/{user_name}/keys/{encoded_key}", use_api_base=True, auth_kind="admin")
res_json: dict[str, Any] = response.json()
if response.status_code == 200:
print(res_json["message"])
@@ -534,23 +494,19 @@ class RAGFlowClient:
var_name = _strip_tree_value(command["var_name"])
var_value = _strip_tree_value(command["var_value"])
response = self.http_client.request("PUT", "/admin/variables",
json_body={"var_name": var_name, "var_value": var_value}, use_api_base=True,
auth_kind="admin")
response = self.http_client.request("PUT", "/admin/variables", json_body={"var_name": var_name, "var_value": var_value}, use_api_base=True, auth_kind="admin")
res_json = response.json()
if response.status_code == 200:
print(res_json["message"])
else:
print(
f"Fail to set variable {var_name} to {var_value}, code: {res_json['code']}, message: {res_json['message']}")
print(f"Fail to set variable {var_name} to {var_value}, code: {res_json['code']}, message: {res_json['message']}")
def show_variable(self, command):
if self.server_type != "admin":
print("This command is only allowed in ADMIN mode")
var_name = _strip_tree_value(command["var_name"])
response = self.http_client.request(method="GET", path="/admin/variables", json_body={"var_name": var_name},
use_api_base=True, auth_kind="admin")
response = self.http_client.request(method="GET", path="/admin/variables", json_body={"var_name": var_name}, use_api_base=True, auth_kind="admin")
res_json = response.json()
if response.status_code == 200:
self._print_table_simple(res_json["data"])
@@ -604,8 +560,7 @@ class RAGFlowClient:
if self.server_type != "admin":
print("This command is only allowed in ADMIN mode")
license = command["license"]
response = self.http_client.request("POST", "/admin/license", json_body={"license": license}, use_api_base=True,
auth_kind="admin")
response = self.http_client.request("POST", "/admin/license", json_body={"license": license}, use_api_base=True, auth_kind="admin")
res_json = response.json()
if response.status_code == 200:
print("Set license successfully")
@@ -617,9 +572,7 @@ class RAGFlowClient:
print("This command is only allowed in ADMIN mode")
value1 = command["value1"]
value2 = command["value2"]
response = self.http_client.request("POST", "/admin/license/config",
json_body={"value1": value1, "value2": value2}, use_api_base=True,
auth_kind="admin")
response = self.http_client.request("POST", "/admin/license/config", json_body={"value1": value1, "value2": value2}, use_api_base=True, auth_kind="admin")
res_json = response.json()
if response.status_code == 200:
print("Set license successfully")
@@ -690,8 +643,7 @@ class RAGFlowClient:
user_name: str = username_tree.children[0].strip("'\"")
print(f"Listing all datasets of user: {user_name}")
response = self.http_client.request("GET", f"/admin/users/{user_name}/datasets", use_api_base=True,
auth_kind="admin")
response = self.http_client.request("GET", f"/admin/users/{user_name}/datasets", use_api_base=True, auth_kind="admin")
res_json = response.json()
if response.status_code == 200:
table_data = res_json["data"]
@@ -708,8 +660,7 @@ class RAGFlowClient:
username_tree: Tree = command["user_name"]
user_name: str = username_tree.children[0].strip("'\"")
print(f"Listing all agents of user: {user_name}")
response = self.http_client.request("GET", f"/admin/users/{user_name}/agents", use_api_base=True,
auth_kind="admin")
response = self.http_client.request("GET", f"/admin/users/{user_name}/agents", use_api_base=True, auth_kind="admin")
res_json = response.json()
if response.status_code == 200:
table_data = res_json["data"]
@@ -733,8 +684,7 @@ class RAGFlowClient:
# Step 1: Add provider
provider_payload = {"provider_name": provider_name}
provider_response = self.http_client.request("PUT", "/providers", json_body=provider_payload,
use_api_base=True, auth_kind="web")
provider_response = self.http_client.request("PUT", "/providers", json_body=provider_payload, use_api_base=True, auth_kind="web")
provider_res = provider_response.json()
if provider_response.status_code == 200 and provider_res.get("code") == 0:
print(f"Success to add provider {provider_name}")
@@ -747,15 +697,8 @@ class RAGFlowClient:
return
# Step 2: Add instance
instance_payload = {
"instance_name": "default",
"api_key": api_key,
"region": "default",
"base_url": ""
}
instance_response = self.http_client.request("POST", f"/providers/{provider_name}/instances",
json_body=instance_payload, use_api_base=True,
auth_kind="web")
instance_payload = {"instance_name": "default", "api_key": api_key, "region": "default", "base_url": ""}
instance_response = self.http_client.request("POST", f"/providers/{provider_name}/instances", json_body=instance_payload, use_api_base=True, auth_kind="web")
instance_res = instance_response.json()
if instance_response.status_code == 200 and instance_res.get("code") == 0:
print(f"Success to add instance for provider {provider_name}")
@@ -771,8 +714,7 @@ class RAGFlowClient:
print("This command is only allowed in USER mode")
return
provider_name: str = command["provider_name"]
response = self.http_client.request("DELETE", f"/providers/{provider_name}", use_api_base=True,
auth_kind="web")
response = self.http_client.request("DELETE", f"/providers/{provider_name}", use_api_base=True, auth_kind="web")
res_json = response.json()
if response.status_code == 200 and res_json.get("code") == 0:
print(f"Success to drop model provider {provider_name}")
@@ -810,8 +752,7 @@ class RAGFlowClient:
"model_type": model_type,
"model_name": model_name,
}
response = self.http_client.request("PATCH", "/models/default", json_body=payload, use_api_base=True,
auth_kind="web")
response = self.http_client.request("PATCH", "/models/default", json_body=payload, use_api_base=True, auth_kind="web")
res_json = response.json()
if response.status_code == 200 and res_json.get("code") == 0:
print(f"Success to set default {model_type} to {model_id}")
@@ -830,8 +771,7 @@ class RAGFlowClient:
return
payload = {"model_type": model_type}
response = self.http_client.request("PATCH", "/models/default", json_body=payload, use_api_base=True,
auth_kind="web")
response = self.http_client.request("PATCH", "/models/default", json_body=payload, use_api_base=True, auth_kind="web")
res_json = response.json()
if response.status_code == 200 and res_json.get("code") == 0:
print(f"Success to reset default {model_type}")
@@ -861,8 +801,7 @@ class RAGFlowClient:
iterations = command.get("iterations", 1)
if iterations > 1:
response = self.http_client.request("GET", "/datasets", use_api_base=True, auth_kind="web",
iterations=iterations)
response = self.http_client.request("GET", "/datasets", use_api_base=True, auth_kind="web", iterations=iterations)
return response
else:
response = self.http_client.request("GET", "/datasets", use_api_base=True, auth_kind="web")
@@ -876,16 +815,12 @@ class RAGFlowClient:
def create_user_dataset(self, command):
if self.server_type != "user":
print("This command is only allowed in USER mode")
payload = {
"name": command["dataset_name"],
"embedding_model": command["embedding"]
}
payload = {"name": command["dataset_name"], "embedding_model": command["embedding"]}
if "parser_id" in command:
payload["chunk_method"] = command["parser"]
if "pipeline" in command:
payload["pipeline_id"] = command["pipeline"]
response = self.http_client.request("POST", "/datasets", json_body=payload, use_api_base=True,
auth_kind="web")
response = self.http_client.request("POST", "/datasets", json_body=payload, use_api_base=True, auth_kind="web")
res_json = response.json()
if response.status_code == 200:
self._print_table_simple(res_json["data"])
@@ -981,8 +916,7 @@ class RAGFlowClient:
dataset_ids = [dataset_id for _, dataset_id in valid_datasets]
kb_ids_param = ",".join(dataset_ids)
response = self.http_client.request("GET", f"/kb/get_meta?kb_ids={kb_ids_param}",
use_api_base=False, auth_kind="web")
response = self.http_client.request("GET", f"/kb/get_meta?kb_ids={kb_ids_param}", use_api_base=False, auth_kind="web")
res_json = response.json()
if response.status_code != 200:
print(f"Fail to get metadata, code: {res_json.get('code')}, message: {res_json.get('message')}")
@@ -996,11 +930,7 @@ class RAGFlowClient:
table_data = []
for field_name, values_dict in meta.items():
for value, docs in values_dict.items():
table_data.append({
"field": field_name,
"value": value,
"doc_ids": ", ".join(docs)
})
table_data.append({"field": field_name, "value": value, "doc_ids": ", ".join(docs)})
self._print_table_simple(table_data)
def list_user_documents_metadata_summary(self, command_dict):
@@ -1018,8 +948,7 @@ class RAGFlowClient:
payload = {"kb_id": kb_id}
if doc_ids:
payload["doc_ids"] = doc_ids
response = self.http_client.request("POST", "/document/metadata/summary", json_body=payload,
use_api_base=False, auth_kind="web")
response = self.http_client.request("POST", "/document/metadata/summary", json_body=payload, use_api_base=False, auth_kind="web")
res_json = response.json()
if response.status_code == 200:
summary = res_json.get("data", {}).get("summary", {})
@@ -1086,16 +1015,11 @@ class RAGFlowClient:
"quote": True,
"keyword": False,
"tts": False,
"system": "You are an intelligent assistant. Your primary function is to answer questions based strictly on the provided knowledge base.\n\n **Essential Rules:**\n - Your answer must be derived **solely** from this knowledge base: `{knowledge}`.\n - **When information is available**: Summarize the content to give a detailed answer.\n - **When information is unavailable**: Your response must contain this exact sentence: \"The answer you are looking for is not found in the knowledge base!\"\n - **Always consider** the entire conversation history.",
"system": 'You are an intelligent assistant. Your primary function is to answer questions based strictly on the provided knowledge base.\n\n **Essential Rules:**\n - Your answer must be derived **solely** from this knowledge base: `{knowledge}`.\n - **When information is available**: Summarize the content to give a detailed answer.\n - **When information is unavailable**: Your response must contain this exact sentence: "The answer you are looking for is not found in the knowledge base!"\n - **Always consider** the entire conversation history.',
"refine_multiturn": False,
"use_kg": False,
"reasoning": False,
"parameters": [
{
"key": "knowledge",
"optional": False
}
],
"parameters": [{"key": "knowledge", "optional": False}],
"toc_enhance": False,
},
"similarity_threshold": 0.2,
@@ -1136,8 +1060,7 @@ class RAGFlowClient:
# Build payload
payload = {"kb_id": dataset_id, "vector_size": vector_size}
# Call API
response = self.http_client.request("POST", "/kb/doc_engine_table", json_body=payload,
use_api_base=False, auth_kind="web")
response = self.http_client.request("POST", "/kb/doc_engine_table", json_body=payload, use_api_base=False, auth_kind="web")
res_json = response.json()
if response.status_code == 200 and res_json.get("code") == 0:
print(f"Success to create table for dataset: {dataset_name}")
@@ -1155,8 +1078,7 @@ class RAGFlowClient:
return
# Call API to delete table
payload = {"kb_id": dataset_id}
response = self.http_client.request("DELETE", "/kb/doc_engine_table", json_body=payload,
use_api_base=False, auth_kind="web")
response = self.http_client.request("DELETE", "/kb/doc_engine_table", json_body=payload, use_api_base=False, auth_kind="web")
res_json = response.json()
if response.status_code == 200 and res_json.get("code") == 0:
print(f"Success to drop table for dataset: {dataset_name}")
@@ -1168,8 +1090,7 @@ class RAGFlowClient:
print("This command is only allowed in USER mode")
return
# Call API to create metadata table
response = self.http_client.request("POST", "/tenant/doc_engine_metadata_table",
use_api_base=False, auth_kind="web")
response = self.http_client.request("POST", "/tenant/doc_engine_metadata_table", use_api_base=False, auth_kind="web")
res_json = response.json()
if response.status_code == 200 and res_json.get("code") == 0:
print("Success to create metadata table")
@@ -1181,8 +1102,7 @@ class RAGFlowClient:
print("This command is only allowed in USER mode")
return
# Call API to delete metadata table
response = self.http_client.request("DELETE", "/tenant/doc_engine_metadata_table",
use_api_base=False, auth_kind="web")
response = self.http_client.request("DELETE", "/tenant/doc_engine_metadata_table", use_api_base=False, auth_kind="web")
res_json = response.json()
if response.status_code == 200 and res_json.get("code") == 0:
print("Success to drop metadata table")
@@ -1225,8 +1145,7 @@ class RAGFlowClient:
def _list_chat_sessions(self, dialog_id):
"""List all sessions (conversations) for a given dialog."""
response = self.http_client.request("GET", f"/chats/{dialog_id}/conversations", use_api_base=True,
auth_kind="web")
response = self.http_client.request("GET", f"/chats/{dialog_id}/conversations", use_api_base=True, auth_kind="web")
res_json = response.json()
if response.status_code == 200 and res_json["code"] == 0:
return res_json["data"]
@@ -1242,14 +1161,12 @@ class RAGFlowClient:
if dialog_id is None:
return
payload = {"name": "New conversation"}
response = self.http_client.request("POST", f"/chats/{dialog_id}/conversations", json_body=payload,
use_api_base=True, auth_kind="web")
response = self.http_client.request("POST", f"/chats/{dialog_id}/conversations", json_body=payload, use_api_base=True, auth_kind="web")
res_json = response.json()
if response.status_code == 200 and res_json["code"] == 0:
print(f"Success to create chat session for chat: {chat_name}")
else:
print(
f"Fail to create chat session for chat {chat_name}, code: {res_json['code']}, message: {res_json['message']}")
print(f"Fail to create chat session for chat {chat_name}, code: {res_json['code']}, message: {res_json['message']}")
def drop_chat_session(self, command):
if self.server_type != "user":
@@ -1270,14 +1187,12 @@ class RAGFlowClient:
print(f"Chat session '{session_id}' not found in chat '{chat_name}'")
return
payload = {"ids": to_drop_session_ids}
response = self.http_client.request("DELETE", f"/chats/{dialog_id}/conversations", json_body=payload,
use_api_base=True, auth_kind="web")
response = self.http_client.request("DELETE", f"/chats/{dialog_id}/conversations", json_body=payload, use_api_base=True, auth_kind="web")
res_json = response.json()
if response.status_code == 200 and res_json["code"] == 0:
print(f"Success to drop chat session '{session_id}' from chat: {chat_name}")
else:
print(
f"Fail to drop chat session '{session_id}' from chat {chat_name}, code: {res_json['code']}, message: {res_json['message']}")
print(f"Fail to drop chat session '{session_id}' from chat {chat_name}, code: {res_json['code']}, message: {res_json['message']}")
def list_chat_sessions(self, command):
if self.server_type != "user":
@@ -1305,13 +1220,9 @@ class RAGFlowClient:
# Prepare payload for completion API
# Note: stream parameter is not sent, server defaults to stream=True
payload = {
"session_id": session_id,
"messages": [{"role": "user", "content": message}]
}
payload = {"session_id": session_id, "messages": [{"role": "user", "content": message}]}
response = self.http_client.request("POST", "/chat/completions", json_body=payload,
use_api_base=True, auth_kind="web", stream=True)
response = self.http_client.request("POST", "/chat/completions", json_body=payload, use_api_base=True, auth_kind="web", stream=True)
if response.status_code != 200:
print(f"Fail to chat on session, status code: {response.status_code}")
@@ -1322,17 +1233,16 @@ class RAGFlowClient:
for line in response.iter_lines():
if not line:
continue
line_str = line.decode('utf-8')
if not line_str.startswith('data:'):
line_str = line.decode("utf-8")
if not line_str.startswith("data:"):
continue
data_str = line_str[5:].strip()
if data_str == '[DONE]':
if data_str == "[DONE]":
break
try:
data_json = json.loads(data_str)
if data_json.get("code") != 0:
print(
f"\nFail to chat on session, code: {data_json.get('code')}, message: {data_json.get('message', '')}")
print(f"\nFail to chat on session, code: {data_json.get('code')}, message: {data_json.get('message', '')}")
return
# Check if it's the final message
if data_json.get("data") is True:
@@ -1416,14 +1326,12 @@ class RAGFlowClient:
print(f"Documents {document_names} not found in {dataset_name}")
payload = {"doc_ids": document_ids, "run": 1}
response = self.http_client.request("POST", "/documents/ingest", json_body=payload, use_api_base=True,
auth_kind="web")
response = self.http_client.request("POST", "/documents/ingest", json_body=payload, use_api_base=True, auth_kind="web")
res_json = response.json()
if response.status_code == 200 and res_json["code"] == 0:
print(f"Success to parse {to_parse_doc_names} of {dataset_name}")
else:
print(
f"Fail to parse documents {res_json["data"]["docs"]}, code: {res_json['code']}, message: {res_json['message']}")
print(f"Fail to parse documents {res_json['data']['docs']}, code: {res_json['code']}, message: {res_json['message']}")
def parse_dataset(self, command_dict):
if self.server_type != "user":
@@ -1442,8 +1350,7 @@ class RAGFlowClient:
document_ids.append(doc["id"])
payload = {"doc_ids": document_ids, "run": 1}
response = self.http_client.request("POST", "/documents/ingest", json_body=payload, use_api_base=True,
auth_kind="web")
response = self.http_client.request("POST", "/documents/ingest", json_body=payload, use_api_base=True, auth_kind="web")
res_json = response.json()
if response.status_code == 200 and res_json["code"] == 0:
pass
@@ -1483,15 +1390,7 @@ class RAGFlowClient:
encoder = MultipartEncoder(fields=fields)
headers = {"Content-Type": encoder.content_type}
response = self.http_client.request(
"POST",
f"/datasets/{dataset_id}/documents?return_raw_files=true",
headers=headers,
data=encoder,
json_body=None,
params=None,
stream=False,
auth_kind="web",
use_api_base=True
"POST", f"/datasets/{dataset_id}/documents?return_raw_files=true", headers=headers, data=encoder, json_body=None, params=None, stream=False, auth_kind="web", use_api_base=True
)
res = response.json()
if res.get("code") == 0:
@@ -1526,22 +1425,18 @@ class RAGFlowClient:
}
iterations = command_dict.get("iterations", 1)
if iterations > 1:
response = self.http_client.request("POST", "/retrieval", json_body=payload, use_api_base=True,
auth_kind="web", iterations=iterations)
response = self.http_client.request("POST", "/retrieval", json_body=payload, use_api_base=True, auth_kind="web", iterations=iterations)
return response
else:
response = self.http_client.request("POST", "/retrieval", json_body=payload, use_api_base=True,
auth_kind="web")
response = self.http_client.request("POST", "/retrieval", json_body=payload, use_api_base=True, auth_kind="web")
res_json = response.json()
if response.status_code == 200:
if res_json["code"] == 0:
self._print_table_simple(res_json["data"]["chunks"])
else:
print(
f"Fail to search datasets: {dataset_names}, code: {res_json['code']}, message: {res_json['message']}")
print(f"Fail to search datasets: {dataset_names}, code: {res_json['code']}, message: {res_json['message']}")
else:
print(
f"Fail to search datasets: {dataset_names}, code: {res_json['code']}, message: {res_json['message']}")
print(f"Fail to search datasets: {dataset_names}, code: {res_json['code']}, message: {res_json['message']}")
def get_chunk(self, command_dict):
if self.server_type != "user":
@@ -1549,8 +1444,7 @@ class RAGFlowClient:
return
chunk_id = command_dict["chunk_id"]
response = self.http_client.request("GET", f"/chunk/get?chunk_id={chunk_id}", use_api_base=False,
auth_kind="web")
response = self.http_client.request("GET", f"/chunk/get?chunk_id={chunk_id}", use_api_base=False, auth_kind="web")
res_json = response.json()
if response.status_code == 200:
if res_json["code"] == 0:
@@ -1568,8 +1462,7 @@ class RAGFlowClient:
file_path = command_dict["file_path"]
payload = {"file_path": file_path}
response = self.http_client.request("POST", "/kb/insert_from_file", json_body=payload,
use_api_base=False, auth_kind="web")
response = self.http_client.request("POST", "/kb/insert_from_file", json_body=payload, use_api_base=False, auth_kind="web")
res_json = response.json()
if response.status_code == 200:
if res_json["code"] == 0:
@@ -1589,8 +1482,7 @@ class RAGFlowClient:
file_path = command_dict["file_path"]
payload = {"file_path": file_path}
response = self.http_client.request("POST", "/tenant/insert_metadata_from_file", json_body=payload,
use_api_base=False, auth_kind="web")
response = self.http_client.request("POST", "/tenant/insert_metadata_from_file", json_body=payload, use_api_base=False, auth_kind="web")
res_json = response.json()
if response.status_code == 200:
if res_json["code"] == 0:
@@ -1617,8 +1509,7 @@ class RAGFlowClient:
return
# Get doc_id from chunk_id via GET /chunk/get
response = self.http_client.request("GET", f"/chunk/get?chunk_id={chunk_id}", use_api_base=False,
auth_kind="web")
response = self.http_client.request("GET", f"/chunk/get?chunk_id={chunk_id}", use_api_base=False, auth_kind="web")
res_json = response.json()
if response.status_code != 200:
print(f"Fail to get chunk info, code: {res_json.get('code')}, message: {res_json.get('message')}")
@@ -1655,14 +1546,8 @@ class RAGFlowClient:
else:
print(f"Fail to update chunk, HTTP {response.status_code}")
def _get_documents_by_ids(self, ids:list[str]):
response = self.http_client.request(
"POST",
"/document/infos",
json_body={"doc_ids": ids},
use_api_base=False,
auth_kind="web"
)
def _get_documents_by_ids(self, ids: list[str]):
response = self.http_client.request("POST", "/document/infos", json_body={"doc_ids": ids}, use_api_base=False, auth_kind="web")
if response.status_code != 200:
return f"Fail to get document info, HTTP {response.status_code}", None
@@ -1687,6 +1572,7 @@ class RAGFlowClient:
# Parse JSON string to dict
import json
try:
meta_fields = json.loads(meta_json_str)
except json.JSONDecodeError as e:
@@ -1713,13 +1599,7 @@ class RAGFlowClient:
"meta_fields": meta_fields,
}
response = self.http_client.request(
"PATCH",
f"/datasets/{dataset_id}/documents/{doc_id}",
json_body=payload,
use_api_base=True,
auth_kind="web"
)
response = self.http_client.request("PATCH", f"/datasets/{dataset_id}/documents/{doc_id}", json_body=payload, use_api_base=True, auth_kind="web")
res_json = response.json()
if response.status_code == 200:
@@ -1747,8 +1627,7 @@ class RAGFlowClient:
"tags": tags,
}
response = self.http_client.request("POST", f"/kb/{dataset_id}/rm_tags", json_body=payload,
use_api_base=False, auth_kind="web")
response = self.http_client.request("POST", f"/kb/{dataset_id}/rm_tags", json_body=payload, use_api_base=False, auth_kind="web")
res_json = response.json()
if response.status_code == 200:
if res_json.get("code") == 0:
@@ -1771,8 +1650,7 @@ class RAGFlowClient:
elif command_dict.get("chunk_ids"):
payload["chunk_ids"] = command_dict["chunk_ids"]
response = self.http_client.request("POST", "/chunk/rm", json_body=payload,
use_api_base=False, auth_kind="web")
response = self.http_client.request("POST", "/chunk/rm", json_body=payload, use_api_base=False, auth_kind="web")
res_json = response.json()
if response.status_code == 200:
if res_json.get("code") == 0:
@@ -1803,15 +1681,14 @@ class RAGFlowClient:
if "available_int" in command_dict:
payload["available_int"] = command_dict["available_int"]
response = self.http_client.request("POST", "/chunk/list", json_body=payload, use_api_base=False,
auth_kind="web")
response = self.http_client.request("POST", "/chunk/list", json_body=payload, use_api_base=False, auth_kind="web")
res_json = response.json()
if response.status_code == 200:
if res_json["code"] == 0:
chunks = res_json["data"]["chunks"]
if chunks:
for i, chunk in enumerate(chunks):
print(f"\n--- Chunk {i+1} ---")
print(f"\n--- Chunk {i + 1} ---")
for key, value in chunk.items():
print(f" {key}: {value}")
else:
@@ -1845,7 +1722,7 @@ class RAGFlowClient:
all_done = True
for doc in docs:
if doc.get("run") != "DONE":
print(f"Document {doc["name"]} is not done, status: {doc.get("run")}")
print(f"Document {doc['name']} is not done, status: {doc.get('run')}")
all_done = False
break
if all_done:
@@ -1856,16 +1733,10 @@ class RAGFlowClient:
def _list_documents(self, dataset_name: str, dataset_id: str):
# Use the new RESTful API: GET /api/v1/datasets/<dataset_id>/documents
response = self.http_client.request(
"GET",
f"/datasets/{dataset_id}/documents",
use_api_base=True,
auth_kind="web"
)
response = self.http_client.request("GET", f"/datasets/{dataset_id}/documents", use_api_base=True, auth_kind="web")
res_json = response.json()
if response.status_code != 200:
print(
f"Fail to list files from dataset {dataset_name}, code: {res_json['code']}, message: {res_json['message']}")
print(f"Fail to list files from dataset {dataset_name}, code: {res_json['code']}, message: {res_json['message']}")
return None
return res_json["data"]["docs"]
@@ -2254,22 +2125,14 @@ def run_benchmark(client: RAGFlowClient, command_dict: dict):
total_duration = result["duration"]
qps = iterations / total_duration if total_duration > 0 else None
print(f"command: {command}, Concurrency: {concurrency}, iterations: {iterations}")
print(
f"total duration: {total_duration:.4f}s, QPS: {qps}, COMMAND_COUNT: {iterations}, SUCCESS: {success_count}, FAILURE: {iterations - success_count}")
print(f"total duration: {total_duration:.4f}s, QPS: {qps}, COMMAND_COUNT: {iterations}, SUCCESS: {success_count}, FAILURE: {iterations - success_count}")
pass
else:
results: List[Optional[dict]] = [None] * concurrency
mp_context = mp.get_context("spawn")
start_time = time.perf_counter()
with ProcessPoolExecutor(max_workers=concurrency, mp_context=mp_context) as executor:
future_map = {
executor.submit(
run_command,
client,
command
): idx
for idx in range(concurrency)
}
future_map = {executor.submit(run_command, client, command): idx for idx in range(concurrency)}
for future in as_completed(future_map):
idx = future_map[future]
results[idx] = future.result()
@@ -2291,7 +2154,6 @@ def run_benchmark(client: RAGFlowClient, command_dict: dict):
total_command_count = iterations * concurrency
qps = total_command_count / total_duration if total_duration > 0 else None
print(f"command: {command}, Concurrency: {concurrency} , iterations: {iterations}")
print(
f"total duration: {total_duration:.4f}s, QPS: {qps}, COMMAND_COUNT: {total_command_count}, SUCCESS: {success_count}, FAILURE: {total_command_count - success_count}")
print(f"total duration: {total_duration:.4f}s, QPS: {qps}, COMMAND_COUNT: {total_command_count}, SUCCESS: {success_count}, FAILURE: {total_command_count - success_count}")
pass
admin/client/user.py
+4 -5
View File
@@ -29,6 +29,7 @@ def encrypt_password(password_plain: str) -> str:
import base64
from Cryptodome.PublicKey import RSA
from Cryptodome.Cipher import PKCS1_v1_5 as Cipher_pkcs1_v1_5
def crypt(line):
"""
decrypt(crypt(input_string)) == base64(input_string), which frontend and ragflow_cli use.
@@ -36,13 +37,11 @@ def encrypt_password(password_plain: str) -> str:
pub = "-----BEGIN PUBLIC KEY-----\nMIIBIjANBgkqhkiG9w0BAQEFAAOCAQ8AMIIBCgKCAQEArq9XTUSeYr2+N1h3Afl/z8Dse/2yD0ZGrKwx+EEEcdsBLca9Ynmx3nIB5obmLlSfmskLpBo0UACBmB5rEjBp2Q2f3AG3Hjd4B+gNCG6BDaawuDlgANIhGnaTLrIqWrrcm4EMzJOnAOI1fgzJRsOOUEfaS318Eq9OVO3apEyCCt0lOQK6PuksduOjVxtltDav+guVAA068NrPYmRNabVKRNLJpL8w4D44sfth5RvZ3q9t+6RTArpEtc5sh5ChzvqPOzKGMXW83C95TxmXqpbK6olN4RevSfVjEAgCydH6HN6OhtOQEcnrU97r9H0iZOWwbw3pVrZiUkuRD1R56Wzs2wIDAQAB\n-----END PUBLIC KEY-----"
rsa_key = RSA.importKey(pub)
cipher = Cipher_pkcs1_v1_5.new(rsa_key)
password_base64 = base64.b64encode(line.encode('utf-8')).decode("utf-8")
password_base64 = base64.b64encode(line.encode("utf-8")).decode("utf-8")
encrypted_password = cipher.encrypt(password_base64.encode())
return base64.b64encode(encrypted_password).decode('utf-8')
return base64.b64encode(encrypted_password).decode("utf-8")
except Exception as exc:
raise AuthException(
"Password encryption unavailable; install pycryptodomex (uv sync --python 3.13 --group test)."
) from exc
raise AuthException("Password encryption unavailable; install pycryptodomex (uv sync --python 3.13 --group test).") from exc
return crypt(password_plain)
admin/server/admin_server.py
+7 -8
View File
@@ -15,6 +15,7 @@
#
import time
start_ts = time.time()
import os
@@ -38,26 +39,24 @@ from common.versions import get_ragflow_version
stop_event = threading.Event()
if __name__ == '__main__':
if __name__ == "__main__":
faulthandler.enable()
init_root_logger("admin_service")
logging.info(r"""
____ ___ ______________ ___ __ _
/ __ \/ | / ____/ ____/ /___ _ __ / | ____/ /___ ___ (_)___
____ ___ ______________ ___ __ _
/ __ \/ | / ____/ ____/ /___ _ __ / | ____/ /___ ___ (_)___
/ /_/ / /| |/ / __/ /_ / / __ \ | /| / / / /| |/ __ / __ `__ \/ / __ \
/ _, _/ ___ / /_/ / __/ / / /_/ / |/ |/ / / ___ / /_/ / / / / / / / / / /
/_/ |_/_/ |_\____/_/ /_/\____/|__/|__/ /_/ |_\__,_/_/ /_/ /_/_/_/ /_/
/_/ |_/_/ |_\____/_/ /_/\____/|__/|__/ /_/ |_\__,_/_/ /_/ /_/_/_/ /_/
""")
app = Flask(__name__)
app.register_blueprint(admin_bp)
app.config["SESSION_PERMANENT"] = False
app.config["SESSION_TYPE"] = "filesystem"
app.config["MAX_CONTENT_LENGTH"] = int(
os.environ.get("MAX_CONTENT_LENGTH", 1024 * 1024 * 1024)
)
app.config["MAX_CONTENT_LENGTH"] = int(os.environ.get("MAX_CONTENT_LENGTH", 1024 * 1024 * 1024))
Session(app)
logging.info(f'RAGFlow admin version: {get_ragflow_version()}')
logging.info(f"RAGFlow admin version: {get_ragflow_version()}")
show_configs()
login_manager = LoginManager()
login_manager.init_app(app)
admin/server/auth.py
+9 -28
View File
@@ -70,9 +70,7 @@ def setup_auth(login_manager):
logging.warning(f"Authentication attempt with invalid token format: {len(access_token)} chars")
return None
user = UserService.query(
access_token=access_token, status=StatusEnum.VALID.value
)
user = UserService.query(access_token=access_token, status=StatusEnum.VALID.value)
if user:
if not user[0].access_token or not user[0].access_token.strip():
logging.warning(f"User {user[0].email} has empty access_token in database")
@@ -126,19 +124,13 @@ def add_tenant_for_admin(user_info: dict, role: str):
"img2txt_id": settings.IMAGE2TEXT_MDL,
"rerank_id": settings.RERANK_MDL,
}
usr_tenant = {
"tenant_id": user_info["id"],
"user_id": user_info["id"],
"invited_by": user_info["id"],
"role": role
}
usr_tenant = {"tenant_id": user_info["id"], "user_id": user_info["id"], "invited_by": user_info["id"], "role": role}
# tenant_llm = get_init_tenant_llm(user_info["id"])
TenantService.insert(**tenant)
UserTenantService.insert(**usr_tenant)
# TenantLLMService.insert_many(tenant_llm)
logging.info(
f"Added tenant for email: {user_info['email']}, A default tenant has been set; changing the default models after login is strongly recommended.")
logging.info(f"Added tenant for email: {user_info['email']}, A default tenant has been set; changing the default models after login is strongly recommended.")
def check_admin_auth(func):
@@ -212,28 +204,17 @@ def login_verify(f):
@wraps(f)
def decorated(*args, **kwargs):
auth = request.authorization
if not auth or 'username' not in auth.parameters or 'password' not in auth.parameters:
return jsonify({
"code": 401,
"message": "Authentication required",
"data": None
}), 200
if not auth or "username" not in auth.parameters or "password" not in auth.parameters:
return jsonify({"code": 401, "message": "Authentication required", "data": None}), 200
username = auth.parameters['username']
password = auth.parameters['password']
username = auth.parameters["username"]
password = auth.parameters["password"]
try:
if not check_admin(username, password):
return jsonify({
"code": 500,
"message": "Access denied",
"data": None
}), 200
return jsonify({"code": 500, "message": "Access denied", "data": None}), 200
except Exception:
logging.exception("An error occurred during admin login verification.")
return jsonify({
"code": 500,
"message": "An internal server error occurred."
}), 200
return jsonify({"code": 500, "message": "An internal server error occurred."}), 200
return f(*args, **kwargs)
admin/server/config.py
+120 -125
View File
@@ -34,8 +34,7 @@ class BaseConfig(BaseModel):
detail_func_name: str
def to_dict(self) -> dict[str, Any]:
return {'id': self.id, 'name': self.name, 'host': self.host, 'port': self.port,
'service_type': self.service_type}
return {"id": self.id, "name": self.name, "host": self.host, "port": self.port, "service_type": self.service_type}
class ServiceConfigs:
@@ -63,11 +62,11 @@ class MetaConfig(BaseConfig):
def to_dict(self) -> dict[str, Any]:
result = super().to_dict()
if 'extra' not in result:
result['extra'] = dict()
extra_dict = result['extra'].copy()
extra_dict['meta_type'] = self.meta_type
result['extra'] = extra_dict
if "extra" not in result:
result["extra"] = dict()
extra_dict = result["extra"].copy()
extra_dict["meta_type"] = self.meta_type
result["extra"] = extra_dict
return result
@@ -77,21 +76,20 @@ class MySQLConfig(MetaConfig):
def to_dict(self) -> dict[str, Any]:
result = super().to_dict()
if 'extra' not in result:
result['extra'] = dict()
extra_dict = result['extra'].copy()
extra_dict['username'] = self.username
extra_dict['password'] = self.password
result['extra'] = extra_dict
if "extra" not in result:
result["extra"] = dict()
extra_dict = result["extra"].copy()
extra_dict["username"] = self.username
extra_dict["password"] = self.password
result["extra"] = extra_dict
return result
class PostgresConfig(MetaConfig):
def to_dict(self) -> dict[str, Any]:
result = super().to_dict()
if 'extra' not in result:
result['extra'] = dict()
if "extra" not in result:
result["extra"] = dict()
return result
@@ -100,11 +98,11 @@ class RetrievalConfig(BaseConfig):
def to_dict(self) -> dict[str, Any]:
result = super().to_dict()
if 'extra' not in result:
result['extra'] = dict()
extra_dict = result['extra'].copy()
extra_dict['retrieval_type'] = self.retrieval_type
result['extra'] = extra_dict
if "extra" not in result:
result["extra"] = dict()
extra_dict = result["extra"].copy()
extra_dict["retrieval_type"] = self.retrieval_type
result["extra"] = extra_dict
return result
@@ -113,11 +111,11 @@ class InfinityConfig(RetrievalConfig):
def to_dict(self) -> dict[str, Any]:
result = super().to_dict()
if 'extra' not in result:
result['extra'] = dict()
extra_dict = result['extra'].copy()
extra_dict['db_name'] = self.db_name
result['extra'] = extra_dict
if "extra" not in result:
result["extra"] = dict()
extra_dict = result["extra"].copy()
extra_dict["db_name"] = self.db_name
result["extra"] = extra_dict
return result
@@ -127,12 +125,12 @@ class ElasticsearchConfig(RetrievalConfig):
def to_dict(self) -> dict[str, Any]:
result = super().to_dict()
if 'extra' not in result:
result['extra'] = dict()
extra_dict = result['extra'].copy()
extra_dict['username'] = self.username
extra_dict['password'] = self.password
result['extra'] = extra_dict
if "extra" not in result:
result["extra"] = dict()
extra_dict = result["extra"].copy()
extra_dict["username"] = self.username
extra_dict["password"] = self.password
result["extra"] = extra_dict
return result
@@ -141,11 +139,11 @@ class MessageQueueConfig(BaseConfig):
def to_dict(self) -> dict[str, Any]:
result = super().to_dict()
if 'extra' not in result:
result['extra'] = dict()
extra_dict = result['extra'].copy()
extra_dict['mq_type'] = self.mq_type
result['extra'] = extra_dict
if "extra" not in result:
result["extra"] = dict()
extra_dict = result["extra"].copy()
extra_dict["mq_type"] = self.mq_type
result["extra"] = extra_dict
return result
@@ -155,30 +153,28 @@ class RedisConfig(MessageQueueConfig):
def to_dict(self) -> dict[str, Any]:
result = super().to_dict()
if 'extra' not in result:
result['extra'] = dict()
extra_dict = result['extra'].copy()
extra_dict['database'] = self.database
extra_dict['password'] = self.password
result['extra'] = extra_dict
if "extra" not in result:
result["extra"] = dict()
extra_dict = result["extra"].copy()
extra_dict["database"] = self.database
extra_dict["password"] = self.password
result["extra"] = extra_dict
return result
class RabbitMQConfig(MessageQueueConfig):
def to_dict(self) -> dict[str, Any]:
result = super().to_dict()
if 'extra' not in result:
result['extra'] = dict()
if "extra" not in result:
result["extra"] = dict()
return result
class RAGFlowServerConfig(BaseConfig):
def to_dict(self) -> dict[str, Any]:
result = super().to_dict()
if 'extra' not in result:
result['extra'] = dict()
if "extra" not in result:
result["extra"] = dict()
return result
@@ -187,9 +183,9 @@ class TaskExecutorConfig(BaseConfig):
def to_dict(self) -> dict[str, Any]:
result = super().to_dict()
if 'extra' not in result:
result['extra'] = dict()
result['extra']['message_queue_type'] = self.message_queue_type
if "extra" not in result:
result["extra"] = dict()
result["extra"]["message_queue_type"] = self.message_queue_type
return result
@@ -198,11 +194,11 @@ class FileStoreConfig(BaseConfig):
def to_dict(self) -> dict[str, Any]:
result = super().to_dict()
if 'extra' not in result:
result['extra'] = dict()
extra_dict = result['extra'].copy()
extra_dict['store_type'] = self.store_type
result['extra'] = extra_dict
if "extra" not in result:
result["extra"] = dict()
extra_dict = result["extra"].copy()
extra_dict["store_type"] = self.store_type
result["extra"] = extra_dict
return result
@@ -212,12 +208,12 @@ class MinioConfig(FileStoreConfig):
def to_dict(self) -> dict[str, Any]:
result = super().to_dict()
if 'extra' not in result:
result['extra'] = dict()
extra_dict = result['extra'].copy()
extra_dict['user'] = self.user
extra_dict['password'] = self.password
result['extra'] = extra_dict
if "extra" not in result:
result["extra"] = dict()
extra_dict = result["extra"].copy()
extra_dict["user"] = self.user
extra_dict["password"] = self.password
result["extra"] = extra_dict
return result
@@ -229,106 +225,105 @@ def load_configurations(config_path: str) -> list[BaseConfig]:
for k, v in raw_configs.items():
match k:
case "ragflow":
name: str = f'ragflow_{ragflow_count}'
host: str = v['host']
http_port: int = v['http_port']
config = RAGFlowServerConfig(id=id_count, name=name, host=host, port=http_port,
service_type="ragflow_server",
detail_func_name="check_ragflow_server_alive")
name: str = f"ragflow_{ragflow_count}"
host: str = v["host"]
http_port: int = v["http_port"]
config = RAGFlowServerConfig(id=id_count, name=name, host=host, port=http_port, service_type="ragflow_server", detail_func_name="check_ragflow_server_alive")
configurations.append(config)
id_count += 1
case "es":
name: str = 'elasticsearch'
url = v['hosts']
name: str = "elasticsearch"
url = v["hosts"]
parsed = urlparse(url)
host: str = parsed.hostname
port: int = parsed.port
username: str = v.get('username')
password: str = v.get('password')
config = ElasticsearchConfig(id=id_count, name=name, host=host, port=port, service_type="retrieval",
retrieval_type="elasticsearch",
username=username, password=password,
detail_func_name="get_es_cluster_stats")
username: str = v.get("username")
password: str = v.get("password")
config = ElasticsearchConfig(
id=id_count,
name=name,
host=host,
port=port,
service_type="retrieval",
retrieval_type="elasticsearch",
username=username,
password=password,
detail_func_name="get_es_cluster_stats",
)
configurations.append(config)
id_count += 1
case "infinity":
name: str = 'infinity'
url = v['uri']
parts = url.split(':', 1)
name: str = "infinity"
url = v["uri"]
parts = url.split(":", 1)
host = parts[0]
port = int(parts[1])
database: str = v.get('db_name', 'default_db')
config = InfinityConfig(id=id_count, name=name, host=host, port=port, service_type="retrieval",
retrieval_type="infinity",
db_name=database, detail_func_name="get_infinity_status")
database: str = v.get("db_name", "default_db")
config = InfinityConfig(id=id_count, name=name, host=host, port=port, service_type="retrieval", retrieval_type="infinity", db_name=database, detail_func_name="get_infinity_status")
configurations.append(config)
id_count += 1
case "minio_0":
name: str = 'minio_0'
url = v['host']
parts = url.split(':', 1)
name: str = "minio_0"
url = v["host"]
parts = url.split(":", 1)
host = parts[0]
port = int(parts[1])
user = v.get('user')
password = v.get('password')
config = MinioConfig(id=id_count, name=name, host=host, port=port, user=user, password=password,
service_type="file_store",
store_type="minio", detail_func_name="check_minio_alive")
user = v.get("user")
password = v.get("password")
config = MinioConfig(id=id_count, name=name, host=host, port=port, user=user, password=password, service_type="file_store", store_type="minio", detail_func_name="check_minio_alive")
configurations.append(config)
id_count += 1
case "minio":
name: str = 'minio'
url = v['host']
parts = url.split(':', 1)
name: str = "minio"
url = v["host"]
parts = url.split(":", 1)
host = parts[0]
port = int(parts[1])
user = v.get('user')
password = v.get('password')
config = MinioConfig(id=id_count, name=name, host=host, port=port, user=user, password=password,
service_type="file_store",
store_type="minio", detail_func_name="check_minio_alive")
user = v.get("user")
password = v.get("password")
config = MinioConfig(id=id_count, name=name, host=host, port=port, user=user, password=password, service_type="file_store", store_type="minio", detail_func_name="check_minio_alive")
configurations.append(config)
id_count += 1
case "redis":
name: str = 'redis'
url = v['host']
parts = url.split(':', 1)
name: str = "redis"
url = v["host"]
parts = url.split(":", 1)
host = parts[0]
port = int(parts[1])
password = v.get('password')
db: int = v.get('db')
config = RedisConfig(id=id_count, name=name, host=host, port=port, password=password, database=db,
service_type="message_queue", mq_type="redis", detail_func_name="get_redis_info")
password = v.get("password")
db: int = v.get("db")
config = RedisConfig(id=id_count, name=name, host=host, port=port, password=password, database=db, service_type="message_queue", mq_type="redis", detail_func_name="get_redis_info")
configurations.append(config)
id_count += 1
case "mysql":
name: str = 'mysql'
host: str = v.get('host')
port: int = v.get('port')
username = v.get('user')
password = v.get('password')
config = MySQLConfig(id=id_count, name=name, host=host, port=port, username=username, password=password,
service_type="meta_data", meta_type="mysql", detail_func_name="get_mysql_status")
name: str = "mysql"
host: str = v.get("host")
port: int = v.get("port")
username = v.get("user")
password = v.get("password")
config = MySQLConfig(
id=id_count, name=name, host=host, port=port, username=username, password=password, service_type="meta_data", meta_type="mysql", detail_func_name="get_mysql_status"
)
configurations.append(config)
id_count += 1
case "admin":
pass
case "task_executor":
name: str = 'task_executor'
host: str = v.get('host', '')
port: int = v.get('port', 0)
message_queue_type: str = v.get('message_queue_type')
config = TaskExecutorConfig(id=id_count, name=name, host=host, port=port, message_queue_type=message_queue_type,
service_type="task_executor", detail_func_name="check_task_executor_alive")
name: str = "task_executor"
host: str = v.get("host", "")
port: int = v.get("port", 0)
message_queue_type: str = v.get("message_queue_type")
config = TaskExecutorConfig(
id=id_count, name=name, host=host, port=port, message_queue_type=message_queue_type, service_type="task_executor", detail_func_name="check_task_executor_alive"
)
configurations.append(config)
id_count += 1
case "rabbitmq":
name: str = 'rabbitmq'
host: str = v.get('host')
port: int = v.get('port')
config = RabbitMQConfig(id=id_count, name=name, host=host, port=port,
service_type="message_queue", mq_type="rabbitmq", detail_func_name="check_rabbitmq_alive")
name: str = "rabbitmq"
host: str = v.get("host")
port: int = v.get("port")
config = RabbitMQConfig(id=id_count, name=name, host=host, port=port, service_type="message_queue", mq_type="rabbitmq", detail_func_name="check_rabbitmq_alive")
configurations.append(config)
id_count += 1
case _:
admin/server/exceptions.py
+4 -1
View File
@@ -4,14 +4,17 @@ class AdminException(Exception):
self.code = code
self.message = message
class UserNotFoundError(AdminException):
def __init__(self, username):
super().__init__(f"User '{username}' not found", 404)
class UserAlreadyExistsError(AdminException):
def __init__(self, username):
super().__init__(f"User '{username}' already exists", 409)
class CannotDeleteAdminError(AdminException):
def __init__(self):
super().__init__("Cannot delete admin account", 403)
super().__init__("Cannot delete admin account", 403)
admin/server/responses.py
+2 -10
View File
@@ -17,16 +17,8 @@ from flask import jsonify
def success_response(data=None, message="Success", code=0):
return jsonify({
"code": code,
"message": message,
"data": data
}), 200
return jsonify({"code": code, "message": message, "data": data}), 200
def error_response(message="Error", code=-1, data=None):
return jsonify({
"code": code,
"message": message,
"data": data
}), 400
return jsonify({"code": code, "message": message, "data": data}), 400
admin/server/services.py
+5 -10
View File
@@ -489,10 +489,7 @@ class SandboxMgr:
"""List all available sandbox providers."""
result = []
for provider_id, metadata in SandboxMgr.PROVIDER_REGISTRY.items():
result.append({
"id": provider_id,
**metadata
})
result.append({"id": provider_id, **metadata})
return result
@staticmethod
@@ -635,6 +632,7 @@ class SandboxMgr:
config_json = json.dumps(config)
SettingsMgr.update_by_name(f"sandbox.{provider_type}", config_json)
from agent.sandbox.client import reload_provider
reload_provider()
return {"provider_type": provider_type, "config": config}
@@ -727,11 +725,7 @@ def main() -> dict:
# Build detailed result message
success = execution_result.exit_code == 0 and "TEST_PASSED" in execution_result.stdout
message_parts = [
f"Test {success and 'PASSED' or 'FAILED'}",
f"Exit code: {execution_result.exit_code}",
f"Execution time: {execution_result.execution_time:.2f}s"
]
message_parts = [f"Test {success and 'PASSED' or 'FAILED'}", f"Exit code: {execution_result.exit_code}", f"Execution time: {execution_result.execution_time:.2f}s"]
if execution_result.stdout.strip():
stdout_preview = execution_result.stdout.strip()[:200]
@@ -751,12 +745,13 @@ def main() -> dict:
"execution_time": execution_result.execution_time,
"stdout": execution_result.stdout,
"stderr": execution_result.stderr,
}
},
}
except AdminException:
raise
except Exception as e:
import traceback
error_details = traceback.format_exc()
raise AdminException(f"Connection test failed: {str(e)}\\n\\nStack trace:\\n{error_details}")
agent/canvas.py
+1
View File
@@ -161,6 +161,7 @@ class Graph:
def close(self):
from common.mcp_tool_call_conn import MCPToolCallSession
seen = set()
for cpn in self.components.values():
obj = cpn.get("obj")
agent/component/__init__.py
+5 -3
View File
@@ -22,8 +22,9 @@ from typing import Dict, Type
_package_path = os.path.dirname(__file__)
__all_classes: Dict[str, Type] = {}
def _import_submodules() -> None:
for filename in os.listdir(_package_path): # noqa: F821
for filename in os.listdir(_package_path): # noqa: F821
if filename.startswith("__") or not filename.endswith(".py") or filename.startswith("base"):
continue
module_name = filename[:-3]
@@ -34,13 +35,14 @@ def _import_submodules() -> None:
except ImportError as e:
print(f"Warning: Failed to import module {module_name}: {str(e)}")
def _extract_classes_from_module(module: ModuleType) -> None:
for name, obj in inspect.getmembers(module):
if (inspect.isclass(obj) and
obj.__module__ == module.__name__ and not name.startswith("_")):
if inspect.isclass(obj) and obj.__module__ == module.__name__ and not name.startswith("_"):
__all_classes[name] = obj
globals()[name] = obj
_import_submodules()
__all__ = list(__all_classes.keys()) + ["__all_classes"]
agent/component/base.py
+23 -52
View File
@@ -28,7 +28,6 @@ from agent import settings
from common.connection_utils import timeout
from common.misc_utils import thread_pool_exec
_logger = logging.getLogger(__name__)
@@ -101,6 +100,7 @@ class ComponentParamBase(ABC):
return None
logging.warning("ComponentParamBase.__str__: JSON fallback via str() for type=%s", type(obj).__name__)
return str(obj)
return json.dumps(self.as_dict(), ensure_ascii=False, default=_serialize_default)
def as_dict(self):
@@ -135,15 +135,11 @@ class ComponentParamBase(ABC):
update_from_raw_conf = conf.get(_IS_RAW_CONF, True)
if update_from_raw_conf:
deprecated_params_set = self._get_or_init_deprecated_params_set()
feeded_deprecated_params_set = (
self._get_or_init_feeded_deprecated_params_set()
)
feeded_deprecated_params_set = self._get_or_init_feeded_deprecated_params_set()
user_feeded_params_set = self._get_or_init_user_feeded_params_set()
setattr(self, _IS_RAW_CONF, False)
else:
feeded_deprecated_params_set = (
self._get_or_init_feeded_deprecated_params_set(conf)
)
feeded_deprecated_params_set = self._get_or_init_feeded_deprecated_params_set(conf)
user_feeded_params_set = self._get_or_init_user_feeded_params_set(conf)
def _recursive_update_param(param, config, depth, prefix):
@@ -179,15 +175,11 @@ class ComponentParamBase(ABC):
else:
# recursive set obj attr
sub_params = _recursive_update_param(
attr, config_value, depth + 1, prefix=f"{prefix}{config_key}."
)
sub_params = _recursive_update_param(attr, config_value, depth + 1, prefix=f"{prefix}{config_key}.")
setattr(param, config_key, sub_params)
if not allow_redundant and redundant_attrs:
raise ValueError(
f"cpn `{getattr(self, '_name', type(self))}` has redundant parameters: `{[redundant_attrs]}`"
)
raise ValueError(f"cpn `{getattr(self, '_name', type(self))}` has redundant parameters: `{[redundant_attrs]}`")
return param
@@ -218,9 +210,7 @@ class ComponentParamBase(ABC):
param_validation_path_prefix = home_dir + "/param_validation/"
param_name = type(self).__name__
param_validation_path = "/".join(
[param_validation_path_prefix, param_name + ".json"]
)
param_validation_path = "/".join([param_validation_path_prefix, param_name + ".json"])
validation_json = None
@@ -253,11 +243,7 @@ class ComponentParamBase(ABC):
break
if not value_legal:
raise ValueError(
"Please check runtime conf, {} = {} does not match user-parameter restriction".format(
variable, value
)
)
raise ValueError("Please check runtime conf, {} = {} does not match user-parameter restriction".format(variable, value))
elif variable in validation_json:
self._validate_param(attr, validation_json)
@@ -335,11 +321,7 @@ class ComponentParamBase(ABC):
def _range(value, ranges):
in_range = False
for left_limit, right_limit in ranges:
if (
left_limit - settings.FLOAT_ZERO
<= value
<= right_limit + settings.FLOAT_ZERO
):
if left_limit - settings.FLOAT_ZERO <= value <= right_limit + settings.FLOAT_ZERO:
in_range = True
break
@@ -355,16 +337,11 @@ class ComponentParamBase(ABC):
def _warn_deprecated_param(self, param_name, description):
if self._deprecated_params_set.get(param_name):
logging.warning(
f"{description} {param_name} is deprecated and ignored in this version."
)
logging.warning(f"{description} {param_name} is deprecated and ignored in this version.")
def _warn_to_deprecate_param(self, param_name, description, new_param):
if self._deprecated_params_set.get(param_name):
logging.warning(
f"{description} {param_name} will be deprecated in future release; "
f"please use {new_param} instead."
)
logging.warning(f"{description} {param_name} will be deprecated in future release; please use {new_param} instead.")
return True
return False
@@ -385,9 +362,7 @@ class ComponentBase(ABC):
return """{{
"component_name": "{}",
"params": {}
}}""".format(self.component_name,
self._param
)
}}""".format(self.component_name, self._param)
def __init__(self, canvas, id, param: ComponentParamBase):
from agent.canvas import Graph # Local import to avoid cyclic dependency
@@ -403,7 +378,7 @@ class ComponentBase(ABC):
def check_if_canceled(self, message: str = "") -> bool:
if self.is_canceled():
task_id = getattr(self._canvas, 'task_id', 'unknown')
task_id = getattr(self._canvas, "task_id", "unknown")
log_message = f"Task {task_id} has been canceled"
if message:
log_message += f" during {message}"
@@ -491,7 +466,9 @@ class ComponentBase(ABC):
input_elements = self.get_input_elements()
_logger.debug(
"[Base] Component '%s' (%s) resolving inputs. Input element keys: %s",
self._id, self.component_name, list(input_elements.keys()),
self._id,
self.component_name,
list(input_elements.keys()),
)
for var, o in input_elements.items():
v = self.get_param(var)
@@ -504,7 +481,7 @@ class ComponentBase(ABC):
_logger.debug("[Base] var '%s': resolved ref '%s' -> %s", var, v, json.dumps(resolved, ensure_ascii=False, default=str)[:200])
elif isinstance(v, str) and re.search(self.variable_ref_patt, v):
elements = self.get_input_elements_from_text(v)
kv = {k: e.get('value', '') for k, e in elements.items()}
kv = {k: e.get("value", "") for k, e in elements.items()}
self.set_input_value(var, self.string_format(v, kv))
_logger.debug("[Base] var '%s': resolved text refs '%s' -> %s", var, v, json.dumps(kv, ensure_ascii=False, default=str)[:200])
else:
@@ -545,7 +522,7 @@ class ComponentBase(ABC):
"name": (self._canvas.get_component_name(cpn_id) + f"@{var_nm}") if cpn_id else exp,
"value": self._canvas.get_variable_value(exp),
"_retrieval": self._canvas.get_variable_value(f"{cpn_id}@_references") if cpn_id else None,
"_cpn_id": cpn_id
"_cpn_id": cpn_id,
}
for r in re.finditer(self.iteration_alias_patt, txt, flags=re.IGNORECASE | re.DOTALL):
exp = r.group(1)
@@ -559,7 +536,7 @@ class ComponentBase(ABC):
"name": (self._canvas.get_component_name(cpn_id) + f"@{var_nm}"),
"value": self._canvas.get_variable_value(ref),
"_retrieval": self._canvas.get_variable_value(f"{cpn_id}@_references"),
"_cpn_id": cpn_id
"_cpn_id": cpn_id,
}
return res
@@ -601,33 +578,27 @@ class ComponentBase(ABC):
return self._canvas.get_component(pid)["obj"]
def get_upstream(self) -> List[str]:
cpn_nms = self._canvas.get_component(self._id)['upstream']
cpn_nms = self._canvas.get_component(self._id)["upstream"]
return cpn_nms
def get_downstream(self) -> List[str]:
cpn_nms = self._canvas.get_component(self._id)['downstream']
cpn_nms = self._canvas.get_component(self._id)["downstream"]
return cpn_nms
@staticmethod
def string_format(content: str, kv: dict[str, str]) -> str:
for n, v in kv.items():
def repl(_match, val=v):
return str(val) if val is not None else ""
content = re.sub(
r"\{%s\}" % re.escape(n),
repl,
content
)
content = re.sub(r"\{%s\}" % re.escape(n), repl, content)
return content
def exception_handler(self):
if not self._param.exception_method:
return None
return {
"goto": self._param.exception_goto,
"default_value": self._param.exception_default_value
}
return {"goto": self._param.exception_goto, "default_value": self._param.exception_default_value}
def get_exception_default_value(self):
if self._param.exception_method != "comment":
agent/component/begin.py
+2 -2
View File
@@ -17,17 +17,17 @@ from agent.component.fillup import UserFillUpParam, UserFillUp
class BeginParam(UserFillUpParam):
"""
Define the Begin component parameters.
"""
def __init__(self):
super().__init__()
self.mode = "conversational"
self.prologue = "Hi! I'm your smart assistant. What can I do for you?"
def check(self):
self.check_valid_value(self.mode, "The 'mode' should be either `conversational` or `task`", ["conversational", "task","Webhook"])
self.check_valid_value(self.mode, "The 'mode' should be either `conversational` or `task`", ["conversational", "task", "Webhook"])
def get_input_form(self) -> dict[str, dict]:
return getattr(self, "inputs")
agent/component/browser.py
+6 -23
View File
@@ -371,9 +371,7 @@ class Browser(ComponentBase, ABC):
sort_keys=True,
ensure_ascii=False,
)
raw_canvas_id = (
f"dsl_{hashlib.sha1(graph_text.encode('utf-8')).hexdigest()[:12]}"
)
raw_canvas_id = f"dsl_{hashlib.sha1(graph_text.encode('utf-8')).hexdigest()[:12]}"
canvas_id = self._safe_path_segment(raw_canvas_id)
node_id = self._safe_path_segment(self._id)
return os.path.join(root, tenant, canvas_id, node_id)
@@ -488,10 +486,7 @@ class Browser(ComponentBase, ABC):
# Keep browser-use watchdog fallback in sync with our resolved path.
os.environ["BROWSER_USE_BROWSER_BINARY_PATH"] = executable_path
else:
logging.warning(
"Browser no local browser executable found. "
"Set BROWSER_USE_EXECUTABLE_PATH or preinstall chromium in image to avoid runtime playwright install."
)
logging.warning("Browser no local browser executable found. Set BROWSER_USE_EXECUTABLE_PATH or preinstall chromium in image to avoid runtime playwright install.")
if profile_dir:
browser_kwargs["user_data_dir"] = profile_dir
# browser-use expects profile_directory to be a profile name
@@ -682,21 +677,13 @@ class Browser(ComponentBase, ABC):
try:
self._prepare_input_values()
user_prompt = self._resolve_text(kwargs.get("prompts", self._param.prompts))
with tempfile.TemporaryDirectory(prefix="browser_use_upload_") as upload_dir, tempfile.TemporaryDirectory(
prefix="browser_use_download_"
) as download_dir:
with tempfile.TemporaryDirectory(prefix="browser_use_upload_") as upload_dir, tempfile.TemporaryDirectory(prefix="browser_use_download_") as download_dir:
uploaded_files = self._prepare_upload_files(upload_dir)
upload_lines = [
f"- file_id={item['file_id']}, name={item['name']}, local_path={item['local_path']}"
for item in uploaded_files
]
upload_lines = [f"- file_id={item['file_id']}, name={item['name']}, local_path={item['local_path']}" for item in uploaded_files]
task_text = user_prompt
if upload_lines:
task_text += (
"\n\nYou can upload files from these local paths when operating web pages:\n"
+ "\n".join(upload_lines)
)
task_text += "\n\nYou can upload files from these local paths when operating web pages:\n" + "\n".join(upload_lines)
upload_local_paths = [item.get("local_path", "") for item in uploaded_files if item.get("local_path")]
if persist_session:
@@ -707,11 +694,7 @@ class Browser(ComponentBase, ABC):
profile_dir = tempfile.mkdtemp(prefix="browser_use_profile_")
except OSError:
profile_dir = None
history = asyncio.run(
self._run_browser_use_async(
task_text, download_dir, upload_local_paths, profile_dir
)
)
history = asyncio.run(self._run_browser_use_async(task_text, download_dir, upload_local_paths, profile_dir))
target_dir_id = FileService.get_root_folder(self._canvas.get_tenant_id())["id"]
downloaded_files = self._save_downloads(download_dir, target_dir_id)
agent/component/categorize.py
+8 -17
View File
@@ -28,10 +28,10 @@ from rag.llm.chat_model import ERROR_PREFIX
class CategorizeParam(LLMParam):
"""
Define the categorize component parameters.
"""
def __init__(self):
super().__init__()
self.category_description = {}
@@ -50,12 +50,7 @@ class CategorizeParam(LLMParam):
raise ValueError(f"[Categorize] 'To' of category {k} can not be empty!")
def get_input_form(self) -> dict[str, dict]:
return {
"query": {
"type": "line",
"name": "Query"
}
}
return {"query": {"type": "line", "name": "Query"}}
def update_prompt(self):
cate_lines = []
@@ -63,13 +58,12 @@ class CategorizeParam(LLMParam):
for line in desc.get("examples", []):
if not line:
continue
cate_lines.append("USER: \"" + re.sub(r"\n", " ", line, flags=re.DOTALL) + "\" → "+c)
cate_lines.append('USER: "' + re.sub(r"\n", " ", line, flags=re.DOTALL) + '" → ' + c)
descriptions = []
for c, desc in self.category_description.items():
if desc.get("description"):
descriptions.append(
"\n------\nCategory: {}\nDescription: {}".format(c, desc["description"]))
descriptions.append("\n------\nCategory: {}\nDescription: {}".format(c, desc["description"]))
self.sys_prompt = """
You are an advanced classification system that categorizes user questions into specific types. Analyze the input question and classify it into ONE of the following categories:
@@ -84,10 +78,7 @@ Here's description of each category:
- Return only the category name without explanations
- Use "Other" only when no other category fits
""".format(
"\n - ".join(list(self.category_description.keys())),
"\n".join(descriptions)
)
""".format("\n - ".join(list(self.category_description.keys())), "\n".join(descriptions))
if cate_lines:
self.sys_prompt += """
@@ -106,7 +97,7 @@ class Categorize(LLM, ABC):
logging.warning(f"[Categorize] input element not detected for query key: {query_key}")
return elements
@timeout(int(os.environ.get("COMPONENT_EXEC_TIMEOUT", 10*60)))
@timeout(int(os.environ.get("COMPONENT_EXEC_TIMEOUT", 10 * 60)))
async def _invoke_async(self, **kwargs):
if self.check_if_canceled("Categorize processing"):
return
@@ -130,7 +121,7 @@ class Categorize(LLM, ABC):
user_prompt = """
---- Real Data ----
{} →
""".format(" | ".join(["{}: \"{}\"".format(c["role"].upper(), re.sub(r"\n", "", c["content"], flags=re.DOTALL)) for c in msg]))
""".format(" | ".join(['{}: "{}"'.format(c["role"].upper(), re.sub(r"\n", "", c["content"], flags=re.DOTALL)) for c in msg]))
if self.check_if_canceled("Categorize processing"):
return
@@ -158,7 +149,7 @@ class Categorize(LLM, ABC):
self.set_output("category_name", max_category)
self.set_output("_next", cpn_ids)
@timeout(int(os.environ.get("COMPONENT_EXEC_TIMEOUT", 10*60)))
@timeout(int(os.environ.get("COMPONENT_EXEC_TIMEOUT", 10 * 60)))
def _invoke(self, **kwargs):
return asyncio.run(self._invoke_async(**kwargs))
agent/component/data_operations.py
+34 -50
View File
@@ -19,55 +19,47 @@ import os
from agent.component.base import ComponentBase, ComponentParamBase
from api.utils.api_utils import timeout
class DataOperationsParam(ComponentParamBase):
"""
Define the Data Operations component parameters.
"""
def __init__(self):
super().__init__()
self.query = []
self.operations = "literal_eval"
self.select_keys = []
self.filter_values=[]
self.updates=[]
self.remove_keys=[]
self.rename_keys=[]
self.outputs = {
"result": {
"value": [],
"type": "Array of Object"
}
}
def check(self):
self.check_valid_value(self.operations, "Support operations", ["select_keys", "literal_eval","combine","filter_values","append_or_update","remove_keys","rename_keys"])
self.filter_values = []
self.updates = []
self.remove_keys = []
self.rename_keys = []
self.outputs = {"result": {"value": [], "type": "Array of Object"}}
class DataOperations(ComponentBase,ABC):
def check(self):
self.check_valid_value(self.operations, "Support operations", ["select_keys", "literal_eval", "combine", "filter_values", "append_or_update", "remove_keys", "rename_keys"])
class DataOperations(ComponentBase, ABC):
component_name = "DataOperations"
def get_input_form(self) -> dict[str, dict]:
return {
k: {"name": o.get("name", ""), "type": "line"}
for input_item in (self._param.query or [])
for k, o in self.get_input_elements_from_text(input_item).items()
}
return {k: {"name": o.get("name", ""), "type": "line"} for input_item in (self._param.query or []) for k, o in self.get_input_elements_from_text(input_item).items()}
@timeout(int(os.environ.get("COMPONENT_EXEC_TIMEOUT", 10*60)))
@timeout(int(os.environ.get("COMPONENT_EXEC_TIMEOUT", 10 * 60)))
def _invoke(self, **kwargs):
self.input_objects=[]
self.input_objects = []
inputs = getattr(self._param, "query", None)
if not isinstance(inputs, (list, tuple)):
inputs = [inputs]
for input_ref in inputs:
input_object=self._canvas.get_variable_value(input_ref)
input_object = self._canvas.get_variable_value(input_ref)
self.set_input_value(input_ref, input_object)
if input_object is None:
continue
if isinstance(input_object,dict):
if isinstance(input_object, dict):
self.input_objects.append(input_object)
elif isinstance(input_object,list):
elif isinstance(input_object, list):
self.input_objects.extend(x for x in input_object if isinstance(x, dict))
else:
continue
@@ -85,13 +77,12 @@ class DataOperations(ComponentBase,ABC):
self._remove_keys()
else:
self._rename_keys()
def _select_keys(self):
filter_criteria: list[str] = self._param.select_keys
results = [{key: value for key, value in data_dict.items() if key in filter_criteria} for data_dict in self.input_objects]
self.set_output("result", results)
def _recursive_eval(self, data):
if isinstance(data, dict):
return {k: self._recursive_eval(v) for k, v in data.items()}
@@ -99,23 +90,19 @@ class DataOperations(ComponentBase,ABC):
return [self._recursive_eval(item) for item in data]
if isinstance(data, str):
try:
if (
data.strip().startswith(("{", "[", "(", "'", '"'))
or data.strip().lower() in ("true", "false", "none")
or data.strip().replace(".", "").isdigit()
):
if data.strip().startswith(("{", "[", "(", "'", '"')) or data.strip().lower() in ("true", "false", "none") or data.strip().replace(".", "").isdigit():
return ast.literal_eval(data)
except (ValueError, SyntaxError, TypeError, MemoryError):
return data
else:
return data
return data
def _literal_eval(self):
self.set_output("result", self._recursive_eval(self.input_objects))
def _combine(self):
result={}
result = {}
for obj in self.input_objects:
for key, value in obj.items():
if key not in result:
@@ -126,15 +113,13 @@ class DataOperations(ComponentBase,ABC):
else:
result[key].append(value)
else:
result[key] = (
[result[key], value] if not isinstance(value, list) else [result[key], *value]
)
result[key] = [result[key], value] if not isinstance(value, list) else [result[key], *value]
self.set_output("result", result)
def norm(self,v):
def norm(self, v):
s = "" if v is None else str(v)
return s
def match_rule(self, obj, rule):
key = rule.get("key")
op = (rule.get("operator") or "equals").lower()
@@ -155,10 +140,10 @@ class DataOperations(ComponentBase,ABC):
if op == "end with":
return v.endswith(target)
return False
def _filter_values(self):
results=[]
rules = (getattr(self._param, "filter_values", None) or [])
results = []
rules = getattr(self._param, "filter_values", None) or []
for obj in self.input_objects:
if not rules:
results.append(obj)
@@ -166,11 +151,10 @@ class DataOperations(ComponentBase,ABC):
if all(self.match_rule(obj, r) for r in rules):
results.append(obj)
self.set_output("result", results)
def _append_or_update(self):
results=[]
updates = getattr(self._param, "updates", []) or []
results = []
updates = getattr(self._param, "updates", []) or []
for obj in self.input_objects:
new_obj = dict(obj)
for item in updates:
@@ -187,7 +171,7 @@ class DataOperations(ComponentBase,ABC):
results = []
remove_keys = getattr(self._param, "remove_keys", []) or []
for obj in (self.input_objects or []):
for obj in self.input_objects or []:
new_obj = dict(obj)
for k in remove_keys:
if not isinstance(k, str):
@@ -200,7 +184,7 @@ class DataOperations(ComponentBase,ABC):
results = []
rename_pairs = getattr(self._param, "rename_keys", []) or []
for obj in (self.input_objects or []):
for obj in self.input_objects or []:
new_obj = dict(obj)
for pair in rename_pairs:
if not isinstance(pair, dict):
agent/component/docs_generator.py
+1
View File
@@ -163,6 +163,7 @@ class DocGenerator(Message, ABC):
logging.info("Starting document generation, content length: %s chars", len(content))
if content:
def _replace_variable(match_obj: re.Match[str]) -> str:
match = match_obj.group(1)
try:
agent/component/excel_processor.py
+58 -80
View File
@@ -39,71 +39,52 @@ class ExcelProcessorParam(ComponentParamBase):
"""
Define the ExcelProcessor component parameters.
"""
def __init__(self):
super().__init__()
# Input configuration
self.input_files = [] # Variable references to uploaded files
self.operation = "read" # read, merge, transform, output
# Processing options
self.sheet_selection = "all" # all, first, or comma-separated sheet names
self.merge_strategy = "concat" # concat, join
self.join_on = "" # Column name for join operations
# Transform options (for LLM-guided transformations)
self.transform_instructions = ""
self.transform_data = "" # Variable reference to transformation data
# Output options
self.output_format = "xlsx" # xlsx, csv
self.output_filename = "output"
# Component outputs
self.outputs = {
"data": {
"type": "object",
"value": {}
},
"summary": {
"type": "str",
"value": ""
},
"markdown": {
"type": "str",
"value": ""
}
}
self.outputs = {"data": {"type": "object", "value": {}}, "summary": {"type": "str", "value": ""}, "markdown": {"type": "str", "value": ""}}
def check(self):
self.check_valid_value(
self.operation,
"[ExcelProcessor] Operation",
["read", "merge", "transform", "output"]
)
self.check_valid_value(
self.output_format,
"[ExcelProcessor] Output format",
["xlsx", "csv"]
)
self.check_valid_value(self.operation, "[ExcelProcessor] Operation", ["read", "merge", "transform", "output"])
self.check_valid_value(self.output_format, "[ExcelProcessor] Output format", ["xlsx", "csv"])
return True
class ExcelProcessor(ComponentBase, ABC):
"""
Excel processing component for RAGFlow agents.
Operations:
- read: Parse Excel files into structured data
- merge: Combine multiple Excel files
- transform: Apply data transformations based on instructions
- output: Generate Excel file output
"""
component_name = "ExcelProcessor"
def get_input_form(self) -> dict[str, dict]:
"""Define input form for the component."""
res = {}
for ref in (self._param.input_files or []):
for ref in self._param.input_files or []:
for k, o in self.get_input_elements_from_text(ref).items():
res[k] = {"name": o.get("name", ""), "type": "file"}
if self._param.transform_data:
@@ -111,13 +92,13 @@ class ExcelProcessor(ComponentBase, ABC):
res[k] = {"name": o.get("name", ""), "type": "object"}
return res
@timeout(int(os.environ.get("COMPONENT_EXEC_TIMEOUT", 10*60)))
@timeout(int(os.environ.get("COMPONENT_EXEC_TIMEOUT", 10 * 60)))
def _invoke(self, **kwargs):
if self.check_if_canceled("ExcelProcessor processing"):
return
operation = self._param.operation.lower()
if operation == "read":
self._read_excels()
elif operation == "merge":
@@ -137,7 +118,7 @@ class ExcelProcessor(ComponentBase, ABC):
value = self._canvas.get_variable_value(file_ref)
if value is None:
return None, None
# Handle different value formats
if isinstance(value, dict):
# File reference from Begin/UserFillUp component
@@ -154,12 +135,13 @@ class ExcelProcessor(ComponentBase, ABC):
# Could be base64 encoded or a path
if value.startswith("data:"):
import base64
# Extract base64 content
_, encoded = value.split(",", 1)
return base64.b64decode(encoded), "uploaded.xlsx"
return None, None
def _get_file_content_from_list(self, item) -> tuple[bytes, str]:
"""Extract file content from a list item."""
if isinstance(item, dict):
@@ -170,15 +152,15 @@ class ExcelProcessor(ComponentBase, ABC):
"""Parse Excel content into a dictionary of DataFrames (one per sheet)."""
try:
excel_file = BytesIO(content)
if filename.lower().endswith(".csv"):
df = pd.read_csv(excel_file)
return {"Sheet1": df}
else:
# Read all sheets
xlsx = pd.ExcelFile(excel_file, engine='openpyxl')
xlsx = pd.ExcelFile(excel_file, engine="openpyxl")
sheet_selection = self._param.sheet_selection
if sheet_selection == "all":
sheets_to_read = xlsx.sheet_names
elif sheet_selection == "first":
@@ -187,12 +169,12 @@ class ExcelProcessor(ComponentBase, ABC):
# Comma-separated sheet names
requested = [s.strip() for s in sheet_selection.split(",")]
sheets_to_read = [s for s in requested if s in xlsx.sheet_names]
dfs = {}
for sheet in sheets_to_read:
dfs[sheet] = pd.read_excel(xlsx, sheet_name=sheet)
return dfs
except Exception as e:
logging.error(f"Error parsing Excel file {filename}: {e}")
return {}
@@ -202,36 +184,36 @@ class ExcelProcessor(ComponentBase, ABC):
all_data = {}
summaries = []
markdown_parts = []
for file_ref in (self._param.input_files or []):
for file_ref in self._param.input_files or []:
if self.check_if_canceled("ExcelProcessor reading"):
return
# Get variable value
value = self._canvas.get_variable_value(file_ref)
self.set_input_value(file_ref, str(value)[:200] if value else "")
if value is None:
continue
# Handle file content
content, filename = self._get_file_content(file_ref)
if content is None:
continue
# Parse Excel
dfs = self._parse_excel_to_dataframes(content, filename)
for sheet_name, df in dfs.items():
key = f"{filename}_{sheet_name}" if len(dfs) > 1 else filename
all_data[key] = df.to_dict(orient="records")
# Build summary
summaries.append(f"**{key}**: {len(df)} rows, {len(df.columns)} columns ({', '.join(df.columns.tolist()[:5])}{'...' if len(df.columns) > 5 else ''})")
# Build markdown table
markdown_parts.append(f"### {key}\n\n{df.head(10).to_markdown(index=False)}\n")
# Set outputs
self.set_output("data", all_data)
self.set_output("summary", "\n".join(summaries) if summaries else "No Excel files found")
@@ -240,29 +222,29 @@ class ExcelProcessor(ComponentBase, ABC):
def _merge_excels(self):
"""Merge multiple Excel files/sheets into one."""
all_dfs = []
for file_ref in (self._param.input_files or []):
for file_ref in self._param.input_files or []:
if self.check_if_canceled("ExcelProcessor merging"):
return
value = self._canvas.get_variable_value(file_ref)
self.set_input_value(file_ref, str(value)[:200] if value else "")
if value is None:
continue
content, filename = self._get_file_content(file_ref)
if content is None:
continue
dfs = self._parse_excel_to_dataframes(content, filename)
all_dfs.extend(dfs.values())
if not all_dfs:
self.set_output("data", {})
self.set_output("summary", "No data to merge")
return
# Merge strategy
if self._param.merge_strategy == "concat":
merged_df = pd.concat(all_dfs, ignore_index=True)
@@ -273,7 +255,7 @@ class ExcelProcessor(ComponentBase, ABC):
merged_df = merged_df.merge(df, on=self._param.join_on, how="outer")
else:
merged_df = pd.concat(all_dfs, ignore_index=True)
self.set_output("data", {"merged": merged_df.to_dict(orient="records")})
self.set_output("summary", f"Merged {len(all_dfs)} sources into {len(merged_df)} rows, {len(merged_df.columns)} columns")
self.set_output("markdown", merged_df.head(20).to_markdown(index=False))
@@ -285,14 +267,14 @@ class ExcelProcessor(ComponentBase, ABC):
if not transform_ref:
self.set_output("summary", "No transform data reference provided")
return
data = self._canvas.get_variable_value(transform_ref)
self.set_input_value(transform_ref, str(data)[:300] if data else "")
if data is None:
self.set_output("summary", "Transform data is empty")
return
# Convert to DataFrame
if isinstance(data, dict):
# Could be {"sheet": [rows]} format
@@ -315,7 +297,7 @@ class ExcelProcessor(ComponentBase, ABC):
else:
self.set_output("data", {"raw": str(data)})
self.set_output("markdown", str(data))
self.set_output("summary", "Transformed data ready for processing")
def _output_excel(self):
@@ -325,14 +307,14 @@ class ExcelProcessor(ComponentBase, ABC):
if not transform_ref:
self.set_output("summary", "No data reference for output")
return
data = self._canvas.get_variable_value(transform_ref)
self.set_input_value(transform_ref, str(data)[:300] if data else "")
if data is None:
self.set_output("summary", "No data to output")
return
try:
# Prepare DataFrames
if isinstance(data, dict):
@@ -346,10 +328,10 @@ class ExcelProcessor(ComponentBase, ABC):
else:
self.set_output("summary", "Invalid data format for Excel output")
return
# Generate output
doc_id = get_uuid()
if self._param.output_format == "csv":
# For CSV, only output first sheet
first_df = list(dfs.values())[0]
@@ -358,7 +340,7 @@ class ExcelProcessor(ComponentBase, ABC):
else:
# Excel output
excel_io = BytesIO()
with pd.ExcelWriter(excel_io, engine='openpyxl') as writer:
with pd.ExcelWriter(excel_io, engine="openpyxl") as writer:
for sheet_name, df in dfs.items():
# Sanitize sheet name (max 31 chars, no special chars)
safe_name = sheet_name[:31].replace("/", "_").replace("\\", "_")
@@ -366,23 +348,19 @@ class ExcelProcessor(ComponentBase, ABC):
excel_io.seek(0)
binary_content = excel_io.read()
filename = f"{self._param.output_filename}.xlsx"
# Store file
settings.STORAGE_IMPL.put(self._canvas._tenant_id, doc_id, binary_content)
# Set attachment output
self.set_output("attachment", {
"doc_id": doc_id,
"format": self._param.output_format,
"file_name": filename
})
self.set_output("attachment", {"doc_id": doc_id, "format": self._param.output_format, "file_name": filename})
total_rows = sum(len(df) for df in dfs.values())
self.set_output("summary", f"Generated {filename} with {len(dfs)} sheet(s), {total_rows} total rows")
self.set_output("data", {k: v.to_dict(orient="records") for k, v in dfs.items()})
logging.info(f"ExcelProcessor: Generated {filename} as {doc_id}")
except Exception as e:
logging.error(f"ExcelProcessor output error: {e}")
self.set_output("summary", f"Error generating output: {str(e)}")
agent/component/exit_loop.py
+1 -1
View File
@@ -29,4 +29,4 @@ class ExitLoop(ComponentBase, ABC):
pass
def thoughts(self) -> str:
return ""
return ""
agent/component/fillup.py
+2 -7
View File
@@ -25,7 +25,6 @@ _INITIAL_USER_INPUT_CONSUMED_KEY = "sys.__initial_user_input_consumed__"
class UserFillUpParam(ComponentParamBase):
def __init__(self):
super().__init__()
self.enable_tips = True
@@ -55,11 +54,7 @@ class UserFillUp(ComponentBase):
return {}
if isinstance(query, dict):
matched = {
key: value if isinstance(value, dict) else {"value": value}
for key, value in query.items()
if key in fields
}
matched = {key: value if isinstance(value, dict) else {"value": value} for key, value in query.items() if key in fields}
if matched:
self._canvas.globals[_INITIAL_USER_INPUT_CONSUMED_KEY] = True
return matched
@@ -108,7 +103,7 @@ class UserFillUp(ComponentBase):
ans = v
if not ans:
ans = ""
content = re.sub(r"\{%s\}"%k, ans, content)
content = re.sub(r"\{%s\}" % k, ans, content)
self.set_output("tips", content)
layout_recognize = self._param.layout_recognize or None
agent/component/iteration.py
+4 -11
View File
@@ -24,6 +24,7 @@ class VariableModel(BaseModel):
model_config = ConfigDict(extra="forbid")
"""
class IterationParam(ComponentParamBase):
"""
Define the Iteration component parameters.
@@ -32,15 +33,10 @@ class IterationParam(ComponentParamBase):
def __init__(self):
super().__init__()
self.items_ref = ""
self.variable={}
self.variable = {}
def get_input_form(self) -> dict[str, dict]:
return {
"items": {
"type": "json",
"name": "Items"
}
}
return {"items": {"type": "json", "name": "Items"}}
def check(self):
return True
@@ -62,10 +58,7 @@ class Iteration(ComponentBase, ABC):
arr = self._canvas.get_variable_value(self._param.items_ref)
if not isinstance(arr, list):
self.set_output("_ERROR", self._param.items_ref + " must be an array, but its type is "+str(type(arr)))
self.set_output("_ERROR", self._param.items_ref + " must be an array, but its type is " + str(type(arr)))
def thoughts(self) -> str:
return "Need to process {} items.".format(len(self._canvas.get_variable_value(self._param.items_ref)))
agent/component/iterationitem.py
+2 -1
View File
@@ -21,6 +21,7 @@ class IterationItemParam(ComponentParamBase):
"""
Define the IterationItem component parameters.
"""
def check(self):
return True
@@ -40,7 +41,7 @@ class IterationItem(ComponentBase, ABC):
arr = self._canvas.get_variable_value(parent._param.items_ref)
if not isinstance(arr, list):
self._idx = -1
raise Exception(parent._param.items_ref + " must be an array, but its type is "+str(type(arr)))
raise Exception(parent._param.items_ref + " must be an array, but its type is " + str(type(arr)))
if self._idx > 0:
if self.check_if_canceled("IterationItem processing"):
agent/component/list_operations.py
+14 -31
View File
@@ -3,10 +3,12 @@ import os
from agent.component.base import ComponentBase, ComponentParamBase
from api.utils.api_utils import timeout
class ListOperationsParam(ComponentParamBase):
"""
Define the List Operations component parameters.
"""
def __init__(self):
super().__init__()
self.query = ""
@@ -20,24 +22,8 @@ class ListOperationsParam(ComponentParamBase):
# first field). Mirrors internal/agent/component/list_operations.go
# parseSortByFieldList + opSort's SortBy path.
self.sort_by = ""
self.filter = {
"operator": "=",
"value": ""
}
self.outputs = {
"result": {
"value": [],
"type": "Array of ?"
},
"first": {
"value": "",
"type": "?"
},
"last": {
"value": "",
"type": "?"
}
}
self.filter = {"operator": "=", "value": ""}
self.outputs = {"result": {"value": [], "type": "Array of ?"}, "first": {"value": "", "type": "?"}, "last": {"value": "", "type": "?"}}
@staticmethod
def _normalize_operation_name(operation):
@@ -45,7 +31,7 @@ class ListOperationsParam(ComponentParamBase):
if op.lower() == "topn":
return "head"
return op or "nth"
def check(self):
self.check_empty(self.query, "query")
self.operations = self._normalize_operation_name(self.operations)
@@ -57,14 +43,14 @@ class ListOperationsParam(ComponentParamBase):
def get_input_form(self) -> dict[str, dict]:
return {}
class ListOperations(ComponentBase,ABC):
class ListOperations(ComponentBase, ABC):
component_name = "ListOperations"
@timeout(int(os.environ.get("COMPONENT_EXEC_TIMEOUT", 10*60)))
@timeout(int(os.environ.get("COMPONENT_EXEC_TIMEOUT", 10 * 60)))
def _invoke(self, **kwargs):
self.input_objects=[]
self.input_objects = []
inputs = getattr(self._param, "query", None)
self.inputs = self._canvas.get_variable_value(inputs)
if not isinstance(self.inputs, list):
@@ -83,7 +69,6 @@ class ListOperations(ComponentBase,ABC):
elif self._param.operations == "drop_duplicates":
self._drop_duplicates()
def _coerce_n(self):
try:
return int(getattr(self._param, "n", 0))
@@ -99,12 +84,10 @@ class ListOperations(ComponentBase,ABC):
def _set_outputs(self, outputs):
self._param.outputs["result"]["value"] = outputs
self._param.outputs["first"]["value"] = outputs[0] if outputs else None
self._param.outputs["last"]["value"] = outputs[-1] if outputs else None
self._param.outputs["last"]["value"] = outputs[-1] if outputs else None
def _raise_strict_range_error(self, operation, n):
raise ValueError(
f"{operation} requires n to be within the valid range in strict mode, got {n}."
)
raise ValueError(f"{operation} requires n to be within the valid range in strict mode, got {n}.")
def _nth(self):
n = self._coerce_n()
@@ -160,9 +143,9 @@ class ListOperations(ComponentBase,ABC):
self._set_outputs(outputs)
def _filter(self):
self._set_outputs([i for i in self.inputs if self._eval(self._norm(i),self._param.filter["operator"],self._param.filter["value"])])
self._set_outputs([i for i in self.inputs if self._eval(self._norm(i), self._param.filter["operator"], self._param.filter["value"])])
def _norm(self,v):
def _norm(self, v):
s = "" if v is None else str(v)
return s
@@ -222,7 +205,7 @@ class ListOperations(ComponentBase,ABC):
outs.append(item)
self._set_outputs(outs)
def _hashable(self,x):
def _hashable(self, x):
if isinstance(x, dict):
return tuple(sorted((k, self._hashable(v)) for k, v in x.items()))
if isinstance(x, (list, tuple)):
agent/component/loop.py
+7 -13
View File
@@ -25,16 +25,11 @@ class LoopParam(ComponentParamBase):
def __init__(self):
super().__init__()
self.loop_variables = []
self.loop_termination_condition=[]
self.loop_termination_condition = []
self.maximum_loop_count = 0
def get_input_form(self) -> dict[str, dict]:
return {
"items": {
"type": "json",
"name": "Items"
}
}
return {"items": {"type": "json", "name": "Items"}}
def check(self):
return True
@@ -83,10 +78,10 @@ class Loop(ComponentBase, ABC):
for item in self._param.loop_variables:
if self._is_incomplete_loop_variable(item):
raise ValueError("Loop Variable is not complete.")
if item["input_mode"]=="variable":
self.set_output(item["variable"],self._canvas.get_variable_value(item["value"]))
elif item["input_mode"]=="constant":
self.set_output(item["variable"],item["value"])
if item["input_mode"] == "variable":
self.set_output(item["variable"], self._canvas.get_variable_value(item["value"]))
elif item["input_mode"] == "constant":
self.set_output(item["variable"], item["value"])
else:
if item["type"] == "number":
self.set_output(item["variable"], 0)
@@ -101,6 +96,5 @@ class Loop(ComponentBase, ABC):
else:
self.set_output(item["variable"], "")
def thoughts(self) -> str:
return "Loop from canvas."
return "Loop from canvas."
agent/component/loopitem.py
+5 -8
View File
@@ -21,9 +21,11 @@ class LoopItemParam(ComponentParamBase):
"""
Define the LoopItem component parameters.
"""
def check(self):
return True
class LoopItem(ComponentBase, ABC):
component_name = "LoopItem"
@@ -31,7 +33,6 @@ class LoopItem(ComponentBase, ABC):
super().__init__(canvas, id, param)
self._idx = 0
def _invoke(self, **kwargs):
if self.check_if_canceled("LoopItem processing"):
return
@@ -45,7 +46,7 @@ class LoopItem(ComponentBase, ABC):
return
self._idx += 1
def evaluate_condition(self,var, operator, value):
def evaluate_condition(self, var, operator, value):
if isinstance(var, str):
if operator == "contains":
return value in var
@@ -140,11 +141,7 @@ class LoopItem(ComponentBase, ABC):
else:
raise ValueError("Invalid input mode.")
conditions.append(self.evaluate_condition(var, operator, value))
should_end = (
all(conditions) if logical_operator == "and"
else any(conditions) if logical_operator == "or"
else None
)
should_end = all(conditions) if logical_operator == "and" else any(conditions) if logical_operator == "or" else None
if should_end is None:
raise ValueError("Invalid logical operator,should be 'and' or 'or'.")
@@ -164,4 +161,4 @@ class LoopItem(ComponentBase, ABC):
return False
def thoughts(self) -> str:
return "Next turn..."
return "Next turn..."
agent/component/message.py
+59 -80
View File
@@ -14,8 +14,10 @@
# limitations under the License.
#
import asyncio
try:
import nest_asyncio
nest_asyncio.apply()
except Exception:
pass
@@ -45,20 +47,14 @@ class MessageParam(ComponentParamBase):
"""
Define the Message component parameters.
"""
def __init__(self):
super().__init__()
self.content = []
self.stream = True
self.output_format = None # default output format
self.auto_play = False
self.outputs = {
"content": {
"type": "str"
},
"downloads": {
"type": "list"
}
}
self.outputs = {"content": {"type": "str"}, "downloads": {"type": "list"}}
def check(self):
self.check_empty(self.content, "[Message] Content")
@@ -71,9 +67,7 @@ class Message(ComponentBase):
@staticmethod
def _is_download_info(value: Any) -> bool:
return isinstance(value, dict) and all(
key in value for key in ("doc_id", "filename", "mime_type")
)
return isinstance(value, dict) and all(key in value for key in ("doc_id", "filename", "mime_type"))
@staticmethod
def _download_info_includes_content(value: Any) -> bool:
@@ -157,7 +151,7 @@ class Message(ComponentBase):
delimiter: str = None,
downloads: list[dict[str, Any]] | None = None,
) -> tuple[str, dict[str, str | list | Any]]:
for k,v in self.get_input_elements_from_text(script).items():
for k, v in self.get_input_elements_from_text(script).items():
if k in kwargs:
continue
v = v["value"]
@@ -191,7 +185,7 @@ class Message(ComponentBase):
buf += t
return buf
async def _stream(self, rand_cnt:str):
async def _stream(self, rand_cnt: str):
s = 0
all_content = ""
cache = {}
@@ -200,8 +194,8 @@ class Message(ComponentBase):
if self.check_if_canceled("Message streaming"):
return
all_content += rand_cnt[s: r.start()]
yield rand_cnt[s: r.start()]
all_content += rand_cnt[s : r.start()]
yield rand_cnt[s : r.start()]
s = r.end()
exp = r.group(1)
if exp in cache:
@@ -235,9 +229,7 @@ class Message(ComponentBase):
continue
elif inspect.isawaitable(v):
v = await v
v = self._stringify_message_value(
v, downloads=downloads, fallback_to_str=True
)
v = self._stringify_message_value(v, downloads=downloads, fallback_to_str=True)
yield v
self.set_input_value(exp, v)
all_content += v
@@ -247,21 +239,19 @@ class Message(ComponentBase):
if self.check_if_canceled("Message streaming"):
return
all_content += rand_cnt[s: ]
yield rand_cnt[s: ]
all_content += rand_cnt[s:]
yield rand_cnt[s:]
self.set_output("downloads", downloads)
self.set_output("content", all_content)
self._convert_content(all_content)
await self._save_to_memory(all_content)
def _is_jinjia2(self, content:str) -> bool:
patt = [
r"\{%.*%\}", "{{", "}}"
]
def _is_jinjia2(self, content: str) -> bool:
patt = [r"\{%.*%\}", "{{", "}}"]
return any([re.search(p, content) for p in patt])
@timeout(int(os.environ.get("COMPONENT_EXEC_TIMEOUT", 10*60)))
@timeout(int(os.environ.get("COMPONENT_EXEC_TIMEOUT", 10 * 60)))
def _invoke(self, **kwargs):
if self.check_if_canceled("Message processing"):
return
@@ -303,19 +293,19 @@ class Message(ComponentBase):
def _parse_markdown_table_lines(self, table_lines: list):
"""
Parse a list of Markdown table lines into a pandas DataFrame.
Args:
table_lines: List of strings, each representing a row in the Markdown table
(excluding separator lines like |---|---|)
Returns:
pandas DataFrame with the table data, or None if parsing fails
"""
import pandas as pd
if not table_lines:
return None
rows = []
headers = None
@@ -350,29 +340,29 @@ class Message(ComponentBase):
return cell
return cell
for line in table_lines:
# Split by | and clean up
cells = [cell.strip() for cell in line.split('|')]
cells = [cell.strip() for cell in line.split("|")]
# Remove empty first and last elements from split (caused by leading/trailing |)
cells = [c for c in cells if c]
if headers is None:
headers = cells
else:
cells = [_coerce_excel_cell_type(c) for c in cells]
rows.append(cells)
if headers and rows:
# Ensure all rows have same number of columns as headers
normalized_rows = []
for row in rows:
while len(row) < len(headers):
row.append('')
normalized_rows.append(row[:len(headers)])
row.append("")
normalized_rows.append(row[: len(headers)])
return pd.DataFrame(normalized_rows, columns=headers)
return None
def _convert_content(self, content):
@@ -380,6 +370,7 @@ class Message(ComponentBase):
return
import pypandoc
doc_id = get_uuid()
if self._param.output_format.lower() not in {"markdown", "html", "pdf", "docx", "xlsx"}:
@@ -408,49 +399,47 @@ class Message(ComponentBase):
# Debug: log the content being parsed
logging.info(f"XLSX Parser: Content length={len(content) if content else 0}, first 500 chars: {content[:500] if content else 'None'}")
# Try to parse ALL Markdown tables from the content
# Each table will be written to a separate sheet
tables = [] # List of (sheet_name, dataframe)
if isinstance(content, str):
lines = content.strip().split('\n')
lines = content.strip().split("\n")
logging.info(f"XLSX Parser: Total lines={len(lines)}, lines starting with '|': {sum(1 for line in lines if line.strip().startswith('|'))}")
current_table_lines = []
current_table_title = None
pending_title = None
in_table = False
table_count = 0
for i, line in enumerate(lines):
stripped = line.strip()
# Check for potential table title (lines before a table)
# Look for patterns like "Table 1:", "## Table", or markdown headers
if not in_table and stripped and not stripped.startswith('|'):
if not in_table and stripped and not stripped.startswith("|"):
# Check if this could be a table title
lower_stripped = stripped.lower()
if (lower_stripped.startswith('table') or
stripped.startswith('#') or
':' in stripped):
pending_title = stripped.lstrip('#').strip()
if stripped.startswith('|') and '|' in stripped[1:]:
if lower_stripped.startswith("table") or stripped.startswith("#") or ":" in stripped:
pending_title = stripped.lstrip("#").strip()
if stripped.startswith("|") and "|" in stripped[1:]:
# Check if this is a separator line (|---|---|)
cleaned = stripped.replace(' ', '').replace('|', '').replace('-', '').replace(':', '')
if cleaned == '':
cleaned = stripped.replace(" ", "").replace("|", "").replace("-", "").replace(":", "")
if cleaned == "":
continue # Skip separator line
if not in_table:
# Starting a new table
in_table = True
current_table_lines = []
current_table_title = pending_title
pending_title = None
current_table_lines.append(stripped)
elif in_table and not stripped.startswith('|'):
elif in_table and not stripped.startswith("|"):
# End of current table - save it
if current_table_lines:
df = self._parse_markdown_table_lines(current_table_lines)
@@ -460,24 +449,22 @@ class Message(ComponentBase):
if current_table_title:
# Clean and truncate title for sheet name
sheet_name = current_table_title[:31]
sheet_name = sheet_name.replace('/', '_').replace('\\', '_').replace('*', '').replace('?', '').replace('[', '').replace(']', '').replace(':', '')
sheet_name = sheet_name.replace("/", "_").replace("\\", "_").replace("*", "").replace("?", "").replace("[", "").replace("]", "").replace(":", "")
else:
sheet_name = f"Table_{table_count}"
tables.append((sheet_name, df))
# Reset for next table
in_table = False
current_table_lines = []
current_table_title = None
# Check if this line could be a title for the next table
if stripped:
lower_stripped = stripped.lower()
if (lower_stripped.startswith('table') or
stripped.startswith('#') or
':' in stripped):
pending_title = stripped.lstrip('#').strip()
if lower_stripped.startswith("table") or stripped.startswith("#") or ":" in stripped:
pending_title = stripped.lstrip("#").strip()
# Don't forget the last table if content ends with a table
if in_table and current_table_lines:
df = self._parse_markdown_table_lines(current_table_lines)
@@ -485,11 +472,11 @@ class Message(ComponentBase):
table_count += 1
if current_table_title:
sheet_name = current_table_title[:31]
sheet_name = sheet_name.replace('/', '_').replace('\\', '_').replace('*', '').replace('?', '').replace('[', '').replace(']', '').replace(':', '')
sheet_name = sheet_name.replace("/", "_").replace("\\", "_").replace("*", "").replace("?", "").replace("[", "").replace("]", "").replace(":", "")
else:
sheet_name = f"Table_{table_count}"
tables.append((sheet_name, df))
# Fallback: if no tables found, create single sheet with content
if not tables:
df = pd.DataFrame({"Content": [content if content else ""]})
@@ -497,7 +484,7 @@ class Message(ComponentBase):
# Write all tables to Excel, each in a separate sheet
excel_io = BytesIO()
with pd.ExcelWriter(excel_io, engine='openpyxl') as writer:
with pd.ExcelWriter(excel_io, engine="openpyxl") as writer:
used_names = set()
for sheet_name, df in tables:
# Ensure unique sheet names
@@ -505,14 +492,14 @@ class Message(ComponentBase):
counter = 1
while sheet_name in used_names:
suffix = f"_{counter}"
sheet_name = original_name[:31-len(suffix)] + suffix
sheet_name = original_name[: 31 - len(suffix)] + suffix
counter += 1
used_names.add(sheet_name)
df.to_excel(writer, sheet_name=sheet_name, index=False)
excel_io.seek(0)
binary_content = excel_io.read()
logging.info(f"Generated Excel with {len(tables)} sheet(s): {[t[0] for t in tables]}")
else: # pdf, docx
@@ -543,10 +530,7 @@ class Message(ComponentBase):
os.remove(tmp_name)
settings.STORAGE_IMPL.put(self._canvas._tenant_id, doc_id, binary_content)
self.set_output("attachment", {
"doc_id":doc_id,
"format":self._param.output_format,
"file_name":f"{doc_id[:8]}.{self._param.output_format}"})
self.set_output("attachment", {"doc_id": doc_id, "format": self._param.output_format, "file_name": f"{doc_id[:8]}.{self._param.output_format}"})
logging.info(f"Converted content uploaded as {doc_id} (format={self._param.output_format})")
@@ -560,15 +544,10 @@ class Message(ComponentBase):
user_id = self._param.user_id if hasattr(self._param, "user_id") else ""
if user_id:
import re
# is variable
if re.match(r"^{.*}$", user_id):
user_id = self._canvas.get_variable_value(user_id)
message_dict = {
"user_id": user_id,
"agent_id": self._canvas._id,
"session_id": self._canvas.task_id,
"user_input": self._canvas.get_sys_query(),
"agent_response": content
}
message_dict = {"user_id": user_id, "agent_id": self._canvas._id, "session_id": self._canvas.task_id, "user_input": self._canvas.get_sys_query(), "agent_response": content}
return await queue_save_to_memory_task(self._param.memory_ids, message_dict)
agent/component/pipeline_chunker.py
+1 -4
View File
@@ -131,10 +131,7 @@ class PipelineChunker(ComponentBase, ABC):
try:
return FileService.get_blob(created_by, file_id), filename
except Exception as e:
logging.exception(
f"[PipelineChunker] FileService.get_blob failed for "
f"file_id={file_id} created_by={created_by} filename={filename}: {e}"
)
logging.exception(f"[PipelineChunker] FileService.get_blob failed for file_id={file_id} created_by={created_by} filename={filename}: {e}")
return None, None
return None, None
agent/component/string_transform.py
+7 -17
View File
@@ -52,18 +52,10 @@ class StringTransform(Message, ABC):
def get_input_form(self) -> dict[str, dict]:
if self._param.method == "split":
return {
"line": {
"name": "String",
"type": "line"
}
}
return {k: {
"name": o["name"],
"type": "line"
} for k, o in self.get_input_elements_from_text(self._param.script).items()}
return {"line": {"name": "String", "type": "line"}}
return {k: {"name": o["name"], "type": "line"} for k, o in self.get_input_elements_from_text(self._param.script).items()}
@timeout(int(os.environ.get("COMPONENT_EXEC_TIMEOUT", 10*60)))
@timeout(int(os.environ.get("COMPONENT_EXEC_TIMEOUT", 10 * 60)))
def _invoke(self, **kwargs):
if self.check_if_canceled("StringTransform processing"):
return
@@ -73,7 +65,7 @@ class StringTransform(Message, ABC):
else:
self._merge(kwargs)
def _split(self, line:str|None = None):
def _split(self, line: str | None = None):
if self.check_if_canceled("StringTransform split processing"):
return
@@ -84,13 +76,13 @@ class StringTransform(Message, ABC):
self.set_input_value(self._param.split_ref, var)
res = []
for i,s in enumerate(re.split(r"(%s)"%("|".join([re.escape(d) for d in self._param.delimiters])), var, flags=re.DOTALL)):
for i, s in enumerate(re.split(r"(%s)" % ("|".join([re.escape(d) for d in self._param.delimiters])), var, flags=re.DOTALL)):
if i % 2 == 1:
continue
res.append(s)
self.set_output("result", res)
def _merge(self, kwargs:dict[str, str] = {}):
def _merge(self, kwargs: dict[str, str] = {}):
if self.check_if_canceled("StringTransform merge processing"):
return
@@ -104,7 +96,7 @@ class StringTransform(Message, ABC):
except Exception:
pass
for k,v in kwargs.items():
for k, v in kwargs.items():
if v is None:
v = ""
script = re.sub(k, lambda match: v, script)
@@ -113,5 +105,3 @@ class StringTransform(Message, ABC):
def thoughts(self) -> str:
return f"It's {self._param.method}ing."
agent/component/switch.py
+4 -9
View File
@@ -40,8 +40,7 @@ class SwitchParam(ComponentParamBase):
"""
self.conditions = []
self.end_cpn_ids = []
self.operators = ['contains', 'not contains', 'start with', 'end with', 'empty', 'not empty', '=', '≠', '>',
'<', '≥', '≤']
self.operators = ["contains", "not contains", "start with", "end with", "empty", "not empty", "=", "≠", ">", "<", "≥", "≤"]
def check(self):
self.check_empty(self.conditions, "[Switch] conditions")
@@ -51,12 +50,8 @@ class SwitchParam(ComponentParamBase):
self.check_empty(self.end_cpn_ids, "[Switch] the ELSE/Other destination can not be empty.")
def get_input_form(self) -> dict[str, dict]:
return {
"urls": {
"name": "URLs",
"type": "line"
}
}
return {"urls": {"name": "URLs", "type": "line"}}
class Switch(ComponentBase, ABC):
component_name = "Switch"
@@ -137,7 +132,7 @@ class Switch(ComponentBase, ABC):
except Exception:
return True if input <= value else False
raise ValueError(f'Not supported operator: {operator}')
raise ValueError(f"Not supported operator: {operator}")
def thoughts(self) -> str:
return "I’m weighing a few options and will pick the next step shortly."
agent/component/variable_aggregator.py
+4 -8
View File
@@ -38,13 +38,9 @@ class VariableAggregatorParam(ComponentParamBase):
if not g.get("group_name"):
raise ValueError("[VariableAggregator] group_name can not be empty!")
if not g.get("variables"):
raise ValueError(
f"[VariableAggregator] variables of group `{g.get('group_name')}` can not be empty"
)
raise ValueError(f"[VariableAggregator] variables of group `{g.get('group_name')}` can not be empty")
if not isinstance(g.get("variables"), list):
raise ValueError(
f"[VariableAggregator] variables of group `{g.get('group_name')}` should be a list of strings"
)
raise ValueError(f"[VariableAggregator] variables of group `{g.get('group_name')}` should be a list of strings")
def get_input_form(self) -> dict[str, dict]:
return {
@@ -67,11 +63,11 @@ class VariableAggregator(ComponentBase):
# record candidate selectors within this group
self.set_input_value(f"{gname}.variables", list(group.get("variables", [])))
for selector in group.get("variables", []):
val = self._canvas.get_variable_value(selector['value'])
val = self._canvas.get_variable_value(selector["value"])
if val:
self.set_output(gname, val)
break
@staticmethod
def _to_object(value: Any) -> Any:
# Try to convert value to serializable object if it has to_object()
agent/component/variable_assigner.py
+59 -61
View File
@@ -19,32 +19,30 @@ import numbers
from agent.component.base import ComponentBase, ComponentParamBase
from api.utils.api_utils import timeout
class VariableAssignerParam(ComponentParamBase):
"""
Define the Variable Assigner component parameters.
"""
def __init__(self):
super().__init__()
self.variables=[]
self.variables = []
def check(self):
return True
def get_input_form(self) -> dict[str, dict]:
return {
"items": {
"type": "json",
"name": "Items"
}
}
class VariableAssigner(ComponentBase,ABC):
def get_input_form(self) -> dict[str, dict]:
return {"items": {"type": "json", "name": "Items"}}
class VariableAssigner(ComponentBase, ABC):
component_name = "VariableAssigner"
_NO_PARAMETER_OPERATORS = {"clear", "remove_first", "remove_last"}
@timeout(int(os.environ.get("COMPONENT_EXEC_TIMEOUT", 10*60)))
@timeout(int(os.environ.get("COMPONENT_EXEC_TIMEOUT", 10 * 60)))
def _invoke(self, **kwargs):
if not isinstance(self._param.variables,list):
if not isinstance(self._param.variables, list):
return
else:
for item in self._param.variables:
@@ -56,105 +54,105 @@ class VariableAssigner(ComponentBase,ABC):
raise ValueError("Variable is not complete.")
if operator not in self._NO_PARAMETER_OPERATORS and parameter is None:
raise ValueError("Variable is not complete.")
variable_value=self._canvas.get_variable_value(variable)
new_variable=self._operate(variable_value,operator,parameter)
variable_value = self._canvas.get_variable_value(variable)
new_variable = self._operate(variable_value, operator, parameter)
self._canvas.set_variable_value(variable, new_variable)
def _operate(self,variable,operator,parameter):
def _operate(self, variable, operator, parameter):
if operator == "overwrite":
return self._overwrite(parameter)
elif operator == "clear":
return self._clear(variable)
elif operator == "set":
return self._set(variable,parameter)
return self._set(variable, parameter)
elif operator == "append":
return self._append(variable,parameter)
return self._append(variable, parameter)
elif operator == "extend":
return self._extend(variable,parameter)
return self._extend(variable, parameter)
elif operator == "remove_first":
return self._remove_first(variable)
elif operator == "remove_last":
return self._remove_last(variable)
elif operator == "+=":
return self._add(variable,parameter)
return self._add(variable, parameter)
elif operator == "-=":
return self._subtract(variable,parameter)
return self._subtract(variable, parameter)
elif operator == "*=":
return self._multiply(variable,parameter)
return self._multiply(variable, parameter)
elif operator == "/=":
return self._divide(variable,parameter)
return self._divide(variable, parameter)
else:
return
def _overwrite(self,parameter):
def _overwrite(self, parameter):
return self._canvas.get_variable_value(parameter)
def _clear(self,variable):
if isinstance(variable,list):
def _clear(self, variable):
if isinstance(variable, list):
return []
elif isinstance(variable,str):
elif isinstance(variable, str):
return ""
elif isinstance(variable,dict):
elif isinstance(variable, dict):
return {}
elif isinstance(variable,bool):
elif isinstance(variable, bool):
return False
elif isinstance(variable,int):
elif isinstance(variable, int):
return 0
elif isinstance(variable,float):
elif isinstance(variable, float):
return 0.0
else:
return None
def _set(self,variable,parameter):
def _set(self, variable, parameter):
if variable is None:
return self._canvas.get_value_with_variable(parameter)
elif isinstance(variable,str):
elif isinstance(variable, str):
return self._canvas.get_value_with_variable(parameter)
elif isinstance(variable,bool):
elif isinstance(variable, bool):
return parameter
elif isinstance(variable,int):
elif isinstance(variable, int):
return parameter
elif isinstance(variable,float):
elif isinstance(variable, float):
return parameter
else:
return parameter
def _append(self,variable,parameter):
parameter=self._canvas.get_variable_value(parameter)
def _append(self, variable, parameter):
parameter = self._canvas.get_variable_value(parameter)
if variable is None:
variable=[]
if not isinstance(variable,list):
variable = []
if not isinstance(variable, list):
return "ERROR:VARIABLE_NOT_LIST"
elif len(variable)!=0 and not isinstance(parameter,type(variable[0])):
elif len(variable) != 0 and not isinstance(parameter, type(variable[0])):
return "ERROR:PARAMETER_NOT_LIST_ELEMENT_TYPE"
else:
variable.append(parameter)
return variable
def _extend(self,variable,parameter):
parameter=self._canvas.get_variable_value(parameter)
def _extend(self, variable, parameter):
parameter = self._canvas.get_variable_value(parameter)
if variable is None:
variable=[]
if not isinstance(variable,list):
variable = []
if not isinstance(variable, list):
return "ERROR:VARIABLE_NOT_LIST"
elif not isinstance(parameter,list):
elif not isinstance(parameter, list):
return "ERROR:PARAMETER_NOT_LIST"
elif len(variable)!=0 and len(parameter)!=0 and not isinstance(parameter[0],type(variable[0])):
elif len(variable) != 0 and len(parameter) != 0 and not isinstance(parameter[0], type(variable[0])):
return "ERROR:PARAMETER_NOT_LIST_ELEMENT_TYPE"
else:
return variable + parameter
def _remove_first(self,variable):
if not isinstance(variable,list):
def _remove_first(self, variable):
if not isinstance(variable, list):
return "ERROR:VARIABLE_NOT_LIST"
if len(variable)==0:
if len(variable) == 0:
return variable
return variable[1:]
def _remove_last(self,variable):
if not isinstance(variable,list):
def _remove_last(self, variable):
if not isinstance(variable, list):
return "ERROR:VARIABLE_NOT_LIST"
if len(variable)==0:
if len(variable) == 0:
return variable
return variable[:-1]
@@ -163,32 +161,32 @@ class VariableAssigner(ComponentBase,ABC):
return False
return isinstance(value, numbers.Number)
def _add(self,variable,parameter):
def _add(self, variable, parameter):
if self.is_number(variable) and self.is_number(parameter):
return variable + parameter
else:
return "ERROR:VARIABLE_NOT_NUMBER or PARAMETER_NOT_NUMBER"
def _subtract(self,variable,parameter):
def _subtract(self, variable, parameter):
if self.is_number(variable) and self.is_number(parameter):
return variable - parameter
else:
return "ERROR:VARIABLE_NOT_NUMBER or PARAMETER_NOT_NUMBER"
def _multiply(self,variable,parameter):
def _multiply(self, variable, parameter):
if self.is_number(variable) and self.is_number(parameter):
return variable * parameter
else:
return "ERROR:VARIABLE_NOT_NUMBER or PARAMETER_NOT_NUMBER"
def _divide(self,variable,parameter):
def _divide(self, variable, parameter):
if self.is_number(variable) and self.is_number(parameter):
if parameter==0:
if parameter == 0:
return "ERROR:DIVIDE_BY_ZERO"
else:
return variable/parameter
return variable / parameter
else:
return "ERROR:VARIABLE_NOT_NUMBER or PARAMETER_NOT_NUMBER"
return "ERROR:VARIABLE_NOT_NUMBER or PARAMETER_NOT_NUMBER"
def thoughts(self) -> str:
return "Assign variables from canvas."
return "Assign variables from canvas."
agent/dsl_migration.py
+5 -19
View File
@@ -56,7 +56,7 @@ def normalize_chunker_dsl(dsl: dict) -> dict:
for old_name, new_name in COMPONENT_RENAMES.items():
prefix = f"{old_name}:"
if component_id.startswith(prefix):
new_component_id = f"{new_name}:{component_id[len(prefix):]}"
new_component_id = f"{new_name}:{component_id[len(prefix) :]}"
break
component_id_map[component_id] = new_component_id
@@ -66,12 +66,7 @@ def normalize_chunker_dsl(dsl: dict) -> dict:
def repl(match: re.Match[str]) -> str:
component_id = match.group(2)
return (
match.group(1)
+ component_id_map.get(component_id, component_id)
+ match.group(3)
+ match.group(4)
)
return match.group(1) + component_id_map.get(component_id, component_id) + match.group(3) + match.group(4)
return VARIABLE_REF_PATTERN.sub(repl, text)
@@ -96,15 +91,9 @@ def normalize_chunker_dsl(dsl: dict) -> dict:
obj["component_name"] = COMPONENT_RENAMES.get(component_name, component_name)
if isinstance(new_component.get("downstream"), list):
new_component["downstream"] = [
component_id_map.get(component_id, component_id)
for component_id in new_component["downstream"]
]
new_component["downstream"] = [component_id_map.get(component_id, component_id) for component_id in new_component["downstream"]]
if isinstance(new_component.get("upstream"), list):
new_component["upstream"] = [
component_id_map.get(component_id, component_id)
for component_id in new_component["upstream"]
]
new_component["upstream"] = [component_id_map.get(component_id, component_id) for component_id in new_component["upstream"]]
parent_id = new_component.get("parent_id")
if isinstance(parent_id, str):
@@ -115,10 +104,7 @@ def normalize_chunker_dsl(dsl: dict) -> dict:
normalized["components"] = rewritten_components
if isinstance(normalized.get("path"), list):
normalized["path"] = [
component_id_map.get(component_id, component_id)
for component_id in normalized["path"]
]
normalized["path"] = [component_id_map.get(component_id, component_id) for component_id in normalized["path"]]
graph = normalized.get("graph")
if isinstance(graph, dict):
agent/plugin/common.py
+1 -1
View File
@@ -1 +1 @@
PLUGIN_TYPE_LLM_TOOLS = "llm_tools"
PLUGIN_TYPE_LLM_TOOLS = "llm_tools"
agent/plugin/embedded_plugins/llm_tools/bad_calculator.py
+4 -13
View File
@@ -7,6 +7,7 @@ class BadCalculatorPlugin(LLMToolPlugin):
A sample LLM tool plugin, will add two numbers with 100.
It only presents for demo purpose. Do not use it in production.
"""
_version_ = "1.0.0"
@classmethod
@@ -17,19 +18,9 @@ class BadCalculatorPlugin(LLMToolPlugin):
"description": "A tool to calculate the sum of two numbers (will give wrong answer)",
"displayDescription": "$t:bad_calculator.description",
"parameters": {
"a": {
"type": "number",
"description": "The first number",
"displayDescription": "$t:bad_calculator.params.a",
"required": True
},
"b": {
"type": "number",
"description": "The second number",
"displayDescription": "$t:bad_calculator.params.b",
"required": True
}
}
"a": {"type": "number", "description": "The first number", "displayDescription": "$t:bad_calculator.params.a", "required": True},
"b": {"type": "number", "description": "The second number", "displayDescription": "$t:bad_calculator.params.b", "required": True},
},
}
def invoke(self, a: int, b: int) -> str:
agent/plugin/llm_tool_plugin.py
+4 -10
View File
@@ -38,14 +38,8 @@ def llm_tool_metadata_to_openai_tool(llm_tool_metadata: LLMToolMetadata) -> dict
"description": llm_tool_metadata["description"],
"parameters": {
"type": "object",
"properties": {
k: {
"type": p["type"],
"description": p["description"]
}
for k, p in llm_tool_metadata["parameters"].items()
},
"required": [k for k, p in llm_tool_metadata["parameters"].items() if p["required"]]
}
}
"properties": {k: {"type": p["type"], "description": p["description"]} for k, p in llm_tool_metadata["parameters"].items()},
"required": [k for k, p in llm_tool_metadata["parameters"].items() if p["required"]],
},
},
}
agent/plugin/plugin_manager.py
+2 -4
View File
@@ -15,10 +15,8 @@ class PluginManager:
self._llm_tool_plugins = {}
def load_plugins(self) -> None:
loader = pluginlib.PluginLoader(
paths=[str(Path(os.path.dirname(__file__), "embedded_plugins"))]
)
loader = pluginlib.PluginLoader(paths=[str(Path(os.path.dirname(__file__), "embedded_plugins"))])
for type, plugins in loader.plugins.items():
for name, plugin in plugins.items():
logging.info(f"Loaded {type} plugin {name} version {plugin.version}")
agent/sandbox/client.py
+6 -16
View File
@@ -111,7 +111,10 @@ def _load_provider_from_settings() -> None:
except Exception as e:
logger.error(f"Failed to load sandbox provider from settings: {e}")
import traceback
traceback.print_exc()
def _load_provider_config_from_settings(provider_type: str) -> Dict[str, Any]:
provider_config_settings = SystemSettingsService.get_by_name(f"sandbox.{provider_type}")
if not provider_config_settings:
@@ -147,12 +150,7 @@ def reload_provider() -> None:
_load_provider_from_settings()
def execute_code(
code: str,
language: str = "python",
timeout: int = 30,
arguments: Optional[Dict[str, Any]] = None
) -> ExecutionResult:
def execute_code(code: str, language: str = "python", timeout: int = 30, arguments: Optional[Dict[str, Any]] = None) -> ExecutionResult:
"""
Execute code in the configured sandbox.
@@ -173,9 +171,7 @@ def execute_code(
provider_manager = get_provider_manager()
if not provider_manager.is_configured():
raise RuntimeError(
"No sandbox provider configured. Please configure sandbox settings in the admin panel."
)
raise RuntimeError("No sandbox provider configured. Please configure sandbox settings in the admin panel.")
provider = provider_manager.get_provider()
provider_name = provider_manager.get_provider_name() or getattr(provider, "__class__", type(provider)).__name__
@@ -192,13 +188,7 @@ def execute_code(
try:
# Execute the code
result = provider.execute_code(
instance_id=instance.instance_id,
code=code,
language=language,
timeout=timeout,
arguments=arguments
)
result = provider.execute_code(instance_id=instance.instance_id, code=code, language=language, timeout=timeout, arguments=arguments)
return result
agent/sandbox/executor_manager/api/routes.py
-1
View File
@@ -22,4 +22,3 @@ router = APIRouter()
router.get("/")(healthz_handler)
router.get("/healthz")(healthz_handler)
router.post("/run")(run_code_handler)
agent/sandbox/executor_manager/services/execution.py
+33 -13
View File
@@ -230,9 +230,7 @@ async def execute_code(req: CodeExecutionRequest):
if returncode != 0:
raise RuntimeError(f"Directory creation failed: {stderr}")
tar_proc = await asyncio.create_subprocess_exec(
"tar", "czf", "-", "-C", workdir, code_name, runner_name, str(bundle["args_name"]), stdout=asyncio.subprocess.PIPE
)
tar_proc = await asyncio.create_subprocess_exec("tar", "czf", "-", "-C", workdir, code_name, runner_name, str(bundle["args_name"]), stdout=asyncio.subprocess.PIPE)
tar_stdout, _ = await tar_proc.communicate()
docker_proc = await asyncio.create_subprocess_exec(
@@ -334,8 +332,16 @@ async def _collect_artifacts(container: str, task_id: str, host_workdir: str) ->
# List files in the artifacts directory inside the container
returncode, stdout, _ = await async_run_command(
"docker", "exec", container, "find", artifacts_path,
"-maxdepth", "1", "-type", "f", timeout=5,
"docker",
"exec",
container,
"find",
artifacts_path,
"-maxdepth",
"1",
"-type",
"f",
timeout=5,
)
if returncode != 0 or not stdout.strip():
return []
@@ -359,7 +365,14 @@ async def _collect_artifacts(container: str, task_id: str, host_workdir: str) ->
# Check file size inside the container
returncode, size_str, _ = await async_run_command(
"docker", "exec", container, "stat", "-c", "%s", file_path, timeout=5,
"docker",
"exec",
container,
"stat",
"-c",
"%s",
file_path,
timeout=5,
)
if returncode != 0:
logger.warning(f"Failed to stat artifact {fname}")
@@ -374,7 +387,12 @@ async def _collect_artifacts(container: str, task_id: str, host_workdir: str) ->
# Read file content via docker exec (docker cp doesn't work with gVisor tmpfs)
returncode, content_b64, stderr = await async_run_command(
"docker", "exec", container, "base64", file_path, timeout=30,
"docker",
"exec",
container,
"base64",
file_path,
timeout=30,
)
if returncode != 0:
logger.warning(f"Failed to read artifact {fname}: {stderr}")
@@ -382,12 +400,14 @@ async def _collect_artifacts(container: str, task_id: str, host_workdir: str) ->
content_b64 = content_b64.replace("\n", "").strip()
items.append(ArtifactItem(
name=fname,
mime_type=mime_type,
size=file_size,
content_b64=content_b64,
))
items.append(
ArtifactItem(
name=fname,
mime_type=mime_type,
size=file_size,
content_b64=content_b64,
)
)
logger.info(f"Collected artifact: {fname} ({file_size} bytes, {mime_type})")
return items
agent/sandbox/providers/aliyun_codeinterpreter.py
+1 -5
View File
@@ -232,11 +232,7 @@ class AliyunCodeInterpreterProvider(SandboxProvider):
# Wrap code to call main() function
# Matches self_managed provider behavior: call main(**arguments)
args_json = json.dumps(arguments or {})
wrapped_code = (
build_python_wrapper(code, args_json)
if normalized_lang == "python"
else build_javascript_wrapper(code, args_json)
)
wrapped_code = build_python_wrapper(code, args_json) if normalized_lang == "python" else build_javascript_wrapper(code, args_json)
logger.debug(f"Aliyun Code Interpreter: Wrapped code (first 200 chars): {wrapped_code[:200]}")
start_time = time.time()
agent/sandbox/providers/base.py
+3 -8
View File
@@ -33,6 +33,7 @@ class SandboxProviderConfigError(Exception):
@dataclass
class SandboxInstance:
"""Represents a sandbox execution instance"""
instance_id: str
provider: str
status: str # running, stopped, error
@@ -46,6 +47,7 @@ class SandboxInstance:
@dataclass
class ExecutionResult:
"""Result of code execution in a sandbox"""
stdout: str
stderr: str
exit_code: int
@@ -96,14 +98,7 @@ class SandboxProvider(ABC):
pass
@abstractmethod
def execute_code(
self,
instance_id: str,
code: str,
language: str,
timeout: int = 10,
arguments: Optional[Dict[str, Any]] = None
) -> ExecutionResult:
def execute_code(self, instance_id: str, code: str, language: str, timeout: int = 10, arguments: Optional[Dict[str, Any]] = None) -> ExecutionResult:
"""
Execute code in a sandbox instance.
agent/sandbox/providers/e2b.py
+4 -12
View File
@@ -97,16 +97,10 @@ class E2BProvider(SandboxProvider):
metadata={
"language": language,
"region": self.region,
}
},
)
def execute_code(
self,
instance_id: str,
code: str,
language: str,
timeout: int = 10
) -> ExecutionResult:
def execute_code(self, instance_id: str, code: str, language: str, timeout: int = 10) -> ExecutionResult:
"""
Execute code in the E2B instance.
@@ -130,9 +124,7 @@ class E2BProvider(SandboxProvider):
# POST /sandbox/{sandboxID}/execute
raise RuntimeError(
"E2B provider is not yet fully implemented. "
"Please use the self-managed provider or implement the E2B API integration. "
"See https://github.com/e2b-dev/e2b for API documentation."
"E2B provider is not yet fully implemented. Please use the self-managed provider or implement the E2B API integration. See https://github.com/e2b-dev/e2b for API documentation."
)
def destroy_instance(self, instance_id: str) -> bool:
@@ -208,7 +200,7 @@ class E2BProvider(SandboxProvider):
"min": 5,
"max": 300,
"description": "API request timeout for code execution",
}
},
}
def _normalize_language(self, language: str) -> str:
agent/sandbox/providers/local.py
+2
View File
@@ -49,6 +49,8 @@ LOCAL_PYTHON_THREAD_ENV_VARS = (
"BLIS_NUM_THREADS",
"VECLIB_MAXIMUM_THREADS",
)
class LocalProvider(SandboxProvider):
"""
Execute code as a local child process.
agent/sandbox/providers/self_managed.py
+9 -28
View File
@@ -108,17 +108,10 @@ class SelfManagedProvider(SandboxProvider):
"language": language,
"endpoint": self.endpoint,
"pool_size": self.pool_size,
}
},
)
def execute_code(
self,
instance_id: str,
code: str,
language: str,
timeout: int = 10,
arguments: Optional[Dict[str, Any]] = None
) -> ExecutionResult:
def execute_code(self, instance_id: str, code: str, language: str, timeout: int = 10, arguments: Optional[Dict[str, Any]] = None) -> ExecutionResult:
"""
Execute code in the sandbox.
@@ -144,11 +137,7 @@ class SelfManagedProvider(SandboxProvider):
# Prepare request
code_b64 = base64.b64encode(code.encode("utf-8")).decode("utf-8")
payload = {
"code_b64": code_b64,
"language": normalized_lang,
"arguments": arguments or {}
}
payload = {"code_b64": code_b64, "language": normalized_lang, "arguments": arguments or {}}
url = f"{self.endpoint}/run"
exec_timeout = timeout or self.timeout
@@ -156,19 +145,12 @@ class SelfManagedProvider(SandboxProvider):
start_time = time.time()
try:
response = requests.post(
url,
json=payload,
timeout=exec_timeout,
headers={"Content-Type": "application/json"}
)
response = requests.post(url, json=payload, timeout=exec_timeout, headers={"Content-Type": "application/json"})
execution_time = time.time() - start_time
if response.status_code != 200:
raise RuntimeError(
f"HTTP {response.status_code}: {response.text}"
)
raise RuntimeError(f"HTTP {response.status_code}: {response.text}")
result = response.json()
structured_result = result.get("result") or {}
@@ -188,14 +170,12 @@ class SelfManagedProvider(SandboxProvider):
"result_present": structured_result.get("present", False),
"result_value": structured_result.get("value"),
"result_type": structured_result.get("type"),
}
},
)
except requests.Timeout:
execution_time = time.time() - start_time
raise TimeoutError(
f"Execution timed out after {exec_timeout} seconds"
)
raise TimeoutError(f"Execution timed out after {exec_timeout} seconds")
except requests.RequestException as e:
raise RuntimeError(f"HTTP request failed: {str(e)}")
@@ -388,7 +368,8 @@ class SelfManagedProvider(SandboxProvider):
if endpoint:
# Check if it's a valid HTTP/HTTPS URL or localhost
import re
url_pattern = r'^(https?://|http://localhost|http://[\d\.]+:[a-z]+:[/]|http://[\w\.]+:)'
url_pattern = r"^(https?://|http://localhost|http://[\d\.]+:[a-z]+:[/]|http://[\w\.]+:)"
if not re.match(url_pattern, endpoint):
return False, f"Invalid endpoint format: {endpoint}. Must start with http:// or https://"
agent/sandbox/providers/ssh.py
+9 -29
View File
@@ -135,9 +135,7 @@ class SSHProvider(SandboxProvider):
timeout=min(self.timeout, 10),
)
if exit_code != 0:
raise RuntimeError(
f"Failed to create remote artifacts directory: {stderr or stdout or 'unknown error'}"
)
raise RuntimeError(f"Failed to create remote artifacts directory: {stderr or stdout or 'unknown error'}")
except Exception:
sftp.close()
client.close()
@@ -211,9 +209,7 @@ class SSHProvider(SandboxProvider):
"status": "ok" if exit_code == 0 else "error",
"timeout": exec_timeout,
"command": command,
"artifacts": self._collect_artifacts(
sftp, posixpath.join(remote_work_dir, "artifacts")
),
"artifacts": self._collect_artifacts(sftp, posixpath.join(remote_work_dir, "artifacts")),
"result_present": structured_result.get("present", False),
"result_value": structured_result.get("value"),
"result_type": structured_result.get("type"),
@@ -269,18 +265,13 @@ class SSHProvider(SandboxProvider):
timeout=min(self.timeout, 10),
)
if exit_code != 0:
raise SandboxProviderConfigError(
f"SSH connectivity check failed on {self.username}@{self.host}:{self.port}: "
f"{stderr or 'remote command returned non-zero exit status'}"
)
raise SandboxProviderConfigError(f"SSH connectivity check failed on {self.username}@{self.host}:{self.port}: {stderr or 'remote command returned non-zero exit status'}")
finally:
client.close()
except SandboxProviderConfigError:
raise
except Exception as exc:
raise SandboxProviderConfigError(
f"Failed to connect to SSH host {self.username}@{self.host}:{self.port}: {exc}"
) from exc
raise SandboxProviderConfigError(f"Failed to connect to SSH host {self.username}@{self.host}:{self.port}: {exc}") from exc
def get_supported_languages(self) -> List[str]:
return ["python", "javascript", "nodejs"]
@@ -470,9 +461,7 @@ class SSHProvider(SandboxProvider):
# Match the Go provider's fail-closed posture (see
# internal/agent/sandbox/ssh.go::hostKeyCallback).
logging.warning("SSH: failed to load configured known_hosts file; refusing connection")
raise SandboxProviderConfigError(
"Failed to load configured SSH known_hosts file."
) from exc
raise SandboxProviderConfigError("Failed to load configured SSH known_hosts file.") from exc
# Reject unknown hosts: this is the default fail-closed posture
# to prevent silent MITM. Operators must either ship a populated
# known_hosts file or accept the warning (paramiko will fail the
@@ -522,9 +511,7 @@ class SSHProvider(SandboxProvider):
except Exception as exc:
errors.append(str(exc))
raise SandboxProviderConfigError(
"Failed to load SSH private key. " + "; ".join(error for error in errors if error)
)
raise SandboxProviderConfigError("Failed to load SSH private key. " + "; ".join(error for error in errors if error))
def _create_remote_workspace(self, client: paramiko.SSHClient) -> str:
base_dir = self.work_dir.rstrip("/") or "/tmp"
@@ -535,9 +522,7 @@ class SSHProvider(SandboxProvider):
timeout=min(self.timeout, 10),
)
if exit_code != 0:
raise RuntimeError(
f"Failed to create remote workspace on {self.host}: {stderr or stdout or 'unknown error'}"
)
raise RuntimeError(f"Failed to create remote workspace on {self.host}: {stderr or stdout or 'unknown error'}")
remote_work_dir = stdout.strip().splitlines()[-1] if stdout.strip() else ""
if not remote_work_dir:
@@ -577,10 +562,7 @@ class SSHProvider(SandboxProvider):
else:
raise RuntimeError(f"Unsupported language for SSH provider: {language}")
return (
f"cd {shlex.quote(remote_work_dir)} && "
f"{shlex.quote(executable)} {shlex.quote(remote_script_path)}"
)
return f"cd {shlex.quote(remote_work_dir)} && {shlex.quote(executable)} {shlex.quote(remote_script_path)}"
def _run_remote_command(
self,
@@ -700,7 +682,5 @@ def _get_paramiko_module():
try:
import paramiko
except ImportError as exc:
raise SandboxProviderConfigError(
"paramiko is required for the SSH sandbox provider. Install the project dependencies to enable it."
) from exc
raise SandboxProviderConfigError("paramiko is required for the SSH sandbox provider. Install the project dependencies to enable it.") from exc
return paramiko
agent/sandbox/result_protocol.py
+2 -2
View File
@@ -36,7 +36,7 @@ if __name__ == "__main__":
def build_javascript_wrapper(code: str, args_json: str) -> str:
return f'''{code}
return f"""{code}
const __ragflowArgs = {args_json};
@@ -55,7 +55,7 @@ const __ragflowArgs = {args_json};
}}
console.log('{RESULT_MARKER_PREFIX}' + Buffer.from(payload, 'utf8').toString('base64'));
}})();
'''
"""
def extract_structured_result(stdout: str) -> tuple[str, dict[str, Any]]:
agent/sandbox/tests/test_aliyun_codeinterpreter_integration.py
+2 -2
View File
@@ -151,7 +151,7 @@ class TestAliyunCodeInterpreterIntegration:
""",
language="python",
timeout=30,
arguments={"name": "World", "count": 2}
arguments={"name": "World", "count": 2},
)
assert result.exit_code == 0
@@ -211,7 +211,7 @@ class TestAliyunCodeInterpreterIntegration:
}""",
language="javascript",
timeout=30,
arguments={"name": "World", "count": 2}
arguments={"name": "World", "count": 2},
)
assert result.exit_code == 0
agent/sandbox/tests/test_providers.py
+26 -83
View File
@@ -32,12 +32,7 @@ class TestSandboxDataclasses:
def test_sandbox_instance_creation(self):
"""Test SandboxInstance dataclass creation."""
instance = SandboxInstance(
instance_id="test-123",
provider="self_managed",
status="running",
metadata={"language": "python"}
)
instance = SandboxInstance(instance_id="test-123", provider="self_managed", status="running", metadata={"language": "python"})
assert instance.instance_id == "test-123"
assert instance.provider == "self_managed"
@@ -46,24 +41,13 @@ class TestSandboxDataclasses:
def test_sandbox_instance_default_metadata(self):
"""Test SandboxInstance with None metadata."""
instance = SandboxInstance(
instance_id="test-123",
provider="self_managed",
status="running",
metadata=None
)
instance = SandboxInstance(instance_id="test-123", provider="self_managed", status="running", metadata=None)
assert instance.metadata == {}
def test_execution_result_creation(self):
"""Test ExecutionResult dataclass creation."""
result = ExecutionResult(
stdout="Hello, World!",
stderr="",
exit_code=0,
execution_time=1.5,
metadata={"status": "success"}
)
result = ExecutionResult(stdout="Hello, World!", stderr="", exit_code=0, execution_time=1.5, metadata={"status": "success"})
assert result.stdout == "Hello, World!"
assert result.stderr == ""
@@ -73,13 +57,7 @@ class TestSandboxDataclasses:
def test_execution_result_default_metadata(self):
"""Test ExecutionResult with None metadata."""
result = ExecutionResult(
stdout="output",
stderr="error",
exit_code=1,
execution_time=0.5,
metadata=None
)
result = ExecutionResult(stdout="output", stderr="error", exit_code=1, execution_time=0.5, metadata=None)
assert result.metadata == {}
@@ -145,7 +123,7 @@ class TestSelfManagedProvider:
assert provider.pool_size == 10
assert not provider._initialized
@patch('requests.get')
@patch("requests.get")
def test_initialize_success(self, mock_get):
"""Test successful initialization."""
mock_response = Mock()
@@ -153,12 +131,7 @@ class TestSelfManagedProvider:
mock_get.return_value = mock_response
provider = SelfManagedProvider()
result = provider.initialize({
"endpoint": "http://test-endpoint:9385",
"timeout": 60,
"max_retries": 5,
"pool_size": 20
})
result = provider.initialize({"endpoint": "http://test-endpoint:9385", "timeout": 60, "max_retries": 5, "pool_size": 20})
assert result is True
assert provider.endpoint == "http://test-endpoint:9385"
@@ -168,7 +141,7 @@ class TestSelfManagedProvider:
assert provider._initialized
mock_get.assert_called_once_with("http://test-endpoint:9385/healthz", timeout=5)
@patch('requests.get')
@patch("requests.get")
def test_initialize_failure(self, mock_get):
"""Test initialization failure."""
mock_get.side_effect = Exception("Connection error")
@@ -181,7 +154,7 @@ class TestSelfManagedProvider:
def test_initialize_default_config(self):
"""Test initialization with default config."""
with patch('requests.get') as mock_get:
with patch("requests.get") as mock_get:
mock_response = Mock()
mock_response.status_code = 200
mock_get.return_value = mock_response
@@ -222,30 +195,18 @@ class TestSelfManagedProvider:
with pytest.raises(RuntimeError, match="Provider not initialized"):
provider.create_instance("python")
@patch('requests.post')
@patch("requests.post")
def test_execute_code_success(self, mock_post):
"""Test successful code execution."""
mock_response = Mock()
mock_response.status_code = 200
mock_response.json.return_value = {
"status": "success",
"stdout": '{"result": 42}',
"stderr": "",
"exit_code": 0,
"time_used_ms": 100.0,
"memory_used_kb": 1024.0
}
mock_response.json.return_value = {"status": "success", "stdout": '{"result": 42}', "stderr": "", "exit_code": 0, "time_used_ms": 100.0, "memory_used_kb": 1024.0}
mock_post.return_value = mock_response
provider = SelfManagedProvider()
provider._initialized = True
result = provider.execute_code(
instance_id="test-123",
code="def main(): return {'result': 42}",
language="python",
timeout=10
)
result = provider.execute_code(instance_id="test-123", code="def main(): return {'result': 42}", language="python", timeout=10)
assert result.stdout == '{"result": 42}'
assert result.stderr == ""
@@ -254,7 +215,7 @@ class TestSelfManagedProvider:
assert result.metadata["status"] == "success"
assert result.metadata["instance_id"] == "test-123"
@patch('requests.post')
@patch("requests.post")
def test_execute_code_maps_structured_result_into_metadata(self, mock_post):
"""Test successful code execution with structured result envelope."""
mock_response = Mock()
@@ -277,19 +238,14 @@ class TestSelfManagedProvider:
provider = SelfManagedProvider()
provider._initialized = True
result = provider.execute_code(
instance_id="test-123",
code="def main(): return {'items': ['a', 'b']}",
language="python",
timeout=10
)
result = provider.execute_code(instance_id="test-123", code="def main(): return {'items': ['a', 'b']}", language="python", timeout=10)
assert result.stdout == "debug line\n"
assert result.metadata["result_present"] is True
assert result.metadata["result_value"] == {"items": ["a", "b"]}
assert result.metadata["result_type"] == "json"
@patch('requests.post')
@patch("requests.post")
def test_execute_code_timeout(self, mock_post):
"""Test code execution timeout."""
mock_post.side_effect = requests.Timeout()
@@ -298,14 +254,9 @@ class TestSelfManagedProvider:
provider._initialized = True
with pytest.raises(TimeoutError, match="Execution timed out"):
provider.execute_code(
instance_id="test-123",
code="while True: pass",
language="python",
timeout=5
)
provider.execute_code(instance_id="test-123", code="while True: pass", language="python", timeout=5)
@patch('requests.post')
@patch("requests.post")
def test_execute_code_http_error(self, mock_post):
"""Test code execution with HTTP error."""
mock_response = Mock()
@@ -317,22 +268,14 @@ class TestSelfManagedProvider:
provider._initialized = True
with pytest.raises(RuntimeError, match="HTTP 500"):
provider.execute_code(
instance_id="test-123",
code="invalid code",
language="python"
)
provider.execute_code(instance_id="test-123", code="invalid code", language="python")
def test_execute_code_not_initialized(self):
"""Test executing code when provider not initialized."""
provider = SelfManagedProvider()
with pytest.raises(RuntimeError, match="Provider not initialized"):
provider.execute_code(
instance_id="test-123",
code="print('hello')",
language="python"
)
provider.execute_code(instance_id="test-123", code="print('hello')", language="python")
def test_destroy_instance(self):
"""Test destroying an instance (no-op for self-managed)."""
@@ -344,7 +287,7 @@ class TestSelfManagedProvider:
assert result is True
@patch('requests.get')
@patch("requests.get")
def test_health_check_success(self, mock_get):
"""Test successful health check."""
mock_response = Mock()
@@ -358,7 +301,7 @@ class TestSelfManagedProvider:
assert result is True
mock_get.assert_called_once_with("http://localhost:9385/healthz", timeout=5)
@patch('requests.get')
@patch("requests.get")
def test_health_check_failure(self, mock_get):
"""Test health check failure."""
mock_get.side_effect = Exception("Connection error")
@@ -439,20 +382,20 @@ class TestProviderInterface:
provider = SelfManagedProvider()
# Check all abstract methods are implemented
assert hasattr(provider, 'initialize')
assert hasattr(provider, "initialize")
assert callable(provider.initialize)
assert hasattr(provider, 'create_instance')
assert hasattr(provider, "create_instance")
assert callable(provider.create_instance)
assert hasattr(provider, 'execute_code')
assert hasattr(provider, "execute_code")
assert callable(provider.execute_code)
assert hasattr(provider, 'destroy_instance')
assert hasattr(provider, "destroy_instance")
assert callable(provider.destroy_instance)
assert hasattr(provider, 'health_check')
assert hasattr(provider, "health_check")
assert callable(provider.health_check)
assert hasattr(provider, 'get_supported_languages')
assert hasattr(provider, "get_supported_languages")
assert callable(provider.get_supported_languages)
agent/sandbox/tests/test_security.py
+2 -6
View File
@@ -76,9 +76,7 @@ def test_python_builtins_import_is_rejected():
assert is_safe is False
# Pin the specific reason: rejection must come from the new ``builtins``
# entry in ``DANGEROUS_IMPORTS``, not from some unrelated parse error.
assert any("builtins" in issue for issue, _ in issues), (
f"expected an issue mentioning 'builtins', got {issues!r}"
)
assert any("builtins" in issue for issue, _ in issues), f"expected an issue mentioning 'builtins', got {issues!r}"
def test_python_attribute_eval_call_is_rejected():
@@ -94,9 +92,7 @@ def test_python_attribute_eval_call_is_rejected():
# not from the ``import builtins`` line above. We assert ``exec`` is in at
# least one finding so the test fails if visit_Call's attribute branch is
# ever reverted.
assert any("exec" in issue for issue, _ in issues), (
f"expected an issue mentioning 'exec', got {issues!r}"
)
assert any("exec" in issue for issue, _ in issues), f"expected an issue mentioning 'exec', got {issues!r}"
def test_javascript_safe_code_still_passes():
agent/sandbox/tests/verify_sdk.py
+2 -5
View File
@@ -36,17 +36,14 @@ print("✓ Provider has all required methods")
print("\n[3/5] Testing SDK imports...")
try:
# Check if agentrun SDK is available using importlib
if (
importlib.util.find_spec("agentrun.sandbox") is None
or importlib.util.find_spec("agentrun.utils.config") is None
or importlib.util.find_spec("agentrun.utils.exception") is None
):
if importlib.util.find_spec("agentrun.sandbox") is None or importlib.util.find_spec("agentrun.utils.config") is None or importlib.util.find_spec("agentrun.utils.exception") is None:
raise ImportError("agentrun SDK not found")
# Verify imports work (assign to _ to indicate they're intentionally unused)
from agentrun.sandbox import CodeInterpreterSandbox, TemplateType, CodeLanguage
from agentrun.utils.config import Config
from agentrun.utils.exception import ServerError
_ = (CodeInterpreterSandbox, TemplateType, CodeLanguage, Config, ServerError)
print("✓ SDK modules imported successfully")
agent/test/client.py
+7 -7
View File
@@ -18,29 +18,29 @@ import os
from agent.canvas import Canvas
from common import settings
if __name__ == '__main__':
if __name__ == "__main__":
parser = argparse.ArgumentParser()
dsl_default_path = os.path.join(
os.path.dirname(os.path.realpath(__file__)),
"dsl_examples",
"retrieval_and_generate.json",
)
parser.add_argument('-s', '--dsl', default=dsl_default_path, help="input dsl", action='store', required=True)
parser.add_argument('-t', '--tenant_id', default=False, help="Tenant ID", action='store', required=True)
parser.add_argument('-m', '--stream', default=False, help="Stream output", action='store_true', required=False)
parser.add_argument("-s", "--dsl", default=dsl_default_path, help="input dsl", action="store", required=True)
parser.add_argument("-t", "--tenant_id", default=False, help="Tenant ID", action="store", required=True)
parser.add_argument("-m", "--stream", default=False, help="Stream output", action="store_true", required=False)
args = parser.parse_args()
settings.init_settings()
canvas = Canvas(open(args.dsl, "r").read(), args.tenant_id)
if canvas.get_prologue():
print(f"==================== Bot =====================\n> {canvas.get_prologue()}", end='')
print(f"==================== Bot =====================\n> {canvas.get_prologue()}", end="")
query = ""
while True:
canvas.reset(True)
query = input("\n==================== User =====================\n> ")
ans = canvas.run(query=query)
print("==================== Bot =====================\n> ", end='')
print("==================== Bot =====================\n> ", end="")
for ans in canvas.run(query=query):
print(ans, end='\n', flush=True)
print(ans, end="\n", flush=True)
print(canvas.path)
agent/tools/__init__.py
+6 -4
View File
@@ -22,8 +22,9 @@ from typing import Dict, Type
_package_path = os.path.dirname(__file__)
__all_classes: Dict[str, Type] = {}
def _import_submodules() -> None:
for filename in os.listdir(_package_path): # noqa: F821
for filename in os.listdir(_package_path): # noqa: F821
if filename.startswith("__") or not filename.endswith(".py") or filename.startswith("base"):
continue
module_name = filename[:-3]
@@ -34,15 +35,16 @@ def _import_submodules() -> None:
except ImportError as e:
print(f"Warning: Failed to import module {module_name}: {str(e)}")
def _extract_classes_from_module(module: ModuleType) -> None:
for name, obj in inspect.getmembers(module):
if (inspect.isclass(obj) and
obj.__module__ == module.__name__ and not name.startswith("_")):
if inspect.isclass(obj) and obj.__module__ == module.__name__ and not name.startswith("_"):
__all_classes[name] = obj
globals()[name] = obj
_import_submodules()
__all__ = list(__all_classes.keys()) + ["__all_classes"]
del _package_path, _import_submodules, _extract_classes_from_module
del _package_path, _import_submodules, _extract_classes_from_module
agent/tools/arxiv.py
+10 -25
View File
@@ -28,7 +28,7 @@ class ArXivParam(ToolParamBase):
"""
def __init__(self):
self.meta:ToolMeta = {
self.meta: ToolMeta = {
"name": "arxiv_search",
"description": """arXiv is a free distribution service and an open-access archive for nearly 2.4 million scholarly articles in the fields of physics, mathematics, computer science, quantitative biology, quantitative finance, statistics, electrical engineering and systems science, and economics. Materials on this site are not peer-reviewed by arXiv.""",
"parameters": {
@@ -36,26 +36,20 @@ class ArXivParam(ToolParamBase):
"type": "string",
"description": "The search keywords to execute with arXiv. The keywords should be the most important words/terms(includes synonyms) from the original request.",
"default": "{sys.query}",
"required": True
"required": True,
}
}
},
}
super().__init__()
self.top_n = 12
self.sort_by = 'submittedDate'
self.sort_by = "submittedDate"
def check(self):
self.check_positive_integer(self.top_n, "Top N")
self.check_valid_value(self.sort_by, "ArXiv Search Sort_by",
['submittedDate', 'lastUpdatedDate', 'relevance'])
self.check_valid_value(self.sort_by, "ArXiv Search Sort_by", ["submittedDate", "lastUpdatedDate", "relevance"])
def get_input_form(self) -> dict[str, dict]:
return {
"query": {
"name": "Query",
"type": "line"
}
}
return {"query": {"name": "Query", "type": "line"}}
class ArXiv(ToolBase, ABC):
@@ -71,29 +65,20 @@ class ArXiv(ToolBase, ABC):
return ""
last_e = ""
for _ in range(self._param.max_retries+1):
for _ in range(self._param.max_retries + 1):
if self.check_if_canceled("ArXiv processing"):
return
try:
sort_choices = {"relevance": arxiv.SortCriterion.Relevance,
"lastUpdatedDate": arxiv.SortCriterion.LastUpdatedDate,
'submittedDate': arxiv.SortCriterion.SubmittedDate}
sort_choices = {"relevance": arxiv.SortCriterion.Relevance, "lastUpdatedDate": arxiv.SortCriterion.LastUpdatedDate, "submittedDate": arxiv.SortCriterion.SubmittedDate}
arxiv_client = arxiv.Client()
search = arxiv.Search(
query=kwargs["query"],
max_results=self._param.top_n,
sort_by=sort_choices[self._param.sort_by]
)
search = arxiv.Search(query=kwargs["query"], max_results=self._param.top_n, sort_by=sort_choices[self._param.sort_by])
results = list(arxiv_client.results(search))
if self.check_if_canceled("ArXiv processing"):
return
self._retrieve_chunks(results,
get_title=lambda r: r.title,
get_url=lambda r: r.pdf_url,
get_content=lambda r: r.summary)
self._retrieve_chunks(results, get_title=lambda r: r.title, get_url=lambda r: r.pdf_url, get_content=lambda r: r.summary)
return self.output("formalized_content")
except Exception as e:
if self.check_if_canceled("ArXiv processing"):
agent/tools/base.py
+12 -30
View File
@@ -28,10 +28,9 @@ from common.mcp_tool_call_conn import MCPToolBinding, MCPToolCallSession, ToolCa
from timeit import default_timer as timer
from common.misc_utils import thread_pool_exec
class ToolParameter(TypedDict):
type: str
description: str
@@ -81,7 +80,9 @@ class LLMToolPluginCallSession(ToolCallSession):
resp = fallback_output
else:
resp = fallback_output
logging.warning(f"[ToolCall] resp is None, fallback to output name={name} output_keys={list(fallback_output.keys()) if isinstance(fallback_output, dict) else type(fallback_output).__name__}")
logging.warning(
f"[ToolCall] resp is None, fallback to output name={name} output_keys={list(fallback_output.keys()) if isinstance(fallback_output, dict) else type(fallback_output).__name__}"
)
except Exception as e:
logging.warning(f"[ToolCall] resp is None and output fallback failed name={name} err={e}")
@@ -96,28 +97,25 @@ class LLMToolPluginCallSession(ToolCallSession):
class ToolParamBase(ComponentParamBase):
def __init__(self):
#self.meta:ToolMeta = None
# self.meta:ToolMeta = None
super().__init__()
self._init_inputs()
self._init_attr_by_meta()
def _init_inputs(self):
self.inputs = {}
for k,p in self.meta["parameters"].items():
for k, p in self.meta["parameters"].items():
self.inputs[k] = deepcopy(p)
def _init_attr_by_meta(self):
for k,p in self.meta["parameters"].items():
for k, p in self.meta["parameters"].items():
if not hasattr(self, k):
setattr(self, k, p.get("default"))
def get_meta(self):
params = {}
for k, p in self.meta["parameters"].items():
params[k] = {
"type": p["type"],
"description": p["description"]
}
params[k] = {"type": p["type"], "description": p["description"]}
if "enum" in p:
params[k]["enum"] = p["enum"]
@@ -129,12 +127,8 @@ class ToolParamBase(ComponentParamBase):
"function": {
"name": function_name,
"description": desc,
"parameters": {
"type": "object",
"properties": params,
"required": [k for k, p in self.meta["parameters"].items() if p["required"]]
}
}
"parameters": {"type": "object", "properties": params, "required": [k for k, p in self.meta["parameters"].items() if p["required"]]},
},
}
@@ -209,20 +203,8 @@ class ToolBase(ComponentBase):
title = get_title(r)
url = get_url(r)
score = get_score(r) if get_score else 1
chunks.append({
"chunk_id": id,
"content": content,
"doc_id": id,
"docnm_kwd": title,
"similarity": score,
"url": url
})
aggs.append({
"doc_name": title,
"doc_id": id,
"count": 1,
"url": url
})
chunks.append({"chunk_id": id, "content": content, "doc_id": id, "docnm_kwd": title, "similarity": score, "url": url})
aggs.append({"doc_name": title, "doc_id": id, "count": 1, "url": url})
self._canvas.add_reference(chunks, aggs)
self.set_output("formalized_content", "\n".join(kb_prompt({"chunks": chunks, "doc_aggs": aggs}, 200000, True)))
agent/tools/deepl.py
+46 -12
View File
@@ -27,18 +27,53 @@ class DeepLParam(ComponentParamBase):
super().__init__()
self.auth_key = "xxx"
self.parameters = []
self.source_lang = 'ZH'
self.target_lang = 'EN-GB'
self.source_lang = "ZH"
self.target_lang = "EN-GB"
def check(self):
self.check_valid_value(self.source_lang, "Source language",
['AR', 'BG', 'CS', 'DA', 'DE', 'EL', 'EN', 'ES', 'ET', 'FI', 'FR', 'HU', 'ID', 'IT',
'JA', 'KO', 'LT', 'LV', 'NB', 'NL', 'PL', 'PT', 'RO', 'RU', 'SK', 'SL', 'SV', 'TR',
'UK', 'ZH'])
self.check_valid_value(self.target_lang, "Target language",
['AR', 'BG', 'CS', 'DA', 'DE', 'EL', 'EN-GB', 'EN-US', 'ES', 'ET', 'FI', 'FR', 'HU',
'ID', 'IT', 'JA', 'KO', 'LT', 'LV', 'NB', 'NL', 'PL', 'PT-BR', 'PT-PT', 'RO', 'RU',
'SK', 'SL', 'SV', 'TR', 'UK', 'ZH'])
self.check_valid_value(
self.source_lang,
"Source language",
["AR", "BG", "CS", "DA", "DE", "EL", "EN", "ES", "ET", "FI", "FR", "HU", "ID", "IT", "JA", "KO", "LT", "LV", "NB", "NL", "PL", "PT", "RO", "RU", "SK", "SL", "SV", "TR", "UK", "ZH"],
)
self.check_valid_value(
self.target_lang,
"Target language",
[
"AR",
"BG",
"CS",
"DA",
"DE",
"EL",
"EN-GB",
"EN-US",
"ES",
"ET",
"FI",
"FR",
"HU",
"ID",
"IT",
"JA",
"KO",
"LT",
"LV",
"NB",
"NL",
"PL",
"PT-BR",
"PT-PT",
"RO",
"RU",
"SK",
"SL",
"SV",
"TR",
"UK",
"ZH",
],
)
class DeepL(ComponentBase, ABC):
@@ -57,8 +92,7 @@ class DeepL(ComponentBase, ABC):
try:
translator = deepl.Translator(self._param.auth_key)
result = translator.translate_text(ans, source_lang=self._param.source_lang,
target_lang=self._param.target_lang)
result = translator.translate_text(ans, source_lang=self._param.source_lang, target_lang=self._param.target_lang)
return DeepL.be_output(result.text)
except Exception as e:
agent/tools/duckduckgo.py
+7 -24
View File
@@ -28,7 +28,7 @@ class DuckDuckGoParam(ToolParamBase):
"""
def __init__(self):
self.meta:ToolMeta = {
self.meta: ToolMeta = {
"name": "duckduckgo_search",
"description": "DuckDuckGo is a search engine focused on privacy. It offers search capabilities for web pages, images, and provides translation services. DuckDuckGo also features a private AI chat interface, providing users with an AI assistant that prioritizes data protection.",
"parameters": {
@@ -36,7 +36,7 @@ class DuckDuckGoParam(ToolParamBase):
"type": "string",
"description": "The search keywords to execute with DuckDuckGo. The keywords should be the most important words/terms(includes synonyms) from the original request.",
"default": "{sys.query}",
"required": True
"required": True,
},
"channel": {
"type": "string",
@@ -45,7 +45,7 @@ class DuckDuckGoParam(ToolParamBase):
"default": "general",
"required": False,
},
}
},
}
super().__init__()
self.top_n = 10
@@ -56,18 +56,7 @@ class DuckDuckGoParam(ToolParamBase):
self.check_valid_value(self.channel, "Web Search or News", ["text", "news"])
def get_input_form(self) -> dict[str, dict]:
return {
"query": {
"name": "Query",
"type": "line"
},
"channel": {
"name": "Channel",
"type": "options",
"value": "general",
"options": ["general", "news"]
}
}
return {"query": {"name": "Query", "type": "line"}, "channel": {"name": "Channel", "type": "options", "value": "general", "options": ["general", "news"]}}
class DuckDuckGo(ToolBase, ABC):
@@ -83,7 +72,7 @@ class DuckDuckGo(ToolBase, ABC):
return ""
last_e = ""
for _ in range(self._param.max_retries+1):
for _ in range(self._param.max_retries + 1):
if self.check_if_canceled("DuckDuckGo processing"):
return
@@ -99,10 +88,7 @@ class DuckDuckGo(ToolBase, ABC):
if self.check_if_canceled("DuckDuckGo processing"):
return
self._retrieve_chunks(duck_res,
get_title=lambda r: r["title"],
get_url=lambda r: r.get("href", r.get("url")),
get_content=lambda r: r["body"])
self._retrieve_chunks(duck_res, get_title=lambda r: r["title"], get_url=lambda r: r.get("href", r.get("url")), get_content=lambda r: r["body"])
self.set_output("json", duck_res)
return self.output("formalized_content")
else:
@@ -116,10 +102,7 @@ class DuckDuckGo(ToolBase, ABC):
if self.check_if_canceled("DuckDuckGo processing"):
return
self._retrieve_chunks(duck_res,
get_title=lambda r: r["title"],
get_url=lambda r: r.get("href", r.get("url")),
get_content=lambda r: r["body"])
self._retrieve_chunks(duck_res, get_title=lambda r: r["title"], get_url=lambda r: r.get("href", r.get("url")), get_content=lambda r: r["body"])
self.set_output("json", duck_res)
return self.output("formalized_content")
except Exception as e:
agent/tools/email.py
+18 -48
View File
@@ -32,36 +32,17 @@ class EmailParam(ToolParamBase):
"""
Define the Email component parameters.
"""
def __init__(self):
self.meta:ToolMeta = {
self.meta: ToolMeta = {
"name": "email",
"description": "The email is a method of electronic communication for sending and receiving information through the Internet. This tool helps users to send emails to one person or to multiple recipients with support for CC, BCC, file attachments, and markdown-to-HTML conversion.",
"parameters": {
"to_email": {
"type": "string",
"description": "The target email address.",
"default": "{sys.query}",
"required": True
},
"cc_email": {
"type": "string",
"description": "The other email addresses needs to be send to. Comma splited.",
"default": "",
"required": False
},
"content": {
"type": "string",
"description": "The content of the email.",
"default": "",
"required": False
},
"subject": {
"type": "string",
"description": "The subject/title of the email.",
"default": "",
"required": False
}
}
"to_email": {"type": "string", "description": "The target email address.", "default": "{sys.query}", "required": True},
"cc_email": {"type": "string", "description": "The other email addresses needs to be send to. Comma splited.", "default": "", "required": False},
"content": {"type": "string", "description": "The content of the email.", "default": "", "required": False},
"subject": {"type": "string", "description": "The subject/title of the email.", "default": "", "required": False},
},
}
super().__init__()
# Fixed configuration parameters
@@ -81,20 +62,9 @@ class EmailParam(ToolParamBase):
def get_input_form(self) -> dict[str, dict]:
return {
"to_email": {
"name": "To ",
"type": "line"
},
"subject": {
"name": "Subject",
"type": "line",
"optional": True
},
"cc_email": {
"name": "CC To",
"type": "line",
"optional": True
},
"to_email": {"name": "To ", "type": "line"},
"subject": {"name": "Subject", "type": "line", "optional": True},
"cc_email": {"name": "CC To", "type": "line", "optional": True},
}
@@ -111,7 +81,7 @@ class Email(ToolBase, ABC):
return ""
last_e = ""
for _ in range(self._param.max_retries+1):
for _ in range(self._param.max_retries + 1):
if self.check_if_canceled("Email processing"):
return
@@ -126,19 +96,19 @@ class Email(ToolBase, ABC):
return False
# Create email object
msg = MIMEMultipart('alternative')
msg = MIMEMultipart("alternative")
# Properly handle sender name encoding
msg['From'] = formataddr((str(Header(self._param.sender_name,'utf-8')), self._param.email))
msg['To'] = email_data["to_email"]
msg["From"] = formataddr((str(Header(self._param.sender_name, "utf-8")), self._param.email))
msg["To"] = email_data["to_email"]
if email_data.get("cc_email"):
msg['Cc'] = email_data["cc_email"]
msg['Subject'] = Header(email_data.get("subject", "No Subject"), 'utf-8').encode()
msg["Cc"] = email_data["cc_email"]
msg["Subject"] = Header(email_data.get("subject", "No Subject"), "utf-8").encode()
# Use content from email_data or default content
email_content = email_data.get("content", "No content provided")
# msg.attach(MIMEText(email_content, 'plain', 'utf-8'))
msg.attach(MIMEText(email_content, 'html', 'utf-8'))
msg.attach(MIMEText(email_content, "html", "utf-8"))
# Connect to SMTP server and send
logging.info(f"Connecting to SMTP server {self._param.smtp_server}:{self._param.smtp_port}")
@@ -160,7 +130,7 @@ class Email(ToolBase, ABC):
# Get all recipient list
recipients = [email_data["to_email"]]
if email_data.get("cc_email"):
recipients.extend(email_data["cc_email"].split(','))
recipients.extend(email_data["cc_email"].split(","))
# Send email
logging.info(f"Sending email to recipients: {recipients}")
agent/tools/exesql.py
+26 -53
View File
@@ -34,17 +34,10 @@ class ExeSQLParam(ToolParamBase):
"""
def __init__(self):
self.meta:ToolMeta = {
self.meta: ToolMeta = {
"name": "execute_sql",
"description": "This is a tool that can execute SQL.",
"parameters": {
"sql": {
"type": "string",
"description": "The SQL needs to be executed.",
"default": "{sys.query}",
"required": True
}
}
"parameters": {"sql": {"type": "string", "description": "The SQL needs to be executed.", "default": "{sys.query}", "required": True}},
}
super().__init__()
self.db_type = "mysql"
@@ -56,7 +49,7 @@ class ExeSQLParam(ToolParamBase):
self.max_records = 1024
def check(self):
self.check_valid_value(self.db_type, "Choose DB type", ['mysql', 'postgres', 'mariadb', 'mssql', 'IBM DB2', 'trino', 'oceanbase'])
self.check_valid_value(self.db_type, "Choose DB type", ["mysql", "postgres", "mariadb", "mssql", "IBM DB2", "trino", "oceanbase"])
self.check_empty(self.database, "Database name")
self.check_empty(self.username, "database username")
self.check_empty(self.host, "IP Address")
@@ -71,12 +64,7 @@ class ExeSQLParam(ToolParamBase):
raise ValueError("For the security reason, it does not support database named rag_flow.")
def get_input_form(self) -> dict[str, dict]:
return {
"sql": {
"name": "SQL",
"type": "line"
}
}
return {"sql": {"name": "SQL", "type": "line"}}
class ExeSQL(ToolBase, ABC):
@@ -90,6 +78,7 @@ class ExeSQL(ToolBase, ABC):
def convert_decimals(obj):
from decimal import Decimal
import math
if isinstance(obj, float):
# Handle NaN and Infinity which are not valid JSON values
if math.isnan(obj) or math.isinf(obj):
@@ -140,24 +129,21 @@ class ExeSQL(ToolBase, ABC):
sqls = sql.split(";")
if self._param.db_type in ["mysql", "mariadb"]:
db = pymysql.connect(db=self._param.database, user=self._param.username, host=safe_host,
port=self._param.port, password=self._param.password)
elif self._param.db_type == 'oceanbase':
db = pymysql.connect(db=self._param.database, user=self._param.username, host=safe_host,
port=self._param.port, password=self._param.password, charset='utf8mb4')
elif self._param.db_type == 'postgres':
db = psycopg2.connect(dbname=self._param.database, user=self._param.username, host=safe_host,
port=self._param.port, password=self._param.password)
elif self._param.db_type == 'mssql':
db = pymysql.connect(db=self._param.database, user=self._param.username, host=safe_host, port=self._param.port, password=self._param.password)
elif self._param.db_type == "oceanbase":
db = pymysql.connect(db=self._param.database, user=self._param.username, host=safe_host, port=self._param.port, password=self._param.password, charset="utf8mb4")
elif self._param.db_type == "postgres":
db = psycopg2.connect(dbname=self._param.database, user=self._param.username, host=safe_host, port=self._param.port, password=self._param.password)
elif self._param.db_type == "mssql":
conn_str = (
r'DRIVER={ODBC Driver 17 for SQL Server};'
r'SERVER=' + safe_host + ',' + str(self._param.port) + ';'
r'DATABASE=' + self._param.database + ';'
r'UID=' + self._param.username + ';'
r'PWD=' + self._param.password
r"DRIVER={ODBC Driver 17 for SQL Server};"
r"SERVER=" + safe_host + "," + str(self._param.port) + ";"
r"DATABASE=" + self._param.database + ";"
r"UID=" + self._param.username + ";"
r"PWD=" + self._param.password
)
db = pyodbc.connect(conn_str)
elif self._param.db_type == 'trino':
elif self._param.db_type == "trino":
try:
import trino
from trino.auth import BasicAuthentication
@@ -186,26 +172,14 @@ class ExeSQL(ToolBase, ABC):
try:
db = trino.dbapi.connect(
host=safe_host,
port=int(self._param.port or 8080),
user=self._param.username or "ragflow",
catalog=catalog,
schema=schema or "default",
http_scheme=http_scheme,
auth=auth
host=safe_host, port=int(self._param.port or 8080), user=self._param.username or "ragflow", catalog=catalog, schema=schema or "default", http_scheme=http_scheme, auth=auth
)
except Exception as e:
raise Exception("Database Connection Failed! \n" + str(e))
elif self._param.db_type == 'IBM DB2':
elif self._param.db_type == "IBM DB2":
import ibm_db
conn_str = (
f"DATABASE={self._param.database};"
f"HOSTNAME={safe_host};"
f"PORT={self._param.port};"
f"PROTOCOL=TCPIP;"
f"UID={self._param.username};"
f"PWD={self._param.password};"
)
conn_str = f"DATABASE={self._param.database};HOSTNAME={safe_host};PORT={self._param.port};PROTOCOL=TCPIP;UID={self._param.username};PWD={self._param.password};"
try:
conn = ibm_db.connect(conn_str, "", "")
except Exception as e:
@@ -275,7 +249,7 @@ class ExeSQL(ToolBase, ABC):
if self.check_if_canceled("ExeSQL processing"):
return
single_sql = single_sql.replace('```', '').strip()
single_sql = single_sql.replace("```", "").strip()
if not single_sql:
continue
single_sql = re.sub(r"\[ID:[0-9]+\]", "", single_sql)
@@ -288,20 +262,19 @@ class ExeSQL(ToolBase, ABC):
if cursor.rowcount == 0:
sql_res.append({"content": "No record in the database!"})
break
if self._param.db_type == 'mssql':
single_res = pd.DataFrame.from_records(cursor.fetchmany(self._param.max_records),
columns=[desc[0] for desc in cursor.description])
if self._param.db_type == "mssql":
single_res = pd.DataFrame.from_records(cursor.fetchmany(self._param.max_records), columns=[desc[0] for desc in cursor.description])
else:
single_res = pd.DataFrame([i for i in cursor.fetchmany(self._param.max_records)])
single_res.columns = [i[0] for i in cursor.description]
for col in single_res.columns:
if pd.api.types.is_datetime64_any_dtype(single_res[col]):
single_res[col] = single_res[col].dt.strftime('%Y-%m-%d')
single_res[col] = single_res[col].dt.strftime("%Y-%m-%d")
single_res = single_res.where(pd.notnull(single_res), None)
sql_res.append(convert_decimals(single_res.to_dict(orient='records')))
sql_res.append(convert_decimals(single_res.to_dict(orient="records")))
formalized_content.append(single_res.to_markdown(index=False, floatfmt=".6f"))
except Exception as e:
# A failing statement must not abort the node: report it and keep
agent/tools/github.py
+10 -18
View File
@@ -29,7 +29,7 @@ class GitHubParam(ToolParamBase):
"""
def __init__(self):
self.meta:ToolMeta = {
self.meta: ToolMeta = {
"name": "github_search",
"description": """GitHub repository search is a feature that enables users to find specific repositories on the GitHub platform. This search functionality allows users to locate projects, codebases, and other content hosted on GitHub based on various criteria.""",
"parameters": {
@@ -37,9 +37,9 @@ class GitHubParam(ToolParamBase):
"type": "string",
"description": "The search keywords to execute with GitHub. The keywords should be the most important words/terms(includes synonyms) from the original request.",
"default": "{sys.query}",
"required": True
"required": True,
}
}
},
}
super().__init__()
self.top_n = 10
@@ -48,12 +48,8 @@ class GitHubParam(ToolParamBase):
self.check_positive_integer(self.top_n, "Top N")
def get_input_form(self) -> dict[str, dict]:
return {
"query": {
"name": "Query",
"type": "line"
}
}
return {"query": {"name": "Query", "type": "line"}}
class GitHub(ToolBase, ABC):
component_name = "GitHub"
@@ -68,24 +64,20 @@ class GitHub(ToolBase, ABC):
return ""
last_e = ""
for _ in range(self._param.max_retries+1):
for _ in range(self._param.max_retries + 1):
if self.check_if_canceled("GitHub processing"):
return
try:
url = 'https://api.github.com/search/repositories?q=' + kwargs["query"] + '&sort=stars&order=desc&per_page=' + str(
self._param.top_n)
headers = {"Content-Type": "application/vnd.github+json", "X-GitHub-Api-Version": '2022-11-28'}
url = "https://api.github.com/search/repositories?q=" + kwargs["query"] + "&sort=stars&order=desc&per_page=" + str(self._param.top_n)
headers = {"Content-Type": "application/vnd.github+json", "X-GitHub-Api-Version": "2022-11-28"}
response = requests.get(url=url, headers=headers, timeout=DEFAULT_TIMEOUT).json()
if self.check_if_canceled("GitHub processing"):
return
self._retrieve_chunks(response['items'],
get_title=lambda r: r["name"],
get_url=lambda r: r["html_url"],
get_content=lambda r: str(r["description"]) + '\n stars:' + str(r['watchers']))
self.set_output("json", response['items'])
self._retrieve_chunks(response["items"], get_title=lambda r: r["name"], get_url=lambda r: r["html_url"], get_content=lambda r: str(r["description"]) + "\n stars:" + str(r["watchers"]))
self.set_output("json", response["items"])
return self.output("formalized_content")
except Exception as e:
if self.check_if_canceled("GitHub processing"):
agent/tools/google.py
+422 -67
View File
@@ -28,7 +28,7 @@ class GoogleParam(ToolParamBase):
"""
def __init__(self):
self.meta:ToolMeta = {
self.meta: ToolMeta = {
"name": "google_search",
"description": """Search the world's information, including webpages, images, videos and more. Google has many special features to help you find exactly what you're looking ...""",
"parameters": {
@@ -36,7 +36,7 @@ class GoogleParam(ToolParamBase):
"type": "string",
"description": "The search keywords to execute with Google. The keywords should be the most important words/terms(includes synonyms) from the original request.",
"default": "{sys.query}",
"required": True
"required": True,
},
"start": {
"type": "integer",
@@ -49,8 +49,8 @@ class GoogleParam(ToolParamBase):
"description": "Parameter defines the maximum number of results to return. (e.g., 10 (default) returns 10 results, 40 returns 40 results, and 100 returns 100 results). The use of num may introduce latency, and/or prevent the inclusion of specialized result types. It is better to omit this parameter unless it is strictly necessary to increase the number of results per page. Results are not guaranteed to have the number of results specified in num.",
"default": "6",
"required": False,
}
}
},
},
}
super().__init__()
self.start = 0
@@ -61,57 +61,418 @@ class GoogleParam(ToolParamBase):
def check(self):
self.check_empty(self.api_key, "SerpApi API key")
self.check_valid_value(self.country, "Google Country",
['af', 'al', 'dz', 'as', 'ad', 'ao', 'ai', 'aq', 'ag', 'ar', 'am', 'aw', 'au', 'at',
'az', 'bs', 'bh', 'bd', 'bb', 'by', 'be', 'bz', 'bj', 'bm', 'bt', 'bo', 'ba', 'bw',
'bv', 'br', 'io', 'bn', 'bg', 'bf', 'bi', 'kh', 'cm', 'ca', 'cv', 'ky', 'cf', 'td',
'cl', 'cn', 'cx', 'cc', 'co', 'km', 'cg', 'cd', 'ck', 'cr', 'ci', 'hr', 'cu', 'cy',
'cz', 'dk', 'dj', 'dm', 'do', 'ec', 'eg', 'sv', 'gq', 'er', 'ee', 'et', 'fk', 'fo',
'fj', 'fi', 'fr', 'gf', 'pf', 'tf', 'ga', 'gm', 'ge', 'de', 'gh', 'gi', 'gr', 'gl',
'gd', 'gp', 'gu', 'gt', 'gn', 'gw', 'gy', 'ht', 'hm', 'va', 'hn', 'hk', 'hu', 'is',
'in', 'id', 'ir', 'iq', 'ie', 'il', 'it', 'jm', 'jp', 'jo', 'kz', 'ke', 'ki', 'kp',
'kr', 'kw', 'kg', 'la', 'lv', 'lb', 'ls', 'lr', 'ly', 'li', 'lt', 'lu', 'mo', 'mk',
'mg', 'mw', 'my', 'mv', 'ml', 'mt', 'mh', 'mq', 'mr', 'mu', 'yt', 'mx', 'fm', 'md',
'mc', 'mn', 'ms', 'ma', 'mz', 'mm', 'na', 'nr', 'np', 'nl', 'an', 'nc', 'nz', 'ni',
'ne', 'ng', 'nu', 'nf', 'mp', 'no', 'om', 'pk', 'pw', 'ps', 'pa', 'pg', 'py', 'pe',
'ph', 'pn', 'pl', 'pt', 'pr', 'qa', 're', 'ro', 'ru', 'rw', 'sh', 'kn', 'lc', 'pm',
'vc', 'ws', 'sm', 'st', 'sa', 'sn', 'rs', 'sc', 'sl', 'sg', 'sk', 'si', 'sb', 'so',
'za', 'gs', 'es', 'lk', 'sd', 'sr', 'sj', 'sz', 'se', 'ch', 'sy', 'tw', 'tj', 'tz',
'th', 'tl', 'tg', 'tk', 'to', 'tt', 'tn', 'tr', 'tm', 'tc', 'tv', 'ug', 'ua', 'ae',
'uk', 'gb', 'us', 'um', 'uy', 'uz', 'vu', 've', 'vn', 'vg', 'vi', 'wf', 'eh', 'ye',
'zm', 'zw'])
self.check_valid_value(self.language, "Google languages",
['af', 'ak', 'sq', 'ws', 'am', 'ar', 'hy', 'az', 'eu', 'be', 'bem', 'bn', 'bh',
'xx-bork', 'bs', 'br', 'bg', 'bt', 'km', 'ca', 'chr', 'ny', 'zh-cn', 'zh-tw', 'co',
'hr', 'cs', 'da', 'nl', 'xx-elmer', 'en', 'eo', 'et', 'ee', 'fo', 'tl', 'fi', 'fr',
'fy', 'gaa', 'gl', 'ka', 'de', 'el', 'kl', 'gn', 'gu', 'xx-hacker', 'ht', 'ha', 'haw',
'iw', 'hi', 'hu', 'is', 'ig', 'id', 'ia', 'ga', 'it', 'ja', 'jw', 'kn', 'kk', 'rw',
'rn', 'xx-klingon', 'kg', 'ko', 'kri', 'ku', 'ckb', 'ky', 'lo', 'la', 'lv', 'ln', 'lt',
'loz', 'lg', 'ach', 'mk', 'mg', 'ms', 'ml', 'mt', 'mv', 'mi', 'mr', 'mfe', 'mo', 'mn',
'sr-me', 'my', 'ne', 'pcm', 'nso', 'no', 'nn', 'oc', 'or', 'om', 'ps', 'fa',
'xx-pirate', 'pl', 'pt', 'pt-br', 'pt-pt', 'pa', 'qu', 'ro', 'rm', 'nyn', 'ru', 'gd',
'sr', 'sh', 'st', 'tn', 'crs', 'sn', 'sd', 'si', 'sk', 'sl', 'so', 'es', 'es-419', 'su',
'sw', 'sv', 'tg', 'ta', 'tt', 'te', 'th', 'ti', 'to', 'lua', 'tum', 'tr', 'tk', 'tw',
'ug', 'uk', 'ur', 'uz', 'vu', 'vi', 'cy', 'wo', 'xh', 'yi', 'yo', 'zu']
)
self.check_valid_value(
self.country,
"Google Country",
[
"af",
"al",
"dz",
"as",
"ad",
"ao",
"ai",
"aq",
"ag",
"ar",
"am",
"aw",
"au",
"at",
"az",
"bs",
"bh",
"bd",
"bb",
"by",
"be",
"bz",
"bj",
"bm",
"bt",
"bo",
"ba",
"bw",
"bv",
"br",
"io",
"bn",
"bg",
"bf",
"bi",
"kh",
"cm",
"ca",
"cv",
"ky",
"cf",
"td",
"cl",
"cn",
"cx",
"cc",
"co",
"km",
"cg",
"cd",
"ck",
"cr",
"ci",
"hr",
"cu",
"cy",
"cz",
"dk",
"dj",
"dm",
"do",
"ec",
"eg",
"sv",
"gq",
"er",
"ee",
"et",
"fk",
"fo",
"fj",
"fi",
"fr",
"gf",
"pf",
"tf",
"ga",
"gm",
"ge",
"de",
"gh",
"gi",
"gr",
"gl",
"gd",
"gp",
"gu",
"gt",
"gn",
"gw",
"gy",
"ht",
"hm",
"va",
"hn",
"hk",
"hu",
"is",
"in",
"id",
"ir",
"iq",
"ie",
"il",
"it",
"jm",
"jp",
"jo",
"kz",
"ke",
"ki",
"kp",
"kr",
"kw",
"kg",
"la",
"lv",
"lb",
"ls",
"lr",
"ly",
"li",
"lt",
"lu",
"mo",
"mk",
"mg",
"mw",
"my",
"mv",
"ml",
"mt",
"mh",
"mq",
"mr",
"mu",
"yt",
"mx",
"fm",
"md",
"mc",
"mn",
"ms",
"ma",
"mz",
"mm",
"na",
"nr",
"np",
"nl",
"an",
"nc",
"nz",
"ni",
"ne",
"ng",
"nu",
"nf",
"mp",
"no",
"om",
"pk",
"pw",
"ps",
"pa",
"pg",
"py",
"pe",
"ph",
"pn",
"pl",
"pt",
"pr",
"qa",
"re",
"ro",
"ru",
"rw",
"sh",
"kn",
"lc",
"pm",
"vc",
"ws",
"sm",
"st",
"sa",
"sn",
"rs",
"sc",
"sl",
"sg",
"sk",
"si",
"sb",
"so",
"za",
"gs",
"es",
"lk",
"sd",
"sr",
"sj",
"sz",
"se",
"ch",
"sy",
"tw",
"tj",
"tz",
"th",
"tl",
"tg",
"tk",
"to",
"tt",
"tn",
"tr",
"tm",
"tc",
"tv",
"ug",
"ua",
"ae",
"uk",
"gb",
"us",
"um",
"uy",
"uz",
"vu",
"ve",
"vn",
"vg",
"vi",
"wf",
"eh",
"ye",
"zm",
"zw",
],
)
self.check_valid_value(
self.language,
"Google languages",
[
"af",
"ak",
"sq",
"ws",
"am",
"ar",
"hy",
"az",
"eu",
"be",
"bem",
"bn",
"bh",
"xx-bork",
"bs",
"br",
"bg",
"bt",
"km",
"ca",
"chr",
"ny",
"zh-cn",
"zh-tw",
"co",
"hr",
"cs",
"da",
"nl",
"xx-elmer",
"en",
"eo",
"et",
"ee",
"fo",
"tl",
"fi",
"fr",
"fy",
"gaa",
"gl",
"ka",
"de",
"el",
"kl",
"gn",
"gu",
"xx-hacker",
"ht",
"ha",
"haw",
"iw",
"hi",
"hu",
"is",
"ig",
"id",
"ia",
"ga",
"it",
"ja",
"jw",
"kn",
"kk",
"rw",
"rn",
"xx-klingon",
"kg",
"ko",
"kri",
"ku",
"ckb",
"ky",
"lo",
"la",
"lv",
"ln",
"lt",
"loz",
"lg",
"ach",
"mk",
"mg",
"ms",
"ml",
"mt",
"mv",
"mi",
"mr",
"mfe",
"mo",
"mn",
"sr-me",
"my",
"ne",
"pcm",
"nso",
"no",
"nn",
"oc",
"or",
"om",
"ps",
"fa",
"xx-pirate",
"pl",
"pt",
"pt-br",
"pt-pt",
"pa",
"qu",
"ro",
"rm",
"nyn",
"ru",
"gd",
"sr",
"sh",
"st",
"tn",
"crs",
"sn",
"sd",
"si",
"sk",
"sl",
"so",
"es",
"es-419",
"su",
"sw",
"sv",
"tg",
"ta",
"tt",
"te",
"th",
"ti",
"to",
"lua",
"tum",
"tr",
"tk",
"tw",
"ug",
"uk",
"ur",
"uz",
"vu",
"vi",
"cy",
"wo",
"xh",
"yi",
"yo",
"zu",
],
)
def get_input_form(self) -> dict[str, dict]:
return {
"q": {
"name": "Query",
"type": "line"
},
"start": {
"name": "From",
"type": "integer",
"value": 0
},
"num": {
"name": "Limit",
"type": "integer",
"value": 12
}
}
return {"q": {"name": "Query", "type": "line"}, "start": {"name": "From", "type": "integer", "value": 0}, "num": {"name": "Limit", "type": "integer", "value": 12}}
class Google(ToolBase, ABC):
component_name = "Google"
@@ -125,16 +486,9 @@ class Google(ToolBase, ABC):
self.set_output("formalized_content", "")
return ""
params = {
"api_key": self._param.api_key,
"engine": "google",
"q": kwargs["q"],
"google_domain": "google.com",
"gl": self._param.country,
"hl": self._param.language
}
params = {"api_key": self._param.api_key, "engine": "google", "q": kwargs["q"], "google_domain": "google.com", "gl": self._param.country, "hl": self._param.language}
last_e = ""
for _ in range(self._param.max_retries+1):
for _ in range(self._param.max_retries + 1):
if self.check_if_canceled("Google processing"):
return
@@ -144,11 +498,12 @@ class Google(ToolBase, ABC):
if self.check_if_canceled("Google processing"):
return
self._retrieve_chunks(search["organic_results"],
get_title=lambda r: r["title"],
get_url=lambda r: r["link"],
get_content=lambda r: r.get("about_this_result", {}).get("source", {}).get("description", r["snippet"])
)
self._retrieve_chunks(
search["organic_results"],
get_title=lambda r: r["title"],
get_url=lambda r: r["link"],
get_content=lambda r: r.get("about_this_result", {}).get("source", {}).get("description", r["snippet"]),
)
self.set_output("json", search["organic_results"])
return self.output("formalized_content")
except Exception as e:
agent/tools/googlescholar.py
+15 -19
View File
@@ -29,7 +29,7 @@ class GoogleScholarParam(ToolParamBase):
"""
def __init__(self):
self.meta:ToolMeta = {
self.meta: ToolMeta = {
"name": "google_scholar_search",
"description": """Google Scholar provides a simple way to broadly search for scholarly literature. From one place, you can search across many disciplines and sources: articles, theses, books, abstracts and court opinions, from academic publishers, professional societies, online repositories, universities and other web sites. Google Scholar helps you find relevant work across the world of scholarly research.""",
"parameters": {
@@ -37,29 +37,25 @@ class GoogleScholarParam(ToolParamBase):
"type": "string",
"description": "The search keyword to execute with Google Scholar. The keywords should be the most important words/terms(includes synonyms) from the original request.",
"default": "{sys.query}",
"required": True
"required": True,
}
}
},
}
super().__init__()
self.top_n = 12
self.sort_by = 'relevance'
self.sort_by = "relevance"
self.year_low = None
self.year_high = None
self.patents = True
def check(self):
self.check_positive_integer(self.top_n, "Top N")
self.check_valid_value(self.sort_by, "GoogleScholar Sort_by", ['date', 'relevance'])
self.check_valid_value(self.sort_by, "GoogleScholar Sort_by", ["date", "relevance"])
self.check_boolean(self.patents, "Whether or not to include patents, defaults to True")
def get_input_form(self) -> dict[str, dict]:
return {
"query": {
"name": "Query",
"type": "line"
}
}
return {"query": {"name": "Query", "type": "line"}}
class GoogleScholar(ToolBase, ABC):
component_name = "GoogleScholar"
@@ -77,13 +73,12 @@ class GoogleScholar(ToolBase, ABC):
return ""
last_e = ""
for _ in range(self._param.max_retries+1):
for _ in range(self._param.max_retries + 1):
if self.check_if_canceled("GoogleScholar processing"):
return
try:
scholar_client = scholarly.search_pubs(kwargs["query"], patents=self._param.patents, year_low=self._param.year_low,
year_high=self._param.year_high, sort_by=self._param.sort_by)
scholar_client = scholarly.search_pubs(kwargs["query"], patents=self._param.patents, year_low=self._param.year_low, year_high=self._param.year_high, sort_by=self._param.sort_by)
if self.check_if_canceled("GoogleScholar processing"):
return
@@ -94,11 +89,12 @@ class GoogleScholar(ToolBase, ABC):
# would otherwise leave json empty).
results = list(itertools.islice(scholar_client, self._param.top_n))
self._retrieve_chunks(results,
get_title=lambda r: r['bib']['title'],
get_url=lambda r: r["pub_url"],
get_content=lambda r: "\n author: " + ",".join(r['bib']['author']) + '\n Abstract: ' + r['bib'].get('abstract', 'no abstract')
)
self._retrieve_chunks(
results,
get_title=lambda r: r["bib"]["title"],
get_url=lambda r: r["pub_url"],
get_content=lambda r: "\n author: " + ",".join(r["bib"]["author"]) + "\n Abstract: " + r["bib"].get("abstract", "no abstract"),
)
self.set_output("json", results)
return self.output("formalized_content")
except Exception as e:
agent/tools/jin10.py
+52 -61
View File
@@ -30,21 +30,21 @@ class Jin10Param(ComponentParamBase):
super().__init__()
self.type = "flash"
self.secret_key = "xxx"
self.flash_type = '1'
self.calendar_type = 'cj'
self.calendar_datatype = 'data'
self.symbols_type = 'GOODS'
self.symbols_datatype = 'symbols'
self.flash_type = "1"
self.calendar_type = "cj"
self.calendar_datatype = "data"
self.symbols_type = "GOODS"
self.symbols_datatype = "symbols"
self.contain = ""
self.filter = ""
def check(self):
self.check_valid_value(self.type, "Type", ['flash', 'calendar', 'symbols', 'news'])
self.check_valid_value(self.flash_type, "Flash Type", ['1', '2', '3', '4', '5'])
self.check_valid_value(self.calendar_type, "Calendar Type", ['cj', 'qh', 'hk', 'us'])
self.check_valid_value(self.calendar_datatype, "Calendar DataType", ['data', 'event', 'holiday'])
self.check_valid_value(self.symbols_type, "Symbols Type", ['GOODS', 'FOREX', 'FUTURE', 'CRYPTO'])
self.check_valid_value(self.symbols_datatype, 'Symbols DataType', ['symbols', 'quotes'])
self.check_valid_value(self.type, "Type", ["flash", "calendar", "symbols", "news"])
self.check_valid_value(self.flash_type, "Flash Type", ["1", "2", "3", "4", "5"])
self.check_valid_value(self.calendar_type, "Calendar Type", ["cj", "qh", "hk", "us"])
self.check_valid_value(self.calendar_datatype, "Calendar DataType", ["data", "event", "holiday"])
self.check_valid_value(self.symbols_type, "Symbols Type", ["GOODS", "FOREX", "FUTURE", "CRYPTO"])
self.check_valid_value(self.symbols_datatype, "Symbols DataType", ["symbols", "quotes"])
class Jin10(ComponentBase, ABC):
@@ -60,86 +60,77 @@ class Jin10(ComponentBase, ABC):
return Jin10.be_output("")
jin10_res = []
headers = {'secret-key': self._param.secret_key}
headers = {"secret-key": self._param.secret_key}
try:
if self.check_if_canceled("Jin10 processing"):
return
if self._param.type == "flash":
params = {
'category': self._param.flash_type,
'contain': self._param.contain,
'filter': self._param.filter
}
response = requests.get(
url='https://open-data-api.jin10.com/data-api/flash?category=' + self._param.flash_type,
headers=headers, data=json.dumps(params), timeout=DEFAULT_TIMEOUT)
params = {"category": self._param.flash_type, "contain": self._param.contain, "filter": self._param.filter}
response = requests.get(url="https://open-data-api.jin10.com/data-api/flash?category=" + self._param.flash_type, headers=headers, data=json.dumps(params), timeout=DEFAULT_TIMEOUT)
response = response.json()
for i in response['data']:
for i in response["data"]:
if self.check_if_canceled("Jin10 processing"):
return
jin10_res.append({"content": i['data']['content']})
jin10_res.append({"content": i["data"]["content"]})
if self._param.type == "calendar":
params = {
'category': self._param.calendar_type
}
params = {"category": self._param.calendar_type}
response = requests.get(
url='https://open-data-api.jin10.com/data-api/calendar/' + self._param.calendar_datatype + '?category=' + self._param.calendar_type,
headers=headers, data=json.dumps(params), timeout=DEFAULT_TIMEOUT)
url="https://open-data-api.jin10.com/data-api/calendar/" + self._param.calendar_datatype + "?category=" + self._param.calendar_type,
headers=headers,
data=json.dumps(params),
timeout=DEFAULT_TIMEOUT,
)
response = response.json()
if self.check_if_canceled("Jin10 processing"):
return
jin10_res.append({"content": pd.DataFrame(response['data']).to_markdown()})
jin10_res.append({"content": pd.DataFrame(response["data"]).to_markdown()})
if self._param.type == "symbols":
params = {
'type': self._param.symbols_type
}
params = {"type": self._param.symbols_type}
if self._param.symbols_datatype == "quotes":
params['codes'] = 'BTCUSD'
params["codes"] = "BTCUSD"
response = requests.get(
url='https://open-data-api.jin10.com/data-api/' + self._param.symbols_datatype + '?type=' + self._param.symbols_type,
headers=headers, data=json.dumps(params), timeout=DEFAULT_TIMEOUT)
url="https://open-data-api.jin10.com/data-api/" + self._param.symbols_datatype + "?type=" + self._param.symbols_type,
headers=headers,
data=json.dumps(params),
timeout=DEFAULT_TIMEOUT,
)
response = response.json()
if self.check_if_canceled("Jin10 processing"):
return
if self._param.symbols_datatype == "symbols":
for i in response['data']:
for i in response["data"]:
if self.check_if_canceled("Jin10 processing"):
return
i['Commodity Code'] = i['c']
i['Stock Exchange'] = i['e']
i['Commodity Name'] = i['n']
i['Commodity Type'] = i['t']
del i['c'], i['e'], i['n'], i['t']
i["Commodity Code"] = i["c"]
i["Stock Exchange"] = i["e"]
i["Commodity Name"] = i["n"]
i["Commodity Type"] = i["t"]
del i["c"], i["e"], i["n"], i["t"]
if self._param.symbols_datatype == "quotes":
for i in response['data']:
for i in response["data"]:
if self.check_if_canceled("Jin10 processing"):
return
i['Selling Price'] = i['a']
i['Buying Price'] = i['b']
i['Commodity Code'] = i['c']
i['Stock Exchange'] = i['e']
i['Highest Price'] = i['h']
i['Yesterday’s Closing Price'] = i['hc']
i['Lowest Price'] = i['l']
i['Opening Price'] = i['o']
i['Latest Price'] = i['p']
i['Market Quote Time'] = i['t']
del i['a'], i['b'], i['c'], i['e'], i['h'], i['hc'], i['l'], i['o'], i['p'], i['t']
jin10_res.append({"content": pd.DataFrame(response['data']).to_markdown()})
i["Selling Price"] = i["a"]
i["Buying Price"] = i["b"]
i["Commodity Code"] = i["c"]
i["Stock Exchange"] = i["e"]
i["Highest Price"] = i["h"]
i["Yesterday’s Closing Price"] = i["hc"]
i["Lowest Price"] = i["l"]
i["Opening Price"] = i["o"]
i["Latest Price"] = i["p"]
i["Market Quote Time"] = i["t"]
del i["a"], i["b"], i["c"], i["e"], i["h"], i["hc"], i["l"], i["o"], i["p"], i["t"]
jin10_res.append({"content": pd.DataFrame(response["data"]).to_markdown()})
if self._param.type == "news":
params = {
'contain': self._param.contain,
'filter': self._param.filter
}
response = requests.get(
url='https://open-data-api.jin10.com/data-api/news',
headers=headers, data=json.dumps(params), timeout=DEFAULT_TIMEOUT)
params = {"contain": self._param.contain, "filter": self._param.filter}
response = requests.get(url="https://open-data-api.jin10.com/data-api/news", headers=headers, data=json.dumps(params), timeout=DEFAULT_TIMEOUT)
response = response.json()
if self.check_if_canceled("Jin10 processing"):
return
jin10_res.append({"content": pd.DataFrame(response['data']).to_markdown()})
jin10_res.append({"content": pd.DataFrame(response["data"]).to_markdown()})
except Exception as e:
if self.check_if_canceled("Jin10 processing"):
return
agent/tools/pubmed.py
+16 -27
View File
@@ -30,7 +30,7 @@ class PubMedParam(ToolParamBase):
"""
def __init__(self):
self.meta:ToolMeta = {
self.meta: ToolMeta = {
"name": "pubmed_search",
"description": """
PubMed is an openly accessible, free database which includes primarily the MEDLINE database of references and abstracts on life sciences and biomedical topics.
@@ -47,9 +47,9 @@ In addition to MEDLINE, PubMed provides access to:
"type": "string",
"description": "The search keywords to execute with PubMed. The keywords should be the most important words/terms(includes synonyms) from the original request.",
"default": "{sys.query}",
"required": True
"required": True,
}
}
},
}
super().__init__()
self.top_n = 12
@@ -59,12 +59,8 @@ In addition to MEDLINE, PubMed provides access to:
self.check_positive_integer(self.top_n, "Top N")
def get_input_form(self) -> dict[str, dict]:
return {
"query": {
"name": "Query",
"type": "line"
}
}
return {"query": {"name": "Query", "type": "line"}}
class PubMed(ToolBase, ABC):
component_name = "PubMed"
@@ -79,27 +75,28 @@ class PubMed(ToolBase, ABC):
return ""
last_e = ""
for _ in range(self._param.max_retries+1):
for _ in range(self._param.max_retries + 1):
if self.check_if_canceled("PubMed processing"):
return
try:
Entrez.email = self._param.email
pubmedids = Entrez.read(Entrez.esearch(db='pubmed', retmax=self._param.top_n, term=kwargs["query"]))['IdList']
pubmedids = Entrez.read(Entrez.esearch(db="pubmed", retmax=self._param.top_n, term=kwargs["query"]))["IdList"]
if self.check_if_canceled("PubMed processing"):
return
pubmedcnt = ET.fromstring(re.sub(r'<(/?)b>|<(/?)i>', '', Entrez.efetch(db='pubmed', id=",".join(pubmedids),
retmode="xml").read().decode("utf-8")))
pubmedcnt = ET.fromstring(re.sub(r"<(/?)b>|<(/?)i>", "", Entrez.efetch(db="pubmed", id=",".join(pubmedids), retmode="xml").read().decode("utf-8")))
if self.check_if_canceled("PubMed processing"):
return
self._retrieve_chunks(pubmedcnt.findall("PubmedArticle"),
get_title=lambda child: child.find("MedlineCitation").find("Article").find("ArticleTitle").text,
get_url=lambda child: "https://pubmed.ncbi.nlm.nih.gov/" + child.find("MedlineCitation").find("PMID").text,
get_content=lambda child: self._format_pubmed_content(child),)
self._retrieve_chunks(
pubmedcnt.findall("PubmedArticle"),
get_title=lambda child: child.find("MedlineCitation").find("Article").find("ArticleTitle").text,
get_url=lambda child: "https://pubmed.ncbi.nlm.nih.gov/" + child.find("MedlineCitation").find("PMID").text,
get_content=lambda child: self._format_pubmed_content(child),
)
return self.output("formalized_content")
except Exception as e:
if self.check_if_canceled("PubMed processing"):
@@ -117,6 +114,7 @@ class PubMed(ToolBase, ABC):
def _format_pubmed_content(self, child):
"""Extract structured reference info from PubMed XML"""
def safe_find(path, base=None):
node = child if base is None else base
for p in path.split("/"):
@@ -149,16 +147,7 @@ class PubMed(ToolBase, ABC):
doi = eid.text
break
return (
f"Title: {title}\n"
f"Authors: {authors_str}\n"
f"Journal: {journal}\n"
f"Volume: {volume}\n"
f"Issue: {issue}\n"
f"Pages: {pages}\n"
f"DOI: {doi or '-'}\n"
f"Abstract: {abstract.strip()}"
)
return f"Title: {title}\nAuthors: {authors_str}\nJournal: {journal}\nVolume: {volume}\nIssue: {issue}\nPages: {pages}\nDOI: {doi or '-'}\nAbstract: {abstract.strip()}"
def thoughts(self) -> str:
return "Looking for scholarly papers on `{}`,” prioritising reputable sources.".format(self.get_input().get("query", "-_-!"))
agent/tools/qweather.py
+45 -15
View File
@@ -30,7 +30,7 @@ class QWeatherParam(ComponentParamBase):
self.web_apikey = "xxx"
self.lang = "zh"
self.type = "weather"
self.user_type = 'free'
self.user_type = "free"
self.error_code = {
"204": "The request was successful, but the region you are querying does not have the data you need at this time.",
"400": "Request error, may contain incorrect request parameters or missing mandatory request parameters.",
@@ -39,20 +39,53 @@ class QWeatherParam(ComponentParamBase):
"403": "No access, may be the binding PackageName, BundleID, domain IP address is inconsistent, or the data that requires additional payment.",
"404": "The queried data or region does not exist.",
"429": "Exceeded the limited QPM (number of accesses per minute), please refer to the QPM description",
"500": "No response or timeout, interface service abnormality please contact us"
}
"500": "No response or timeout, interface service abnormality please contact us",
}
# Weather
self.time_period = 'now'
self.time_period = "now"
def check(self):
self.check_empty(self.web_apikey, "BaiduFanyi APPID")
self.check_valid_value(self.type, "Type", ["weather", "indices", "airquality"])
self.check_valid_value(self.user_type, "Free subscription or paid subscription", ["free", "paid"])
self.check_valid_value(self.lang, "Use language",
['zh', 'zh-hant', 'en', 'de', 'es', 'fr', 'it', 'ja', 'ko', 'ru', 'hi', 'th', 'ar', 'pt',
'bn', 'ms', 'nl', 'el', 'la', 'sv', 'id', 'pl', 'tr', 'cs', 'et', 'vi', 'fil', 'fi',
'he', 'is', 'nb'])
self.check_valid_value(self.time_period, "Time period", ['now', '3d', '7d', '10d', '15d', '30d'])
self.check_valid_value(
self.lang,
"Use language",
[
"zh",
"zh-hant",
"en",
"de",
"es",
"fr",
"it",
"ja",
"ko",
"ru",
"hi",
"th",
"ar",
"pt",
"bn",
"ms",
"nl",
"el",
"la",
"sv",
"id",
"pl",
"tr",
"cs",
"et",
"vi",
"fil",
"fi",
"he",
"is",
"nb",
],
)
self.check_valid_value(self.time_period, "Time period", ["now", "3d", "7d", "10d", "15d", "30d"])
class QWeather(ComponentBase, ABC):
@@ -71,9 +104,7 @@ class QWeather(ComponentBase, ABC):
if self.check_if_canceled("Qweather processing"):
return
response = requests.get(
url="https://geoapi.qweather.com/v2/city/lookup?location=" + ans + "&key=" + self._param.web_apikey,
timeout=DEFAULT_TIMEOUT).json()
response = requests.get(url="https://geoapi.qweather.com/v2/city/lookup?location=" + ans + "&key=" + self._param.web_apikey, timeout=DEFAULT_TIMEOUT).json()
if response["code"] == "200":
location_id = response["location"][0]["id"]
else:
@@ -82,7 +113,7 @@ class QWeather(ComponentBase, ABC):
if self.check_if_canceled("Qweather processing"):
return
base_url = "https://api.qweather.com/v7/" if self._param.user_type == 'paid' else "https://devapi.qweather.com/v7/"
base_url = "https://api.qweather.com/v7/" if self._param.user_type == "paid" else "https://devapi.qweather.com/v7/"
if self._param.type == "weather":
url = base_url + "weather/" + self._param.time_period + "?location=" + location_id + "&key=" + self._param.web_apikey + "&lang=" + self._param.lang
@@ -110,8 +141,7 @@ class QWeather(ComponentBase, ABC):
if self.check_if_canceled("Qweather processing"):
return
if response["code"] == "200":
indices_res = response["daily"][0]["date"] + "\n" + "\n".join(
[i["name"] + ": " + i["category"] + ", " + i["text"] for i in response["daily"]])
indices_res = response["daily"][0]["date"] + "\n" + "\n".join([i["name"] + ": " + i["category"] + ", " + i["text"] for i in response["daily"]])
return QWeather.be_output(indices_res)
else:
agent/tools/retrieval.py
+16 -30
View File
@@ -40,7 +40,7 @@ class RetrievalParam(ToolParamBase):
"""
def __init__(self):
self.meta:ToolMeta = {
self.meta: ToolMeta = {
"name": "search_my_dateset",
"description": "This tool can be utilized for relevant content searching in the datasets.",
"parameters": {
@@ -48,9 +48,9 @@ class RetrievalParam(ToolParamBase):
"type": "string",
"description": "The keywords to search the dataset. The keywords should be the most important words/terms(includes synonyms) from the original request.",
"default": "",
"required": True
"required": True,
}
}
},
}
super().__init__()
self.function_name = "search_my_dateset"
@@ -68,7 +68,7 @@ class RetrievalParam(ToolParamBase):
self.use_kg = False
self.cross_languages = []
self.toc_enhance = False
self.meta_data_filter={}
self.meta_data_filter = {}
def check(self):
self.check_decimal_float(self.similarity_threshold, "[Retrieval] Similarity threshold")
@@ -76,12 +76,8 @@ class RetrievalParam(ToolParamBase):
self.check_positive_number(self.top_n, "[Retrieval] Top N")
def get_input_form(self) -> dict[str, dict]:
return {
"query": {
"name": "Query",
"type": "line"
}
}
return {"query": {"name": "Query", "type": "line"}}
class Retrieval(ToolBase, ABC):
component_name = "Retrieval"
@@ -101,8 +97,7 @@ class Retrieval(ToolBase, ABC):
# if kb_nm is a list
kb_nm_list = kb_nm if isinstance(kb_nm, list) else [kb_nm]
for nm_or_id in kb_nm_list:
e, kb = KnowledgebaseService.get_by_name(nm_or_id,
self._canvas._tenant_id)
e, kb = KnowledgebaseService.get_by_name(nm_or_id, self._canvas._tenant_id)
if not e:
e, kb = KnowledgebaseService.get_by_id(nm_or_id)
if not e:
@@ -153,7 +148,7 @@ class Retrieval(ToolBase, ABC):
last = 0
for m in pat.finditer(s):
out_parts.append(s[last:m.start()])
out_parts.append(s[last : m.start()])
key = m.group(1)
v = self._canvas.get_variable_value(key)
if v is None:
@@ -220,22 +215,16 @@ class Retrieval(ToolBase, ABC):
tenant_id = self._canvas._tenant_id
chat_model_config = get_tenant_default_model_by_type(tenant_id, LLMType.CHAT)
chat_mdl = LLMBundle(tenant_id, chat_model_config)
cks = await settings.retriever.retrieval_by_toc(query, kbinfos["chunks"], [kb.tenant_id for kb in kbs],
chat_mdl, self._param.top_n)
cks = await settings.retriever.retrieval_by_toc(query, kbinfos["chunks"], [kb.tenant_id for kb in kbs], chat_mdl, self._param.top_n)
if self.check_if_canceled("Retrieval processing"):
return
if cks:
kbinfos["chunks"] = cks
kbinfos["chunks"] = settings.retriever.retrieval_by_children(kbinfos["chunks"],
[kb.tenant_id for kb in kbs])
kbinfos["chunks"] = settings.retriever.retrieval_by_children(kbinfos["chunks"], [kb.tenant_id for kb in kbs])
if self._param.use_kg:
tenant_id = self._canvas.get_tenant_id()
chat_model_config = get_tenant_default_model_by_type(tenant_id, LLMType.CHAT)
ck = await settings.kg_retriever.retrieval(query,
[kb.tenant_id for kb in kbs],
kb_ids,
embd_mdl,
LLMBundle(tenant_id, chat_model_config))
ck = await settings.kg_retriever.retrieval(query, [kb.tenant_id for kb in kbs], kb_ids, embd_mdl, LLMBundle(tenant_id, chat_model_config))
if self.check_if_canceled("Retrieval processing"):
return
if ck["content_with_weight"]:
@@ -245,8 +234,7 @@ class Retrieval(ToolBase, ABC):
if self._param.use_kg and kbs:
chat_model_config = get_tenant_default_model_by_type(kbs[0].tenant_id, LLMType.CHAT)
ck = await settings.kg_retriever.retrieval(query, [kb.tenant_id for kb in kbs], filtered_kb_ids, embd_mdl,
LLMBundle(kbs[0].tenant_id, chat_model_config))
ck = await settings.kg_retriever.retrieval(query, [kb.tenant_id for kb in kbs], filtered_kb_ids, embd_mdl, LLMBundle(kbs[0].tenant_id, chat_model_config))
if self.check_if_canceled("Retrieval processing"):
return
if ck["content_with_weight"]:
@@ -293,16 +281,14 @@ class Retrieval(ToolBase, ABC):
filter_dict: dict = {"memory_id": memory_ids}
if user_id:
import re
# is variable
if re.match(r"^{.*}$", user_id):
user_id = self._canvas.get_variable_value(user_id)
filter_dict["user_id"] = user_id
message_list = memory_message_service.query_message(filter_dict, {
"query": query,
"similarity_threshold": self._param.similarity_threshold,
"keywords_similarity_weight": self._param.keywords_similarity_weight,
"top_n": self._param.top_n
})
message_list = memory_message_service.query_message(
filter_dict, {"query": query, "similarity_threshold": self._param.similarity_threshold, "keywords_similarity_weight": self._param.keywords_similarity_weight, "top_n": self._param.top_n}
)
if not message_list:
self.set_output("formalized_content", self._param.empty_response)
return ""
agent/tools/tavily.py
+28 -43
View File
@@ -28,7 +28,7 @@ class TavilySearchParam(ToolParamBase):
"""
def __init__(self):
self.meta:ToolMeta = {
self.meta: ToolMeta = {
"name": "tavily_search",
"description": """
Tavily is a search engine optimized for LLMs, aimed at efficient, quick and persistent search results.
@@ -43,7 +43,7 @@ When searching:
"type": "string",
"description": "The search keywords to execute with Tavily. The keywords should be the most important words/terms(includes synonyms) from the original request.",
"default": "{sys.query}",
"required": True
"required": True,
},
"topic": {
"type": "string",
@@ -56,27 +56,21 @@ When searching:
"type": "array",
"description": "default:[]. A list of domains only from which the search results can be included.",
"default": [],
"items": {
"type": "string",
"description": "Domain name that must be included, e.g. www.yahoo.com"
},
"required": False
"items": {"type": "string", "description": "Domain name that must be included, e.g. www.yahoo.com"},
"required": False,
},
"exclude_domains": {
"type": "array",
"description": "default:[]. A list of domains from which the search results can not be included",
"default": [],
"items": {
"type": "string",
"description": "Domain name that must be excluded, e.g. www.yahoo.com"
},
"required": False
"items": {"type": "string", "description": "Domain name that must be excluded, e.g. www.yahoo.com"},
"required": False,
},
}
},
}
super().__init__()
self.api_key = ""
self.search_depth = "basic" # basic/advanced
self.search_depth = "basic" # basic/advanced
self.max_results = 6
self.days = 14
self.include_answer = False
@@ -91,12 +85,8 @@ When searching:
self.check_positive_integer(self.days, "Tavily days should be greater than 1")
def get_input_form(self) -> dict[str, dict]:
return {
"query": {
"name": "Query",
"type": "line"
}
}
return {"query": {"name": "Query", "type": "line"}}
class TavilySearch(ToolBase, ABC):
component_name = "TavilySearch"
@@ -115,7 +105,7 @@ class TavilySearch(ToolBase, ABC):
for fld in ["search_depth", "topic", "max_results", "days", "include_answer", "include_raw_content", "include_images", "include_image_descriptions", "include_domains", "exclude_domains"]:
if fld not in kwargs:
kwargs[fld] = getattr(self._param, fld)
for _ in range(self._param.max_retries+1):
for _ in range(self._param.max_retries + 1):
if self.check_if_canceled("TavilySearch processing"):
return
@@ -126,11 +116,13 @@ class TavilySearch(ToolBase, ABC):
if self.check_if_canceled("TavilySearch processing"):
return
self._retrieve_chunks(res["results"],
get_title=lambda r: r["title"],
get_url=lambda r: r["url"],
get_content=lambda r: r["raw_content"] if r["raw_content"] else r["content"],
get_score=lambda r: r["score"])
self._retrieve_chunks(
res["results"],
get_title=lambda r: r["title"],
get_url=lambda r: r["url"],
get_content=lambda r: r["raw_content"] if r["raw_content"] else r["content"],
get_score=lambda r: r["score"],
)
self.set_output("json", res["results"])
return self.output("formalized_content")
except Exception as e:
@@ -159,7 +151,7 @@ class TavilyExtractParam(ToolParamBase):
"""
def __init__(self):
self.meta:ToolMeta = {
self.meta: ToolMeta = {
"name": "tavily_extract",
"description": "Extract web page content from one or more specified URLs using Tavily Extract.",
"parameters": {
@@ -167,11 +159,8 @@ class TavilyExtractParam(ToolParamBase):
"type": "array",
"description": "The URLs to extract content from.",
"default": "",
"items": {
"type": "string",
"description": "The URL to extract content from, e.g. www.yahoo.com"
},
"required": True
"items": {"type": "string", "description": "The URL to extract content from, e.g. www.yahoo.com"},
"required": True,
},
"extract_depth": {
"type": "string",
@@ -186,12 +175,12 @@ class TavilyExtractParam(ToolParamBase):
"enum": ["markdown", "text"],
"default": "markdown",
"required": False,
}
}
},
},
}
super().__init__()
self.api_key = ""
self.extract_depth = "basic" # basic/advanced
self.extract_depth = "basic" # basic/advanced
self.urls = []
self.format = "markdown"
self.include_images = False
@@ -201,17 +190,13 @@ class TavilyExtractParam(ToolParamBase):
self.check_valid_value(self.format, "Tavily extract format should be in 'markdown/text'", ["markdown", "text"])
def get_input_form(self) -> dict[str, dict]:
return {
"urls": {
"name": "URLs",
"type": "line"
}
}
return {"urls": {"name": "URLs", "type": "line"}}
class TavilyExtract(ToolBase, ABC):
component_name = "TavilyExtract"
@timeout(int(os.environ.get("COMPONENT_EXEC_TIMEOUT", 10*60)))
@timeout(int(os.environ.get("COMPONENT_EXEC_TIMEOUT", 10 * 60)))
def _invoke(self, **kwargs):
if self.check_if_canceled("TavilyExtract processing"):
return
@@ -223,7 +208,7 @@ class TavilyExtract(ToolBase, ABC):
kwargs[fld] = getattr(self._param, fld)
if kwargs.get("urls") and isinstance(kwargs["urls"], str):
kwargs["urls"] = kwargs["urls"].split(",")
for _ in range(self._param.max_retries+1):
for _ in range(self._param.max_retries + 1):
if self.check_if_canceled("TavilyExtract processing"):
return
agent/tools/tushare.py
+8 -20
View File
@@ -37,8 +37,7 @@ class TuShareParam(ComponentParamBase):
self.keyword = ""
def check(self):
self.check_valid_value(self.src, "Quick News Source",
["sina", "wallstreetcn", "10jqka", "eastmoney", "yuncaijing", "fenghuang", "jinrongjie"])
self.check_valid_value(self.src, "Quick News Source", ["sina", "wallstreetcn", "10jqka", "eastmoney", "yuncaijing", "fenghuang", "jinrongjie"])
class TuShare(ComponentBase, ABC):
@@ -58,20 +57,15 @@ class TuShare(ComponentBase, ABC):
return
tus_res = []
params = {
"api_name": "news",
"token": self._param.token,
"params": {"src": self._param.src, "start_date": self._param.start_date,
"end_date": self._param.end_date}
}
response = requests.post(url="http://api.tushare.pro", data=json.dumps(params).encode('utf-8'), timeout=DEFAULT_TIMEOUT)
params = {"api_name": "news", "token": self._param.token, "params": {"src": self._param.src, "start_date": self._param.start_date, "end_date": self._param.end_date}}
response = requests.post(url="http://api.tushare.pro", data=json.dumps(params).encode("utf-8"), timeout=DEFAULT_TIMEOUT)
response = response.json()
if self.check_if_canceled("TuShare processing"):
return
if response['code'] != 0:
return TuShare.be_output(response['msg'])
df = pd.DataFrame(response['data']['items'])
df.columns = response['data']['fields']
if response["code"] != 0:
return TuShare.be_output(response["msg"])
df = pd.DataFrame(response["data"]["items"])
df.columns = response["data"]["fields"]
if self.check_if_canceled("TuShare processing"):
return
keyword = self._param.keyword or ans
@@ -79,13 +73,7 @@ class TuShare(ComponentBase, ABC):
"TuShare news filter keyword source=%s",
"param.keyword" if self._param.keyword else "upstream_input",
)
tus_res.append(
{
"content": (
df[df["content"].str.contains(keyword, case=False, na=False, regex=False)]
).to_markdown()
}
)
tus_res.append({"content": (df[df["content"].str.contains(keyword, case=False, na=False, regex=False)]).to_markdown()})
except Exception as e:
if self.check_if_canceled("TuShare processing"):
return
agent/tools/wencai.py
+7 -21
View File
@@ -30,21 +30,14 @@ class WenCaiParam(ToolParamBase):
"""
def __init__(self):
self.meta:ToolMeta = {
self.meta: ToolMeta = {
"name": "iwencai",
"description": """
iwencai search: search platform is committed to providing hundreds of millions of investors with the most timely, accurate and comprehensive information, covering news, announcements, research reports, blogs, forums, Weibo, characters, etc.
robo-advisor intelligent stock selection platform: through AI technology, is committed to providing investors with intelligent stock selection, quantitative investment, main force tracking, value investment, technical analysis and other types of stock selection technologies.
fund selection platform: through AI technology, is committed to providing excellent fund, value investment, quantitative analysis and other fund selection technologies for foundation citizens.
""",
"parameters": {
"query": {
"type": "string",
"description": "The question/conditions to select stocks.",
"default": "{sys.query}",
"required": True
}
}
"parameters": {"query": {"type": "string", "description": "The question/conditions to select stocks.", "default": "{sys.query}", "required": True}},
}
super().__init__()
self.top_n = 10
@@ -52,18 +45,11 @@ fund selection platform: through AI technology, is committed to providing excell
def check(self):
self.check_positive_integer(self.top_n, "Top N")
self.check_valid_value(self.query_type, "Query type",
['stock', 'zhishu', 'fund', 'hkstock', 'usstock', 'threeboard', 'conbond', 'insurance',
'futures', 'lccp',
'foreign_exchange'])
self.check_valid_value(self.query_type, "Query type", ["stock", "zhishu", "fund", "hkstock", "usstock", "threeboard", "conbond", "insurance", "futures", "lccp", "foreign_exchange"])
def get_input_form(self) -> dict[str, dict]:
return {
"query": {
"name": "Query",
"type": "line"
}
}
return {"query": {"name": "Query", "type": "line"}}
class WenCai(ToolBase, ABC):
component_name = "WenCai"
@@ -78,7 +64,7 @@ class WenCai(ToolBase, ABC):
return ""
last_e = ""
for _ in range(self._param.max_retries+1):
for _ in range(self._param.max_retries + 1):
if self.check_if_canceled("WenCai processing"):
return
@@ -103,7 +89,7 @@ class WenCai(ToolBase, ABC):
elif isinstance(item[1], dict):
if "meta" in item[1].keys():
continue
wencai_res.append(pd.DataFrame.from_dict(item[1], orient='index').to_markdown())
wencai_res.append(pd.DataFrame.from_dict(item[1], orient="index").to_markdown())
elif isinstance(item[1], pd.DataFrame):
if "image_url" in item[1].columns:
continue
agent/tools/wikipedia.py
+83 -20
View File
@@ -28,7 +28,7 @@ class WikipediaParam(ToolParamBase):
"""
def __init__(self):
self.meta:ToolMeta = {
self.meta: ToolMeta = {
"name": "wikipedia_search",
"description": """A wide range of how-to and information pages are made available in wikipedia. Since 2001, it has grown rapidly to become the world's largest reference website. From Wikipedia, the free encyclopedia.""",
"parameters": {
@@ -36,9 +36,9 @@ class WikipediaParam(ToolParamBase):
"type": "string",
"description": "The search keyword to execute with wikipedia. The keyword MUST be a specific subject that can match the title.",
"default": "{sys.query}",
"required": True
"required": True,
}
}
},
}
super().__init__()
self.top_n = 10
@@ -46,20 +46,86 @@ class WikipediaParam(ToolParamBase):
def check(self):
self.check_positive_integer(self.top_n, "Top N")
self.check_valid_value(self.language, "Wikipedia languages",
['af', 'pl', 'ar', 'ast', 'az', 'bg', 'nan', 'bn', 'be', 'ca', 'cs', 'cy', 'da', 'de',
'et', 'el', 'en', 'es', 'eo', 'eu', 'fa', 'fr', 'gl', 'ko', 'hy', 'hi', 'hr', 'id',
'it', 'he', 'ka', 'lld', 'la', 'lv', 'lt', 'hu', 'mk', 'arz', 'ms', 'min', 'my', 'nl',
'ja', 'nb', 'nn', 'ce', 'uz', 'pt', 'kk', 'ro', 'ru', 'ceb', 'sk', 'sl', 'sr', 'sh',
'fi', 'sv', 'ta', 'tt', 'th', 'tg', 'azb', 'tr', 'uk', 'ur', 'vi', 'war', 'zh', 'yue'])
self.check_valid_value(
self.language,
"Wikipedia languages",
[
"af",
"pl",
"ar",
"ast",
"az",
"bg",
"nan",
"bn",
"be",
"ca",
"cs",
"cy",
"da",
"de",
"et",
"el",
"en",
"es",
"eo",
"eu",
"fa",
"fr",
"gl",
"ko",
"hy",
"hi",
"hr",
"id",
"it",
"he",
"ka",
"lld",
"la",
"lv",
"lt",
"hu",
"mk",
"arz",
"ms",
"min",
"my",
"nl",
"ja",
"nb",
"nn",
"ce",
"uz",
"pt",
"kk",
"ro",
"ru",
"ceb",
"sk",
"sl",
"sr",
"sh",
"fi",
"sv",
"ta",
"tt",
"th",
"tg",
"azb",
"tr",
"uk",
"ur",
"vi",
"war",
"zh",
"yue",
],
)
def get_input_form(self) -> dict[str, dict]:
return {
"query": {
"name": "Query",
"type": "line"
}
}
return {"query": {"name": "Query", "type": "line"}}
class Wikipedia(ToolBase, ABC):
"""Wikipedia search tool that retrieves and processes Wikipedia articles."""
@@ -84,7 +150,7 @@ class Wikipedia(ToolBase, ABC):
return ""
last_e = ""
for _ in range(self._param.max_retries+1):
for _ in range(self._param.max_retries + 1):
if self.check_if_canceled("Wikipedia processing"):
return
@@ -104,10 +170,7 @@ class Wikipedia(ToolBase, ABC):
logging.info(f"Wikipedia page not found: '{p}'")
except Exception as e:
logging.exception(f"Unexpected error fetching Wikipedia page '{p}': {e}")
self._retrieve_chunks(pages,
get_title=lambda r: r.title,
get_url=lambda r: r.url,
get_content=lambda r: r.summary)
self._retrieve_chunks(pages, get_title=lambda r: r.title, get_url=lambda r: r.url, get_content=lambda r: r.summary)
return self.output("formalized_content")
except Exception as e:
if self.check_if_canceled("Wikipedia processing"):
agent/tools/yahoofinance.py
+5 -16
View File
@@ -29,17 +29,10 @@ class YahooFinanceParam(ToolParamBase):
"""
def __init__(self):
self.meta:ToolMeta = {
self.meta: ToolMeta = {
"name": "yahoo_finance",
"description": "The Yahoo Finance is a service that provides access to real-time and historical stock market data. It enables users to fetch various types of stock information, such as price quotes, historical prices, company profiles, and financial news. The API offers structured data, allowing developers to integrate market data into their applications and analysis tools.",
"parameters": {
"stock_code": {
"type": "string",
"description": "The stock code or company name.",
"default": "{sys.query}",
"required": True
}
}
"parameters": {"stock_code": {"type": "string", "description": "The stock code or company name.", "default": "{sys.query}", "required": True}},
}
super().__init__()
self.info = True
@@ -62,12 +55,8 @@ class YahooFinanceParam(ToolParamBase):
self.check_boolean(self.news, "show news")
def get_input_form(self) -> dict[str, dict]:
return {
"stock_code": {
"name": "Stock code/Company name",
"type": "line"
}
}
return {"stock_code": {"name": "Stock code/Company name", "type": "line"}}
class YahooFinance(ToolBase, ABC):
component_name = "YahooFinance"
@@ -82,7 +71,7 @@ class YahooFinance(ToolBase, ABC):
return ""
last_e = ""
for _ in range(self._param.max_retries+1):
for _ in range(self._param.max_retries + 1):
if self.check_if_canceled("YahooFinance processing"):
return None
api/apps/auth/__init__.py
+2 -6
View File
@@ -19,14 +19,10 @@ from .oidc import OIDCClient
from .github import GithubOAuthClient
CLIENT_TYPES = {
"oauth2": OAuthClient,
"oidc": OIDCClient,
"github": GithubOAuthClient
}
CLIENT_TYPES = {"oauth2": OAuthClient, "oidc": OIDCClient, "github": GithubOAuthClient}
def get_auth_client(config)->OAuthClient:
def get_auth_client(config) -> OAuthClient:
channel_type = str(config.get("type", "")).lower()
if channel_type == "":
if config.get("issuer"):
api/apps/auth/github.py
+9 -11
View File
@@ -23,15 +23,16 @@ class GithubOAuthClient(OAuthClient):
"""
Initialize the GithubOAuthClient with the provider's configuration.
"""
config.update({
"authorization_url": "https://github.com/login/oauth/authorize",
"token_url": "https://github.com/login/oauth/access_token",
"userinfo_url": "https://api.github.com/user",
"scope": "user:email"
})
config.update(
{
"authorization_url": "https://github.com/login/oauth/authorize",
"token_url": "https://github.com/login/oauth/access_token",
"userinfo_url": "https://api.github.com/user",
"scope": "user:email",
}
)
super().__init__(config)
def fetch_user_info(self, access_token, **kwargs):
"""
Fetch GitHub user info (synchronous).
@@ -42,9 +43,7 @@ class GithubOAuthClient(OAuthClient):
response = sync_request("GET", self.userinfo_url, headers=headers, timeout=self.http_request_timeout)
response.raise_for_status()
user_info.update(response.json())
email_response = sync_request(
"GET", self.userinfo_url + "/emails", headers=headers, timeout=self.http_request_timeout
)
email_response = sync_request("GET", self.userinfo_url + "/emails", headers=headers, timeout=self.http_request_timeout)
email_response.raise_for_status()
email_info = email_response.json()
user_info["email"] = next((email for email in email_info if email["primary"]), None)["email"]
@@ -79,7 +78,6 @@ class GithubOAuthClient(OAuthClient):
except Exception as e:
raise ValueError(f"Failed to fetch github user info: {e}")
def normalize_user_info(self, user_info):
email = user_info.get("email")
username = user_info.get("login", str(email).split("@")[0])
api/apps/auth/oauth.py
+2 -12
View File
@@ -24,7 +24,7 @@ class UserInfo:
self.username = username
self.nickname = nickname
self.avatar_url = avatar_url
def to_dict(self):
return {key: value for key, value in self.__dict__.items()}
@@ -44,7 +44,6 @@ class OAuthClient:
self.http_request_timeout = 7
def get_authorization_url(self, state=None):
"""
Generate the authorization URL for user login.
@@ -61,19 +60,12 @@ class OAuthClient:
authorization_url = f"{self.authorization_url}?{urllib.parse.urlencode(params)}"
return authorization_url
def exchange_code_for_token(self, code):
"""
Exchange authorization code for access token.
"""
try:
payload = {
"client_id": self.client_id,
"client_secret": self.client_secret,
"code": code,
"redirect_uri": self.redirect_uri,
"grant_type": "authorization_code"
}
payload = {"client_id": self.client_id, "client_secret": self.client_secret, "code": code, "redirect_uri": self.redirect_uri, "grant_type": "authorization_code"}
response = sync_request(
"POST",
self.token_url,
@@ -110,7 +102,6 @@ class OAuthClient:
except Exception as e:
raise ValueError(f"Failed to exchange authorization code for token: {e}")
def fetch_user_info(self, access_token, **kwargs):
"""
Fetch user information using access token.
@@ -140,7 +131,6 @@ class OAuthClient:
except Exception as e:
raise ValueError(f"Failed to fetch user info: {e}")
def normalize_user_info(self, user_info):
email = user_info.get("email")
username = user_info.get("username", str(email).split("@")[0])
api/apps/auth/oidc.py
+25 -19
View File
@@ -26,12 +26,20 @@ from .oauth import OAuthClient
# forge tokens by HMAC-signing them with the public key bytes
# (RSA/HMAC algorithm-confusion attack, CWE-347). "none" is excluded for the
# obvious reason that it disables signature verification entirely.
_ALLOWED_OIDC_SIGNING_ALGS = frozenset({
"RS256", "RS384", "RS512",
"ES256", "ES384", "ES512",
"PS256", "PS384", "PS512",
"EdDSA",
})
_ALLOWED_OIDC_SIGNING_ALGS = frozenset(
{
"RS256",
"RS384",
"RS512",
"ES256",
"ES384",
"ES512",
"PS256",
"PS384",
"PS512",
"EdDSA",
}
)
# OIDC Core 1.0 § 2 makes RS256 the spec-default ``id_token_signing_alg``,
# so this is the safe fallback when a provider's discovery document does not
@@ -69,24 +77,25 @@ class OIDCClient(OAuthClient):
raise ValueError("Missing issuer in configuration.")
oidc_metadata = self._load_oidc_metadata(self.issuer)
config.update({
'issuer': oidc_metadata['issuer'],
'jwks_uri': oidc_metadata['jwks_uri'],
'authorization_url': oidc_metadata['authorization_endpoint'],
'token_url': oidc_metadata['token_endpoint'],
'userinfo_url': oidc_metadata['userinfo_endpoint']
})
config.update(
{
"issuer": oidc_metadata["issuer"],
"jwks_uri": oidc_metadata["jwks_uri"],
"authorization_url": oidc_metadata["authorization_endpoint"],
"token_url": oidc_metadata["token_endpoint"],
"userinfo_url": oidc_metadata["userinfo_endpoint"],
}
)
super().__init__(config)
self.issuer = config['issuer']
self.jwks_uri = config['jwks_uri']
self.issuer = config["issuer"]
self.jwks_uri = config["jwks_uri"]
# Pin the accepted ID-token signing algorithms at construction time
# from a trusted source (provider metadata + safe allowlist) so the
# JWT verification step in :meth:`parse_id_token` cannot be tricked
# by attacker-controlled JWT headers (CWE-345 / CWE-347).
self.id_token_signing_algs = _resolve_id_token_signing_algs(oidc_metadata)
@staticmethod
def _load_oidc_metadata(issuer):
"""
@@ -100,7 +109,6 @@ class OIDCClient(OAuthClient):
except Exception as e:
raise ValueError(f"Failed to fetch OIDC metadata: {e}")
def parse_id_token(self, id_token):
"""
Parse and validate OIDC ID Token (JWT format) with signature verification.
@@ -134,7 +142,6 @@ class OIDCClient(OAuthClient):
except Exception as e:
raise ValueError(f"Error parsing ID Token: {e}")
def fetch_user_info(self, access_token, id_token=None, **kwargs):
"""
Fetch user info.
@@ -152,6 +159,5 @@ class OIDCClient(OAuthClient):
user_info.update((await super().async_fetch_user_info(access_token)).to_dict())
return self.normalize_user_info(user_info)
def normalize_user_info(self, user_info):
return super().normalize_user_info(user_info)
api/apps/backward_compat.py
+77 -103
View File
@@ -44,6 +44,7 @@ Deprecated APIs and their replacements:
- POST /api/v1/sessions/related_questions -> POST /api/v1/chat/recommandation
- PUT (chunk update) -> PATCH (chunk update)
"""
import logging
from quart import Blueprint, jsonify, request
@@ -68,6 +69,7 @@ def _index_result(success, result):
# System APIs
# =============================================================================
@legacy_v1_manager.route("/system/healthz", methods=["GET"])
async def deprecated_system_healthz():
"""
@@ -76,17 +78,16 @@ async def deprecated_system_healthz():
Old path: GET /v1/system/healthz
New path: GET /api/v1/system/healthz
"""
logging.warning(
"API endpoint /v1/system/healthz is deprecated. "
"Please use /api/v1/system/healthz instead."
)
logging.warning("API endpoint /v1/system/healthz is deprecated. Please use /api/v1/system/healthz instead.")
result, all_ok = run_health_checks()
return jsonify(result), (200 if all_ok else 500)
# =============================================================================
# Chat Completion APIs
# =============================================================================
@manager.route("/chats/<chat_id>/completions", methods=["POST"])
@login_required
async def deprecated_chat_completions(chat_id):
@@ -97,8 +98,7 @@ async def deprecated_chat_completions(chat_id):
New path: POST /api/v1/chat/completions
"""
logging.warning(
"API endpoint /api/v1/chats/%s/completions is deprecated. "
"Please use /api/v1/chat/completions instead.",
"API endpoint /api/v1/chats/%s/completions is deprecated. Please use /api/v1/chat/completions instead.",
chat_id,
)
# Forward to the new API implementation
@@ -115,9 +115,9 @@ async def deprecated_openai_chat_completions(chat_id):
New path: POST /api/v1/openai/{chat_id}/chat/completions
"""
logging.warning(
"API endpoint /api/v1/chats_openai/%s/chat/completions is deprecated. "
"Please use /api/v1/openai/%s/chat/completions instead.",
chat_id, chat_id,
"API endpoint /api/v1/chats_openai/%s/chat/completions is deprecated. Please use /api/v1/openai/%s/chat/completions instead.",
chat_id,
chat_id,
)
# Forward to the new API implementation
return await openai_api.openai_chat_completions(chat_id)
@@ -134,8 +134,7 @@ async def deprecated_agents_openai_chat_completions(agent_id, tenant_id=None):
New path: POST /api/v1/agents/chat/completions
"""
logging.warning(
"API endpoint /api/v1/agents_openai/%s/chat/completions is deprecated. "
"Please use /api/v1/agents/chat/completions with `openai-compatible` instead.",
"API endpoint /api/v1/agents_openai/%s/chat/completions is deprecated. Please use /api/v1/agents/chat/completions with `openai-compatible` instead.",
agent_id,
)
req = dict(await get_request_json())
@@ -148,6 +147,7 @@ async def deprecated_agents_openai_chat_completions(agent_id, tenant_id=None):
# Dataset Graph and Index APIs
# =============================================================================
@manager.route("/datasets/<dataset_id>/knowledge_graph", methods=["GET"])
@login_required
async def deprecated_get_knowledge_graph(dataset_id):
@@ -158,9 +158,9 @@ async def deprecated_get_knowledge_graph(dataset_id):
New path: GET /api/v1/datasets/{dataset_id}/graph
"""
logging.warning(
"API endpoint /api/v1/datasets/%s/knowledge_graph is deprecated. "
"Please use /api/v1/datasets/%s/graph instead.",
dataset_id, dataset_id,
"API endpoint /api/v1/datasets/%s/knowledge_graph is deprecated. Please use /api/v1/datasets/%s/graph instead.",
dataset_id,
dataset_id,
)
return await dataset_api.get_knowledge_graph(dataset_id=dataset_id)
@@ -175,9 +175,9 @@ async def deprecated_delete_knowledge_graph(dataset_id):
New path: DELETE /api/v1/datasets/{dataset_id}/graph
"""
logging.warning(
"API endpoint DELETE /api/v1/datasets/%s/knowledge_graph is deprecated. "
"Please use DELETE /api/v1/datasets/%s/graph instead.",
dataset_id, dataset_id,
"API endpoint DELETE /api/v1/datasets/%s/knowledge_graph is deprecated. Please use DELETE /api/v1/datasets/%s/graph instead.",
dataset_id,
dataset_id,
)
return await dataset_api.delete_knowledge_graph(dataset_id=dataset_id)
@@ -193,9 +193,9 @@ async def deprecated_run_graphrag(dataset_id, tenant_id=None):
New path: POST /api/v1/datasets/{dataset_id}/index?type=graph
"""
logging.warning(
"API endpoint /api/v1/datasets/%s/run_graphrag is deprecated. "
"Please use /api/v1/datasets/%s/index?type=graph instead.",
dataset_id, dataset_id,
"API endpoint /api/v1/datasets/%s/run_graphrag is deprecated. Please use /api/v1/datasets/%s/index?type=graph instead.",
dataset_id,
dataset_id,
)
return _index_result(*dataset_api_service.run_index(dataset_id, tenant_id, "graph"))
@@ -211,9 +211,9 @@ async def deprecated_trace_graphrag(dataset_id, tenant_id=None):
New path: GET /api/v1/datasets/{dataset_id}/index?type=graph
"""
logging.warning(
"API endpoint /api/v1/datasets/%s/trace_graphrag is deprecated. "
"Please use /api/v1/datasets/%s/index?type=graph instead.",
dataset_id, dataset_id,
"API endpoint /api/v1/datasets/%s/trace_graphrag is deprecated. Please use /api/v1/datasets/%s/index?type=graph instead.",
dataset_id,
dataset_id,
)
return _index_result(*dataset_api_service.trace_index(dataset_id, tenant_id, "graph"))
@@ -229,9 +229,9 @@ async def deprecated_run_raptor(dataset_id, tenant_id=None):
New path: POST /api/v1/datasets/{dataset_id}/index?type=raptor
"""
logging.warning(
"API endpoint /api/v1/datasets/%s/run_raptor is deprecated. "
"Please use /api/v1/datasets/%s/index?type=raptor instead.",
dataset_id, dataset_id,
"API endpoint /api/v1/datasets/%s/run_raptor is deprecated. Please use /api/v1/datasets/%s/index?type=raptor instead.",
dataset_id,
dataset_id,
)
return _index_result(*dataset_api_service.run_index(dataset_id, tenant_id, "raptor"))
@@ -247,9 +247,9 @@ async def deprecated_trace_raptor(dataset_id, tenant_id=None):
New path: GET /api/v1/datasets/{dataset_id}/index?type=raptor
"""
logging.warning(
"API endpoint /api/v1/datasets/%s/trace_raptor is deprecated. "
"Please use /api/v1/datasets/%s/index?type=raptor instead.",
dataset_id, dataset_id,
"API endpoint /api/v1/datasets/%s/trace_raptor is deprecated. Please use /api/v1/datasets/%s/index?type=raptor instead.",
dataset_id,
dataset_id,
)
return _index_result(*dataset_api_service.trace_index(dataset_id, tenant_id, "raptor"))
@@ -258,6 +258,7 @@ async def deprecated_trace_raptor(dataset_id, tenant_id=None):
# Chat Session APIs
# =============================================================================
@manager.route("/chats/<chat_id>/sessions/<session_id>", methods=["PUT"])
@login_required
async def deprecated_update_session(chat_id, session_id):
@@ -268,9 +269,11 @@ async def deprecated_update_session(chat_id, session_id):
New path: PATCH /api/v1/chats/{chat_id}/sessions/{session_id}
"""
logging.warning(
"API endpoint PUT /api/v1/chats/%s/sessions/%s is deprecated. "
"Please use PATCH /api/v1/chats/%s/sessions/%s instead.",
chat_id, session_id, chat_id, session_id,
"API endpoint PUT /api/v1/chats/%s/sessions/%s is deprecated. Please use PATCH /api/v1/chats/%s/sessions/%s instead.",
chat_id,
session_id,
chat_id,
session_id,
)
# Forward to the new API implementation
return await chat_api.update_session(chat_id, session_id)
@@ -280,6 +283,7 @@ async def deprecated_update_session(chat_id, session_id):
# File APIs (Old /api/v1/file/* -> New /api/v1/files*)
# =============================================================================
@manager.route("/file/get/<file_id>", methods=["GET"])
@login_required
async def deprecated_file_get(file_id):
@@ -290,9 +294,9 @@ async def deprecated_file_get(file_id):
New path: GET /api/v1/files/{file_id}
"""
logging.warning(
"API endpoint /api/v1/file/get/%s is deprecated. "
"Please use /api/v1/files/%s instead.",
file_id, file_id,
"API endpoint /api/v1/file/get/%s is deprecated. Please use /api/v1/files/%s instead.",
file_id,
file_id,
)
# Forward to the new API implementation (download)
return await file_api.download(file_id=file_id)
@@ -307,10 +311,7 @@ async def deprecated_file_list():
Old path: GET /api/v1/file/list?...
New path: GET /api/v1/files?...
"""
logging.warning(
"API endpoint /api/v1/file/list is deprecated. "
"Please use /api/v1/files instead."
)
logging.warning("API endpoint /api/v1/file/list is deprecated. Please use /api/v1/files instead.")
# Forward to the new API implementation
return await file_api.list_files()
@@ -328,8 +329,7 @@ async def deprecated_file_all_parent_folder():
if not file_id:
return get_data_error_result(message="`file_id` query parameter is required")
logging.warning(
"API endpoint /api/v1/file/all_parent_folder is deprecated. "
"Please use /api/v1/files/%s/ancestors instead.",
"API endpoint /api/v1/file/all_parent_folder is deprecated. Please use /api/v1/files/%s/ancestors instead.",
file_id,
)
# Forward to the new API implementation
@@ -349,8 +349,7 @@ async def deprecated_file_parent_folder():
if not file_id:
return get_data_error_result(message="`file_id` query parameter is required")
logging.warning(
"API endpoint /api/v1/file/parent_folder is deprecated. "
"Please use /api/v1/files/%s/parent instead.",
"API endpoint /api/v1/file/parent_folder is deprecated. Please use /api/v1/files/%s/parent instead.",
file_id,
)
# Forward to the new API implementation
@@ -366,10 +365,7 @@ async def deprecated_file_root_folder():
Old path: GET /api/v1/file/root_folder
New path: GET /api/v1/files?parent_id=<root_id>
"""
logging.warning(
"API endpoint /api/v1/file/root_folder is deprecated. "
"Please use /api/v1/files with appropriate parent_id instead."
)
logging.warning("API endpoint /api/v1/file/root_folder is deprecated. Please use /api/v1/files with appropriate parent_id instead.")
# Forward to the new API implementation with empty parent_id to get root
return await file_api.list_files()
@@ -384,10 +380,7 @@ async def deprecated_file_create(tenant_id=None):
Old path: POST /api/v1/file/create
New path: POST /api/v1/files
"""
logging.warning(
"API endpoint /api/v1/file/create is deprecated. "
"Please use POST /api/v1/files instead."
)
logging.warning("API endpoint /api/v1/file/create is deprecated. Please use POST /api/v1/files instead.")
# Forward to the new API implementation
return await file_api.create_or_upload(tenant_id=tenant_id)
@@ -402,10 +395,7 @@ async def deprecated_file_upload(tenant_id=None):
Old path: POST /api/v1/file/upload
New path: POST /api/v1/files
"""
logging.warning(
"API endpoint /api/v1/file/upload is deprecated. "
"Please use POST /api/v1/files with multipart/form-data instead."
)
logging.warning("API endpoint /api/v1/file/upload is deprecated. Please use POST /api/v1/files with multipart/form-data instead.")
# Forward to the new API implementation
return await file_api.create_or_upload(tenant_id=tenant_id)
@@ -419,10 +409,7 @@ async def deprecated_file_convert():
Old path: POST /api/v1/file/convert
New path: POST /api/v1/files/link-to-datasets
"""
logging.warning(
"API endpoint /api/v1/file/convert is deprecated. "
"Please use POST /api/v1/files/link-to-datasets instead."
)
logging.warning("API endpoint /api/v1/file/convert is deprecated. Please use POST /api/v1/files/link-to-datasets instead.")
return await file2document_api.convert()
@@ -436,10 +423,7 @@ async def deprecated_file_mv(tenant_id=None):
Old path: POST /api/v1/file/mv
New path: POST /api/v1/files/move
"""
logging.warning(
"API endpoint /api/v1/file/mv is deprecated. "
"Please use POST /api/v1/files/move instead."
)
logging.warning("API endpoint /api/v1/file/mv is deprecated. Please use POST /api/v1/files/move instead.")
# Forward to the new API implementation
return await file_api.move(tenant_id=tenant_id)
@@ -454,10 +438,7 @@ async def deprecated_file_rename(tenant_id=None):
Old path: POST /api/v1/file/rename
New path: POST /api/v1/files/move
"""
logging.warning(
"API endpoint /api/v1/file/rename is deprecated. "
"Please use POST /api/v1/files/move with `new_name` instead."
)
logging.warning("API endpoint /api/v1/file/rename is deprecated. Please use POST /api/v1/files/move with `new_name` instead.")
# Transform the old API format to new format
req = await request.get_json()
# Old API used `file_id` and `name`, new API uses `src_file_ids` and `new_name`
@@ -465,9 +446,7 @@ async def deprecated_file_rename(tenant_id=None):
new_name = req.get("name")
# Call the underlying service directly with transformed data
try:
success, result = await file_api_service.move_files(
tenant_id, src_file_ids, None, new_name
)
success, result = await file_api_service.move_files(tenant_id, src_file_ids, None, new_name)
if success:
return get_json_result(data=result)
else:
@@ -487,10 +466,7 @@ async def deprecated_file_rm(tenant_id=None):
Old path: POST /api/v1/file/rm
New path: DELETE /api/v1/files
"""
logging.warning(
"API endpoint /api/v1/file/rm is deprecated. "
"Please use DELETE /api/v1/files instead."
)
logging.warning("API endpoint /api/v1/file/rm is deprecated. Please use DELETE /api/v1/files instead.")
# Transform POST with body to DELETE behavior
# The new API expects a JSON body with `ids`
return await file_api.delete(tenant_id=tenant_id)
@@ -500,6 +476,7 @@ async def deprecated_file_rm(tenant_id=None):
# Related Questions API
# =============================================================================
@manager.route("/sessions/related_questions", methods=["POST"])
@login_required
async def deprecated_related_questions():
@@ -509,10 +486,7 @@ async def deprecated_related_questions():
Old path: POST /api/v1/sessions/related_questions
New path: POST /api/v1/chat/recommendation
"""
logging.warning(
"API endpoint /api/v1/sessions/related_questions is deprecated. "
"Please use /api/v1/chat/recommendation instead."
)
logging.warning("API endpoint /api/v1/sessions/related_questions is deprecated. Please use /api/v1/chat/recommendation instead.")
# Forward to the new API implementation
return await chat_api.recommendation()
@@ -521,6 +495,7 @@ async def deprecated_related_questions():
# Chunk Update API (PUT -> PATCH)
# =============================================================================
@manager.route("/datasets/<dataset_id>/documents/<document_id>/chunks/<chunk_id>", methods=["PUT"])
@login_required
async def deprecated_update_chunk(dataset_id, document_id, chunk_id):
@@ -531,9 +506,10 @@ async def deprecated_update_chunk(dataset_id, document_id, chunk_id):
New path: PATCH /api/v1/datasets/{dataset_id}/documents/{document_id}/chunks/{chunk_id}
"""
logging.warning(
"API endpoint PUT /api/v1/datasets/%s/documents/%s/chunks/%s is deprecated. "
"Please use PATCH instead.",
dataset_id, document_id, chunk_id,
"API endpoint PUT /api/v1/datasets/%s/documents/%s/chunks/%s is deprecated. Please use PATCH instead.",
dataset_id,
document_id,
chunk_id,
)
# Forward to the new API implementation
return await chunk_api.update_chunk(dataset_id=dataset_id, document_id=document_id, chunk_id=chunk_id)
@@ -543,6 +519,7 @@ async def deprecated_update_chunk(dataset_id, document_id, chunk_id):
# File Upload Info API
# =============================================================================
@manager.route("/file/upload_info", methods=["POST"])
@login_required
async def deprecated_file_upload_info():
@@ -554,10 +531,7 @@ async def deprecated_file_upload_info():
"""
from api.apps import current_user
logging.warning(
"API endpoint /api/v1/file/upload_info is deprecated. "
"Please use POST /api/v1/documents/upload instead."
)
logging.warning("API endpoint /api/v1/file/upload_info is deprecated. Please use POST /api/v1/documents/upload instead.")
# Forward to the new API implementation
# Need to pass tenant_id explicitly since we're calling the function directly
tenant_id = current_user.id
@@ -575,10 +549,7 @@ async def deprecated_legacy_document_upload_info():
"""
from api.apps import current_user
logging.warning(
"API endpoint /v1/document/upload_info is deprecated. "
"Please use POST /api/v1/documents/upload instead."
)
logging.warning("API endpoint /v1/document/upload_info is deprecated. Please use POST /api/v1/documents/upload instead.")
tenant_id = current_user.id
return await document_api.upload_info(tenant_id=tenant_id)
@@ -587,6 +558,7 @@ async def deprecated_legacy_document_upload_info():
# Document APIs
# =============================================================================
@manager.route("/datasets/<dataset_id>/documents/<document_id>", methods=["PUT"])
@login_required
async def deprecated_update_document(dataset_id, document_id):
@@ -597,9 +569,9 @@ async def deprecated_update_document(dataset_id, document_id):
New path: PATCH /api/v1/datasets/{dataset_id}/documents/{document_id}
"""
logging.warning(
"API endpoint PUT /api/v1/datasets/%s/documents/%s is deprecated. "
"Please use PATCH instead.",
dataset_id, document_id,
"API endpoint PUT /api/v1/datasets/%s/documents/%s is deprecated. Please use PATCH instead.",
dataset_id,
document_id,
)
# Forward to the new API implementation
return await document_api.update_document(dataset_id=dataset_id, document_id=document_id)
@@ -615,9 +587,9 @@ async def deprecated_document_get(doc_id):
New path: GET /api/v1/documents/{doc_id}/preview
"""
logging.warning(
"API endpoint /api/v1/document/get/%s is deprecated. "
"Please use /api/v1/documents/%s/preview instead.",
doc_id, doc_id,
"API endpoint /api/v1/document/get/%s is deprecated. Please use /api/v1/documents/%s/preview instead.",
doc_id,
doc_id,
)
return await document_api.get(doc_id)
@@ -632,9 +604,9 @@ async def deprecated_document_download(doc_id):
New path: GET /api/v1/agents/attachments/{doc_id}/download
"""
logging.warning(
"API endpoint /api/v1/document/download/%s is deprecated. "
"Please use /api/v1/agents/attachments/%s/download instead.",
doc_id, doc_id,
"API endpoint /api/v1/document/download/%s is deprecated. Please use /api/v1/agents/attachments/%s/download instead.",
doc_id,
doc_id,
)
return await agent_api.download_attachment(attachment_id=doc_id)
@@ -649,16 +621,18 @@ async def document_download_v1(attachment_id):
New path: GET /api/v1/agents/attachments/{attachment_id}/download
"""
logging.warning(
"API endpoint /v1/document/download/%s is deprecated. "
"Please use /api/v1/agents/attachments/%s/download instead.",
attachment_id, attachment_id,
"API endpoint /v1/document/download/%s is deprecated. Please use /api/v1/agents/attachments/%s/download instead.",
attachment_id,
attachment_id,
)
return await agent_api.download_attachment(attachment_id=attachment_id)
# =============================================================================
# Agent Chat API
# =============================================================================
@manager.route("/agents/<agent_id>/completions", methods=["POST"])
@login_required
@add_tenant_id_to_kwargs
@@ -670,12 +644,12 @@ async def deprecated_agent_completions(agent_id, tenant_id=None):
New path: POST /api/v1/agents/chat/completions
"""
logging.warning(
"API endpoint /api/v1/agents/%s/completions is deprecated. "
"Please use /api/v1/agents/chat/completions instead.",
"API endpoint /api/v1/agents/%s/completions is deprecated. Please use /api/v1/agents/chat/completions instead.",
agent_id,
)
return await agent_api.agent_chat_completion(tenant_id=tenant_id, agent_id=agent_id)
def register_backward_compat_routes(app_instance):
"""
Register all backward compatibility routes with the app.
api/apps/canvas_app.py
+78 -139
View File
@@ -50,27 +50,25 @@ from api.apps.services.canvas_replica_service import CanvasReplicaService
from api.db.services.canvas_service import completion as agent_completion
@manager.route('/templates', methods=['GET']) # noqa: F821
@manager.route("/templates", methods=["GET"]) # noqa: F821
@login_required
def templates():
return get_json_result(data=[c.to_dict() for c in CanvasTemplateService.get_all()])
@manager.route('/rm', methods=['POST']) # noqa: F821
@manager.route("/rm", methods=["POST"]) # noqa: F821
@validate_request("canvas_ids")
@login_required
async def rm():
req = await get_request_json()
for i in req["canvas_ids"]:
if not UserCanvasService.accessible(i, current_user.id):
return get_json_result(
data=False, message='Only owner of canvas authorized for this operation.',
code=RetCode.OPERATING_ERROR)
return get_json_result(data=False, message="Only owner of canvas authorized for this operation.", code=RetCode.OPERATING_ERROR)
UserCanvasService.delete_by_id(i)
return get_json_result(data=True)
@manager.route('/set', methods=['POST']) # noqa: F821
@manager.route("/set", methods=["POST"]) # noqa: F821
@validate_request("dsl", "title")
@login_required
async def save():
@@ -89,9 +87,7 @@ async def save():
return get_data_error_result(message="Fail to save canvas.")
else:
if not UserCanvasService.accessible(req["id"], current_user.id):
return get_json_result(
data=False, message='Only owner of canvas authorized for this operation.',
code=RetCode.OPERATING_ERROR)
return get_json_result(data=False, message="Only owner of canvas authorized for this operation.", code=RetCode.OPERATING_ERROR)
UserCanvasService.update_by_id(req["id"], req)
# save version
UserCanvasVersionService.save_or_replace_latest(
@@ -112,7 +108,7 @@ async def save():
return get_json_result(data=req)
@manager.route('/get/<canvas_id>', methods=['GET']) # noqa: F821
@manager.route("/get/<canvas_id>", methods=["GET"]) # noqa: F821
@login_required
def get(canvas_id):
if not UserCanvasService.accessible(canvas_id, current_user.id):
@@ -135,29 +131,25 @@ def get(canvas_id):
return get_json_result(data=c)
@manager.route('/getsse/<canvas_id>', methods=['GET']) # type: ignore # noqa: F821
@manager.route("/getsse/<canvas_id>", methods=["GET"]) # type: ignore # noqa: F821
def getsse(canvas_id):
token = request.headers.get('Authorization').split()
token = request.headers.get("Authorization").split()
if len(token) != 2:
return get_data_error_result(message='Authorization is not valid!')
return get_data_error_result(message="Authorization is not valid!")
token = token[1]
objs = APIToken.query(beta=token)
if not objs:
return get_data_error_result(message='Authentication error: API key is invalid!"')
tenant_id = objs[0].tenant_id
if not UserCanvasService.query(user_id=tenant_id, id=canvas_id):
return get_json_result(
data=False,
message='Only owner of canvas authorized for this operation.',
code=RetCode.OPERATING_ERROR
)
return get_json_result(data=False, message="Only owner of canvas authorized for this operation.", code=RetCode.OPERATING_ERROR)
e, c = UserCanvasService.get_by_id(canvas_id)
if not e or c.user_id != tenant_id:
return get_data_error_result(message="canvas not found.")
return get_json_result(data=c.to_dict())
@manager.route('/completion', methods=['POST']) # noqa: F821
@manager.route("/completion", methods=["POST"]) # noqa: F821
@validate_request("id")
@login_required
async def run():
@@ -169,9 +161,7 @@ async def run():
runtime_user_id = req.get("user_id") or tenant_id
user_id = str(runtime_user_id)
if not await thread_pool_exec(UserCanvasService.accessible, req["id"], tenant_id):
return get_json_result(
data=False, message='Only owner of canvas authorized for this operation.',
code=RetCode.OPERATING_ERROR)
return get_json_result(data=False, message="Only owner of canvas authorized for this operation.", code=RetCode.OPERATING_ERROR)
replica_payload = CanvasReplicaService.load_for_run(
canvas_id=req["id"],
@@ -234,7 +224,7 @@ async def run():
resp.headers.add_header("Connection", "keep-alive")
resp.headers.add_header("X-Accel-Buffering", "no")
resp.headers.add_header("Content-Type", "text/event-stream; charset=utf-8")
#resp.call_on_close(lambda: canvas.cancel_task())
# resp.call_on_close(lambda: canvas.cancel_task())
return resp
@@ -244,6 +234,7 @@ async def exp_agent_completion(canvas_id):
tenant_id = current_user.id
req = await get_request_json()
return_trace = bool(req.get("return_trace", False))
async def generate():
trace_items = []
async for answer in agent_completion(tenant_id=tenant_id, agent_id=canvas_id, **req):
@@ -280,9 +271,9 @@ async def exp_agent_completion(canvas_id):
resp.headers.add_header("X-Accel-Buffering", "no")
resp.headers.add_header("Content-Type", "text/event-stream; charset=utf-8")
return resp
@manager.route('/rerun', methods=['POST']) # noqa: F821
@manager.route("/rerun", methods=["POST"]) # noqa: F821
@validate_request("id", "dsl", "component_id")
@login_required
async def rerun():
@@ -310,7 +301,7 @@ async def rerun():
return get_json_result(data=True)
@manager.route('/cancel/<task_id>', methods=['PUT']) # noqa: F821
@manager.route("/cancel/<task_id>", methods=["PUT"]) # noqa: F821
@login_required
def cancel(task_id):
try:
@@ -320,15 +311,13 @@ def cancel(task_id):
return get_json_result(data=True)
@manager.route('/reset', methods=['POST']) # noqa: F821
@manager.route("/reset", methods=["POST"]) # noqa: F821
@validate_request("id")
@login_required
async def reset():
req = await get_request_json()
if not UserCanvasService.accessible(req["id"], current_user.id):
return get_json_result(
data=False, message='Only owner of canvas authorized for this operation.',
code=RetCode.OPERATING_ERROR)
return get_json_result(data=False, message="Only owner of canvas authorized for this operation.", code=RetCode.OPERATING_ERROR)
try:
e, user_canvas = UserCanvasService.get_by_id(req["id"])
if not e:
@@ -361,7 +350,7 @@ async def upload(canvas_id):
return server_error_response(e)
@manager.route('/input_form', methods=['GET']) # noqa: F821
@manager.route("/input_form", methods=["GET"]) # noqa: F821
@login_required
def input_form():
cvs_id = request.args.get("id")
@@ -371,9 +360,7 @@ def input_form():
if not e:
return get_data_error_result(message="canvas not found.")
if not UserCanvasService.query(user_id=current_user.id, id=cvs_id):
return get_json_result(
data=False, message='Only owner of canvas authorized for this operation.',
code=RetCode.OPERATING_ERROR)
return get_json_result(data=False, message="Only owner of canvas authorized for this operation.", code=RetCode.OPERATING_ERROR)
canvas = Canvas(json.dumps(user_canvas.dsl), current_user.id, canvas_id=user_canvas.id)
return get_json_result(data=canvas.get_component_input_form(cpn_id))
@@ -381,15 +368,13 @@ def input_form():
return server_error_response(e)
@manager.route('/debug', methods=['POST']) # noqa: F821
@manager.route("/debug", methods=["POST"]) # noqa: F821
@validate_request("id", "component_id", "params")
@login_required
async def debug():
req = await get_request_json()
if not UserCanvasService.accessible(req["id"], current_user.id):
return get_json_result(
data=False, message='Only owner of canvas authorized for this operation.',
code=RetCode.OPERATING_ERROR)
return get_json_result(data=False, message="Only owner of canvas authorized for this operation.", code=RetCode.OPERATING_ERROR)
try:
e, user_canvas = UserCanvasService.get_by_id(req["id"])
canvas = Canvas(json.dumps(user_canvas.dsl), current_user.id, canvas_id=user_canvas.id)
@@ -400,7 +385,7 @@ async def debug():
if isinstance(component, LLM):
component.set_debug_inputs(req["params"])
component.invoke(**{k: o["value"] for k,o in req["params"].items()})
component.invoke(**{k: o["value"] for k, o in req["params"].items()})
outputs = component.output()
for k in outputs.keys():
if isinstance(outputs[k], partial):
@@ -418,59 +403,39 @@ async def debug():
return server_error_response(e)
@manager.route('/test_db_connect', methods=['POST']) # noqa: F821
@manager.route("/test_db_connect", methods=["POST"]) # noqa: F821
@validate_request("db_type", "database", "username", "host", "port", "password")
@login_required
async def test_db_connect():
req = await get_request_json()
try:
if req["db_type"] in ["mysql", "mariadb"]:
db = MySQLDatabase(req["database"], user=req["username"], host=req["host"], port=req["port"],
password=req["password"])
db = MySQLDatabase(req["database"], user=req["username"], host=req["host"], port=req["port"], password=req["password"])
elif req["db_type"] == "oceanbase":
db = MySQLDatabase(req["database"], user=req["username"], host=req["host"], port=req["port"],
password=req["password"], charset="utf8mb4")
elif req["db_type"] == 'postgres':
db = PostgresqlDatabase(req["database"], user=req["username"], host=req["host"], port=req["port"],
password=req["password"])
elif req["db_type"] == 'mssql':
db = MySQLDatabase(req["database"], user=req["username"], host=req["host"], port=req["port"], password=req["password"], charset="utf8mb4")
elif req["db_type"] == "postgres":
db = PostgresqlDatabase(req["database"], user=req["username"], host=req["host"], port=req["port"], password=req["password"])
elif req["db_type"] == "mssql":
import pyodbc
connection_string = (
f"DRIVER={{ODBC Driver 17 for SQL Server}};"
f"SERVER={req['host']},{req['port']};"
f"DATABASE={req['database']};"
f"UID={req['username']};"
f"PWD={req['password']};"
)
connection_string = f"DRIVER={{ODBC Driver 17 for SQL Server}};SERVER={req['host']},{req['port']};DATABASE={req['database']};UID={req['username']};PWD={req['password']};"
db = pyodbc.connect(connection_string)
cursor = db.cursor()
cursor.execute("SELECT 1")
cursor.close()
elif req["db_type"] == 'IBM DB2':
elif req["db_type"] == "IBM DB2":
import ibm_db
conn_str = (
f"DATABASE={req['database']};"
f"HOSTNAME={req['host']};"
f"PORT={req['port']};"
f"PROTOCOL=TCPIP;"
f"UID={req['username']};"
f"PWD={req['password']};"
)
redacted_conn_str = (
f"DATABASE={req['database']};"
f"HOSTNAME={req['host']};"
f"PORT={req['port']};"
f"PROTOCOL=TCPIP;"
f"UID={req['username']};"
f"PWD=****;"
)
conn_str = f"DATABASE={req['database']};HOSTNAME={req['host']};PORT={req['port']};PROTOCOL=TCPIP;UID={req['username']};PWD={req['password']};"
redacted_conn_str = f"DATABASE={req['database']};HOSTNAME={req['host']};PORT={req['port']};PROTOCOL=TCPIP;UID={req['username']};PWD=****;"
logging.info(redacted_conn_str)
conn = ibm_db.connect(conn_str, "", "")
stmt = ibm_db.exec_immediate(conn, "SELECT 1 FROM sysibm.sysdummy1")
ibm_db.fetch_assoc(stmt)
ibm_db.close(conn)
return get_json_result(data="Database Connection Successful!")
elif req["db_type"] == 'trino':
elif req["db_type"] == "trino":
def _parse_catalog_schema(db_name: str):
if not db_name:
return None, None
@@ -481,6 +446,7 @@ async def test_db_connect():
else:
catalog_name, schema_name = db_name, "default"
return catalog_name, schema_name
try:
import trino
import os
@@ -498,13 +464,7 @@ async def test_db_connect():
auth = trino.BasicAuthentication(req.get("username") or "ragflow", req["password"])
conn = trino.dbapi.connect(
host=req["host"],
port=int(req["port"] or 8080),
user=req["username"] or "ragflow",
catalog=catalog,
schema=schema or "default",
http_scheme=http_scheme,
auth=auth
host=req["host"], port=int(req["port"] or 8080), user=req["username"] or "ragflow", catalog=catalog, schema=schema or "default", http_scheme=http_scheme, auth=auth
)
cur = conn.cursor()
cur.execute("SELECT 1")
@@ -514,7 +474,7 @@ async def test_db_connect():
return get_json_result(data="Database Connection Successful!")
else:
return server_error_response("Unsupported database type.")
if req["db_type"] != 'mssql':
if req["db_type"] != "mssql":
db.connect()
db.close()
@@ -523,21 +483,21 @@ async def test_db_connect():
return server_error_response(e)
#api get list version dsl of canvas
@manager.route('/getlistversion/<canvas_id>', methods=['GET']) # noqa: F821
# api get list version dsl of canvas
@manager.route("/getlistversion/<canvas_id>", methods=["GET"]) # noqa: F821
@login_required
def getlistversion(canvas_id):
try:
versions =sorted([c.to_dict() for c in UserCanvasVersionService.list_by_canvas_id(canvas_id)], key=lambda x: x["update_time"]*-1)
versions = sorted([c.to_dict() for c in UserCanvasVersionService.list_by_canvas_id(canvas_id)], key=lambda x: x["update_time"] * -1)
return get_json_result(data=versions)
except Exception as e:
return get_data_error_result(message=f"Error getting history files: {e}")
#api get version dsl of canvas
@manager.route('/getversion/<version_id>', methods=['GET']) # noqa: F821
# api get version dsl of canvas
@manager.route("/getversion/<version_id>", methods=["GET"]) # noqa: F821
@login_required
def getversion( version_id):
def getversion(version_id):
try:
e, version = UserCanvasVersionService.get_by_id(version_id)
if version:
@@ -546,7 +506,7 @@ def getversion( version_id):
return get_json_result(data=f"Error getting history file: {e}")
@manager.route('/list', methods=['GET']) # noqa: F821
@manager.route("/list", methods=["GET"]) # noqa: F821
@login_required
def list_canvas():
keywords = request.args.get("keywords", "")
@@ -563,18 +523,14 @@ def list_canvas():
tenants = TenantService.get_joined_tenants_by_user_id(current_user.id)
tenants = [m["tenant_id"] for m in tenants]
tenants.append(current_user.id)
canvas, total = UserCanvasService.get_by_tenant_ids(
tenants, current_user.id, page_number,
items_per_page, orderby, desc, keywords, canvas_category)
canvas, total = UserCanvasService.get_by_tenant_ids(tenants, current_user.id, page_number, items_per_page, orderby, desc, keywords, canvas_category)
else:
tenants = owner_ids
canvas, total = UserCanvasService.get_by_tenant_ids(
tenants, current_user.id, 0,
0, orderby, desc, keywords, canvas_category)
canvas, total = UserCanvasService.get_by_tenant_ids(tenants, current_user.id, 0, 0, orderby, desc, keywords, canvas_category)
return get_json_result(data={"canvas": canvas, "total": total})
@manager.route('/setting', methods=['POST']) # noqa: F821
@manager.route("/setting", methods=["POST"]) # noqa: F821
@validate_request("id", "title", "permission")
@login_required
async def setting():
@@ -582,11 +538,9 @@ async def setting():
req["user_id"] = current_user.id
if not UserCanvasService.accessible(req["id"], current_user.id):
return get_json_result(
data=False, message='Only owner of canvas authorized for this operation.',
code=RetCode.OPERATING_ERROR)
return get_json_result(data=False, message="Only owner of canvas authorized for this operation.", code=RetCode.OPERATING_ERROR)
e,flow = UserCanvasService.get_by_id(req["id"])
e, flow = UserCanvasService.get_by_id(req["id"])
if not e:
return get_data_error_result(message="canvas not found.")
flow = flow.to_dict()
@@ -596,11 +550,11 @@ async def setting():
if value := req.get(key):
flow[key] = value
num= UserCanvasService.update_by_id(req["id"], flow)
num = UserCanvasService.update_by_id(req["id"], flow)
return get_json_result(data=num)
@manager.route('/trace', methods=['GET']) # noqa: F821
@manager.route("/trace", methods=["GET"]) # noqa: F821
def trace():
cvs_id = request.args.get("canvas_id")
msg_id = request.args.get("message_id")
@@ -614,14 +568,12 @@ def trace():
logging.exception(e)
@manager.route('/<canvas_id>/sessions', methods=['GET']) # noqa: F821
@manager.route("/<canvas_id>/sessions", methods=["GET"]) # noqa: F821
@login_required
def sessions(canvas_id):
tenant_id = current_user.id
if not UserCanvasService.accessible(canvas_id, tenant_id):
return get_json_result(
data=False, message='Only owner of canvas authorized for this operation.',
code=RetCode.OPERATING_ERROR)
return get_json_result(data=False, message="Only owner of canvas authorized for this operation.", code=RetCode.OPERATING_ERROR)
user_id = request.args.get("user_id")
page_number = int(request.args.get("page", 1))
@@ -639,18 +591,17 @@ def sessions(canvas_id):
if exp_user_id:
sess = API4ConversationService.get_names(canvas_id, exp_user_id)
return get_json_result(data={"total": len(sess), "sessions": sess})
# dsl defaults to True in all cases except for False and false
include_dsl = request.args.get("dsl") != "False" and request.args.get("dsl") != "false"
total, sess = API4ConversationService.get_list(canvas_id, tenant_id, page_number, items_per_page, orderby, desc,
None, user_id, include_dsl, keywords, from_date, to_date, exp_user_id=exp_user_id)
total, sess = API4ConversationService.get_list(canvas_id, tenant_id, page_number, items_per_page, orderby, desc, None, user_id, include_dsl, keywords, from_date, to_date, exp_user_id=exp_user_id)
try:
return get_json_result(data={"total": total, "sessions": sess})
except Exception as e:
return server_error_response(e)
@manager.route('/<canvas_id>/sessions', methods=['PUT']) # noqa: F821
@manager.route("/<canvas_id>/sessions", methods=["PUT"]) # noqa: F821
@login_required
async def set_session(canvas_id):
req = await get_request_json()
@@ -659,63 +610,51 @@ async def set_session(canvas_id):
assert e, "Agent not found."
if not isinstance(cvs.dsl, str):
cvs.dsl = json.dumps(cvs.dsl, ensure_ascii=False)
session_id=get_uuid()
session_id = get_uuid()
canvas = Canvas(cvs.dsl, tenant_id, canvas_id, canvas_id=cvs.id)
canvas.reset()
conv = {
"id": session_id,
"name": req.get("name", ""),
"dialog_id": cvs.id,
"user_id": tenant_id,
"exp_user_id": tenant_id,
"message": [],
"source": "agent",
"dsl": cvs.dsl,
"reference": []
}
conv = {"id": session_id, "name": req.get("name", ""), "dialog_id": cvs.id, "user_id": tenant_id, "exp_user_id": tenant_id, "message": [], "source": "agent", "dsl": cvs.dsl, "reference": []}
API4ConversationService.save(**conv)
return get_json_result(data=conv)
@manager.route('/<canvas_id>/sessions/<session_id>', methods=['GET']) # noqa: F821
@manager.route("/<canvas_id>/sessions/<session_id>", methods=["GET"]) # noqa: F821
@login_required
def get_session(canvas_id, session_id):
tenant_id = current_user.id
if not UserCanvasService.accessible(canvas_id, tenant_id):
return get_json_result(
data=False, message='Only owner of canvas authorized for this operation.',
code=RetCode.OPERATING_ERROR)
return get_json_result(data=False, message="Only owner of canvas authorized for this operation.", code=RetCode.OPERATING_ERROR)
_, conv = API4ConversationService.get_by_id(session_id)
return get_json_result(data=conv.to_dict())
@manager.route('/<canvas_id>/sessions/<session_id>', methods=['DELETE']) # noqa: F821
@manager.route("/<canvas_id>/sessions/<session_id>", methods=["DELETE"]) # noqa: F821
@login_required
def del_session(canvas_id, session_id):
tenant_id = current_user.id
if not UserCanvasService.accessible(canvas_id, tenant_id):
return get_json_result(
data=False, message='Only owner of canvas authorized for this operation.',
code=RetCode.OPERATING_ERROR)
return get_json_result(data=False, message="Only owner of canvas authorized for this operation.", code=RetCode.OPERATING_ERROR)
return get_json_result(data=API4ConversationService.delete_by_id(session_id))
@manager.route('/prompts', methods=['GET']) # noqa: F821
@manager.route("/prompts", methods=["GET"]) # noqa: F821
@login_required
def prompts():
from rag.prompts.generator import ANALYZE_TASK_SYSTEM, ANALYZE_TASK_USER, NEXT_STEP, REFLECT, CITATION_PROMPT_TEMPLATE
return get_json_result(data={
"task_analysis": ANALYZE_TASK_SYSTEM +"\n\n"+ ANALYZE_TASK_USER,
"plan_generation": NEXT_STEP,
"reflection": REFLECT,
#"context_summary": SUMMARY4MEMORY,
#"context_ranking": RANK_MEMORY,
"citation_guidelines": CITATION_PROMPT_TEMPLATE
})
return get_json_result(
data={
"task_analysis": ANALYZE_TASK_SYSTEM + "\n\n" + ANALYZE_TASK_USER,
"plan_generation": NEXT_STEP,
"reflection": REFLECT,
# "context_summary": SUMMARY4MEMORY,
# "context_ranking": RANK_MEMORY,
"citation_guidelines": CITATION_PROMPT_TEMPLATE,
}
)
@manager.route('/download', methods=['GET']) # noqa: F821
@manager.route("/download", methods=["GET"]) # noqa: F821
async def download():
id = request.args.get("id")
created_by = request.args.get("created_by")
api/apps/llm_app.py
+18 -8
View File
@@ -111,6 +111,7 @@ async def set_api_key():
assert factory in ChatModel, f"Chat model from {factory} is not supported yet."
mdl = ChatModel[factory](req["api_key"], llm.llm_name, base_url=base_url, **extra)
try:
async def check_streamly():
async for chunk in mdl.async_chat_streamly(
None,
@@ -149,7 +150,7 @@ async def set_api_key():
break
if req.get("verify", False):
return get_json_result(data={"message": msg, "success": len(msg.strip())==0})
return get_json_result(data={"message": msg, "success": len(msg.strip()) == 0})
if msg:
return get_data_error_result(message=msg)
@@ -207,7 +208,10 @@ async def add_llm():
saved_llm_name = llm_name + _LLM_NAME_SUFFIX.get(factory, "")
logging.debug(
"add_llm: attempting api_key recovery factory=%s llm_name=%s saved_llm_name=%s tenant_id=%s",
factory, llm_name, saved_llm_name, current_user.id,
factory,
llm_name,
saved_llm_name,
current_user.id,
)
existing_llms = TenantLLMService.query(
tenant_id=current_user.id,
@@ -216,21 +220,25 @@ async def add_llm():
)
logging.debug(
"add_llm: api_key recovery query matched=%d factory=%s saved_llm_name=%s",
len(existing_llms) if existing_llms else 0, factory, saved_llm_name,
len(existing_llms) if existing_llms else 0,
factory,
saved_llm_name,
)
if existing_llms:
existing_api_key, _, _ = TenantLLMService._decode_api_key_config(
existing_llms[0].api_key
)
existing_api_key, _, _ = TenantLLMService._decode_api_key_config(existing_llms[0].api_key)
logging.debug(
"add_llm: api_key recovery decoded=%s factory=%s saved_llm_name=%s",
"present" if existing_api_key else "absent", factory, saved_llm_name,
"present" if existing_api_key else "absent",
factory,
saved_llm_name,
)
if existing_api_key:
req["api_key"] = existing_api_key
logging.info(
"add_llm: recovered saved api_key from existing record factory=%s saved_llm_name=%s tenant_id=%s",
factory, saved_llm_name, current_user.id,
factory,
saved_llm_name,
current_user.id,
)
api_key = req.get("api_key", "x")
@@ -335,6 +343,7 @@ async def add_llm():
**extra,
)
try:
async def check_streamly():
async for chunk in mdl.async_chat_streamly(
None,
@@ -387,6 +396,7 @@ async def add_llm():
assert factory in TTSModel, f"TTS model from {factory} is not supported yet."
mdl = TTSModel[factory](key=model_api_key, model_name=mdl_nm, base_url=model_base_url)
try:
def drain_tts():
for _ in mdl.tts("Hello~ RAGFlower!"):
pass
api/apps/restful_apis/bot_api.py
+48 -42
View File
@@ -39,8 +39,7 @@ from api.db.services.search_service import SearchService
from api.db.services.user_service import UserTenantService
from api.db.joint_services.tenant_model_service import get_tenant_default_model_by_type, get_model_config_from_provider_instance
from common.misc_utils import thread_pool_exec
from api.utils.api_utils import get_error_data_result, get_json_result, \
add_tenant_id_to_kwargs, get_result, get_request_json, server_error_response, validate_request
from api.utils.api_utils import get_error_data_result, get_json_result, add_tenant_id_to_kwargs, get_result, get_request_json, server_error_response, validate_request
from rag.app.tag import label_question
from rag.prompts.template import load_prompt
from rag.prompts.generator import cross_languages, keyword_extraction
@@ -58,7 +57,7 @@ def _get_sdk_authorization_token():
auth_header = request.headers.get("Authorization", "")
if not auth_header.startswith("Bearer "):
return ""
return auth_header[len("Bearer "):].strip()
return auth_header[len("Bearer ") :].strip()
@manager.route("/chatbots/<dialog_id>/completions", methods=["POST"]) # noqa: F821
@@ -68,9 +67,7 @@ async def chatbot_completions(dialog_id, tenant_id=None):
req = await get_request_json()
exists, dialog = DialogService.get_by_id(dialog_id)
if (not exists
or getattr(dialog, "tenant_id", None) != tenant_id
or str(getattr(dialog, "status", "")) != StatusEnum.VALID.value):
if not exists or getattr(dialog, "tenant_id", None) != tenant_id or str(getattr(dialog, "status", "")) != StatusEnum.VALID.value:
logger.warning(
"Denied chatbot access: reason=%s tenant_id=%s dialog_id=%s user_id=%s session_id=%s",
"no access to this chatbot",
@@ -132,14 +129,13 @@ async def chatbot_completions(dialog_id, tenant_id=None):
return None
@manager.route("/chatbots/<dialog_id>/info", methods=["GET"]) # noqa: F821
@login_required(auth_types=AUTH_BETA)
@add_tenant_id_to_kwargs
async def chatbots_inputs(dialog_id, tenant_id=None):
exists, dialog = await thread_pool_exec(DialogService.get_by_id, dialog_id)
if (not exists
or getattr(dialog, "tenant_id", None) != tenant_id
or str(getattr(dialog, "status", "")) != StatusEnum.VALID.value):
if not exists or getattr(dialog, "tenant_id", None) != tenant_id or str(getattr(dialog, "status", "")) != StatusEnum.VALID.value:
request_args = getattr(request, "args", {}) or {}
request_user_id = request_args.get("user_id") if hasattr(request_args, "get") else None
request_session_id = request_args.get("session_id") if hasattr(request_args, "get") else None
@@ -178,6 +174,7 @@ async def agent_bot_completions(agent_id, tenant_id=None):
return get_error_data_result(message=f"Can't find agent by ID: {agent_id}")
if req.get("stream", True):
async def stream():
try:
async for answer in agent_completion(tenant_id, agent_id, **req):
@@ -185,14 +182,18 @@ async def agent_bot_completions(agent_id, tenant_id=None):
except Exception as e:
logging.exception(e)
error_result = get_error_data_result(message=str(e) or "Unknown error")
yield "data:" + json.dumps(
{
"event": "message",
"data": {"content": f"Error {error_result['code']}: {error_result['message']}\n\n"},
**error_result,
},
ensure_ascii=False,
) + "\n\n"
yield (
"data:"
+ json.dumps(
{
"event": "message",
"data": {"content": f"Error {error_result['code']}: {error_result['message']}\n\n"},
**error_result,
},
ensure_ascii=False,
)
+ "\n\n"
)
resp = Response(stream(), mimetype="text/event-stream")
resp.headers.add_header("Cache-control", "no-cache")
@@ -217,7 +218,7 @@ async def agent_bot_completions(agent_id, tenant_id=None):
line = line.strip()
if not line.startswith("data:"):
continue
payload = line[len("data:"):].strip()
payload = line[len("data:") :].strip()
if not payload:
continue
try:
@@ -270,9 +271,7 @@ async def begin_inputs(agent_id, tenant_id=None):
return get_error_data_result(f"Can't find agent by ID: {agent_id}")
canvas = Canvas(json.dumps(cvs.dsl), tenant_id, canvas_id=cvs.id)
return get_result(
data={"title": cvs.title, "avatar": cvs.avatar, "inputs": canvas.get_component_input_form("begin"),
"prologue": canvas.get_prologue(), "mode": canvas.get_mode()})
return get_result(data={"title": cvs.title, "avatar": cvs.avatar, "inputs": canvas.get_component_input_form("begin"), "prologue": canvas.get_prologue(), "mode": canvas.get_mode()})
@manager.route("/agentbots/<shared_id>/logs/<message_id>", methods=["GET"]) # noqa: F821
@@ -289,9 +288,10 @@ async def agent_bot_logs(shared_id, message_id):
if not token:
logger.warning(
"agent_bot_logs: missing Authorization header (shared_id=%s message_id=%s)",
shared_id, message_id,
shared_id,
message_id,
)
return get_error_data_result(message='Authorization is not valid!')
return get_error_data_result(message="Authorization is not valid!")
# Non-reversible fingerprint of the share token: lets operators correlate
# auth-failure log lines for the same token without leaking a guessable
# substring of the secret itself.
@@ -300,7 +300,8 @@ async def agent_bot_logs(shared_id, message_id):
if not objs:
logger.warning(
"agent_bot_logs: invalid beta token (fingerprint=%s shared_id=%s)",
token_fp, shared_id,
token_fp,
shared_id,
)
return get_error_data_result(message='Authentication error: API key is invalid!"')
@@ -308,9 +309,10 @@ async def agent_bot_logs(shared_id, message_id):
if not agent_id:
logger.warning(
"agent_bot_logs: APIToken has no dialog_id (tenant_id=%s fingerprint=%s)",
objs[0].tenant_id, token_fp,
objs[0].tenant_id,
token_fp,
)
return get_error_data_result(message='API token is not bound to an agent.')
return get_error_data_result(message="API token is not bound to an agent.")
try:
binary = await thread_pool_exec(REDIS_CONN.get, f"{agent_id}-{message_id}-logs")
@@ -348,9 +350,7 @@ async def ask_about_embedded(tenant_id=None):
async for ans in async_ask(req["question"], req["kb_ids"], uid, chat_llm_name=chat_llm_name, search_config=search_config):
yield "data:" + json.dumps({"code": 0, "message": "", "data": ans}, ensure_ascii=False) + "\n\n"
except Exception as e:
yield "data:" + json.dumps(
{"code": 500, "message": str(e), "data": {"answer": "**ERROR**: " + str(e), "reference": []}},
ensure_ascii=False) + "\n\n"
yield "data:" + json.dumps({"code": 500, "message": str(e), "data": {"answer": "**ERROR**: " + str(e), "reference": []}}, ensure_ascii=False) + "\n\n"
yield "data:" + json.dumps({"code": 0, "message": "", "data": True}, ensure_ascii=False) + "\n\n"
resp = Response(stream(), mimetype="text/event-stream")
@@ -374,8 +374,7 @@ async def retrieval_test_embedded(tenant_id=None):
if isinstance(kb_ids, str):
kb_ids = [kb_ids]
if not kb_ids:
return get_json_result(data=False, message='Please specify dataset firstly.',
code=RetCode.DATA_ERROR)
return get_json_result(data=False, message="Please specify dataset firstly.", code=RetCode.DATA_ERROR)
doc_ids = req.get("doc_ids", [])
similarity_threshold = float(req.get("similarity_threshold", 0.0))
vector_similarity_weight = float(req.get("vector_similarity_weight", 0.3))
@@ -443,8 +442,7 @@ async def retrieval_test_embedded(tenant_id=None):
tenant_ids.append(tenant.tenant_id)
break
else:
return get_json_result(data=False, message="Only owner of dataset authorized for this operation.",
code=RetCode.OPERATING_ERROR)
return get_json_result(data=False, message="Only owner of dataset authorized for this operation.", code=RetCode.OPERATING_ERROR)
e, kb = await thread_pool_exec(KnowledgebaseService.get_by_id, kb_ids[0])
if not e:
@@ -467,13 +465,23 @@ async def retrieval_test_embedded(tenant_id=None):
labels = label_question(_question, [kb])
ranks = await settings.retriever.retrieval(
_question, embd_mdl, tenant_ids, kb_ids, page, size, similarity_threshold, vector_similarity_weight, top,
local_doc_ids, rerank_mdl=rerank_mdl, highlight=req.get("highlight"), rank_feature=labels
_question,
embd_mdl,
tenant_ids,
kb_ids,
page,
size,
similarity_threshold,
vector_similarity_weight,
top,
local_doc_ids,
rerank_mdl=rerank_mdl,
highlight=req.get("highlight"),
rank_feature=labels,
)
if use_kg:
default_chat_model = await thread_pool_exec(get_tenant_default_model_by_type, kb.tenant_id, LLMType.CHAT)
ck = await settings.kg_retriever.retrieval(_question, tenant_ids, kb_ids, embd_mdl,
LLMBundle(kb.tenant_id, default_chat_model))
ck = await settings.kg_retriever.retrieval(_question, tenant_ids, kb_ids, embd_mdl, LLMBundle(kb.tenant_id, default_chat_model))
if ck["content_with_weight"]:
ranks["chunks"].insert(0, ck)
@@ -492,8 +500,7 @@ async def retrieval_test_embedded(tenant_id=None):
return await _retrieval()
except Exception as e:
if "not_found" in str(e):
return get_json_result(data=False, message="No chunk found! Check the chunk status please!",
code=RetCode.DATA_ERROR)
return get_json_result(data=False, message="No chunk found! Check the chunk status please!", code=RetCode.DATA_ERROR)
return server_error_response(e)
@@ -552,8 +559,7 @@ async def detail_share_embedded(tenant_id=None):
if await thread_pool_exec(SearchService.query, tenant_id=tenant.tenant_id, id=search_id):
break
else:
return get_json_result(data=False, message="Has no permission for this operation.",
code=RetCode.OPERATING_ERROR)
return get_json_result(data=False, message="Has no permission for this operation.", code=RetCode.OPERATING_ERROR)
search = await thread_pool_exec(SearchService.get_detail, search_id)
if not search:
@@ -573,7 +579,7 @@ async def mindmap(tenant_id=None):
search_id = req.get("search_id", "")
search_app = await thread_pool_exec(SearchService.get_detail, search_id) if search_id else {}
mind_map =await gen_mindmap(req["question"], req["kb_ids"], tenant_id, search_app.get("search_config", {}))
mind_map = await gen_mindmap(req["question"], req["kb_ids"], tenant_id, search_app.get("search_config", {}))
if "error" in mind_map:
return server_error_response(Exception(mind_map["error"]))
return get_json_result(data=mind_map)
api/apps/restful_apis/chat_api.py
+43 -41
View File
@@ -27,9 +27,7 @@ from quart import Response, request
from api.apps import current_user, login_required
from api.apps.restful_apis._generation_params import merge_generation_config, pop_generation_config
from api.db.joint_services.tenant_model_service import (
get_tenant_default_model_by_type, get_model_config_from_provider_instance, get_api_key, split_model_name
)
from api.db.joint_services.tenant_model_service import get_tenant_default_model_by_type, get_model_config_from_provider_instance, get_api_key, split_model_name
from api.db.services.chunk_feedback_service import ChunkFeedbackService
from api.db.services.conversation_service import ConversationService, structure_answer
from api.db.services.dialog_service import DialogService, async_chat, gen_mindmap
@@ -52,6 +50,7 @@ from common.misc_utils import get_uuid, thread_pool_exec
from rag.prompts.generator import chunks_format
from rag.prompts.template import load_prompt
def _sanitize_json_floats(obj):
"""Replace NaN/Infinity floats with None so the result is RFC 8259 JSON.
@@ -86,8 +85,8 @@ def _sanitize_json_floats(obj):
_DEFAULT_PROMPT_CONFIG = {
"system": (
'You are an intelligent assistant. Please summarize the content of the dataset to answer the question. '
'Please list the data in the dataset and answer in detail. When all dataset content is irrelevant to the '
"You are an intelligent assistant. Please summarize the content of the dataset to answer the question. "
"Please list the data in the dataset and answer in detail. When all dataset content is irrelevant to the "
'question, your answer must include the sentence "The answer you are looking for is not found in the dataset!" '
"Answers need to consider chat history.\n"
" Here is the knowledge base:\n"
@@ -162,10 +161,7 @@ def _build_session_response(conv: dict) -> dict:
async def _ensure_owned_chat(chat_id):
return await thread_pool_exec(
DialogService.query,
tenant_id=current_user.id, id=chat_id, status=StatusEnum.VALID.value
)
return await thread_pool_exec(DialogService.query, tenant_id=current_user.id, id=chat_id, status=StatusEnum.VALID.value)
def _build_default_completion_dialog():
@@ -292,10 +288,11 @@ async def _validate_llm_id(llm_id, tenant_id, llm_setting=None):
return None
async def _validate_rerank_id(rerank_id, tenant_id):
if not rerank_id:
return None
parts = rerank_id.split('@')
parts = rerank_id.split("@")
llm_name = parts[0]
if llm_name in _DEFAULT_RERANK_MODELS:
return None
@@ -342,7 +339,7 @@ async def _validate_dataset_ids(dataset_ids, tenant_id):
embd_ids = [split_model_name(kb.embd_id)[0] for kb in kbs]
if len(set(embd_ids)) > 1:
return f'Datasets use different embedding models: {[kb.embd_id for kb in kbs]}'
return f"Datasets use different embedding models: {[kb.embd_id for kb in kbs]}"
return normalized_ids
@@ -462,7 +459,14 @@ async def list_chats():
if owner_ids:
chats, total = await thread_pool_exec(
DialogService.get_by_tenant_ids,
owner_ids, current_user.id, 0, 0, orderby, desc, keywords, **exact_filters,
owner_ids,
current_user.id,
0,
0,
orderby,
desc,
keywords,
**exact_filters,
)
chats = [chat for chat in chats if chat["tenant_id"] in owner_ids]
total = len(chats)
@@ -472,12 +476,17 @@ async def list_chats():
else:
chats, total = await thread_pool_exec(
DialogService.get_by_tenant_ids,
[], current_user.id, page_number, items_per_page, orderby, desc, keywords, **exact_filters,
[],
current_user.id,
page_number,
items_per_page,
orderby,
desc,
keywords,
**exact_filters,
)
return get_json_result(
data={"chats": [_build_chat_response(chat) for chat in chats], "total": total}
)
return get_json_result(data={"chats": [_build_chat_response(chat) for chat in chats], "total": total})
except Exception as ex:
return server_error_response(ex)
@@ -490,7 +499,9 @@ async def get_chat(chat_id):
for tenant in tenants:
if await thread_pool_exec(
DialogService.query,
tenant_id=tenant.tenant_id, id=chat_id, status=StatusEnum.VALID.value,
tenant_id=tenant.tenant_id,
id=chat_id,
status=StatusEnum.VALID.value,
):
break
else:
@@ -512,9 +523,7 @@ async def get_chat(chat_id):
@login_required
async def update_chat(chat_id):
if not await _ensure_owned_chat(chat_id):
return get_json_result(
data=False, message="No authorization.", code=RetCode.AUTHENTICATION_ERROR
)
return get_json_result(data=False, message="No authorization.", code=RetCode.AUTHENTICATION_ERROR)
try:
req = await get_request_json()
@@ -596,9 +605,7 @@ async def update_chat(chat_id):
@login_required
async def patch_chat(chat_id):
if not await _ensure_owned_chat(chat_id):
return get_json_result(
data=False, message="No authorization.", code=RetCode.AUTHENTICATION_ERROR
)
return get_json_result(data=False, message="No authorization.", code=RetCode.AUTHENTICATION_ERROR)
try:
req = await get_request_json()
@@ -686,9 +693,7 @@ async def patch_chat(chat_id):
@login_required
async def delete_chat(chat_id):
if not await _ensure_owned_chat(chat_id):
return get_json_result(
data=False, message="No authorization.", code=RetCode.AUTHENTICATION_ERROR
)
return get_json_result(data=False, message="No authorization.", code=RetCode.AUTHENTICATION_ERROR)
try:
if not DialogService.update_by_id(chat_id, {"status": StatusEnum.INVALID.value}):
@@ -708,12 +713,7 @@ async def bulk_delete_chats():
ids = req.get("ids")
if not ids:
if req.get("delete_all") is True:
ids = [
chat.id
for chat in DialogService.query(
tenant_id=current_user.id, status=StatusEnum.VALID.value
)
]
ids = [chat.id for chat in DialogService.query(tenant_id=current_user.id, status=StatusEnum.VALID.value)]
if not ids:
return get_json_result(data={})
else:
@@ -799,9 +799,7 @@ async def list_sessions(chat_id):
session_id = request.args.get("id")
name = request.args.get("name")
user_id = request.args.get("user_id")
convs = ConversationService.get_list(
chat_id, page_number, items_per_page, orderby, desc, session_id, name, user_id
)
convs = ConversationService.get_list(chat_id, page_number, items_per_page, orderby, desc, session_id, name, user_id)
if items_per_page == 0:
convs = []
return get_json_result(data=[_build_session_response(c) for c in convs])
@@ -1045,17 +1043,21 @@ async def transcription():
uploaded = files["file"]
ALLOWED_EXTS = {
".wav", ".mp3", ".m4a", ".aac",
".flac", ".ogg", ".webm",
".opus", ".wma",
".wav",
".mp3",
".m4a",
".aac",
".flac",
".ogg",
".webm",
".opus",
".wma",
}
filename = uploaded.filename or ""
suffix = os.path.splitext(filename)[-1].lower()
if suffix not in ALLOWED_EXTS:
return get_data_error_result(
message=f"Unsupported audio format: {suffix}. Allowed: {', '.join(sorted(ALLOWED_EXTS))}"
)
return get_data_error_result(message=f"Unsupported audio format: {suffix}. Allowed: {', '.join(sorted(ALLOWED_EXTS))}")
fd, temp_audio_path = tempfile.mkstemp(suffix=suffix)
os.close(fd)
api/apps/restful_apis/chat_channel_api.py
+1 -8
View File
@@ -37,14 +37,7 @@ def _chat_channel_auth_error(channel_id: str, user_id: str):
async def create_chat_channel():
"""Create a chat channel bot owned by the current tenant."""
req = await get_request_json()
channel = {
"id": get_uuid(),
"tenant_id": current_user.id,
"name": req["name"],
"channel": req["channel"],
"config": req.get("config") or {},
"chat_id": req.get("chat_id") or None
}
channel = {"id": get_uuid(), "tenant_id": current_user.id, "name": req["name"], "channel": req["channel"], "config": req.get("config") or {}, "chat_id": req.get("chat_id") or None}
ChatChannelService.insert(**channel)
e, conn = ChatChannelService.get_by_id(channel["id"])
api/apps/restful_apis/connector_api.py
+11 -6
View File
@@ -551,6 +551,7 @@ async def google_drive_web_oauth_callback():
return await _render_web_oauth_popup(state_id, True, "Authorization completed successfully.", source)
@manager.route("/connectors/google/oauth/web/result", methods=["POST"]) # noqa: F821
@login_required
@validate_request("flow_id")
@@ -571,13 +572,14 @@ async def poll_google_web_result():
REDIS_CONN.delete(_web_result_cache_key(flow_id, source))
return get_json_result(data={"credentials": result.get("credentials")})
@manager.route("/connectors/box/oauth/web/start", methods=["POST"]) # noqa: F821
@login_required
async def start_box_web_oauth():
req = await get_request_json()
client_id = req.get("client_id")
client_secret = req.get("client_secret")
client_secret = req.get("client_secret")
redirect_uri = req.get("redirect_uri", BOX_WEB_OAUTH_REDIRECT_URI)
if not client_id or not client_secret:
@@ -608,18 +610,20 @@ async def start_box_web_oauth():
}
REDIS_CONN.set_obj(_web_state_cache_key(flow_id, "box"), cache_payload, WEB_FLOW_TTL_SECS)
return get_json_result(
data = {
data={
"flow_id": flow_id,
"authorization_url": auth_url,
"expires_in": WEB_FLOW_TTL_SECS,}
"expires_in": WEB_FLOW_TTL_SECS,
}
)
@manager.route("/connectors/box/oauth/web/callback", methods=["GET"]) # noqa: F821
async def box_web_oauth_callback():
flow_id = request.args.get("state")
if not flow_id:
return await _render_web_oauth_popup("", False, "Missing OAuth parameters.", "box")
code = request.args.get("code")
if not code:
return await _render_web_oauth_popup(flow_id, False, "Missing authorization code from Box.", "box")
@@ -633,7 +637,7 @@ async def box_web_oauth_callback():
if error:
REDIS_CONN.delete(_web_state_cache_key(flow_id, "box"))
return await _render_web_oauth_popup(flow_id, False, error_description or "Authorization failed.", "box")
auth = BoxOAuth(
OAuthConfig(
client_id=cache_payload.get("client_id"),
@@ -656,6 +660,7 @@ async def box_web_oauth_callback():
return await _render_web_oauth_popup(flow_id, True, "Authorization completed successfully.", "box")
@manager.route("/connectors/box/oauth/web/result", methods=["POST"]) # noqa: F821
@login_required
@validate_request("flow_id")
@@ -670,7 +675,7 @@ async def poll_box_web_result():
cache_raw = json.loads(cache_blob)
if cache_raw.get("user_id") != current_user.id:
return get_json_result(code=RetCode.PERMISSION_ERROR, message="You are not allowed to access this authorization result.")
REDIS_CONN.delete(_web_result_cache_key(flow_id, "box"))
return get_json_result(data={"credentials": cache_raw})
api/apps/restful_apis/dify_retrieval_api.py
+9 -23
View File
@@ -108,7 +108,7 @@ def _parse_retrieval_options(retrieval_setting):
return retrieval_setting, similarity_threshold, top
@manager.route('/dify/retrieval', methods=['POST', 'GET']) # noqa: F821
@manager.route("/dify/retrieval", methods=["POST", "GET"]) # noqa: F821
@login_required
@add_tenant_id_to_kwargs
async def retrieval(tenant_id):
@@ -251,7 +251,6 @@ async def retrieval(tenant_id):
doc_ids = []
try:
e, kb = KnowledgebaseService.get_by_id(kb_id)
if not e:
return build_error_result(message="Knowledgebase not found!", code=RetCode.NOT_FOUND)
@@ -279,17 +278,13 @@ async def retrieval(tenant_id):
vector_similarity_weight=0.3,
top=top,
doc_ids=doc_ids,
rank_feature=label_question(question, [kb])
rank_feature=label_question(question, [kb]),
)
ranks["chunks"] = settings.retriever.retrieval_by_children(ranks["chunks"], [tenant_id])
if use_kg:
model_config = get_tenant_default_model_by_type(kb.tenant_id, LLMType.CHAT)
ck = await settings.kg_retriever.retrieval(question,
[tenant_id],
[kb_id],
embd_mdl,
LLMBundle(kb.tenant_id, model_config))
ck = await settings.kg_retriever.retrieval(question, [tenant_id], [kb_id], embd_mdl, LLMBundle(kb.tenant_id, model_config))
if ck["content_with_weight"]:
ranks["chunks"].insert(0, ck)
@@ -303,30 +298,21 @@ async def retrieval(tenant_id):
if not doc:
continue
c.pop("vector", None)
meta = getattr(doc, 'meta_fields', {})
meta = getattr(doc, "meta_fields", {})
meta["doc_id"] = c["doc_id"]
# Dify expects metadata.document_id for external retrieval sources.
meta["document_id"] = c["doc_id"]
records.append({
"content": c["content_with_weight"],
"score": c["similarity"],
"title": c["docnm_kwd"],
"metadata": meta
})
records.append({"content": c["content_with_weight"], "score": c["similarity"], "title": c["docnm_kwd"], "metadata": meta})
return jsonify({"records": records})
except Exception as e:
if "not_found" in str(e):
return build_error_result(
message='No chunk found! Check the chunk status please!',
code=RetCode.NOT_FOUND
)
return build_error_result(message="No chunk found! Check the chunk status please!", code=RetCode.NOT_FOUND)
logging.exception(e)
return build_error_result(message=str(e), code=RetCode.SERVER_ERROR)
@manager.route('/dify/retrieval/health', methods=['GET']) # noqa: F821
@manager.route("/dify/retrieval/health", methods=["GET"]) # noqa: F821
async def retrieval_health_check():
"""Health check endpoint for Dify external knowledge base connectivity verification."""
return get_json_result(data=True)
api/apps/restful_apis/file2document_api.py
+24 -22
View File
@@ -56,27 +56,31 @@ def _convert_files(file_ids, kb_ids, user_id):
e, kb = KnowledgebaseService.get_by_id(kb_id)
if not e:
continue
doc = DocumentService.insert({
"id": get_uuid(),
"kb_id": kb.id,
"parser_id": FileService.get_parser(file.type, file.name, kb.parser_id),
"pipeline_id": kb.pipeline_id,
"parser_config": kb.parser_config,
"created_by": user_id,
"type": file.type,
"name": file.name,
"suffix": Path(file.name).suffix.lstrip("."),
"location": file.location,
"size": file.size
})
File2DocumentService.insert({
"id": get_uuid(),
"file_id": id,
"document_id": doc.id,
})
doc = DocumentService.insert(
{
"id": get_uuid(),
"kb_id": kb.id,
"parser_id": FileService.get_parser(file.type, file.name, kb.parser_id),
"pipeline_id": kb.pipeline_id,
"parser_config": kb.parser_config,
"created_by": user_id,
"type": file.type,
"name": file.name,
"suffix": Path(file.name).suffix.lstrip("."),
"location": file.location,
"size": file.size,
}
)
File2DocumentService.insert(
{
"id": get_uuid(),
"file_id": id,
"document_id": doc.id,
}
)
@manager.route('/files/link-to-datasets', methods=['POST']) # noqa: F821
@manager.route("/files/link-to-datasets", methods=["POST"]) # noqa: F821
@login_required
@validate_request("file_ids", "kb_ids")
async def convert():
@@ -162,9 +166,7 @@ async def convert():
# soon as the background task is scheduled.
loop = asyncio.get_running_loop()
future = loop.run_in_executor(None, _convert_files, all_file_ids, kb_ids, user_id)
future.add_done_callback(
lambda f: logging.error("_convert_files failed: %s", f.exception()) if f.exception() else None
)
future.add_done_callback(lambda f: logging.error("_convert_files failed: %s", f.exception()) if f.exception() else None)
logger.info(
"user_id=%s resource_type=file_to_dataset_link resource_id=batch action=schedule_convert result=scheduled file_ids=%s kb_ids=%s",
user_id,
api/apps/restful_apis/file_api.py
+6 -13
View File
@@ -69,11 +69,11 @@ async def create_or_upload(tenant_id: str = None):
form = await request.form
pf_id = form.get("parent_id")
files = await request.files
if 'file' not in files:
if "file" not in files:
return get_error_argument_result("No file part!")
file_objs = files.getlist('file')
file_objs = files.getlist("file")
for file_obj in file_objs:
if file_obj.filename == '':
if file_obj.filename == "":
return get_error_argument_result("No file selected!")
success, result = await file_api_service.upload_file(tenant_id, pf_id, file_objs)
@@ -86,9 +86,7 @@ async def create_or_upload(tenant_id: str = None):
if err is not None:
return get_error_argument_result(err)
success, result = await file_api_service.create_folder(
tenant_id, req["name"], req.get("parent_id"), req.get("type")
)
success, result = await file_api_service.create_folder(tenant_id, req["name"], req.get("parent_id"), req.get("type"))
if success:
return get_result(data=result)
else:
@@ -198,9 +196,7 @@ async def delete(tenant_id: str = None):
errors = result.get("errors", [])
return get_json_result(
code=RetCode.DATA_ERROR,
message=f"Partially deleted {success_count} files with {len(errors)} errors"
if success_count > 0
else f"Deleted files failed with {len(errors)} errors",
message=f"Partially deleted {success_count} files with {len(errors)} errors" if success_count > 0 else f"Deleted files failed with {len(errors)} errors",
data=result,
)
return get_error_data_result(message=result)
@@ -209,7 +205,6 @@ async def delete(tenant_id: str = None):
return get_error_data_result(message="Internal server error")
@manager.route("/files/move", methods=["POST"]) # noqa: F821
@login_required
@add_tenant_id_to_kwargs
@@ -254,9 +249,7 @@ async def move(tenant_id: str = None):
return get_error_argument_result(err)
try:
success, result = await file_api_service.move_files(
tenant_id, req["src_file_ids"], req.get("dest_file_id"), req.get("new_name")
)
success, result = await file_api_service.move_files(tenant_id, req["src_file_ids"], req.get("dest_file_id"), req.get("new_name"))
if success:
return get_result(data=result)
else:
api/apps/restful_apis/memory_api.py
+38 -40
View File
@@ -35,12 +35,7 @@ async def create_memory():
req = await get_request_json()
t_parsed = time.perf_counter() if timing_enabled else None
try:
memory_info = {
"name": req["name"],
"memory_type": req["memory_type"],
"embd_id": req["embd_id"],
"llm_id": req["llm_id"]
}
memory_info = {"name": req["name"], "memory_type": req["memory_type"], "embd_id": req["embd_id"], "llm_id": req["llm_id"]}
success, res = await memory_api_service.create_memory(memory_info)
if timing_enabled:
logging.info(
@@ -84,10 +79,26 @@ async def create_memory():
@login_required
async def update_memory(memory_id):
req = await get_request_json()
new_settings = {k: req[k] for k in [
"name", "permissions", "llm_id", "embd_id", "memory_type", "memory_size", "forgetting_policy", "temperature",
"avatar", "description", "system_prompt", "user_prompt", "tenant_llm_id", "tenant_embd_id"
] if k in req}
new_settings = {
k: req[k]
for k in [
"name",
"permissions",
"llm_id",
"embd_id",
"memory_type",
"memory_size",
"forgetting_policy",
"temperature",
"avatar",
"description",
"system_prompt",
"user_prompt",
"tenant_llm_id",
"tenant_embd_id",
]
if k in req
}
try:
success, res = await memory_api_service.update_memory(memory_id, new_settings)
if success:
@@ -122,9 +133,7 @@ async def delete_memory(memory_id):
@manager.route("/memories", methods=["GET"]) # noqa: F821
@login_required
async def list_memory():
filter_params = {
k: request.args.get(k) for k in ["memory_type", "tenant_id", "owner_ids", "storage_type"] if k in request.args
}
filter_params = {k: request.args.get(k) for k in ["memory_type", "tenant_id", "owner_ids", "storage_type"] if k in request.args}
keywords = request.args.get("keywords")
page = int(request.args.get("page", 1))
page_size = validate_rest_api_page_size(int(request.args.get("page_size", 50)))
@@ -155,16 +164,14 @@ async def get_memory_config(memory_id):
async def get_memory_messages(memory_id):
args = request.args
agent_ids = args.getlist("agent_id")
if len(agent_ids) == 1 and ',' in agent_ids[0]:
agent_ids = agent_ids[0].split(',')
if len(agent_ids) == 1 and "," in agent_ids[0]:
agent_ids = agent_ids[0].split(",")
keywords = args.get("keywords", "")
keywords = keywords.strip()
page = int(args.get("page", 1))
page_size = validate_rest_api_page_size(int(args.get("page_size", 50)))
try:
res = await memory_api_service.get_memory_messages(
memory_id, agent_ids, keywords, page, page_size
)
res = await memory_api_service.get_memory_messages(memory_id, agent_ids, keywords, page, page_size)
return get_json_result(message=True, data=res)
except NotFoundException as not_found_exception:
logging.error(not_found_exception)
@@ -174,7 +181,7 @@ async def get_memory_messages(memory_id):
return get_json_result(code=RetCode.SERVER_ERROR, message="Internal server error")
@manager.route("/messages", methods=["POST"]) # noqa: F821
@manager.route("/messages", methods=["POST"]) # noqa: F821
@login_required
@validate_request("memory_id", "agent_id", "session_id", "user_input", "agent_response")
async def add_message():
@@ -206,7 +213,7 @@ async def add_message():
return get_json_result(message="Some messages failed to add. Detail:" + msg, code=RetCode.SERVER_ERROR)
@manager.route("/messages/<memory_id>:<message_id>", methods=["DELETE"]) # noqa: F821
@manager.route("/messages/<memory_id>:<message_id>", methods=["DELETE"]) # noqa: F821
@login_required
async def forget_message(memory_id: str, message_id: int):
try:
@@ -220,7 +227,7 @@ async def forget_message(memory_id: str, message_id: int):
return get_json_result(code=RetCode.SERVER_ERROR, message="Internal server error")
@manager.route("/messages/<memory_id>:<message_id>", methods=["PUT"]) # noqa: F821
@manager.route("/messages/<memory_id>:<message_id>", methods=["PUT"]) # noqa: F821
@login_required
@validate_request("status")
async def update_message(memory_id: str, message_id: int):
@@ -243,13 +250,13 @@ async def update_message(memory_id: str, message_id: int):
return get_json_result(code=RetCode.SERVER_ERROR, message="Internal server error")
@manager.route("/messages/search", methods=["GET"]) # noqa: F821
@manager.route("/messages/search", methods=["GET"]) # noqa: F821
@login_required
async def search_message():
args = request.args
memory_ids = args.getlist("memory_id")
if len(memory_ids) == 1 and ',' in memory_ids[0]:
memory_ids = memory_ids[0].split(',')
if len(memory_ids) == 1 and "," in memory_ids[0]:
memory_ids = memory_ids[0].split(",")
query = args.get("query")
similarity_threshold = float(args.get("similarity_threshold", 0.2))
keywords_similarity_weight = float(args.get("keywords_similarity_weight", 0.7))
@@ -258,28 +265,19 @@ async def search_message():
session_id = args.get("session_id", "")
user_id = args.get("user_id", "")
filter_dict = {
"memory_id": memory_ids,
"agent_id": agent_id,
"session_id": session_id,
"user_id": user_id
}
params = {
"query": query,
"similarity_threshold": similarity_threshold,
"keywords_similarity_weight": keywords_similarity_weight,
"top_n": top_n
}
filter_dict = {"memory_id": memory_ids, "agent_id": agent_id, "session_id": session_id, "user_id": user_id}
params = {"query": query, "similarity_threshold": similarity_threshold, "keywords_similarity_weight": keywords_similarity_weight, "top_n": top_n}
res = await memory_api_service.search_message(filter_dict, params)
return get_json_result(message=True, data=res)
@manager.route("/messages", methods=["GET"]) # noqa: F821
@manager.route("/messages", methods=["GET"]) # noqa: F821
@login_required
async def get_messages():
args = request.args
memory_ids = args.getlist("memory_id")
if len(memory_ids) == 1 and ',' in memory_ids[0]:
memory_ids = memory_ids[0].split(',')
if len(memory_ids) == 1 and "," in memory_ids[0]:
memory_ids = memory_ids[0].split(",")
agent_id = args.get("agent_id", "")
session_id = args.get("session_id", "")
limit = int(args.get("limit", 10))
@@ -293,7 +291,7 @@ async def get_messages():
return get_json_result(code=RetCode.SERVER_ERROR, message="Internal server error")
@manager.route("/messages/<memory_id>:<message_id>/content", methods=["GET"]) # noqa: F821
@manager.route("/messages/<memory_id>:<message_id>/content", methods=["GET"]) # noqa: F821
@login_required
async def get_message_content(memory_id: str, message_id: int):
try:
api/apps/restful_apis/models_api.py
+1 -3
View File
@@ -188,9 +188,7 @@ async def set_default_models(tenant_id: str):
model_type = data["model_type"]
try:
success, msg = models_api_service.set_tenant_default_models(
tenant_id, model_provider, model_instance, model_name, model_type
)
success, msg = models_api_service.set_tenant_default_models(tenant_id, model_provider, model_instance, model_name, model_type)
if success:
logging.info(f"success: {success}, msg: {msg}")
return get_result(message=msg)
api/apps/restful_apis/openai_api.py
+2
View File
@@ -54,6 +54,7 @@ def _validate_llm_id(llm_id, tenant_id, llm_setting=None):
import logging
from api.utils.reference_metadata_utils import enrich_chunks_with_document_metadata
def _build_reference_chunks(reference, include_metadata=False, metadata_fields=None):
chunks = chunks_format(reference)
if not include_metadata:
@@ -195,6 +196,7 @@ async def _stream_chat_completion_sse(
yield f"data:{json.dumps(response, ensure_ascii=False)}\n\n"
yield "data:[DONE]\n\n"
def _normalize_message_content(content):
"""Convert OpenAI message content to a string for the dialog layer.
api/apps/restful_apis/plugin_api.py
+1 -1
View File
@@ -21,7 +21,7 @@ from api.utils.api_utils import get_json_result
from agent.plugin import GlobalPluginManager
@manager.route('/plugin/tools', methods=['GET']) # noqa: F821
@manager.route("/plugin/tools", methods=["GET"]) # noqa: F821
@login_required
def llm_tools() -> Response:
tools = GlobalPluginManager.get_llm_tools()
api/apps/restful_apis/provider_api.py
+14 -19
View File
@@ -618,9 +618,7 @@ def list_instance_models(tenant_id: str = None, provider_id_or_name: str = None,
"""
supported_only = request.args.get("supported", "").lower() == "true"
try:
success, result = provider_api_service.list_instance_models(
tenant_id, provider_id_or_name, instance_id_or_name, supported_only
)
success, result = provider_api_service.list_instance_models(tenant_id, provider_id_or_name, instance_id_or_name, supported_only)
if success:
return get_result(data=result)
else:
@@ -680,9 +678,7 @@ async def update_instance_models(tenant_id: str, provider_id_or_name: str, insta
model_name = data["model_name"]
model_type = data["model_type"]
try:
success, msg = provider_api_service.update_instance_models(
tenant_id, provider_id_or_name, instance_id_or_name, model_name, model_type
)
success, msg = provider_api_service.update_instance_models(tenant_id, provider_id_or_name, instance_id_or_name, model_name, model_type)
if success:
return get_result(message=msg)
else:
@@ -755,9 +751,7 @@ async def add_model_to_instance(tenant_id: str, provider_id_or_name: str, instan
extra = data.get("extra", {})
try:
success, result = provider_api_service.add_model_to_instance(
tenant_id, provider_id_or_name, instance_id_or_name, model_name, model_type, max_tokens, extra
)
success, result = provider_api_service.add_model_to_instance(tenant_id, provider_id_or_name, instance_id_or_name, model_name, model_type, max_tokens, extra)
if success:
return get_result(message=result)
else:
@@ -827,9 +821,7 @@ async def enable_or_disable_model(tenant_id: str = None, provider_id_or_name: st
return get_error_argument_result(message="status must be 'active' or 'inactive'")
try:
success, msg = provider_api_service.update_model_status(
tenant_id, provider_id_or_name, instance_id_or_name, model_name, status
)
success, msg = provider_api_service.update_model_status(tenant_id, provider_id_or_name, instance_id_or_name, model_name, status)
if success:
return get_result(message=msg)
else:
@@ -904,15 +896,14 @@ async def chat_to_model(tenant_id: str = None, provider_id_or_name: str = None,
thinking = data.get("thinking", False)
try:
success, result = await provider_api_service.chat_to_model(
tenant_id, provider_id_or_name, instance_id_or_name, model_name, message, stream, thinking
)
success, result = await provider_api_service.chat_to_model(tenant_id, provider_id_or_name, instance_id_or_name, model_name, message, stream, thinking)
if not success:
return get_error_data_result(message=result)
if stream and isinstance(result, dict) and result.get("type") == "stream":
# Streaming response using SSE
from quart import Response
llm = result["llm"]
async def generate():
@@ -925,10 +916,14 @@ async def chat_to_model(tenant_id: str = None, provider_id_or_name: str = None,
yield f"data: [MESSAGE]{chunk}\n\n"
yield "data: [DONE]\n\n"
return Response(generate(), mimetype="text/event-stream", headers={
"Cache-Control": "no-cache",
"Connection": "keep-alive",
})
return Response(
generate(),
mimetype="text/event-stream",
headers={
"Cache-Control": "no-cache",
"Connection": "keep-alive",
},
)
# Non-streaming response
return get_result(data=result)
api/apps/restful_apis/search_api.py
+11 -10
View File
@@ -92,7 +92,7 @@ def list_searches():
search_apps = [s for s in search_apps if s["tenant_id"] in owner_ids]
total = len(search_apps)
if page_number and items_per_page:
search_apps = search_apps[(page_number - 1) * items_per_page: page_number * items_per_page]
search_apps = search_apps[(page_number - 1) * items_per_page : page_number * items_per_page]
return get_json_result(data={"search_apps": search_apps, "total": total})
except Exception as e:
return server_error_response(e)
@@ -151,9 +151,7 @@ async def update(search_id):
return get_data_error_result(message="search_config must be a JSON object")
req["search_config"] = {**current_config, **new_config}
logging.debug(
"Search update weight: search_id=%s user_id=%s "
"incoming_vector_similarity_weight=%s stored_vector_similarity_weight=%s "
"stored_full_text_weight=%s",
"Search update weight: search_id=%s user_id=%s incoming_vector_similarity_weight=%s stored_vector_similarity_weight=%s stored_full_text_weight=%s",
search_id,
current_user.id,
new_config.get("vector_similarity_weight"),
@@ -212,8 +210,7 @@ async def completion(search_id):
search_config = search_app.get("search_config", {})
logging.debug(
"Search completion loaded weight: search_id=%s user_id=%s "
"stored_vector_similarity_weight=%s stored_full_text_weight=%s",
"Search completion loaded weight: search_id=%s user_id=%s stored_vector_similarity_weight=%s stored_full_text_weight=%s",
search_id,
uid,
search_config.get("vector_similarity_weight", 0.3),
@@ -229,10 +226,14 @@ async def completion(search_id):
async for ans in async_ask(req["question"], kb_ids, uid, search_config=search_config, search_id=search_id):
yield "data:" + json.dumps({"code": 0, "message": "", "data": ans}, ensure_ascii=False) + "\n\n"
except Exception as ex:
yield "data:" + json.dumps(
{"code": 500, "message": str(ex), "data": {"answer": "**ERROR**: " + str(ex), "reference": []}},
ensure_ascii=False,
) + "\n\n"
yield (
"data:"
+ json.dumps(
{"code": 500, "message": str(ex), "data": {"answer": "**ERROR**: " + str(ex), "reference": []}},
ensure_ascii=False,
)
+ "\n\n"
)
yield "data:" + json.dumps({"code": 0, "message": "", "data": True}, ensure_ascii=False) + "\n\n"
resp = Response(stream(), mimetype="text/event-stream")
api/apps/restful_apis/stats_api.py
+8 -12
View File
@@ -20,7 +20,8 @@ from api.db.services.user_service import UserTenantService
from api.utils.api_utils import get_data_error_result, get_json_result, server_error_response
from api.apps import login_required, current_user
@manager.route('/system/stats', methods=['GET']) # noqa: F821
@manager.route("/system/stats", methods=["GET"]) # noqa: F821
@login_required
def stats():
try:
@@ -29,15 +30,10 @@ def stats():
return get_data_error_result(message="Tenant not found!")
objs = API4ConversationService.stats(
tenants[0].tenant_id,
request.args.get(
"from_date",
(datetime.now() -
timedelta(
days=7)).strftime("%Y-%m-%d 00:00:00")),
request.args.get(
"to_date",
datetime.now().strftime("%Y-%m-%d %H:%M:%S")),
"agent" if "canvas_id" in request.args else None)
request.args.get("from_date", (datetime.now() - timedelta(days=7)).strftime("%Y-%m-%d 00:00:00")),
request.args.get("to_date", datetime.now().strftime("%Y-%m-%d %H:%M:%S")),
"agent" if "canvas_id" in request.args else None,
)
res = {"pv": [], "uv": [], "speed": [], "tokens": [], "round": [], "thumb_up": []}
@@ -45,8 +41,8 @@ def stats():
dt = obj["dt"]
res["pv"].append((dt, obj["pv"]))
res["uv"].append((dt, obj["uv"]))
res["speed"].append((dt, float(obj["tokens"]) / (float(obj["duration"]) + 0.1))) # +0.1 to avoid division by zero
res["tokens"].append((dt, float(obj["tokens"]) / 1000.0)) # convert to thousands
res["speed"].append((dt, float(obj["tokens"]) / (float(obj["duration"]) + 0.1))) # +0.1 to avoid division by zero
res["tokens"].append((dt, float(obj["tokens"]) / 1000.0)) # convert to thousands
res["round"].append((dt, obj["round"]))
res["thumb_up"].append((dt, obj["thumb_up"]))
api/apps/restful_apis/system_api.py
+12 -11
View File
@@ -34,10 +34,12 @@ from common.log_utils import get_log_levels, set_log_level
from common import settings
from rag.utils.redis_conn import REDIS_CONN
@manager.route("/system/ping", methods=["GET"]) # noqa: F821
async def ping():
return "pong", 200
@manager.route("/system/version", methods=["GET"]) # noqa: F821
def version():
"""
@@ -196,13 +198,7 @@ def oceanbase_status():
status_info = get_oceanbase_status()
return get_json_result(data=status_info)
except Exception as e:
return get_json_result(
data={
"status": "error",
"message": f"Failed to get OceanBase status: {str(e)}"
},
code=500
)
return get_json_result(data={"status": "error", "message": f"Failed to get OceanBase status: {str(e)}"}, code=500)
@manager.route("/system/config", methods=["GET"]) # noqa: F821
@@ -222,16 +218,20 @@ def get_config():
type: integer 0 means disabled, 1 means enabled
description: Whether user registration is enabled
"""
return get_json_result(data={
"registerEnabled": settings.REGISTER_ENABLED,
"disablePasswordLogin": settings.DISABLE_PASSWORD_LOGIN,
})
return get_json_result(
data={
"registerEnabled": settings.REGISTER_ENABLED,
"disablePasswordLogin": settings.DISABLE_PASSWORD_LOGIN,
}
)
@manager.route("/system/healthz", methods=["GET"]) # noqa: F821
def healthz():
result, all_ok = run_health_checks()
return jsonify(result), (200 if all_ok else 500)
@manager.route("/system/tokens", methods=["GET"]) # noqa: F821
@login_required
def token_list():
@@ -409,6 +409,7 @@ async def set_logger_level():
description: Log level updated successfully
"""
from quart import request
data = await request.get_json()
if not data or "pkg_name" not in data or "level" not in data:
return get_data_error_result(message="pkg_name and level are required")
api/apps/restful_apis/task_api.py
+3 -7
View File
@@ -30,8 +30,7 @@ from rag.utils.redis_conn import REDIS_CONN
@manager.route("/tasks/<task_id>/cancel", methods=["POST"]) # noqa: F821
@login_required
async def cancel_task(task_id):
"""Cancel a running task.
"""
"""Cancel a running task."""
return await _cancel_task(task_id)
@@ -77,11 +76,7 @@ async def _cancel_task(task_id):
TaskService.model.update(
progress_msg=TaskService.model.progress_msg + cancel_msg,
progress=-1,
).where(
(TaskService.model.id == task_id)
& (TaskService.model.progress >= 0)
& (TaskService.model.progress < 1)
).execute()
).where((TaskService.model.id == task_id) & (TaskService.model.progress >= 0) & (TaskService.model.progress < 1)).execute()
except Exception as e:
logging.warning("Failed to update task %s progress after cancellation: %s", task_id, str(e))
@@ -89,6 +84,7 @@ async def _cancel_task(task_id):
# cancelled so that the UI reflects the state correctly.
try:
from api.db.services.document_service import DocumentService
doc_id = task.doc_id
if doc_id and doc_id not in (CANVAS_DEBUG_DOC_ID, GRAPH_RAPTOR_FAKE_DOC_ID):
_, doc = DocumentService.get_by_id(doc_id)
api/apps/restful_apis/tenant_api.py
+1 -3
View File
@@ -82,9 +82,7 @@ async def create(tenant_id):
return get_data_error_result(message=f"{invite_user_email} is already in the team.")
if user_tenant_role == UserTenantRole.OWNER:
return get_data_error_result(message=f"{invite_user_email} is the owner of the team.")
return get_data_error_result(
message=f"{invite_user_email} is in the team, but the role: {user_tenant_role} is invalid."
)
return get_data_error_result(message=f"{invite_user_email} is in the team, but the role: {user_tenant_role} is invalid.")
UserTenantService.save(
id=get_uuid(),
api/apps/services/canvas_replica_service.py
-12
View File
@@ -42,7 +42,6 @@ class CanvasReplicaService:
LOCK_RETRY_ATTEMPTS = 3
LOCK_RETRY_SLEEP_SECS = 0.2
@classmethod
def normalize_dsl(cls, dsl):
"""Normalize DSL to a JSON-serializable dict. Raise ValueError on invalid input."""
@@ -61,17 +60,14 @@ class CanvasReplicaService:
except Exception as e:
raise ValueError("DSL is not JSON-serializable.") from e
@classmethod
def _replica_key(cls, canvas_id: str, tenant_id: str, runtime_user_id: str) -> str:
return f"{cls.REPLICA_KEY_PREFIX}:{canvas_id}:{tenant_id}:{runtime_user_id}"
@classmethod
def _lock_key(cls, canvas_id: str, tenant_id: str, runtime_user_id: str) -> str:
return f"{cls.LOCK_KEY_PREFIX}:{canvas_id}:{tenant_id}:{runtime_user_id}"
@classmethod
def _read_payload(cls, replica_key: str):
"""Read replica payload from Redis; return None on missing/invalid content."""
@@ -88,14 +84,12 @@ class CanvasReplicaService:
logging.warning("Failed to parse canvas replica %s: %s", replica_key, e)
return None
@classmethod
def _write_payload(cls, replica_key: str, payload: dict):
"""Write payload and refresh TTL."""
payload["updated_at"] = int(time.time())
REDIS_CONN.set_obj(replica_key, payload, cls.TTL_SECS)
@classmethod
def _build_payload(
cls,
@@ -116,7 +110,6 @@ class CanvasReplicaService:
"updated_at": int(time.time()),
}
@classmethod
def create_if_absent(
cls,
@@ -136,7 +129,6 @@ class CanvasReplicaService:
cls._write_payload(replica_key, payload)
return payload
@classmethod
def bootstrap(
cls,
@@ -157,14 +149,12 @@ class CanvasReplicaService:
title=title,
)
@classmethod
def load_for_run(cls, canvas_id: str, tenant_id: str, runtime_user_id: str):
"""Load current runtime replica used by /completions."""
replica_key = cls._replica_key(canvas_id, str(tenant_id), str(runtime_user_id))
return cls._read_payload(replica_key)
@classmethod
def replace_for_set(
cls,
@@ -203,7 +193,6 @@ class CanvasReplicaService:
except Exception:
logging.exception("Failed to release canvas replica lock: %s", lock_key)
@classmethod
def _acquire_lock_with_retry(cls, lock_key: str):
"""Acquire distributed lock with bounded retries; return lock object or None."""
@@ -219,7 +208,6 @@ class CanvasReplicaService:
time.sleep(cls.LOCK_RETRY_SLEEP_SECS + random.uniform(0, 0.1))
return None
@classmethod
def commit_after_run(
cls,
api/apps/services/file_api_service.py
+50 -71
View File
@@ -48,15 +48,15 @@ async def upload_file(tenant_id: str, pf_id: str, file_objs: list):
file_res = []
for file_obj in file_objs:
MAX_FILE_NUM_PER_USER = int(os.environ.get('MAX_FILE_NUM_PER_USER', 0))
MAX_FILE_NUM_PER_USER = int(os.environ.get("MAX_FILE_NUM_PER_USER", 0))
if 0 < MAX_FILE_NUM_PER_USER <= await thread_pool_exec(DocumentService.get_doc_count, tenant_id):
return False, "Exceed the maximum file number of a free user!"
if not file_obj.filename:
file_obj_names = [pf_folder.name, file_obj.filename]
else:
full_path = '/' + file_obj.filename
file_obj_names = full_path.split('/')
full_path = "/" + file_obj.filename
file_obj_names = full_path.split("/")
file_len = len(file_obj_names)
file_id_list = await thread_pool_exec(FileService.get_id_list_by_id, pf_id, file_obj_names, 1, [pf_id])
@@ -66,25 +66,19 @@ async def upload_file(tenant_id: str, pf_id: str, file_objs: list):
e, file = await thread_pool_exec(FileService.get_by_id, file_id_list[len_id_list - 1])
if not e:
return False, "Folder not found!"
last_folder = await thread_pool_exec(
FileService.create_folder, file, file_id_list[len_id_list - 1], file_obj_names, len_id_list, tenant_id, tenant_id
)
last_folder = await thread_pool_exec(FileService.create_folder, file, file_id_list[len_id_list - 1], file_obj_names, len_id_list, tenant_id, tenant_id)
else:
e, file = await thread_pool_exec(FileService.get_by_id, file_id_list[len_id_list - 2])
if not e:
return False, "Folder not found!"
last_folder = await thread_pool_exec(
FileService.create_folder, file, file_id_list[len_id_list - 2], file_obj_names, len_id_list, tenant_id, tenant_id
)
last_folder = await thread_pool_exec(FileService.create_folder, file, file_id_list[len_id_list - 2], file_obj_names, len_id_list, tenant_id, tenant_id)
filetype = filename_type(file_obj_names[file_len - 1])
location = file_obj_names[file_len - 1]
while await thread_pool_exec(settings.STORAGE_IMPL.obj_exist, last_folder.id, location):
location += "_"
blob = await thread_pool_exec(file_obj.read)
filename = await thread_pool_exec(
duplicate_name, FileService.query, name=file_obj_names[file_len - 1], parent_id=last_folder.id
)
filename = await thread_pool_exec(duplicate_name, FileService.query, name=file_obj_names[file_len - 1], parent_id=last_folder.id)
await thread_pool_exec(settings.STORAGE_IMPL.put, last_folder.id, location, blob)
file_data = {
"id": get_uuid(),
@@ -126,16 +120,18 @@ async def create_folder(tenant_id: str, name: str, pf_id: str = None, file_type:
else:
ft = FileType.VIRTUAL.value
file = FileService.insert({
"id": get_uuid(),
"parent_id": pf_id,
"tenant_id": tenant_id,
"created_by": tenant_id,
"name": name,
"location": "",
"size": 0,
"type": ft,
})
file = FileService.insert(
{
"id": get_uuid(),
"parent_id": pf_id,
"tenant_id": tenant_id,
"created_by": tenant_id,
"name": name,
"location": "",
"size": 0,
"type": ft,
}
)
return True, file.to_json()
@@ -173,7 +169,6 @@ def list_files(tenant_id: str, args: dict):
return True, {"total": total, "files": files, "parent_folder": parent_folder.to_json()}
def get_parent_folder(file_id: str, user_id: str = None):
"""
Get parent folder of a file with permission check.
@@ -235,9 +230,9 @@ async def delete_files(uid: str, file_ids: list, auth_header: str = ""):
try:
import requests
host = getattr(settings, 'HOST_IP', '127.0.0.1')
host = getattr(settings, "HOST_IP", "127.0.0.1")
# Go service runs on port+4 (9384 by default)
port = getattr(settings, 'HOST_PORT', 9380) + 4
port = getattr(settings, "HOST_PORT", 9380) + 4
service_url = f"http://{host}:{port}"
# List all spaces and find the one matching the name
@@ -270,9 +265,9 @@ async def delete_files(uid: str, file_ids: list, auth_header: str = ""):
from urllib.parse import quote
# Construct service URL from settings
host = getattr(settings, 'HOST_IP', '127.0.0.1')
host = getattr(settings, "HOST_IP", "127.0.0.1")
# Go service runs on port+4 (9384 by default)
port = getattr(settings, 'HOST_PORT', 9380) + 4
port = getattr(settings, "HOST_PORT", 9380) + 4
service_url = f"http://{host}:{port}"
# Get space UUID from space name
@@ -289,37 +284,24 @@ async def delete_files(uid: str, file_ids: list, auth_header: str = ""):
try:
data = response.json()
if data.get("code") == 0:
logging.info(
f"Successfully deleted skill index: space={space_name}, skill={skill_name}, "
f"status={response.status_code}, code=0"
)
logging.info(f"Successfully deleted skill index: space={space_name}, skill={skill_name}, status={response.status_code}, code=0")
return True
else:
app_code = data.get("code", "unknown")
app_msg = data.get("message", "no message")
logging.error(
f"Failed to delete skill index: space={space_name}, skill={skill_name}, "
f"status={response.status_code}, app_code={app_code}, app_msg={app_msg}, "
f"response={response.text}"
f"Failed to delete skill index: space={space_name}, skill={skill_name}, status={response.status_code}, app_code={app_code}, app_msg={app_msg}, response={response.text}"
)
return False
except ValueError as json_err:
# JSON decode error - treat as failure
logging.error(
f"Failed to parse delete response JSON: space={space_name}, skill={skill_name}, "
f"error={json_err}, raw_response={response.text}"
)
logging.error(f"Failed to parse delete response JSON: space={space_name}, skill={skill_name}, error={json_err}, raw_response={response.text}")
return False
else:
logging.error(
f"Failed to delete skill index: space={space_name}, skill={skill_name}, "
f"status={response.status_code}, response={response.text}"
)
logging.error(f"Failed to delete skill index: space={space_name}, skill={skill_name}, status={response.status_code}, response={response.text}")
return False
except Exception as e:
logging.error(
f"Exception deleting skill index: space={space_name}, skill={skill_name}, error={e}"
)
logging.error(f"Exception deleting skill index: space={space_name}, skill={skill_name}, error={e}")
return False
def _delete_single_file(file) -> int:
@@ -408,18 +390,12 @@ async def delete_files(uid: str, file_ids: list, auth_header: str = ""):
logging.info(f"Deleting skill index for skill '{folder.name}' in space '{current_space_name}'")
index_deleted = _delete_skill_index(tenant_id, current_space_name, folder.name, auth_header)
if not index_deleted:
logging.error(
f"Aborting folder deletion due to index deletion failure: "
f"folder={folder.name}, space={current_space_name}"
)
errors.append(
f"Failed to delete skill index for folder '{folder.name}' in space '{current_space_name}'. "
f"Folder deletion aborted to prevent orphaned indexes."
)
logging.error(f"Aborting folder deletion due to index deletion failure: folder={folder.name}, space={current_space_name}")
errors.append(f"Failed to delete skill index for folder '{folder.name}' in space '{current_space_name}'. Folder deletion aborted to prevent orphaned indexes.")
return deleted
sub_files = FileService.list_all_files_by_parent_id(folder.id)
logging.info(f"Folder '{folder.name}': found {len(sub_files)} children to delete")
for sub_file in sub_files:
if sub_file.type == FileType.FOLDER.value:
deleted += _delete_folder_recursive(sub_file, tenant_id)
@@ -432,16 +408,16 @@ async def delete_files(uid: str, file_ids: list, auth_header: str = ""):
errors.append(f"Failed to delete folder record {folder.id}: {e}")
else:
deleted += 1
try:
if hasattr(settings.STORAGE_IMPL, 'remove_bucket'):
if hasattr(settings.STORAGE_IMPL, "remove_bucket"):
logging.info(f"Removing storage bucket for folder '{folder.name}' (id={folder.id})")
settings.STORAGE_IMPL.remove_bucket(folder.id)
else:
logging.debug(f"Storage implementation does not support remove_bucket, skipping for folder '{folder.name}'")
except Exception as e:
logging.warning(f"Failed to remove storage bucket for folder '{folder.name}' (id={folder.id}): {e}")
return deleted
def _rm_sync():
@@ -513,8 +489,7 @@ async def move_files(uid: str, src_file_ids: list, dest_file_id: str = None, new
if new_name:
file = files_dict[src_file_ids[0]]
if file.type != FileType.FOLDER.value and \
pathlib.Path(new_name.lower()).suffix != pathlib.Path(file.name.lower()).suffix:
if file.type != FileType.FOLDER.value and pathlib.Path(new_name.lower()).suffix != pathlib.Path(file.name.lower()).suffix:
return False, "The extension of file can't be changed"
target_parent_id = dest_folder.id if dest_folder else file.parent_id
for f in FileService.query(name=new_name, parent_id=target_parent_id):
@@ -541,16 +516,18 @@ async def move_files(uid: str, src_file_ids: list, dest_file_id: str = None, new
if existing_folder:
new_folder = existing_folder[0]
else:
new_folder = FileService.insert({
"id": get_uuid(),
"parent_id": dest_folder_entry.id,
"tenant_id": source_file_entry.tenant_id,
"created_by": source_file_entry.tenant_id,
"name": effective_name,
"location": "",
"size": 0,
"type": FileType.FOLDER.value,
})
new_folder = FileService.insert(
{
"id": get_uuid(),
"parent_id": dest_folder_entry.id,
"tenant_id": source_file_entry.tenant_id,
"created_by": source_file_entry.tenant_id,
"name": effective_name,
"location": "",
"size": 0,
"type": FileType.FOLDER.value,
}
)
sub_files = FileService.list_all_files_by_parent_id(source_file_entry.id)
for sub_file in sub_files:
@@ -569,8 +546,10 @@ async def move_files(uid: str, src_file_ids: list, dest_file_id: str = None, new
new_location += "_"
try:
moved = settings.STORAGE_IMPL.move(
source_file_entry.parent_id, source_file_entry.location,
dest_folder_entry.id, new_location,
source_file_entry.parent_id,
source_file_entry.location,
dest_folder_entry.id,
new_location,
)
except Exception as storage_err:
raise RuntimeError(f"Move file failed at storage layer: {str(storage_err)}")
api/apps/services/memory_api_service.py
+9 -30
View File
@@ -98,13 +98,7 @@ async def create_memory(memory_info: dict):
if invalid_type:
raise ArgumentException(f"Memory type '{invalid_type}' is not supported.")
memory_type = list(memory_type)
success, res = MemoryService.create_memory(
tenant_id=current_user.id,
name=memory_name,
memory_type=memory_type,
embd_id=memory_info["embd_id"],
llm_id=memory_info["llm_id"]
)
success, res = MemoryService.create_memory(tenant_id=current_user.id, name=memory_name, memory_type=memory_type, embd_id=memory_info["embd_id"], llm_id=memory_info["llm_id"])
if success:
return True, format_ret_data_from_memory(res)
else:
@@ -243,7 +237,7 @@ async def delete_memory(memory_id):
return True
async def list_memory(filter_params: dict, keywords: str, page: int=1, page_size: int = 50):
async def list_memory(filter_params: dict, keywords: str, page: int = 1, page_size: int = 50):
"""
:param filter_params: {
"memory_type": list[str],
@@ -268,9 +262,7 @@ async def list_memory(filter_params: dict, keywords: str, page: int=1, page_size
memory_list, count = MemoryService.get_by_filter(filter_dict, keywords, page, page_size)
[memory.update({"memory_type": get_memory_type_human(memory["memory_type"])}) for memory in memory_list]
return {
"memory_list": memory_list, "total_count": count
}
return {"memory_list": memory_list, "total_count": count}
async def get_memory_config(memory_id):
@@ -280,10 +272,9 @@ async def get_memory_config(memory_id):
return format_ret_data_from_memory(memory)
async def get_memory_messages(memory_id, agent_ids: list[str], keywords: str, page: int=1, page_size: int = 50):
async def get_memory_messages(memory_id, agent_ids: list[str], keywords: str, page: int = 1, page_size: int = 50):
memory = _require_memory_access(memory_id)
messages = MessageService.list_message(
memory.tenant_id, memory_id, agent_ids, keywords, page, page_size)
messages = MessageService.list_message(memory.tenant_id, memory_id, agent_ids, keywords, page, page_size)
agent_name_mapping = {}
extract_task_mapping = {}
if messages["message_list"]:
@@ -291,7 +282,7 @@ async def get_memory_messages(memory_id, agent_ids: list[str], keywords: str, pa
agent_name_mapping = {agent["id"]: agent["title"] for agent in agent_list}
task_list = TaskService.get_tasks_progress_by_doc_ids([memory_id])
if task_list:
task_list.sort(key=lambda t: t["create_time"]) # asc, use newer when exist more than one task
task_list.sort(key=lambda t: t["create_time"]) # asc, use newer when exist more than one task
for task in task_list:
# the 'digest' field carries the source_id when a task is created, so use 'digest' as key
extract_task_mapping.update({int(task["digest"]): task})
@@ -324,10 +315,7 @@ async def forget_message(memory_id: str, message_id: int):
memory = _require_memory_access(memory_id)
forget_time = timestamp_to_date(current_timestamp())
update_succeed = MessageService.update_message(
{"memory_id": memory_id, "message_id": int(message_id)},
{"forget_at": forget_time},
memory.tenant_id, memory_id)
update_succeed = MessageService.update_message({"memory_id": memory_id, "message_id": int(message_id)}, {"forget_at": forget_time}, memory.tenant_id, memory_id)
if update_succeed:
return True
raise Exception(f"Failed to forget message '{message_id}' in memory '{memory_id}'.")
@@ -336,10 +324,7 @@ async def forget_message(memory_id: str, message_id: int):
async def update_message_status(memory_id: str, message_id: int, status: bool):
memory = _require_memory_access(memory_id)
update_succeed = MessageService.update_message(
{"memory_id": memory_id, "message_id": int(message_id)},
{"status": status},
memory.tenant_id, memory_id)
update_succeed = MessageService.update_message({"memory_id": memory_id, "message_id": int(message_id)}, {"status": status}, memory.tenant_id, memory_id)
if update_succeed:
return True
raise Exception(f"Failed to set status for message '{message_id}' in memory '{memory_id}'.")
@@ -383,13 +368,7 @@ async def get_messages(memory_ids: list[str], agent_id: str = "", session_id: st
return []
uids = [memory.tenant_id for memory in memory_list]
accessible_memory_ids = [memory.id for memory in memory_list]
res = MessageService.get_recent_messages(
uids,
accessible_memory_ids,
agent_id,
session_id,
limit
)
res = MessageService.get_recent_messages(uids, accessible_memory_ids, agent_id, session_id, limit)
return res
api/apps/services/models_api_service.py
+42 -52
View File
@@ -98,11 +98,7 @@ def _get_model_info(tenant_id: str, default_model: str, model_type: str):
# Special case: TEI Builtin embedding model
compose_profiles = os.getenv("COMPOSE_PROFILES", "")
tei_model = os.getenv("TEI_MODEL", "")
if (model_type == "embedding"
and "tei-" in compose_profiles
and tei_model
and model_name == tei_model
and (not provider_name or provider_name == "Builtin")):
if model_type == "embedding" and "tei-" in compose_profiles and tei_model and model_name == tei_model and (not provider_name or provider_name == "Builtin"):
return {
"model_provider": "Builtin",
"model_instance": "default",
@@ -124,9 +120,7 @@ def _get_model_info(tenant_id: str, default_model: str, model_type: str):
return None
# Check if model is enabled (no TenantModel record or status != inactive means enabled)
model_entity = TenantModelService.get_by_provider_id_and_instance_id_and_model_type_and_model_name(
provider_obj.id, instance_obj.id, model_type, model_name
)
model_entity = TenantModelService.get_by_provider_id_and_instance_id_and_model_type_and_model_name(provider_obj.id, instance_obj.id, model_type, model_name)
enable = model_entity is None or model_entity.status == ActiveStatusEnum.ACTIVE.value
if not enable:
@@ -134,12 +128,12 @@ def _get_model_info(tenant_id: str, default_model: str, model_type: str):
if model_entity:
return {
"model_provider": provider_name,
"model_instance": instance_name,
"model_name": model_name,
"model_type": model_type,
"enable": enable,
}
"model_provider": provider_name,
"model_instance": instance_name,
"model_name": model_name,
"model_type": model_type,
"enable": enable,
}
# Check if model is in the LLM factory info
factory_info = [f for f in (FACTORY_LLM_INFOS or []) if f["name"] == provider_name]
@@ -185,12 +179,7 @@ def _check_model_available(tenant_id: str, provider_name: str, instance_name: st
return True, None
compose_profiles = os.getenv("COMPOSE_PROFILES", "")
is_tei_builtin_embedding = (
model_type == LLMType.EMBEDDING.value
and "tei-" in compose_profiles
and model_name == os.getenv("TEI_MODEL", "")
and (provider_name == "Builtin" or not provider_name)
)
is_tei_builtin_embedding = model_type == LLMType.EMBEDDING.value and "tei-" in compose_profiles and model_name == os.getenv("TEI_MODEL", "") and (provider_name == "Builtin" or not provider_name)
if is_tei_builtin_embedding:
return True, None
@@ -210,9 +199,7 @@ def _check_model_available(tenant_id: str, provider_name: str, instance_name: st
return False, f"Provider '{provider_name}' not found in factory info"
model_type = MODEL_TAG_TO_TYPE.get(model_type, model_type)
# Check if model is disabled
model_entity = TenantModelService.get_by_provider_id_and_instance_id_and_model_type_and_model_name(
provider_obj.id, instance_obj.id, model_type, model_name
)
model_entity = TenantModelService.get_by_provider_id_and_instance_id_and_model_type_and_model_name(provider_obj.id, instance_obj.id, model_type, model_name)
if model_entity:
if model_entity.status != ActiveStatusEnum.ACTIVE.value:
return False, f"Model '{model_name}' isn't available"
@@ -297,7 +284,7 @@ def set_tenant_default_models(tenant_id: str, model_provider: str, model_instanc
return True, "success"
def list_tenant_added_models(tenant_id: str, model_type_filter: str=None):
def list_tenant_added_models(tenant_id: str, model_type_filter: str = None):
"""
List all added models for a tenant.
@@ -369,14 +356,16 @@ def list_tenant_added_models(tenant_id: str, model_type_filter: str=None):
if not model_types:
continue
added_models.append({
"model_type": model_types,
"name": llm["llm_name"],
"provider_id": factory_instance.provider_id,
"provider_name": provider_info_map[factory_instance.provider_id].provider_name if provider_info_map.get(factory_instance.provider_id) else "",
"instance_id": factory_instance.id,
"instance_name": factory_instance.instance_name
})
added_models.append(
{
"model_type": model_types,
"name": llm["llm_name"],
"provider_id": factory_instance.provider_id,
"provider_name": provider_info_map[factory_instance.provider_id].provider_name if provider_info_map.get(factory_instance.provider_id) else "",
"instance_id": factory_instance.id,
"instance_name": factory_instance.instance_name,
}
)
manual_added_model_record_keys = list(set(model_record_map.keys()) - set(model_key_in_factory))
if manual_added_model_record_keys:
@@ -390,33 +379,34 @@ def list_tenant_added_models(tenant_id: str, model_type_filter: str=None):
if not model_types:
continue
added_models.append({
"model_type": model_types,
"name": model_name,
"provider_id": provider_id,
"provider_name": provider_info_map[provider_id].provider_name if provider_info_map.get(provider_id) else "",
"instance_id": instance_id,
"instance_name": instance_info_map[instance_id].instance_name if instance_info_map.get(instance_id) else ""
})
added_models.append(
{
"model_type": model_types,
"name": model_name,
"provider_id": provider_id,
"provider_name": provider_info_map[provider_id].provider_name if provider_info_map.get(provider_id) else "",
"instance_id": instance_id,
"instance_name": instance_info_map[instance_id].instance_name if instance_info_map.get(instance_id) else "",
}
)
# Add TEI Builtin embedding model if configured
compose_profiles = os.getenv("COMPOSE_PROFILES", "")
tei_model = os.getenv("TEI_MODEL", "")
if "tei-" in compose_profiles and tei_model:
if not model_type_filter or model_type_filter == "embedding":
tei_already_added = any(
m["provider_name"] == "Builtin" and m["name"] == tei_model
for m in added_models
)
tei_already_added = any(m["provider_name"] == "Builtin" and m["name"] == tei_model for m in added_models)
if not tei_already_added:
added_models.append({
"model_type": ["embedding"],
"name": tei_model,
"provider_id": "",
"provider_name": "Builtin",
"instance_id": "",
"instance_name": "default",
})
added_models.append(
{
"model_type": ["embedding"],
"name": tei_model,
"provider_id": "",
"provider_name": "Builtin",
"instance_id": "",
"instance_name": "default",
}
)
added_models.sort(key=lambda x: (factory_rank_mapping.get(x["provider_name"]), x["provider_name"], x["instance_name"]))

Some files were not shown because too many files have changed in this diff Show More