diff --git a/autogpt_platform/backend/backend/data/block.py b/autogpt_platform/backend/backend/data/block.py index 096195e2c97c..5bc16a68bc79 100644 --- a/autogpt_platform/backend/backend/data/block.py +++ b/autogpt_platform/backend/backend/data/block.py @@ -22,10 +22,10 @@ from backend.util.settings import Config from .model import ( - CREDENTIALS_FIELD_NAME, ContributorDetails, Credentials, CredentialsMetaInput, + is_credentials_field_name, ) app_config = Config() @@ -140,17 +140,38 @@ def get_required_fields(cls) -> set[str]: @classmethod def __pydantic_init_subclass__(cls, **kwargs): """Validates the schema definition. Rules: - - Only one `CredentialsMetaInput` field may be present. - - This field MUST be called `credentials`. - - A field that is called `credentials` MUST be a `CredentialsMetaInput`. + - Fields with annotation `CredentialsMetaInput` MUST be + named `credentials` or `*_credentials` + - Fields named `credentials` or `*_credentials` MUST be + of type `CredentialsMetaInput` """ super().__pydantic_init_subclass__(**kwargs) # Reset cached JSON schema to prevent inheriting it from parent class cls.cached_jsonschema = {} - credentials_fields = [ - field_name + credentials_fields = cls.get_credentials_fields() + + for field_name in cls.get_fields(): + if is_credentials_field_name(field_name): + if field_name not in credentials_fields: + raise TypeError( + f"Credentials field '{field_name}' on {cls.__qualname__} " + f"is not of type {CredentialsMetaInput.__name__}" + ) + + credentials_fields[field_name].validate_credentials_field_schema(cls) + + elif field_name in credentials_fields: + raise KeyError( + f"Credentials field '{field_name}' on {cls.__qualname__} " + "has invalid name: must be 'credentials' or *_credentials" + ) + + @classmethod + def get_credentials_fields(cls) -> dict[str, type[CredentialsMetaInput]]: + return { + field_name: info.annotation for field_name, info in cls.model_fields.items() if ( inspect.isclass(info.annotation) @@ -159,32 +180,7 @@ def __pydantic_init_subclass__(cls, **kwargs): CredentialsMetaInput, ) ) - ] - if len(credentials_fields) > 1: - raise ValueError( - f"{cls.__qualname__} can only have one CredentialsMetaInput field" - ) - elif ( - len(credentials_fields) == 1 - and credentials_fields[0] != CREDENTIALS_FIELD_NAME - ): - raise ValueError( - f"CredentialsMetaInput field on {cls.__qualname__} " - "must be named 'credentials'" - ) - elif ( - len(credentials_fields) == 0 - and CREDENTIALS_FIELD_NAME in cls.model_fields.keys() - ): - raise TypeError( - f"Field 'credentials' on {cls.__qualname__} " - f"must be of type {CredentialsMetaInput.__name__}" - ) - if credentials_field := cls.model_fields.get(CREDENTIALS_FIELD_NAME): - credentials_input_type = cast( - CredentialsMetaInput, credentials_field.annotation - ) - credentials_input_type.validate_credentials_field_schema(cls) + } BlockSchemaInputType = TypeVar("BlockSchemaInputType", bound=BlockSchema) @@ -242,7 +238,7 @@ def __init__( test_input: BlockInput | list[BlockInput] | None = None, test_output: BlockData | list[BlockData] | None = None, test_mock: dict[str, Any] | None = None, - test_credentials: Optional[Credentials] = None, + test_credentials: Optional[Credentials | dict[str, Credentials]] = None, disabled: bool = False, static_output: bool = False, block_type: BlockType = BlockType.STANDARD, @@ -299,6 +295,12 @@ def __init__( "field must be a BaseModel and all its fields must be boolean" ) + # Disallow multiple credentials inputs on webhook blocks + if len(self.input_schema.get_credentials_fields()) > 1: + raise ValueError( + "Multiple credentials input fields not supported on webhook blocks" + ) + # Enforce presence of 'payload' input if "payload" not in self.input_schema.model_fields: raise TypeError( diff --git a/autogpt_platform/backend/backend/data/model.py b/autogpt_platform/backend/backend/data/model.py index 9e79fd7da394..aedd8fa5aa95 100644 --- a/autogpt_platform/backend/backend/data/model.py +++ b/autogpt_platform/backend/backend/data/model.py @@ -238,7 +238,8 @@ class UserIntegrations(BaseModel): CT = TypeVar("CT", bound=CredentialsType) -CREDENTIALS_FIELD_NAME = "credentials" +def is_credentials_field_name(field_name: str) -> bool: + return field_name == "credentials" or field_name.endswith("_credentials") class CredentialsMetaInput(BaseModel, Generic[CP, CT]): @@ -247,21 +248,21 @@ class CredentialsMetaInput(BaseModel, Generic[CP, CT]): provider: CP type: CT - @staticmethod - def _add_json_schema_extra(schema, cls: CredentialsMetaInput): - schema["credentials_provider"] = get_args( - cls.model_fields["provider"].annotation - ) - schema["credentials_types"] = get_args(cls.model_fields["type"].annotation) + @classmethod + def allowed_providers(cls) -> tuple[ProviderName, ...]: + return get_args(cls.model_fields["provider"].annotation) - model_config = ConfigDict( - json_schema_extra=_add_json_schema_extra, # type: ignore - ) + @classmethod + def allowed_cred_types(cls) -> tuple[CredentialsType, ...]: + return get_args(cls.model_fields["type"].annotation) @classmethod def validate_credentials_field_schema(cls, model: type["BlockSchema"]): - """Validates the schema of a `credentials` field""" - field_schema = model.jsonschema()["properties"][CREDENTIALS_FIELD_NAME] + """Validates the schema of a credentials input field""" + field_name = next( + name for name, type in model.get_credentials_fields().items() if type is cls + ) + field_schema = model.jsonschema()["properties"][field_name] try: schema_extra = _CredentialsFieldSchemaExtra[CP, CT].model_validate( field_schema @@ -275,11 +276,20 @@ def validate_credentials_field_schema(cls, model: type["BlockSchema"]): f"{field_schema}" ) from e - if ( - len(schema_extra.credentials_provider) > 1 - and not schema_extra.discriminator - ): - raise TypeError("Multi-provider CredentialsField requires discriminator!") + if len(cls.allowed_providers()) > 1 and not schema_extra.discriminator: + raise TypeError( + f"Multi-provider CredentialsField '{field_name}' " + "requires discriminator!" + ) + + @staticmethod + def _add_json_schema_extra(schema, cls: CredentialsMetaInput): + schema["credentials_provider"] = cls.allowed_providers() + schema["credentials_types"] = cls.allowed_cred_types() + + model_config = ConfigDict( + json_schema_extra=_add_json_schema_extra, # type: ignore + ) class _CredentialsFieldSchemaExtra(BaseModel, Generic[CP, CT]): diff --git a/autogpt_platform/backend/backend/executor/manager.py b/autogpt_platform/backend/backend/executor/manager.py index 0d6ed816f156..c6097a1b50c2 100644 --- a/autogpt_platform/backend/backend/executor/manager.py +++ b/autogpt_platform/backend/backend/executor/manager.py @@ -10,7 +10,6 @@ from multiprocessing.pool import AsyncResult, Pool from typing import TYPE_CHECKING, Any, Generator, TypeVar, cast -from pydantic import BaseModel from redis.lock import Lock as RedisLock if TYPE_CHECKING: @@ -20,7 +19,14 @@ from backend.blocks.agent import AgentExecutorBlock from backend.data import redis -from backend.data.block import Block, BlockData, BlockInput, BlockType, get_block +from backend.data.block import ( + Block, + BlockData, + BlockInput, + BlockSchema, + BlockType, + get_block, +) from backend.data.execution import ( ExecutionQueue, ExecutionResult, @@ -31,7 +37,6 @@ parse_execution_output, ) from backend.data.graph import GraphModel, Link, Node -from backend.data.model import CREDENTIALS_FIELD_NAME, CredentialsMetaInput from backend.integrations.creds_manager import IntegrationCredentialsManager from backend.util import json from backend.util.decorator import error_logged, time_measured @@ -170,10 +175,11 @@ def update_execution(status: ExecutionStatus) -> ExecutionResult: # one (running) block at a time; simultaneous execution of blocks using same # credentials is not supported. creds_lock = None - if CREDENTIALS_FIELD_NAME in input_data: - credentials_meta = CredentialsMetaInput(**input_data[CREDENTIALS_FIELD_NAME]) + input_model = cast(type[BlockSchema], node_block.input_schema) + for field_name, input_type in input_model.get_credentials_fields().items(): + credentials_meta = input_type(**input_data[field_name]) credentials, creds_lock = creds_manager.acquire(user_id, credentials_meta.id) - extra_exec_kwargs["credentials"] = credentials + extra_exec_kwargs[field_name] = credentials output_size = 0 end_status = ExecutionStatus.COMPLETED @@ -890,41 +896,39 @@ def _validate_node_input_credentials(self, graph: GraphModel, user_id: str): raise ValueError(f"Unknown block {node.block_id} for node #{node.id}") # Find any fields of type CredentialsMetaInput - model_fields = cast(type[BaseModel], block.input_schema).model_fields - if CREDENTIALS_FIELD_NAME not in model_fields: + credentials_fields = cast( + type[BlockSchema], block.input_schema + ).get_credentials_fields() + if not credentials_fields: continue - field = model_fields[CREDENTIALS_FIELD_NAME] - - # The BlockSchema class enforces that a `credentials` field is always a - # `CredentialsMetaInput`, so we can safely assume this here. - credentials_meta_type = cast(CredentialsMetaInput, field.annotation) - credentials_meta = credentials_meta_type.model_validate( - node.input_default[CREDENTIALS_FIELD_NAME] - ) - # Fetch the corresponding Credentials and perform sanity checks - credentials = self.credentials_store.get_creds_by_id( - user_id, credentials_meta.id - ) - if not credentials: - raise ValueError( - f"Unknown credentials #{credentials_meta.id} " - f"for node #{node.id}" - ) - if ( - credentials.provider != credentials_meta.provider - or credentials.type != credentials_meta.type - ): - logger.warning( - f"Invalid credentials #{credentials.id} for node #{node.id}: " - "type/provider mismatch: " - f"{credentials_meta.type}<>{credentials.type};" - f"{credentials_meta.provider}<>{credentials.provider}" + for field_name, credentials_meta_type in credentials_fields.items(): + credentials_meta = credentials_meta_type.model_validate( + node.input_default[field_name] ) - raise ValueError( - f"Invalid credentials #{credentials.id} for node #{node.id}: " - "type/provider mismatch" + # Fetch the corresponding Credentials and perform sanity checks + credentials = self.credentials_store.get_creds_by_id( + user_id, credentials_meta.id ) + if not credentials: + raise ValueError( + f"Unknown credentials #{credentials_meta.id} " + f"for node #{node.id} input '{field_name}'" + ) + if ( + credentials.provider != credentials_meta.provider + or credentials.type != credentials_meta.type + ): + logger.warning( + f"Invalid credentials #{credentials.id} for node #{node.id}: " + "type/provider mismatch: " + f"{credentials_meta.type}<>{credentials.type};" + f"{credentials_meta.provider}<>{credentials.provider}" + ) + raise ValueError( + f"Invalid credentials #{credentials.id} for node #{node.id}: " + "type/provider mismatch" + ) # ------- UTILITIES ------- # diff --git a/autogpt_platform/backend/backend/integrations/webhooks/graph_lifecycle_hooks.py b/autogpt_platform/backend/backend/integrations/webhooks/graph_lifecycle_hooks.py index c241bc3a4a41..d22a0f22f090 100644 --- a/autogpt_platform/backend/backend/integrations/webhooks/graph_lifecycle_hooks.py +++ b/autogpt_platform/backend/backend/integrations/webhooks/graph_lifecycle_hooks.py @@ -1,9 +1,8 @@ import logging from typing import TYPE_CHECKING, Callable, Optional, cast -from backend.data.block import get_block +from backend.data.block import BlockSchema, get_block from backend.data.graph import set_node_webhook -from backend.data.model import CREDENTIALS_FIELD_NAME from backend.integrations.webhooks import WEBHOOK_MANAGERS_BY_NAME if TYPE_CHECKING: @@ -30,14 +29,28 @@ async def on_graph_activate( # Compare nodes in new_graph_version with previous_graph_version updated_nodes = [] for new_node in graph.nodes: + block = get_block(new_node.block_id) + if not block: + raise ValueError( + f"Node #{new_node.id} is instance of unknown block #{new_node.block_id}" + ) + block_input_schema = cast(BlockSchema, block.input_schema) + node_credentials = None - if creds_meta := new_node.input_default.get(CREDENTIALS_FIELD_NAME): - node_credentials = get_credentials(creds_meta["id"]) - if not node_credentials: - raise ValueError( - f"Node #{new_node.id} updated with non-existent " - f"credentials #{node_credentials}" + if ( + # Webhook-triggered blocks are only allowed to have 1 credentials input + ( + creds_field_name := next( + iter(block_input_schema.get_credentials_fields()), None ) + ) + and (creds_meta := new_node.input_default.get(creds_field_name)) + and not (node_credentials := get_credentials(creds_meta["id"])) + ): + raise ValueError( + f"Node #{new_node.id} input '{creds_field_name}' updated with " + f"non-existent credentials #{creds_meta['id']}" + ) updated_node = await on_node_activate( graph.user_id, new_node, credentials=node_credentials @@ -62,14 +75,28 @@ async def on_graph_deactivate( """ updated_nodes = [] for node in graph.nodes: + block = get_block(node.block_id) + if not block: + raise ValueError( + f"Node #{node.id} is instance of unknown block #{node.block_id}" + ) + block_input_schema = cast(BlockSchema, block.input_schema) + node_credentials = None - if creds_meta := node.input_default.get(CREDENTIALS_FIELD_NAME): - node_credentials = get_credentials(creds_meta["id"]) - if not node_credentials: - logger.error( - f"Node #{node.id} referenced non-existent " - f"credentials #{creds_meta['id']}" + if ( + # Webhook-triggered blocks are only allowed to have 1 credentials input + ( + creds_field_name := next( + iter(block_input_schema.get_credentials_fields()), None ) + ) + and (creds_meta := node.input_default.get(creds_field_name)) + and not (node_credentials := get_credentials(creds_meta["id"])) + ): + logger.error( + f"Node #{node.id} input '{creds_field_name}' referenced non-existent " + f"credentials #{creds_meta['id']}" + ) updated_node = await on_node_deactivate(node, credentials=node_credentials) updated_nodes.append(updated_node) @@ -116,10 +143,12 @@ async def on_node_activate( f"Constructed resource string {resource} from input {node.input_default}" ) + block_input_schema = cast(BlockSchema, block.input_schema) + credentials_field_name = next(iter(block_input_schema.get_credentials_fields())) event_filter_input_name = block.webhook_config.event_filter_input has_everything_for_webhook = ( resource is not None - and CREDENTIALS_FIELD_NAME in node.input_default + and credentials_field_name in node.input_default and event_filter_input_name in node.input_default and any(is_on for is_on in node.input_default[event_filter_input_name].values()) ) @@ -127,7 +156,7 @@ async def on_node_activate( if has_everything_for_webhook and resource: logger.debug(f"Node #{node} has everything for a webhook!") if not credentials: - credentials_meta = node.input_default[CREDENTIALS_FIELD_NAME] + credentials_meta = node.input_default[credentials_field_name] raise ValueError( f"Cannot set up webhook for node #{node.id}: " f"credentials #{credentials_meta['id']} not available" diff --git a/autogpt_platform/backend/backend/util/test.py b/autogpt_platform/backend/backend/util/test.py index 5433fd90226c..89ebbb4977de 100644 --- a/autogpt_platform/backend/backend/util/test.py +++ b/autogpt_platform/backend/backend/util/test.py @@ -1,11 +1,11 @@ import logging import time -from typing import Sequence +from typing import Sequence, cast from backend.data import db -from backend.data.block import Block, initialize_blocks +from backend.data.block import Block, BlockSchema, initialize_blocks from backend.data.execution import ExecutionResult, ExecutionStatus -from backend.data.model import CREDENTIALS_FIELD_NAME +from backend.data.model import _BaseCredentials from backend.data.user import create_default_user from backend.executor import DatabaseManager, ExecutionManager, ExecutionScheduler from backend.server.rest_api import AgentServer @@ -100,14 +100,22 @@ def execute_block_test(block: Block): else: log.info(f"{prefix} mock {mock_name} not found in block") + # Populate credentials argument(s) extra_exec_kwargs = {} - - if CREDENTIALS_FIELD_NAME in block.input_schema.model_fields: - if not block.test_credentials: - raise ValueError( - f"{prefix} requires credentials but has no test_credentials" - ) - extra_exec_kwargs[CREDENTIALS_FIELD_NAME] = block.test_credentials + input_model = cast(type[BlockSchema], block.input_schema) + credentials_input_fields = input_model.get_credentials_fields() + if len(credentials_input_fields) == 1 and isinstance( + block.test_credentials, _BaseCredentials + ): + field_name = next(iter(credentials_input_fields)) + extra_exec_kwargs[field_name] = block.test_credentials + elif credentials_input_fields and block.test_credentials: + if not isinstance(block.test_credentials, dict): + log.warning(f"Block {block.name} has no usable test credentials") + else: + for field_name in credentials_input_fields: + if field_name in block.test_credentials: + extra_exec_kwargs[field_name] = block.test_credentials[field_name] for input_data in block.test_input: log.info(f"{prefix} in: {input_data}") diff --git a/autogpt_platform/frontend/src/components/CustomNode.tsx b/autogpt_platform/frontend/src/components/CustomNode.tsx index 6541c351f5b0..e9b2782c5008 100644 --- a/autogpt_platform/frontend/src/components/CustomNode.tsx +++ b/autogpt_platform/frontend/src/components/CustomNode.tsx @@ -240,6 +240,7 @@ export function CustomNode({ ![BlockUIType.INPUT, BlockUIType.WEBHOOK].includes(nodeType) && // No input connection handles for credentials propKey !== "credentials" && + !propKey.endsWith("_credentials") && // For OUTPUT blocks, only show the 'value' (hides 'name') input connection handle !(nodeType == BlockUIType.OUTPUT && propKey == "name"); const isConnected = isInputHandleConnected(propKey); @@ -256,7 +257,8 @@ export function CustomNode({ side="left" /> ) : ( - propKey != "credentials" && ( + propKey != "credentials" && + !propKey.endsWith("_credentials") && (
{propSchema.title || beautifyString(propKey)} diff --git a/autogpt_platform/frontend/src/components/integrations/credentials-input.tsx b/autogpt_platform/frontend/src/components/integrations/credentials-input.tsx index 879f70aff203..c787bcebf204 100644 --- a/autogpt_platform/frontend/src/components/integrations/credentials-input.tsx +++ b/autogpt_platform/frontend/src/components/integrations/credentials-input.tsx @@ -1,5 +1,5 @@ import { z } from "zod"; -import { cn } from "@/lib/utils"; +import { beautifyString, cn } from "@/lib/utils"; import { useForm } from "react-hook-form"; import { Input } from "@/components/ui/input"; import { Button } from "@/components/ui/button"; @@ -87,12 +87,13 @@ export type OAuthPopupResultMessage = { message_type: "oauth_popup_result" } & ( ); export const CredentialsInput: FC<{ + selfKey: string; className?: string; selectedCredentials?: CredentialsMetaInput; onSelectCredentials: (newValue?: CredentialsMetaInput) => void; -}> = ({ className, selectedCredentials, onSelectCredentials }) => { +}> = ({ selfKey, className, selectedCredentials, onSelectCredentials }) => { const api = useMemo(() => new AutoGPTServerAPI(), []); - const credentials = useCredentials(); + const credentials = useCredentials(selfKey); const [isAPICredentialsModalOpen, setAPICredentialsModalOpen] = useState(false); const [isOAuth2FlowInProgress, setOAuth2FlowInProgress] = useState(false); @@ -209,6 +210,7 @@ export const CredentialsInput: FC<{ <> {supportsApiKey && ( setAPICredentialsModalOpen(false)} onCredentialsCreate={(credsMeta) => { @@ -242,7 +244,9 @@ export const CredentialsInput: FC<{ return ( <>
- Credentials + + {providerName} Credentials +
@@ -310,7 +314,12 @@ export const CredentialsInput: FC<{ // Saved credentials exist return ( <> - Credentials +
+ + {providerName} Credentials + + +