WindyMadman/Auto-GPT
1
0
Fork 0
mirror of https://github.com/Significant-Gravitas/Auto-GPT.git synced 2025-01-08 11:57:32 +08:00
Code Issues Packages Projects Releases Wiki Activity

feat(platform): Support multiple credentials inputs on blocks (#8932)

Browse Source 
- Resolves #8930
- Depends on #8725

### Changes 🏗️

- feat(platform): Support multiple credentials inputs on blocks

Aside from `credentials`, fields within the name pattern `*_credentials`
are now also supported!

- Update docs with info on multi credentials support

### Checklist 📋

#### For code changes:
- [x] I have clearly listed my changes in the PR description
- [x] I have made a test plan
- [x] I have tested my changes according to the test plan:
  - [x] Ask @aarushik93 to test
This commit is contained in:
Reinier van der Leer 2025-01-03 16:18:57 +01:00 committed by GitHub
parent 84af37a27a
commit 1375a0fdbc
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
10 changed files with 207 additions and 127 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

70
autogpt_platform/backend/backend/data/block.py
View File

@ -22,10 +22,10 @@ from backend.util import json
from backend.util.settings import Config
from .model import (
CREDENTIALS_FIELD_NAME,
ContributorDetails,
Credentials,
CredentialsMetaInput,
is_credentials_field_name,
)
app_config = Config()
@ -138,17 +138,38 @@ class BlockSchema(BaseModel):
@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)
@ -157,32 +178,7 @@ class BlockSchema(BaseModel):
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)
@ -255,7 +251,7 @@ class Block(ABC, Generic[BlockSchemaInputType, BlockSchemaOutputType]):
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,
@ -297,10 +293,16 @@ class Block(ABC, Generic[BlockSchemaInputType, BlockSchemaOutputType]):
if self.webhook_config:
if isinstance(self.webhook_config, BlockWebhookConfig):
# Enforce presence of credentials field on auto-setup webhook blocks
if CREDENTIALS_FIELD_NAME not in self.input_schema.model_fields:
if not (cred_fields := self.input_schema.get_credentials_fields()):
raise TypeError(
"credentials field is required on auto-setup webhook blocks"
)
# Disallow multiple credentials inputs on webhook blocks
elif len(cred_fields) > 1:
raise ValueError(
"Multiple credentials inputs not supported on webhook blocks"
)
self.block_type = BlockType.WEBHOOK
else:
self.block_type = BlockType.WEBHOOK_MANUAL

44
autogpt_platform/backend/backend/data/model.py
View File

@ -245,7 +245,8 @@ CP = TypeVar("CP", bound=ProviderName)
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]):
@ -254,21 +255,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
@ -282,11 +283,20 @@ class CredentialsMetaInput(BaseModel, Generic[CP, CT]):
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]):

78
autogpt_platform/backend/backend/executor/manager.py
View File

@ -10,7 +10,6 @@ from contextlib import contextmanager
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 autogpt_libs.utils.cache import thread_cached
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 @@ from backend.data.execution import (
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 execute_node(
# 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
@ -893,41 +899,39 @@ class ExecutionManager(AppService):
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}"
for field_name, credentials_meta_type in credentials_fields.items():
credentials_meta = credentials_meta_type.model_validate(
node.input_default[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"
# 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 ------- #

66
autogpt_platform/backend/backend/integrations/webhooks/graph_lifecycle_hooks.py
View File

@ -1,9 +1,8 @@
import logging
from typing import TYPE_CHECKING, Callable, Optional, cast
from backend.data.block import BlockWebhookConfig, get_block
from backend.data.block import BlockSchema, BlockWebhookConfig, 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)
@ -119,14 +146,17 @@ async def on_node_activate(
else:
resource = "" # not relevant for manual webhooks
needs_credentials = CREDENTIALS_FIELD_NAME in block.input_schema.model_fields
block_input_schema = cast(BlockSchema, block.input_schema)
credentials_field_name = next(iter(block_input_schema.get_credentials_fields()), "")
credentials_meta = (
node.input_default.get(CREDENTIALS_FIELD_NAME) if needs_credentials else None
node.input_default.get(credentials_field_name)
if credentials_field_name
else None
)
event_filter_input_name = block.webhook_config.event_filter_input
has_everything_for_webhook = (
resource is not None
and (credentials_meta or not needs_credentials)
and (credentials_meta or not credentials_field_name)
and (
not event_filter_input_name
or (
@ -230,7 +260,7 @@ async def on_node_deactivate(
)
await webhooks_manager.prune_webhook_if_dangling(webhook.id, credentials)
if (
CREDENTIALS_FIELD_NAME in block.input_schema.model_fields
cast(BlockSchema, block.input_schema).get_credentials_fields()
and not credentials
):
logger.warning(

28
autogpt_platform/backend/backend/util/test.py
View File

@ -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):
raise TypeError(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}")

4
autogpt_platform/frontend/src/components/CustomNode.tsx
View File

@ -245,6 +245,7 @@ export function CustomNode({
].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);
@ -261,7 +262,8 @@ export function CustomNode({
side="left"
/>
) : (
propKey != "credentials" && (
propKey !== "credentials" &&
!propKey.endsWith("_credentials") && (
<div className="flex gap-1">
<span className="text-m green mb-0 text-gray-900 dark:text-gray-100">
{propSchema.title || beautifyString(propKey)}

24
autogpt_platform/frontend/src/components/integrations/credentials-input.tsx
View File

@ -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 = useBackendAPI();
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 && (
<APIKeyCredentialsModal
credentialsFieldName={selfKey}
open={isAPICredentialsModalOpen}
onClose={() => setAPICredentialsModalOpen(false)}
onCredentialsCreate={(credsMeta) => {
@ -242,7 +244,9 @@ export const CredentialsInput: FC<{
return (
<>
<div className="mb-2 flex gap-1">
<span className="text-m green text-gray-900">Credentials</span>
<span className="text-m green text-gray-900">
{providerName} Credentials
</span>
<SchemaTooltip description={schema.description} />
</div>
<div className={cn("flex flex-row space-x-2", className)}>
@ -310,7 +314,12 @@ export const CredentialsInput: FC<{
// Saved credentials exist
return (
<>
<span className="text-m green mb-0 text-gray-900">Credentials</span>
<div className="flex gap-1">
<span className="text-m green mb-0 text-gray-900">
{providerName} Credentials
</span>
<SchemaTooltip description={schema.description} />
</div>
<Select value={selectedCredentials?.id} onValueChange={handleValueChange}>
<SelectTrigger>
<SelectValue placeholder={schema.placeholder} />
@ -353,11 +362,12 @@ export const CredentialsInput: FC<{
};
export const APIKeyCredentialsModal: FC<{
credentialsFieldName: string;
open: boolean;
onClose: () => void;
onCredentialsCreate: (creds: CredentialsMetaInput) => void;
}> = ({ open, onClose, onCredentialsCreate }) => {
const credentials = useCredentials();
}> = ({ credentialsFieldName, open, onClose, onCredentialsCreate }) => {
const credentials = useCredentials(credentialsFieldName);
const formSchema = z.object({
apiKey: z.string().min(1, "API Key is required"),

1
autogpt_platform/frontend/src/components/node-input-components.tsx
View File

@ -653,6 +653,7 @@ const NodeCredentialsInput: FC<{
return (
<div className={cn("flex flex-col", className)}>
<CredentialsInput
selfKey={selfKey}
onSelectCredentials={(credsMeta) =>
handleInputChange(selfKey, credsMeta)
}

9
autogpt_platform/frontend/src/hooks/useCredentials.ts
View File

@ -26,7 +26,9 @@ export type CredentialsData =
isLoading: false;
});
export default function useCredentials(): CredentialsData | null {
export default function useCredentials(
inputFieldName: string,
): CredentialsData | null {
const nodeId = useNodeId();
const allProviders = useContext(CredentialsProvidersContext);
@ -35,8 +37,9 @@ export default function useCredentials(): CredentialsData | null {
}
const data = useNodesData<Node<CustomNodeData>>(nodeId)!.data;
const credentialsSchema = data.inputSchema.properties
.credentials as BlockIOCredentialsSubSchema;
const credentialsSchema = data.inputSchema.properties[
inputFieldName
] as BlockIOCredentialsSubSchema;
const discriminatorValue: CredentialsProviderName | null =
(credentialsSchema.discriminator &&

10
docs/content/platform/new_blocks.md
View File

@ -238,6 +238,16 @@ Naturally, to add an authenticated block for a new provider, you'll have to add
```
</details>
#### Multiple credentials inputs
Multiple credentials inputs are supported, under the following conditions:
- The name of each of the credentials input fields must end with `_credentials`.
- The names of the credentials input fields must match the names of the corresponding
parameters on the `run(..)` method of the block.
- If more than one of the credentials parameters are required, `test_credentials`
is a `dict[str, Credentials]`, with for each required credentials input the
parameter name as the key and suitable test credentials as the value.
#### Adding an OAuth2 service integration
To add support for a new OAuth2-authenticated service, you'll need to add an `OAuthHandler`.