diff --git a/node/server/grpc_server.py b/node/server/grpc_server.py index a0dc241..67b70c3 100644 --- a/node/server/grpc_server.py +++ b/node/server/grpc_server.py @@ -34,9 +34,12 @@ ToolRunInput, EnvironmentRunInput, KBRunInput, - ModuleExecutionType + ModuleExecutionType, + SecretInput ) +from node.secret import Secret from node.module_manager import setup_module_deployment +from node.utils import get_logger load_dotenv() @@ -44,7 +47,7 @@ root_dir = file_path.parent.parent.parent MODULES_SOURCE_DIR = root_dir / os.getenv("MODULES_SOURCE_DIR") -logger = logging.getLogger(__name__) +logger = get_logger(__name__) class GrpcServerServicer(grpc_server_pb2_grpc.GrpcServerServicer): async def CheckUser(self, request, context): @@ -112,6 +115,17 @@ async def create_module(self, module_deployment: Union[AgentDeployment, MemoryDe async def RunModule(self, request, context): try: module_type = request.module_type + if len(request.secrets) > 0: + secrets = [SecretInput( + user_id=secret.user_id, + secret_value=secret.secret_value, + key_name=secret.key_name + ) for secret in request.secrets] + else: + secrets = [] + + request_dict = MessageToDict(request, preserving_proto_field_name=True) + module_configs = { "agent": { "input_class": AgentRunInput, @@ -196,8 +210,19 @@ async def RunModule(self, request, context): else: execution_type = run_input.deployment.module.execution_type + # Process secrets for environment variables + user_env_data = {} + secret_obj = Secret() + if secrets: + for secret in secrets: + user_env_data[secret.key_name] = secret_obj.decrypt_rsa(secret.secret_value) + if execution_type == ModuleExecutionType.package or execution_type == "package": - task = config["worker"].delay(module_run_data) + task = config["worker"].delay( + module_run_data, + user_env_data, + [secret.model_dump() for secret in secrets] if secrets else [] + ) elif execution_type == ModuleExecutionType.docker or execution_type == "docker": if config["docker_support"]: task = execute_docker_agent.delay(module_run_data) diff --git a/node/server/grpc_server_pb2.py b/node/server/grpc_server_pb2.py index de459e4..051f56c 100644 --- a/node/server/grpc_server_pb2.py +++ b/node/server/grpc_server_pb2.py @@ -2,7 +2,7 @@ # Generated by the protocol buffer compiler. DO NOT EDIT! # NO CHECKED-IN PROTOBUF GENCODE # source: grpc_server.proto -# Protobuf Python Version: 5.27.2 +# Protobuf Python Version: 5.29.0 """Generated protocol buffer code.""" from google.protobuf import descriptor as _descriptor from google.protobuf import descriptor_pool as _descriptor_pool @@ -12,8 +12,8 @@ _runtime_version.ValidateProtobufRuntimeVersion( _runtime_version.Domain.PUBLIC, 5, - 27, - 2, + 29, + 0, '', 'grpc_server.proto' ) @@ -26,7 +26,7 @@ from google.protobuf import empty_pb2 as google_dot_protobuf_dot_empty__pb2 -DESCRIPTOR = _descriptor_pool.Default().AddSerializedFile(b'\n\x11grpc_server.proto\x12\x05\x61gent\x1a\x1cgoogle/protobuf/struct.proto\x1a\x1bgoogle/protobuf/empty.proto\".\n\x0fGeneralResponse\x12\n\n\x02ok\x18\x01 \x01(\x08\x12\x0f\n\x07message\x18\x02 \x01(\t\"\xe9\x01\n\tLLMConfig\x12\x18\n\x0b\x63onfig_name\x18\x01 \x01(\tH\x00\x88\x01\x01\x12\x13\n\x06\x63lient\x18\x02 \x01(\tH\x01\x88\x01\x01\x12\x12\n\x05model\x18\x03 \x01(\tH\x02\x88\x01\x01\x12\x17\n\nmax_tokens\x18\x04 \x01(\x05H\x03\x88\x01\x01\x12\x18\n\x0btemperature\x18\x05 \x01(\x02H\x04\x88\x01\x01\x12\x15\n\x08\x61pi_base\x18\x06 \x01(\tH\x05\x88\x01\x01\x42\x0e\n\x0c_config_nameB\t\n\x07_clientB\x08\n\x06_modelB\r\n\x0b_max_tokensB\x0e\n\x0c_temperatureB\x0b\n\t_api_base\"K\n\nNodeServer\x12\x1e\n\x16\x63ommunication_protocol\x18\x01 \x01(\t\x12\x0c\n\x04port\x18\x02 \x01(\x05\x12\x0f\n\x07node_id\x18\x03 \x01(\t\"\xa9\x01\n\x0fNodeConfigInput\x12\n\n\x02ip\x18\x01 \x01(\t\x12$\n\x17user_communication_port\x18\x02 \x01(\x05H\x00\x88\x01\x01\x12(\n\x1buser_communication_protocol\x18\x03 \x01(\tH\x01\x88\x01\x01\x42\x1a\n\x18_user_communication_portB\x1e\n\x1c_user_communication_protocol\"\xc5\x04\n\nNodeConfig\x12\n\n\x02id\x18\x01 \x01(\t\x12\r\n\x05owner\x18\x02 \x01(\t\x12\x12\n\npublic_key\x18\x03 \x01(\t\x12\n\n\x02ip\x18\x04 \x01(\t\x12#\n\x1buser_communication_protocol\x18\x05 \x01(\t\x12#\n\x1bnode_communication_protocol\x18\x06 \x01(\t\x12\x1f\n\x17user_communication_port\x18\x07 \x01(\x05\x12&\n\x1enum_node_communication_servers\x18\x08 \x01(\x05\x12\"\n\x07servers\x18\t \x03(\x0b\x32\x11.agent.NodeServer\x12\x15\n\rollama_models\x18\n \x03(\t\x12\x13\n\x0b\x64ocker_jobs\x18\x0b \x01(\x08\x12\r\n\x05ports\x18\x0c \x03(\x05\x12\x19\n\x0crouting_type\x18\r \x01(\tH\x00\x88\x01\x01\x12\x18\n\x0brouting_url\x18\x0e \x01(\tH\x01\x88\x01\x01\x12\x15\n\x08num_gpus\x18\x0f \x01(\x05H\x02\x88\x01\x01\x12\x11\n\x04\x61rch\x18\x10 \x01(\tH\x03\x88\x01\x01\x12\x0f\n\x02os\x18\x11 \x01(\tH\x04\x88\x01\x01\x12\x10\n\x03ram\x18\x12 \x01(\x03H\x05\x88\x01\x01\x12\x11\n\x04vram\x18\x13 \x01(\x03H\x06\x88\x01\x01\x12\x16\n\x0eprovider_types\x18\x14 \x03(\t\x12\x0e\n\x06models\x18\x15 \x03(\tB\x0f\n\r_routing_typeB\x0e\n\x0c_routing_urlB\x0b\n\t_num_gpusB\x07\n\x05_archB\x05\n\x03_osB\x06\n\x04_ramB\x07\n\x05_vram\"\x9b\x02\n\x06Module\x12\n\n\x02id\x18\x01 \x01(\t\x12\x0c\n\x04name\x18\x02 \x01(\t\x12\x13\n\x0b\x64\x65scription\x18\x03 \x01(\t\x12\x0e\n\x06\x61uthor\x18\x04 \x01(\t\x12\x12\n\nmodule_url\x18\x05 \x01(\t\x12\x18\n\x0bmodule_type\x18\x06 \x01(\tH\x00\x88\x01\x01\x12\x1b\n\x0emodule_version\x18\x07 \x01(\tH\x01\x88\x01\x01\x12\x1e\n\x11module_entrypoint\x18\x08 \x01(\tH\x02\x88\x01\x01\x12\x1b\n\x0e\x65xecution_type\x18\t \x01(\tH\x03\x88\x01\x01\x42\x0e\n\x0c_module_typeB\x11\n\x0f_module_versionB\x14\n\x12_module_entrypointB\x11\n\x0f_execution_type\"\xd0\x02\n\x14\x44\x61taGenerationConfig\x12\x19\n\x0csave_outputs\x18\x01 \x01(\x08H\x00\x88\x01\x01\x12\"\n\x15save_outputs_location\x18\x02 \x01(\tH\x01\x88\x01\x01\x12\x1e\n\x11save_outputs_path\x18\x03 \x01(\tH\x02\x88\x01\x01\x12\x18\n\x0bsave_inputs\x18\x04 \x01(\x08H\x03\x88\x01\x01\x12!\n\x14save_inputs_location\x18\x05 \x01(\tH\x04\x88\x01\x01\x12\x1d\n\x10\x64\x65\x66\x61ult_filename\x18\x06 \x01(\tH\x05\x88\x01\x01\x42\x0f\n\r_save_outputsB\x18\n\x16_save_outputs_locationB\x14\n\x12_save_outputs_pathB\x0e\n\x0c_save_inputsB\x17\n\x15_save_inputs_locationB\x13\n\x11_default_filename\"\xdb\x01\n\x0e\x42\x61seDeployment\x12(\n\x0bnode_config\x18\x01 \x01(\x0b\x32\x11.agent.NodeConfigH\x00\x12,\n\nnode_input\x18\x02 \x01(\x0b\x32\x16.agent.NodeConfigInputH\x00\x12\x0c\n\x04name\x18\x03 \x01(\t\x12\x1d\n\x06module\x18\x04 \x01(\x0b\x32\r.agent.Module\x12\'\n\x06\x63onfig\x18\x05 \x01(\x0b\x32\x17.google.protobuf.Struct\x12\x13\n\x0binitialized\x18\x06 \x01(\x08\x42\x06\n\x04node\"\xb8\x02\n\x0eToolDeployment\x12(\n\x0bnode_config\x18\x01 \x01(\x0b\x32\x11.agent.NodeConfigH\x00\x12,\n\nnode_input\x18\x02 \x01(\x0b\x32\x16.agent.NodeConfigInputH\x00\x12\x0c\n\x04name\x18\x03 \x01(\t\x12\x1d\n\x06module\x18\x04 \x01(\x0b\x32\r.agent.Module\x12\'\n\x06\x63onfig\x18\x05 \x01(\x0b\x32\x17.google.protobuf.Struct\x12@\n\x16\x64\x61ta_generation_config\x18\x06 \x01(\x0b\x32\x1b.agent.DataGenerationConfigH\x01\x88\x01\x01\x12\x13\n\x0binitialized\x18\x07 \x01(\x08\x42\x06\n\x04nodeB\x19\n\x17_data_generation_config\"\xd1\x03\n\x0f\x41gentDeployment\x12(\n\x0bnode_config\x18\x01 \x01(\x0b\x32\x11.agent.NodeConfigH\x00\x12,\n\nnode_input\x18\x02 \x01(\x0b\x32\x16.agent.NodeConfigInputH\x00\x12\x0c\n\x04name\x18\x03 \x01(\t\x12\x1d\n\x06module\x18\x04 \x01(\x0b\x32\r.agent.Module\x12\'\n\x06\x63onfig\x18\x05 \x01(\x0b\x32\x17.google.protobuf.Struct\x12@\n\x16\x64\x61ta_generation_config\x18\x06 \x01(\x0b\x32\x1b.agent.DataGenerationConfigH\x01\x88\x01\x01\x12/\n\x10tool_deployments\x18\x07 \x03(\x0b\x32\x15.agent.ToolDeployment\x12-\n\x0ekb_deployments\x18\x08 \x03(\x0b\x32\x15.agent.BaseDeployment\x12\x36\n\x17\x65nvironment_deployments\x18\t \x03(\x0b\x32\x15.agent.BaseDeployment\x12\x13\n\x0binitialized\x18\n \x01(\x08\x42\x06\n\x04nodeB\x19\n\x17_data_generation_config\"\xc9\x03\n\x10ModuleRunRequest\x12\x13\n\x0bmodule_type\x18\x01 \x01(\t\x12\x13\n\x0b\x63onsumer_id\x18\x02 \x01(\t\x12\'\n\x06inputs\x18\x03 \x01(\x0b\x32\x17.google.protobuf.Struct\x12\x32\n\x10\x61gent_deployment\x18\x04 \x01(\x0b\x32\x16.agent.AgentDeploymentH\x00\x12\x30\n\x0ftool_deployment\x18\x05 \x01(\x0b\x32\x15.agent.ToolDeploymentH\x00\x12\x32\n\x11memory_deployment\x18\x06 \x01(\x0b\x32\x15.agent.BaseDeploymentH\x00\x12.\n\rkb_deployment\x18\x07 \x01(\x0b\x32\x15.agent.BaseDeploymentH\x00\x12\x37\n\x16\x65nvironment_deployment\x18\x08 \x01(\x0b\x32\x15.agent.BaseDeploymentH\x00\x12+\n\x11orchestrator_runs\x18\t \x03(\x0b\x32\x10.agent.ModuleRun\x12\x16\n\tsignature\x18\n \x01(\tH\x01\x88\x01\x01\x42\x0c\n\ndeploymentB\x0c\n\n_signature\"\xb6\x06\n\tModuleRun\x12\x13\n\x0bmodule_type\x18\x01 \x01(\t\x12\x13\n\x0b\x63onsumer_id\x18\x02 \x01(\t\x12\'\n\x06inputs\x18\x03 \x01(\x0b\x32\x17.google.protobuf.Struct\x12\x32\n\x10\x61gent_deployment\x18\x04 \x01(\x0b\x32\x16.agent.AgentDeploymentH\x00\x12\x30\n\x0ftool_deployment\x18\x05 \x01(\x0b\x32\x15.agent.ToolDeploymentH\x00\x12\x32\n\x11memory_deployment\x18\x06 \x01(\x0b\x32\x15.agent.BaseDeploymentH\x00\x12.\n\rkb_deployment\x18\x07 \x01(\x0b\x32\x15.agent.BaseDeploymentH\x00\x12\x37\n\x16\x65nvironment_deployment\x18\x08 \x01(\x0b\x32\x15.agent.BaseDeploymentH\x00\x12+\n\x11orchestrator_runs\x18\t \x03(\x0b\x32\x10.agent.ModuleRun\x12\x0e\n\x06status\x18\n \x01(\t\x12\r\n\x05\x65rror\x18\x0b \x01(\x08\x12\x0f\n\x02id\x18\x0c \x01(\tH\x01\x88\x01\x01\x12\x0f\n\x07results\x18\r \x03(\t\x12\x1a\n\rerror_message\x18\x0e \x01(\tH\x02\x88\x01\x01\x12\x19\n\x0c\x63reated_time\x18\x0f \x01(\tH\x03\x88\x01\x01\x12\"\n\x15start_processing_time\x18\x10 \x01(\tH\x04\x88\x01\x01\x12\x1b\n\x0e\x63ompleted_time\x18\x11 \x01(\tH\x05\x88\x01\x01\x12\x15\n\x08\x64uration\x18\x12 \x01(\x01H\x06\x88\x01\x01\x12#\n\x16input_schema_ipfs_hash\x18\x13 \x01(\tH\x07\x88\x01\x01\x12\x16\n\tsignature\x18\x14 \x01(\tH\x08\x88\x01\x01\x42\x0c\n\ndeploymentB\x05\n\x03_idB\x10\n\x0e_error_messageB\x0f\n\r_created_timeB\x18\n\x16_start_processing_timeB\x11\n\x0f_completed_timeB\x0b\n\t_durationB\x19\n\x17_input_schema_ipfs_hashB\x0c\n\n_signature\"5\n\x0eModuleRunCheck\x12\x13\n\x0bmodule_type\x18\x01 \x01(\t\x12\x0e\n\x06run_id\x18\x02 \x01(\t\"H\n\x10\x43heckUserRequest\x12\x14\n\x07user_id\x18\x01 \x01(\tH\x00\x88\x01\x01\x12\x12\n\npublic_key\x18\x02 \x01(\tB\n\n\x08_user_id\"J\n\x11\x43heckUserResponse\x12\x15\n\ris_registered\x18\x01 \x01(\x08\x12\n\n\x02id\x18\x02 \x01(\t\x12\x12\n\npublic_key\x18\x03 \x01(\t\")\n\x13RegisterUserRequest\x12\x12\n\npublic_key\x18\x01 \x01(\t\"6\n\x14RegisterUserResponse\x12\n\n\x02id\x18\x01 \x01(\t\x12\x12\n\npublic_key\x18\x02 \x01(\t2\x8a\x03\n\nGrpcServer\x12<\n\x08is_alive\x12\x16.google.protobuf.Empty\x1a\x16.agent.GeneralResponse\"\x00\x12\x38\n\x04stop\x12\x16.google.protobuf.Empty\x1a\x16.agent.GeneralResponse\"\x00\x12@\n\tCheckUser\x12\x17.agent.CheckUserRequest\x1a\x18.agent.CheckUserResponse\"\x00\x12I\n\x0cRegisterUser\x12\x1a.agent.RegisterUserRequest\x1a\x1b.agent.RegisterUserResponse\"\x00\x12:\n\tRunModule\x12\x17.agent.ModuleRunRequest\x1a\x10.agent.ModuleRun\"\x00\x30\x01\x12;\n\x0e\x43heckModuleRun\x12\x15.agent.ModuleRunCheck\x1a\x10.agent.ModuleRun\"\x00\x62\x06proto3') +DESCRIPTOR = _descriptor_pool.Default().AddSerializedFile(b'\n\x11grpc_server.proto\x12\x05\x61gent\x1a\x1cgoogle/protobuf/struct.proto\x1a\x1bgoogle/protobuf/empty.proto\".\n\x0fGeneralResponse\x12\n\n\x02ok\x18\x01 \x01(\x08\x12\x0f\n\x07message\x18\x02 \x01(\t\"\xe9\x01\n\tLLMConfig\x12\x18\n\x0b\x63onfig_name\x18\x01 \x01(\tH\x00\x88\x01\x01\x12\x13\n\x06\x63lient\x18\x02 \x01(\tH\x01\x88\x01\x01\x12\x12\n\x05model\x18\x03 \x01(\tH\x02\x88\x01\x01\x12\x17\n\nmax_tokens\x18\x04 \x01(\x05H\x03\x88\x01\x01\x12\x18\n\x0btemperature\x18\x05 \x01(\x02H\x04\x88\x01\x01\x12\x15\n\x08\x61pi_base\x18\x06 \x01(\tH\x05\x88\x01\x01\x42\x0e\n\x0c_config_nameB\t\n\x07_clientB\x08\n\x06_modelB\r\n\x0b_max_tokensB\x0e\n\x0c_temperatureB\x0b\n\t_api_base\"K\n\nNodeServer\x12\x1e\n\x16\x63ommunication_protocol\x18\x01 \x01(\t\x12\x0c\n\x04port\x18\x02 \x01(\x05\x12\x0f\n\x07node_id\x18\x03 \x01(\t\"\xa9\x01\n\x0fNodeConfigInput\x12\n\n\x02ip\x18\x01 \x01(\t\x12$\n\x17user_communication_port\x18\x02 \x01(\x05H\x00\x88\x01\x01\x12(\n\x1buser_communication_protocol\x18\x03 \x01(\tH\x01\x88\x01\x01\x42\x1a\n\x18_user_communication_portB\x1e\n\x1c_user_communication_protocol\"\xc5\x04\n\nNodeConfig\x12\n\n\x02id\x18\x01 \x01(\t\x12\r\n\x05owner\x18\x02 \x01(\t\x12\x12\n\npublic_key\x18\x03 \x01(\t\x12\n\n\x02ip\x18\x04 \x01(\t\x12#\n\x1buser_communication_protocol\x18\x05 \x01(\t\x12#\n\x1bnode_communication_protocol\x18\x06 \x01(\t\x12\x1f\n\x17user_communication_port\x18\x07 \x01(\x05\x12&\n\x1enum_node_communication_servers\x18\x08 \x01(\x05\x12\"\n\x07servers\x18\t \x03(\x0b\x32\x11.agent.NodeServer\x12\x15\n\rollama_models\x18\n \x03(\t\x12\x13\n\x0b\x64ocker_jobs\x18\x0b \x01(\x08\x12\r\n\x05ports\x18\x0c \x03(\x05\x12\x19\n\x0crouting_type\x18\r \x01(\tH\x00\x88\x01\x01\x12\x18\n\x0brouting_url\x18\x0e \x01(\tH\x01\x88\x01\x01\x12\x15\n\x08num_gpus\x18\x0f \x01(\x05H\x02\x88\x01\x01\x12\x11\n\x04\x61rch\x18\x10 \x01(\tH\x03\x88\x01\x01\x12\x0f\n\x02os\x18\x11 \x01(\tH\x04\x88\x01\x01\x12\x10\n\x03ram\x18\x12 \x01(\x03H\x05\x88\x01\x01\x12\x11\n\x04vram\x18\x13 \x01(\x03H\x06\x88\x01\x01\x12\x16\n\x0eprovider_types\x18\x14 \x03(\t\x12\x0e\n\x06models\x18\x15 \x03(\tB\x0f\n\r_routing_typeB\x0e\n\x0c_routing_urlB\x0b\n\t_num_gpusB\x07\n\x05_archB\x05\n\x03_osB\x06\n\x04_ramB\x07\n\x05_vram\"A\n\x06Secret\x12\x0f\n\x07user_id\x18\x01 \x01(\t\x12\x14\n\x0csecret_value\x18\x02 \x01(\t\x12\x10\n\x08key_name\x18\x03 \x01(\t\"\x9b\x02\n\x06Module\x12\n\n\x02id\x18\x01 \x01(\t\x12\x0c\n\x04name\x18\x02 \x01(\t\x12\x13\n\x0b\x64\x65scription\x18\x03 \x01(\t\x12\x0e\n\x06\x61uthor\x18\x04 \x01(\t\x12\x12\n\nmodule_url\x18\x05 \x01(\t\x12\x18\n\x0bmodule_type\x18\x06 \x01(\tH\x00\x88\x01\x01\x12\x1b\n\x0emodule_version\x18\x07 \x01(\tH\x01\x88\x01\x01\x12\x1e\n\x11module_entrypoint\x18\x08 \x01(\tH\x02\x88\x01\x01\x12\x1b\n\x0e\x65xecution_type\x18\t \x01(\tH\x03\x88\x01\x01\x42\x0e\n\x0c_module_typeB\x11\n\x0f_module_versionB\x14\n\x12_module_entrypointB\x11\n\x0f_execution_type\"\xd0\x02\n\x14\x44\x61taGenerationConfig\x12\x19\n\x0csave_outputs\x18\x01 \x01(\x08H\x00\x88\x01\x01\x12\"\n\x15save_outputs_location\x18\x02 \x01(\tH\x01\x88\x01\x01\x12\x1e\n\x11save_outputs_path\x18\x03 \x01(\tH\x02\x88\x01\x01\x12\x18\n\x0bsave_inputs\x18\x04 \x01(\x08H\x03\x88\x01\x01\x12!\n\x14save_inputs_location\x18\x05 \x01(\tH\x04\x88\x01\x01\x12\x1d\n\x10\x64\x65\x66\x61ult_filename\x18\x06 \x01(\tH\x05\x88\x01\x01\x42\x0f\n\r_save_outputsB\x18\n\x16_save_outputs_locationB\x14\n\x12_save_outputs_pathB\x0e\n\x0c_save_inputsB\x17\n\x15_save_inputs_locationB\x13\n\x11_default_filename\"\xdb\x01\n\x0e\x42\x61seDeployment\x12(\n\x0bnode_config\x18\x01 \x01(\x0b\x32\x11.agent.NodeConfigH\x00\x12,\n\nnode_input\x18\x02 \x01(\x0b\x32\x16.agent.NodeConfigInputH\x00\x12\x0c\n\x04name\x18\x03 \x01(\t\x12\x1d\n\x06module\x18\x04 \x01(\x0b\x32\r.agent.Module\x12\'\n\x06\x63onfig\x18\x05 \x01(\x0b\x32\x17.google.protobuf.Struct\x12\x13\n\x0binitialized\x18\x06 \x01(\x08\x42\x06\n\x04node\"\xb8\x02\n\x0eToolDeployment\x12(\n\x0bnode_config\x18\x01 \x01(\x0b\x32\x11.agent.NodeConfigH\x00\x12,\n\nnode_input\x18\x02 \x01(\x0b\x32\x16.agent.NodeConfigInputH\x00\x12\x0c\n\x04name\x18\x03 \x01(\t\x12\x1d\n\x06module\x18\x04 \x01(\x0b\x32\r.agent.Module\x12\'\n\x06\x63onfig\x18\x05 \x01(\x0b\x32\x17.google.protobuf.Struct\x12@\n\x16\x64\x61ta_generation_config\x18\x06 \x01(\x0b\x32\x1b.agent.DataGenerationConfigH\x01\x88\x01\x01\x12\x13\n\x0binitialized\x18\x07 \x01(\x08\x42\x06\n\x04nodeB\x19\n\x17_data_generation_config\"\xd1\x03\n\x0f\x41gentDeployment\x12(\n\x0bnode_config\x18\x01 \x01(\x0b\x32\x11.agent.NodeConfigH\x00\x12,\n\nnode_input\x18\x02 \x01(\x0b\x32\x16.agent.NodeConfigInputH\x00\x12\x0c\n\x04name\x18\x03 \x01(\t\x12\x1d\n\x06module\x18\x04 \x01(\x0b\x32\r.agent.Module\x12\'\n\x06\x63onfig\x18\x05 \x01(\x0b\x32\x17.google.protobuf.Struct\x12@\n\x16\x64\x61ta_generation_config\x18\x06 \x01(\x0b\x32\x1b.agent.DataGenerationConfigH\x01\x88\x01\x01\x12/\n\x10tool_deployments\x18\x07 \x03(\x0b\x32\x15.agent.ToolDeployment\x12-\n\x0ekb_deployments\x18\x08 \x03(\x0b\x32\x15.agent.BaseDeployment\x12\x36\n\x17\x65nvironment_deployments\x18\t \x03(\x0b\x32\x15.agent.BaseDeployment\x12\x13\n\x0binitialized\x18\n \x01(\x08\x42\x06\n\x04nodeB\x19\n\x17_data_generation_config\"\xe9\x03\n\x10ModuleRunRequest\x12\x13\n\x0bmodule_type\x18\x01 \x01(\t\x12\x13\n\x0b\x63onsumer_id\x18\x02 \x01(\t\x12\'\n\x06inputs\x18\x03 \x01(\x0b\x32\x17.google.protobuf.Struct\x12\x32\n\x10\x61gent_deployment\x18\x04 \x01(\x0b\x32\x16.agent.AgentDeploymentH\x00\x12\x30\n\x0ftool_deployment\x18\x05 \x01(\x0b\x32\x15.agent.ToolDeploymentH\x00\x12\x32\n\x11memory_deployment\x18\x06 \x01(\x0b\x32\x15.agent.BaseDeploymentH\x00\x12.\n\rkb_deployment\x18\x07 \x01(\x0b\x32\x15.agent.BaseDeploymentH\x00\x12\x37\n\x16\x65nvironment_deployment\x18\x08 \x01(\x0b\x32\x15.agent.BaseDeploymentH\x00\x12+\n\x11orchestrator_runs\x18\t \x03(\x0b\x32\x10.agent.ModuleRun\x12\x16\n\tsignature\x18\n \x01(\tH\x01\x88\x01\x01\x12\x1e\n\x07secrets\x18\x0b \x03(\x0b\x32\r.agent.SecretB\x0c\n\ndeploymentB\x0c\n\n_signature\"\xb6\x06\n\tModuleRun\x12\x13\n\x0bmodule_type\x18\x01 \x01(\t\x12\x13\n\x0b\x63onsumer_id\x18\x02 \x01(\t\x12\'\n\x06inputs\x18\x03 \x01(\x0b\x32\x17.google.protobuf.Struct\x12\x32\n\x10\x61gent_deployment\x18\x04 \x01(\x0b\x32\x16.agent.AgentDeploymentH\x00\x12\x30\n\x0ftool_deployment\x18\x05 \x01(\x0b\x32\x15.agent.ToolDeploymentH\x00\x12\x32\n\x11memory_deployment\x18\x06 \x01(\x0b\x32\x15.agent.BaseDeploymentH\x00\x12.\n\rkb_deployment\x18\x07 \x01(\x0b\x32\x15.agent.BaseDeploymentH\x00\x12\x37\n\x16\x65nvironment_deployment\x18\x08 \x01(\x0b\x32\x15.agent.BaseDeploymentH\x00\x12+\n\x11orchestrator_runs\x18\t \x03(\x0b\x32\x10.agent.ModuleRun\x12\x0e\n\x06status\x18\n \x01(\t\x12\r\n\x05\x65rror\x18\x0b \x01(\x08\x12\x0f\n\x02id\x18\x0c \x01(\tH\x01\x88\x01\x01\x12\x0f\n\x07results\x18\r \x03(\t\x12\x1a\n\rerror_message\x18\x0e \x01(\tH\x02\x88\x01\x01\x12\x19\n\x0c\x63reated_time\x18\x0f \x01(\tH\x03\x88\x01\x01\x12\"\n\x15start_processing_time\x18\x10 \x01(\tH\x04\x88\x01\x01\x12\x1b\n\x0e\x63ompleted_time\x18\x11 \x01(\tH\x05\x88\x01\x01\x12\x15\n\x08\x64uration\x18\x12 \x01(\x01H\x06\x88\x01\x01\x12#\n\x16input_schema_ipfs_hash\x18\x13 \x01(\tH\x07\x88\x01\x01\x12\x16\n\tsignature\x18\x14 \x01(\tH\x08\x88\x01\x01\x42\x0c\n\ndeploymentB\x05\n\x03_idB\x10\n\x0e_error_messageB\x0f\n\r_created_timeB\x18\n\x16_start_processing_timeB\x11\n\x0f_completed_timeB\x0b\n\t_durationB\x19\n\x17_input_schema_ipfs_hashB\x0c\n\n_signature\"5\n\x0eModuleRunCheck\x12\x13\n\x0bmodule_type\x18\x01 \x01(\t\x12\x0e\n\x06run_id\x18\x02 \x01(\t\"H\n\x10\x43heckUserRequest\x12\x14\n\x07user_id\x18\x01 \x01(\tH\x00\x88\x01\x01\x12\x12\n\npublic_key\x18\x02 \x01(\tB\n\n\x08_user_id\"J\n\x11\x43heckUserResponse\x12\x15\n\ris_registered\x18\x01 \x01(\x08\x12\n\n\x02id\x18\x02 \x01(\t\x12\x12\n\npublic_key\x18\x03 \x01(\t\")\n\x13RegisterUserRequest\x12\x12\n\npublic_key\x18\x01 \x01(\t\"6\n\x14RegisterUserResponse\x12\n\n\x02id\x18\x01 \x01(\t\x12\x12\n\npublic_key\x18\x02 \x01(\t2\x8a\x03\n\nGrpcServer\x12<\n\x08is_alive\x12\x16.google.protobuf.Empty\x1a\x16.agent.GeneralResponse\"\x00\x12\x38\n\x04stop\x12\x16.google.protobuf.Empty\x1a\x16.agent.GeneralResponse\"\x00\x12@\n\tCheckUser\x12\x17.agent.CheckUserRequest\x1a\x18.agent.CheckUserResponse\"\x00\x12I\n\x0cRegisterUser\x12\x1a.agent.RegisterUserRequest\x1a\x1b.agent.RegisterUserResponse\"\x00\x12:\n\tRunModule\x12\x17.agent.ModuleRunRequest\x1a\x10.agent.ModuleRun\"\x00\x30\x01\x12;\n\x0e\x43heckModuleRun\x12\x15.agent.ModuleRunCheck\x1a\x10.agent.ModuleRun\"\x00\x62\x06proto3') _globals = globals() _builder.BuildMessageAndEnumDescriptors(DESCRIPTOR, _globals) @@ -43,30 +43,32 @@ _globals['_NODECONFIGINPUT']._serialized_end=618 _globals['_NODECONFIG']._serialized_start=621 _globals['_NODECONFIG']._serialized_end=1202 - _globals['_MODULE']._serialized_start=1205 - _globals['_MODULE']._serialized_end=1488 - _globals['_DATAGENERATIONCONFIG']._serialized_start=1491 - _globals['_DATAGENERATIONCONFIG']._serialized_end=1827 - _globals['_BASEDEPLOYMENT']._serialized_start=1830 - _globals['_BASEDEPLOYMENT']._serialized_end=2049 - _globals['_TOOLDEPLOYMENT']._serialized_start=2052 - _globals['_TOOLDEPLOYMENT']._serialized_end=2364 - _globals['_AGENTDEPLOYMENT']._serialized_start=2367 - _globals['_AGENTDEPLOYMENT']._serialized_end=2832 - _globals['_MODULERUNREQUEST']._serialized_start=2835 - _globals['_MODULERUNREQUEST']._serialized_end=3292 - _globals['_MODULERUN']._serialized_start=3295 - _globals['_MODULERUN']._serialized_end=4117 - _globals['_MODULERUNCHECK']._serialized_start=4119 - _globals['_MODULERUNCHECK']._serialized_end=4172 - _globals['_CHECKUSERREQUEST']._serialized_start=4174 - _globals['_CHECKUSERREQUEST']._serialized_end=4246 - _globals['_CHECKUSERRESPONSE']._serialized_start=4248 - _globals['_CHECKUSERRESPONSE']._serialized_end=4322 - _globals['_REGISTERUSERREQUEST']._serialized_start=4324 - _globals['_REGISTERUSERREQUEST']._serialized_end=4365 - _globals['_REGISTERUSERRESPONSE']._serialized_start=4367 - _globals['_REGISTERUSERRESPONSE']._serialized_end=4421 - _globals['_GRPCSERVER']._serialized_start=4424 - _globals['_GRPCSERVER']._serialized_end=4818 + _globals['_SECRET']._serialized_start=1204 + _globals['_SECRET']._serialized_end=1269 + _globals['_MODULE']._serialized_start=1272 + _globals['_MODULE']._serialized_end=1555 + _globals['_DATAGENERATIONCONFIG']._serialized_start=1558 + _globals['_DATAGENERATIONCONFIG']._serialized_end=1894 + _globals['_BASEDEPLOYMENT']._serialized_start=1897 + _globals['_BASEDEPLOYMENT']._serialized_end=2116 + _globals['_TOOLDEPLOYMENT']._serialized_start=2119 + _globals['_TOOLDEPLOYMENT']._serialized_end=2431 + _globals['_AGENTDEPLOYMENT']._serialized_start=2434 + _globals['_AGENTDEPLOYMENT']._serialized_end=2899 + _globals['_MODULERUNREQUEST']._serialized_start=2902 + _globals['_MODULERUNREQUEST']._serialized_end=3391 + _globals['_MODULERUN']._serialized_start=3394 + _globals['_MODULERUN']._serialized_end=4216 + _globals['_MODULERUNCHECK']._serialized_start=4218 + _globals['_MODULERUNCHECK']._serialized_end=4271 + _globals['_CHECKUSERREQUEST']._serialized_start=4273 + _globals['_CHECKUSERREQUEST']._serialized_end=4345 + _globals['_CHECKUSERRESPONSE']._serialized_start=4347 + _globals['_CHECKUSERRESPONSE']._serialized_end=4421 + _globals['_REGISTERUSERREQUEST']._serialized_start=4423 + _globals['_REGISTERUSERREQUEST']._serialized_end=4464 + _globals['_REGISTERUSERRESPONSE']._serialized_start=4466 + _globals['_REGISTERUSERRESPONSE']._serialized_end=4520 + _globals['_GRPCSERVER']._serialized_start=4523 + _globals['_GRPCSERVER']._serialized_end=4917 # @@protoc_insertion_point(module_scope) diff --git a/node/server/http_server.py b/node/server/http_server.py index 972ace3..235e605 100644 --- a/node/server/http_server.py +++ b/node/server/http_server.py @@ -123,13 +123,20 @@ async def agent_create_endpoint(agent_input: AgentDeployment) -> AgentDeployment return await self.agent_create(agent_input) @router.post("/agent/run") - async def agent_run_endpoint(agent_run_input: AgentRunInput, secrets: List[SecretInput] = []) -> AgentRun: + async def agent_run_endpoint(payload: dict) -> AgentRun: """ Run an agent - :param agent_run_input: Agent run specifications - :param secrets: Optional list of secrets to pass to the agent + :param payload: Agent run specifications :return: Status """ + if payload.get("agent_run_input"): + agent_run_input = AgentRunInput(**payload.get("agent_run_input", {})) + secrets_dict = payload.get("secrets", []) + secrets = [SecretInput(**secret) for secret in secrets_dict] + else: + agent_run_input = AgentRunInput(**payload) + secrets = [] + return await self.agent_run(agent_run_input, secrets) @router.post("/agent/check") @@ -151,13 +158,20 @@ async def tool_create_endpoint(tool_input: ToolDeployment) -> ToolDeployment: return await self.tool_create(tool_input) @router.post("/tool/run") - async def tool_run_endpoint(tool_run_input: ToolRunInput, secrets: List[SecretInput] = []) -> ToolRun: + async def tool_run_endpoint(payload: dict) -> ToolRun: """ Run a tool - :param tool_run_input: Tool run specifications - :param secrets: Optional list of secrets to pass to the tool + :param payload: Tool run specifications :return: Status """ + if payload.get("tool_run_input"): + tool_run_input = ToolRunInput(**payload.get("tool_run_input", {})) + secrets_dict = payload.get("secrets", []) + secrets = [SecretInput(**secret) for secret in secrets_dict] + else: + tool_run_input = ToolRunInput(**payload) + secrets = [] + return await self.tool_run(tool_run_input, secrets) @router.post("/tool/check") @@ -179,13 +193,21 @@ async def orchestrator_create_endpoint(orchestrator_input: OrchestratorDeploymen return await self.orchestrator_create(orchestrator_input) @router.post("/orchestrator/run") - async def orchestrator_run_endpoint(orchestrator_run_input: OrchestratorRunInput, secrets: List[SecretInput] = []) -> OrchestratorRun: + async def orchestrator_run_endpoint(payload: dict) -> OrchestratorRun: """ Run an agent orchestrator - :param orchestrator_run_input: Orchestrator run specifications - :param secrets: Optional list of secrets to pass to the orchestrator + :param payload: Orchestrator run specifications :return: Status """ + logger.info(f"Orchestrator run payload: {payload}") + if payload.get("orchestrator_run_input"): + orchestrator_run_input = OrchestratorRunInput(**payload.get("orchestrator_run_input", {})) + secrets_dict = payload.get("secrets", []) + secrets = [SecretInput(**secret) for secret in secrets_dict] + else: + orchestrator_run_input = OrchestratorRunInput(**payload) + secrets = [] + return await self.orchestrator_run(orchestrator_run_input, secrets) @router.post("/orchestrator/check") @@ -207,13 +229,19 @@ async def environment_create_endpoint(environment_input: EnvironmentDeployment) return await self.environment_create(environment_input) @router.post("/environment/run") - async def environment_run_endpoint(environment_run_input: EnvironmentRunInput, secrets: List[SecretInput] = []) -> EnvironmentRun: + async def environment_run_endpoint(payload: dict) -> EnvironmentRun: """ Run an environment - :param environment_run_input: Environment run specifications - :param secrets: Optional list of secrets to pass to the environment + :param payload: Environment run specifications :return: Status """ + if payload.get("environment_run_input"): + environment_run_input = EnvironmentRunInput(**payload.get("environment_run_input", {})) + secrets_dict = payload.get("secrets", []) + secrets = [SecretInput(**secret) for secret in secrets_dict] + else: + environment_run_input = EnvironmentRunInput(**payload) + secrets = [] return await self.environment_run(environment_run_input, secrets) @router.post("/environment/check") @@ -235,13 +263,19 @@ async def kb_create_endpoint(kb_input: KBDeployment) -> KBDeployment: return await self.kb_create(kb_input) @router.post("/kb/run") - async def kb_run_endpoint(kb_run_input: KBRunInput, secrets: List[SecretInput] = []) -> KBRun: + async def kb_run_endpoint(payload: dict) -> KBRun: """ Run a knowledge base - :param kb_run_input: KBRunInput - :param secrets: Optional list of secrets to pass to the knowledge base + :param payload: KB run specifications :return: KBRun """ + if payload.get("kb_run_input"): + kb_run_input = KBRunInput(**payload.get("kb_run_input", {})) + secrets_dict = payload.get("secrets", []) + secrets = [SecretInput(**secret) for secret in secrets_dict] + else: + kb_run_input = KBRunInput(**payload) + secrets = [] return await self.kb_run(kb_run_input, secrets) @router.post("/kb/check") @@ -263,13 +297,19 @@ async def memory_create_endpoint(memory_input: MemoryDeployment) -> MemoryDeploy return await self.memory_create(memory_input) @router.post("/memory/run") - async def memory_run_endpoint(memory_run_input: MemoryRunInput, secrets: List[SecretInput] = [] ) -> MemoryRun: + async def memory_run_endpoint(payload: dict) -> MemoryRun: """ Run a memory module - :param memory_run_input: Memory run specifications - :param secrets: Optional list of secrets to pass to the memory + :param payload: Memory run specifications :return: Status """ + if payload.get("memory_run_input"): + memory_run_input = MemoryRunInput(**payload.get("memory_run_input", {})) + secrets_dict = payload.get("secrets", []) + secrets = [SecretInput(**secret) for secret in secrets_dict] + else: + memory_run_input = MemoryRunInput(**payload) + secrets = [] return await self.memory_run(memory_run_input, secrets) @router.post("/memory/check") @@ -297,35 +337,6 @@ async def user_register_endpoint(user_input: dict): if '_sa_instance_state' in response: response.pop('_sa_instance_state') return response - - @router.post("/user/secret/create") - async def user_secret_create_endpoint( - existing_secrets: List[SecretInput] = [], - secrets: List[SecretInput] = [], - signature: str = Query(...), - is_update: Optional[bool] = Query(False) - ): - user_id = secrets[0].user_id.replace("", "") if secrets else "" - - if not user_id: - raise HTTPException(status_code=400, detail=f"Data cannot be empty") - - if not verify_signature(user_id, signature, user_id.split(":")[1]): - raise HTTPException(status_code=401, detail="Unauthorized: Invalid signature") - - try: - aes_secret = base64.b64decode(os.getenv("AES_SECRET")) - for data in secrets: - rsa_decrypted_secret_value = self.secret.decrypt_rsa(data.secret_value) - encrypted_secret_value = self.secret.encrypt_with_aes(rsa_decrypted_secret_value, aes_secret) - data.secret_value = encrypted_secret_value - - async with HubDBSurreal() as hub: - result = await hub.create_secret(secrets, is_update, existing_secrets) - - return result - except Exception as e: - raise HTTPException(status_code=400, detail=f"Error processing secrets: {str(e)}") async def get_server_connection(self, server_id: str): """ @@ -461,7 +472,7 @@ async def run_module( secret = Secret() for record in secrets: - decrypted_value = secret.decrypt_with_aes(record.secret_value, base64.b64decode(os.getenv("AES_SECRET"))) + decrypted_value = secret.decrypt_rsa(record.secret_value) user_env_data[record.key_name] = decrypted_value # Create module run record in DB @@ -478,17 +489,17 @@ async def run_module( # Execute the task based on module type if module_run_input.deployment.module.execution_type == ModuleExecutionType.package: if module_type == "agent": - _ = run_agent.delay(module_run_data, user_env_data) + _ = run_agent.delay(module_run_data, user_env_data, [secret.model_dump() for secret in secrets] if secrets else []) elif module_type == "tool": - _ = run_tool.delay(module_run_data, user_env_data) + _ = run_tool.delay(module_run_data, user_env_data, [secret.model_dump() for secret in secrets] if secrets else []) elif module_type == "orchestrator": - _ = run_orchestrator.delay(module_run_data, user_env_data) + _ = run_orchestrator.delay(module_run_data, user_env_data, [secret.model_dump() for secret in secrets] if secrets else []) elif module_type == "environment": - _ = run_environment.delay(module_run_data, user_env_data) + _ = run_environment.delay(module_run_data, user_env_data, [secret.model_dump() for secret in secrets] if secrets else []) elif module_type == "kb": - _ = run_kb.delay(module_run_data, user_env_data) + _ = run_kb.delay(module_run_data, user_env_data, [secret.model_dump() for secret in secrets] if secrets else []) elif module_type == "memory": - _ = run_memory.delay(module_run_data, user_env_data) + _ = run_memory.delay(module_run_data, user_env_data, [secret.model_dump() for secret in secrets] if secrets else []) elif module_run_input.deployment.module.execution_type == ModuleExecutionType.docker and module_type == "agent": # validate docker params try: diff --git a/node/server/protos/grpc_server.proto b/node/server/protos/grpc_server.proto index 84b8b70..b85e817 100644 --- a/node/server/protos/grpc_server.proto +++ b/node/server/protos/grpc_server.proto @@ -65,6 +65,11 @@ message NodeConfig { repeated string models = 21; } +message Secret { + string user_id = 1; + string secret_value = 2; + string key_name = 3; +} message Module { string id = 1; @@ -140,6 +145,7 @@ message ModuleRunRequest { } repeated ModuleRun orchestrator_runs = 9; optional string signature = 10; + repeated Secret secrets = 11; } message ModuleRun { diff --git a/node/server/server.py b/node/server/server.py index 7e6a794..0255c75 100644 --- a/node/server/server.py +++ b/node/server/server.py @@ -279,7 +279,6 @@ def signal_handler(sig): secret = Secret(server_private_key_path, server_public_key_path, env_file_path) secret.check_and_generate_keys() - secret.check_and_generate_aes_secret() asyncio.run(run_server( communication_protocol=args.communication_protocol, diff --git a/node/server/ws_server.py b/node/server/ws_server.py index 152319c..67d0b34 100644 --- a/node/server/ws_server.py +++ b/node/server/ws_server.py @@ -7,7 +7,7 @@ from datetime import datetime import traceback import logging -from typing import Dict, Any, Union +from typing import Dict, Any, Union, Optional, List from pathlib import Path from fastapi import FastAPI, HTTPException, WebSocket, WebSocketDisconnect from fastapi.middleware.cors import CORSMiddleware @@ -20,6 +20,7 @@ KBRunInput, KBDeployment, OrchestratorRunInput, OrchestratorDeployment, ModuleExecutionType, + SecretInput ) from node.storage.db.db import LocalDBPostgres from node.user import register_user, check_user @@ -34,6 +35,7 @@ ) from node.module_manager import setup_module_deployment +from node.secret import Secret logger = logging.getLogger(__name__) load_dotenv() @@ -252,8 +254,13 @@ async def run_module_endpoint(self, websocket: WebSocket, module_type: str, clie try: data = await websocket.receive_json() logger.info(f"Endpoint: run_{module_type} :: Received data: {data}") + if data.get('secrets'): + secrets = data.pop('secrets') + logger.info(f"Endpoint: run_{module_type} :: Received1 secrets: {secrets}") + else: + secrets = [] - result = await self.run_module(module_type, data, client_id) + result = await self.run_module(module_type, data, client_id, [SecretInput(**secret) for secret in secrets] if secrets else []) logger.info(f"Endpoint: run_{module_type} :: Sending result: {result}") await self.manager.send_message(result, client_id, f"run_{module_type}") @@ -283,7 +290,7 @@ async def run_module_endpoint(self, websocket: WebSocket, module_type: str, clie logger.warning(f"Client {client_id} disconnected (outer).") self.manager.disconnect(client_id, f"run_{module_type}") - async def run_module(self, module_type: str, data: dict, client_id: str) -> str: + async def run_module(self, module_type: str, data: dict, client_id: str, secrets: List[SecretInput] = []) -> str: try: # Map module types to their corresponding input and run classes module_configs = { @@ -355,8 +362,13 @@ async def run_module(self, module_type: str, data: dict, client_id: str) -> str: else: execution_type = module_run.deployment.module.execution_type + secret_obj = Secret() + user_env_data = {} + for secret in secrets: + user_env_data[secret.key_name] = secret_obj.decrypt_rsa(secret.secret_value) + if execution_type == ModuleExecutionType.package or execution_type == 'package': - task = config["worker"].delay(module_run_data) + task = config["worker"].delay(module_run_data, user_env_data, [secret.model_dump() for secret in secrets] if secrets else []) elif execution_type == ModuleExecutionType.docker or execution_type == 'docker': task = execute_docker_agent.delay(module_run_data) else: diff --git a/node/storage/hub/hub.py b/node/storage/hub/hub.py index be4edc7..14802c3 100644 --- a/node/storage/hub/hub.py +++ b/node/storage/hub/hub.py @@ -8,7 +8,6 @@ import traceback from typing import Dict, List, Optional, Tuple from node.schemas import SecretInput -from contextlib import asynccontextmanager load_dotenv() logger = logging.getLogger(__name__) @@ -51,29 +50,6 @@ def _decode_token(self, token: str) -> str: except jwt.PyJWTError as e: logger.error(f"Token decoding failed: {e}") return None - - @asynccontextmanager - async def root_user_context(self): - try: - # Sign in as regular user if local hub is false - if os.getenv("LOCAL_HUB").lower() == "false": - await self.surrealdb.signin( - { - "username": os.getenv("HUB_USERNAME"), - "password": os.getenv("HUB_PASSWORD"), - "NS": self.ns, - "DB": self.db, - "AC": "user" - } - ) - else: - # Sign in as root user if local hub is true - await self.surrealdb.signin({"user": os.getenv("HUB_DB_SURREAL_ROOT_USER"), "pass": os.getenv("HUB_DB_SURREAL_ROOT_PASS")}) - yield - finally: - if os.getenv("LOCAL_HUB").lower() == "true": - logger.info("Signing out from root user") - await self.close() async def signin( self, username: str, password: str @@ -276,104 +252,6 @@ async def list_memory_modules(self, memory_module_name=None) -> List: async def create_agent(self, agent_config: Dict) -> Tuple[bool, Optional[Dict]]: return await self.surrealdb.create("agent", agent_config) - - def prepare_batch_query(self, secret_config: List[SecretInput], existing_secrets: List[SecretInput], update:bool = False) -> str: - existing_secrets_dict = {secret.key_name for secret in existing_secrets} - records_to_insert = [] - records_to_update = [] - - for secret in secret_config: - user_id = secret.user_id.replace("", "").strip() - key_name = secret.key_name - key_value = secret.secret_value - - if not existing_secrets_dict or key_name not in existing_secrets_dict: - records_to_insert.append({ - "user_id": user_id, - "key_name": key_name, - "secret_value": key_value - }) - else: - if update: - records_to_update.append({ - "secret_value": key_value, - "user_id": user_id, - "key_name": key_name - }) - - insert_query = "" - if records_to_insert: - insert_query = "INSERT INTO api_secrets $records;" - - update_query = "" - if records_to_update: - update_query = "UPDATE api_secrets SET secret_value = $secret_value WHERE user_id = $user_id AND key_name = $key_name;" - - return { - "insert_query": insert_query, - "insert_params": {"records": records_to_insert}, - "update_query": update_query, - "update_params": records_to_update - } - - async def create_secret(self, secret_config: List[SecretInput], update: bool = False, existing_secrets: List[SecretInput] = []) -> str: - try: - user_id = secret_config[0].user_id.replace("", "").strip() - if not user_id: - return "Invalid user ID" - - query_data = self.prepare_batch_query(secret_config, existing_secrets, update) - - if not (query_data["insert_query"] or query_data["update_query"]): - return "Records already exist" - - async with self.root_user_context(): - try: - transaction_query = "BEGIN TRANSACTION;" - - if query_data["insert_query"]: - transaction_query += "\n" + query_data["insert_query"] - - if query_data["update_query"]: - for i, _ in enumerate(query_data["update_params"]): - parameterized_query = query_data["update_query"].replace( - "$secret_value", f"$secret_value_{i}" - ).replace( - "$user_id", f"$user_id_{i}" - ).replace( - "$key_name", f"$key_name_{i}" - ) - transaction_query += f"\n{parameterized_query}" - - transaction_query += "\nCOMMIT TRANSACTION;" - - params = {} - if query_data["insert_params"]: - params.update(query_data["insert_params"]) - - for i, update_params in enumerate(query_data["update_params"]): - params.update({ - f"secret_value_{i}": update_params["secret_value"], - f"user_id_{i}": update_params["user_id"], - f"key_name_{i}": update_params["key_name"] - }) - - results = await self.surrealdb.query(transaction_query, params) - - logger.debug(f"Results: {results}") - - if all(result.get('status') == 'OK' for result in results) and any(result.get('result') for result in results): - return "Records updated successfully" - else: - return "Operation failed: Database error" - - except Exception as e: - logger.error(f"Secret creation failed: {str(e)}") - return f"Operation failed: {type(e).__name__}" - - except Exception as e: - logger.error(f"Secret creation failed: {str(e)}") - return "Operation failed: Invalid input" async def close(self): """Close the database connection""" diff --git a/node/worker/package_worker.py b/node/worker/package_worker.py index fb689b1..47dbdf6 100644 --- a/node/worker/package_worker.py +++ b/node/worker/package_worker.py @@ -37,66 +37,66 @@ sys.path.append(MODULES_SOURCE_DIR) @app.task(bind=True, acks_late=True) -def run_agent(self, agent_run, user_env_data = {}): +def run_agent(self, agent_run, user_env_data = {}, secrets = []): try: agent_run = AgentRun(**agent_run) loop = asyncio.get_event_loop() - return loop.run_until_complete(_run_module_async(agent_run, user_env_data)) + return loop.run_until_complete(_run_module_async(agent_run, user_env_data, secrets)) finally: # Force cleanup of channels app.backend.cleanup() @app.task(bind=True, acks_late=True) -def run_memory(self, memory_run, user_env_data = {}): +def run_memory(self, memory_run, user_env_data = {}, secrets = []): try: memory_run = MemoryRun(**memory_run) loop = asyncio.get_event_loop() - return loop.run_until_complete(_run_module_async(memory_run, user_env_data)) + return loop.run_until_complete(_run_module_async(memory_run, user_env_data, secrets)) finally: # Force cleanup of channels app.backend.cleanup() @app.task(bind=True, acks_late=True) -def run_tool(self, tool_run, user_env_data = {}): +def run_tool(self, tool_run, user_env_data = {}, secrets = []): try: tool_run = ToolRun(**tool_run) loop = asyncio.get_event_loop() - return loop.run_until_complete(_run_module_async(tool_run, user_env_data)) + return loop.run_until_complete(_run_module_async(tool_run, user_env_data, secrets)) finally: # Force cleanup of channels app.backend.cleanup() @app.task(bind=True, acks_late=True) -def run_orchestrator(self, orchestrator_run, user_env_data = {}): +def run_orchestrator(self, orchestrator_run, user_env_data = {}, secrets = []): try: orchestrator_run = OrchestratorRun(**orchestrator_run) loop = asyncio.get_event_loop() - return loop.run_until_complete(_run_module_async(orchestrator_run, user_env_data)) + return loop.run_until_complete(_run_module_async(orchestrator_run, user_env_data, secrets)) finally: # Force cleanup of channels app.backend.cleanup() @app.task(bind=True, acks_late=True) -def run_environment(self, environment_run, user_env_data = {}): +def run_environment(self, environment_run, user_env_data = {}, secrets = []): try: environment_run = EnvironmentRun(**environment_run) loop = asyncio.get_event_loop() - return loop.run_until_complete(_run_module_async(environment_run, user_env_data)) + return loop.run_until_complete(_run_module_async(environment_run, user_env_data, secrets)) finally: # Force cleanup of channels app.backend.cleanup() @app.task(bind=True, acks_late=True) -def run_kb(self, kb_run, user_env_data = {}): +def run_kb(self, kb_run, user_env_data = {}, secrets = []): try: kb_run = KBRun(**kb_run) loop = asyncio.get_event_loop() - return loop.run_until_complete(_run_module_async(kb_run, user_env_data)) + return loop.run_until_complete(_run_module_async(kb_run, user_env_data, secrets)) finally: # Force cleanup of channels app.backend.cleanup() -async def _run_module_async(module_run: Union[AgentRun, MemoryRun, ToolRun, OrchestratorRun, EnvironmentRun, KBRun], user_env_data = {}) -> None: +async def _run_module_async(module_run: Union[AgentRun, MemoryRun, ToolRun, OrchestratorRun, EnvironmentRun, KBRun], user_env_data = {}, secrets = []) -> None: """Handles execution of agent, memory, orchestrator, and environment runs. Args: @@ -123,7 +123,7 @@ async def _run_module_async(module_run: Union[AgentRun, MemoryRun, ToolRun, Orch return await module_run_engine.init_run() - await module_run_engine.start_run(user_env_data) + await module_run_engine.start_run(user_env_data, secrets) if module_run_engine.module_run.status == "completed": await module_run_engine.complete() @@ -238,7 +238,7 @@ def package_context(self, env_vars): os.environ.clear() os.environ.update(old_env) - async def load_and_run(self, module_path: Path, entrypoint: str, module_run, user_env_data = {}): + async def load_and_run(self, module_path: Path, entrypoint: str, module_run, user_env_data = {}, secrets = []): with self.package_context(user_env_data): try: # Remove any existing module references @@ -266,12 +266,21 @@ async def load_and_run(self, module_path: Path, entrypoint: str, module_run, use module_run_dict = module_run.model_dump() else: module_run_dict = module_run + + # Check if run_func accepts variable keyword arguments + sig = inspect.signature(run_func) + has_kwargs = any(p.kind == inspect.Parameter.VAR_KEYWORD for p in sig.parameters.values()) if inspect.iscoroutinefunction(run_func): - result = await run_func(module_run=module_run_dict) + if has_kwargs: + result = await run_func(module_run=module_run_dict, secrets=secrets) + else: + result = await run_func(module_run=module_run_dict) else: - result = await maybe_async_call(run_func, module_run=module_run_dict) - + if has_kwargs: + result = await maybe_async_call(run_func, module_run=module_run_dict, secrets=secrets) + else: + result = await maybe_async_call(run_func, module_run=module_run_dict) return result except Exception as e: @@ -313,7 +322,7 @@ async def init_run(self): # Load the module self.module_run = await load_and_validate_input_schema(self.module_run) - async def start_run(self, env_data = {}): + async def start_run(self, env_data = {}, secrets = []): """Executes the module run""" logger.info(f"Starting {self.module_type} run") self.module_run.status = "running" @@ -343,7 +352,8 @@ async def start_run(self, env_data = {}): module_path=module_path, entrypoint=entrypoint, module_run=self.module_run, - user_env_data=env_data + user_env_data=env_data, + secrets=secrets ) # Handle response