From 43f736b969476bedba34f9c77672587f2361819d Mon Sep 17 00:00:00 2001 From: Reinier van der Leer Date: Fri, 6 Feb 2026 01:05:00 +0100 Subject: [PATCH] Make `credentials_input_schema` cheaper to compute --- .../backend/backend/data/block.py | 4 +- .../backend/backend/data/graph.py | 102 +++++++++++------- .../backend/backend/data/model.py | 17 ++- 3 files changed, 74 insertions(+), 49 deletions(-) diff --git a/autogpt_platform/backend/backend/data/block.py b/autogpt_platform/backend/backend/data/block.py index eb9360b037..f67134ceb3 100644 --- a/autogpt_platform/backend/backend/data/block.py +++ b/autogpt_platform/backend/backend/data/block.py @@ -246,7 +246,9 @@ class BlockSchema(BaseModel): f"is not of type {CredentialsMetaInput.__name__}" ) - credentials_fields[field_name].validate_credentials_field_schema(cls) + CredentialsMetaInput.validate_credentials_field_schema( + cls.get_field_schema(field_name), field_name + ) elif field_name in credentials_fields: raise KeyError( diff --git a/autogpt_platform/backend/backend/data/graph.py b/autogpt_platform/backend/backend/data/graph.py index e52d04b4fe..d040a1b851 100644 --- a/autogpt_platform/backend/backend/data/graph.py +++ b/autogpt_platform/backend/backend/data/graph.py @@ -20,7 +20,7 @@ from prisma.types import ( AgentNodeLinkCreateInput, StoreListingVersionWhereInput, ) -from pydantic import BaseModel, BeforeValidator, Field, create_model +from pydantic import BaseModel, BeforeValidator, Field from pydantic.fields import computed_field from backend.blocks.agent import AgentExecutorBlock @@ -30,7 +30,6 @@ from backend.data.db import prisma as db from backend.data.dynamic_fields import is_tool_pin, sanitize_pin_name from backend.data.includes import MAX_GRAPH_VERSIONS_FETCH from backend.data.model import ( - CredentialsField, CredentialsFieldInfo, CredentialsMetaInput, is_credentials_field_name, @@ -45,7 +44,6 @@ from .block import ( AnyBlockSchema, Block, BlockInput, - BlockSchema, BlockType, EmptySchema, get_block, @@ -392,42 +390,70 @@ class Graph(BaseGraph): f"keys: {keys} <> {other_keys}." ) - # Build the Pydantic model for the credentials input schema - fields: dict[str, tuple[type[CredentialsMetaInput], CredentialsMetaInput]] = { - agg_field_key: ( - ( - CMI - if ( - ( - CMI := CredentialsMetaInput[ - Literal[tuple(field_info.provider)], # type: ignore - Literal[tuple(field_info.supported_types)], # type: ignore - ] - ) - or True - ) - and is_required - else CMI | None - ), - CredentialsField( - required_scopes=set(field_info.required_scopes or []), - discriminator=field_info.discriminator, - discriminator_mapping=field_info.discriminator_mapping, - discriminator_values=field_info.discriminator_values, - ), - ) - for agg_field_key, ( - field_info, - _, - is_required, - ) in graph_credentials_inputs.items() - } + # Build JSON schema directly to avoid expensive create_model + validation overhead + properties = {} + required_fields = [] - return create_model( - self.name.replace(" ", "") + "CredentialsInputSchema", - __base__=BlockSchema, - **fields, # type: ignore - ).jsonschema() + for agg_field_key, (field_info, _, is_required) in graph_credentials_inputs.items(): + providers = list(field_info.provider) + cred_types = list(field_info.supported_types) + + field_schema: dict[str, Any] = { + "credentials_provider": providers, + "credentials_types": cred_types, + "type": "object", + "properties": { + "id": {"title": "Id", "type": "string"}, + "title": { + "anyOf": [{"type": "string"}, {"type": "null"}], + "default": None, + "title": "Title", + }, + "provider": { + "title": "Provider", + "type": "string", + **( + {"enum": providers} + if len(providers) > 1 + else {"const": providers[0]} + ), + }, + "type": { + "title": "Type", + "type": "string", + **( + {"enum": cred_types} + if len(cred_types) > 1 + else {"const": cred_types[0]} + ), + }, + }, + "required": ["id", "provider", "type"], + } + + # Add other (optional) field info items + field_schema.update( + **field_info.model_dump( + by_alias=True, + exclude_defaults=True, + exclude={"provider", "supported_types"}, # already included above + ) + ) + + # Ensure field schema is well-formed + CredentialsMetaInput.validate_credentials_field_schema( + field_schema, agg_field_key + ) + + properties[agg_field_key] = field_schema + if is_required: + required_fields.append(agg_field_key) + + return { + "type": "object", + "properties": properties, + "required": required_fields, + } def aggregate_credentials_inputs( self, diff --git a/autogpt_platform/backend/backend/data/model.py b/autogpt_platform/backend/backend/data/model.py index 5a09c591c9..7bdfef059b 100644 --- a/autogpt_platform/backend/backend/data/model.py +++ b/autogpt_platform/backend/backend/data/model.py @@ -163,7 +163,6 @@ class User(BaseModel): if TYPE_CHECKING: from prisma.models import User as PrismaUser - from backend.data.block import BlockSchema T = TypeVar("T") logger = logging.getLogger(__name__) @@ -508,15 +507,13 @@ class CredentialsMetaInput(BaseModel, Generic[CP, CT]): 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"]): + @staticmethod + def validate_credentials_field_schema( + field_schema: dict[str, Any], field_name: str + ): """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 = CredentialsFieldInfo[CP, CT].model_validate(field_schema) + field_info = CredentialsFieldInfo[CP, CT].model_validate(field_schema) except ValidationError as e: if "Field required [type=missing" not in str(e): raise @@ -526,11 +523,11 @@ class CredentialsMetaInput(BaseModel, Generic[CP, CT]): f"{field_schema}" ) from e - providers = cls.allowed_providers() + providers = field_info.provider if ( providers is not None and len(providers) > 1 - and not schema_extra.discriminator + and not field_info.discriminator ): raise TypeError( f"Multi-provider CredentialsField '{field_name}' "