refactor(backend): Simplify CredentialsField usage + use `ProviderN…

…ame` globally (#8725) - Resolves #8931 - Follow-up to #8358 ### Changes 🏗️ - Avoid double specifying provider and cred types on `credentials` inputs - Move `credentials` sub-schema validation from `CredentialsField` to `CredentialsMetaInput.validate_credentials_field_schema(..)`, which is called in `BlockSchema.__pydantic_init_subclass__` - Use `ProviderName` enum globally
Significant-Gravitas · Dec 11, 2024 · 33b9eef · 33b9eef
1 parent b8a3ffc
commit 33b9eef
Show file tree

Hide file tree

Showing 38 changed files with 314 additions and 192 deletions.
diff --git a/autogpt_platform/backend/backend/blocks/ai_image_generator_block.py b/autogpt_platform/backend/backend/blocks/ai_image_generator_block.py
@@ -12,6 +12,7 @@
     CredentialsMetaInput,
     SchemaField,
 )
+from backend.integrations.providers import ProviderName
 
 
 class ImageSize(str, Enum):
@@ -101,12 +102,10 @@ class ImageGenModel(str, Enum):
 
 class AIImageGeneratorBlock(Block):
     class Input(BlockSchema):
-        credentials: CredentialsMetaInput[Literal["replicate"], Literal["api_key"]] = (
-            CredentialsField(
-                provider="replicate",
-                supported_credential_types={"api_key"},
-                description="Enter your Replicate API key to access the image generation API. You can obtain an API key from https://replicate.com/account/api-tokens.",
-            )
+        credentials: CredentialsMetaInput[
+            Literal[ProviderName.REPLICATE], Literal["api_key"]
+        ] = CredentialsField(
+            description="Enter your Replicate API key to access the image generation API. You can obtain an API key from https://replicate.com/account/api-tokens.",
         )
         prompt: str = SchemaField(
             description="Text prompt for image generation",

diff --git a/autogpt_platform/backend/backend/blocks/ai_music_generator.py b/autogpt_platform/backend/backend/blocks/ai_music_generator.py
@@ -13,6 +13,7 @@
     CredentialsMetaInput,
     SchemaField,
 )
+from backend.integrations.providers import ProviderName
 
 logger = logging.getLogger(__name__)
 
@@ -54,13 +55,11 @@ class NormalizationStrategy(str, Enum):
 
 class AIMusicGeneratorBlock(Block):
     class Input(BlockSchema):
-        credentials: CredentialsMetaInput[Literal["replicate"], Literal["api_key"]] = (
-            CredentialsField(
-                provider="replicate",
-                supported_credential_types={"api_key"},
-                description="The Replicate integration can be used with "
-                "any API key with sufficient permissions for the blocks it is used on.",
-            )
+        credentials: CredentialsMetaInput[
+            Literal[ProviderName.REPLICATE], Literal["api_key"]
+        ] = CredentialsField(
+            description="The Replicate integration can be used with "
+            "any API key with sufficient permissions for the blocks it is used on.",
         )
         prompt: str = SchemaField(
             description="A description of the music you want to generate",

diff --git a/autogpt_platform/backend/backend/blocks/ai_shortform_video_block.py b/autogpt_platform/backend/backend/blocks/ai_shortform_video_block.py
@@ -12,6 +12,7 @@
     CredentialsMetaInput,
     SchemaField,
 )
+from backend.integrations.providers import ProviderName
 from backend.util.request import requests
 
 TEST_CREDENTIALS = APIKeyCredentials(
@@ -140,13 +141,11 @@ class VisualMediaType(str, Enum):
 
 class AIShortformVideoCreatorBlock(Block):
     class Input(BlockSchema):
-        credentials: CredentialsMetaInput[Literal["revid"], Literal["api_key"]] = (
-            CredentialsField(
-                provider="revid",
-                supported_credential_types={"api_key"},
-                description="The revid.ai integration can be used with "
-                "any API key with sufficient permissions for the blocks it is used on.",
-            )
+        credentials: CredentialsMetaInput[
+            Literal[ProviderName.REVID], Literal["api_key"]
+        ] = CredentialsField(
+            description="The revid.ai integration can be used with "
+            "any API key with sufficient permissions for the blocks it is used on.",
         )
         script: str = SchemaField(
             description="""1. Use short and punctuated sentences\n\n2. Use linebreaks to create a new clip\n\n3. Text outside of brackets is spoken by the AI, and [text between brackets] will be used to guide the visual generation. For example, [close-up of a cat] will show a close-up of a cat.""",

diff --git a/autogpt_platform/backend/backend/blocks/code_executor.py b/autogpt_platform/backend/backend/blocks/code_executor.py
@@ -11,6 +11,7 @@
     CredentialsMetaInput,
     SchemaField,
 )
+from backend.integrations.providers import ProviderName
 
 TEST_CREDENTIALS = APIKeyCredentials(
     id="01234567-89ab-cdef-0123-456789abcdef",
@@ -39,12 +40,10 @@ class CodeExecutionBlock(Block):
     # TODO : Add support to upload and download files
     # Currently, You can customized the CPU and Memory, only by creating a pre customized sandbox template
     class Input(BlockSchema):
-        credentials: CredentialsMetaInput[Literal["e2b"], Literal["api_key"]] = (
-            CredentialsField(
-                provider="e2b",
-                supported_credential_types={"api_key"},
-                description="Enter your api key for the E2B Sandbox. You can get it in here - https://e2b.dev/docs",
-            )
+        credentials: CredentialsMetaInput[
+            Literal[ProviderName.E2B], Literal["api_key"]
+        ] = CredentialsField(
+            description="Enter your api key for the E2B Sandbox. You can get it in here - https://e2b.dev/docs",
         )
 
         # Todo : Option to run commond in background

diff --git a/autogpt_platform/backend/backend/blocks/discord.py b/autogpt_platform/backend/backend/blocks/discord.py
@@ -12,16 +12,15 @@
     CredentialsMetaInput,
     SchemaField,
 )
+from backend.integrations.providers import ProviderName
 
-DiscordCredentials = CredentialsMetaInput[Literal["discord"], Literal["api_key"]]
+DiscordCredentials = CredentialsMetaInput[
+    Literal[ProviderName.DISCORD], Literal["api_key"]
+]
 
 
 def DiscordCredentialsField() -> DiscordCredentials:
-    return CredentialsField(
-        description="Discord bot token",
-        provider="discord",
-        supported_credential_types={"api_key"},
-    )
+    return CredentialsField(description="Discord bot token")
 
 
 TEST_CREDENTIALS = APIKeyCredentials(

diff --git a/autogpt_platform/backend/backend/blocks/exa/_auth.py b/autogpt_platform/backend/backend/blocks/exa/_auth.py
@@ -3,10 +3,11 @@
 from pydantic import SecretStr
 
 from backend.data.model import APIKeyCredentials, CredentialsField, CredentialsMetaInput
+from backend.integrations.providers import ProviderName
 
 ExaCredentials = APIKeyCredentials
 ExaCredentialsInput = CredentialsMetaInput[
-    Literal["exa"],
+    Literal[ProviderName.EXA],
     Literal["api_key"],
 ]
 
@@ -28,8 +29,4 @@
 
 def ExaCredentialsField() -> ExaCredentialsInput:
     """Creates an Exa credentials input on a block."""
-    return CredentialsField(
-        provider="exa",
-        supported_credential_types={"api_key"},
-        description="The Exa integration requires an API Key.",
-    )
+    return CredentialsField(description="The Exa integration requires an API Key.")
diff --git a/autogpt_platform/backend/backend/blocks/fal/_auth.py b/autogpt_platform/backend/backend/blocks/fal/_auth.py
@@ -3,10 +3,11 @@
 from pydantic import SecretStr
 
 from backend.data.model import APIKeyCredentials, CredentialsField, CredentialsMetaInput
+from backend.integrations.providers import ProviderName
 
 FalCredentials = APIKeyCredentials
 FalCredentialsInput = CredentialsMetaInput[
-    Literal["fal"],
+    Literal[ProviderName.FAL],
     Literal["api_key"],
 ]
 
@@ -30,7 +31,5 @@ def FalCredentialsField() -> FalCredentialsInput:
     Creates a FAL credentials input on a block.
     """
     return CredentialsField(
-        provider="fal",
-        supported_credential_types={"api_key"},
         description="The FAL integration can be used with an API Key.",
     )
diff --git a/autogpt_platform/backend/backend/blocks/github/_auth.py b/autogpt_platform/backend/backend/blocks/github/_auth.py
@@ -8,6 +8,7 @@
     CredentialsMetaInput,
     OAuth2Credentials,
 )
+from backend.integrations.providers import ProviderName
 from backend.util.settings import Secrets
 
 secrets = Secrets()
@@ -17,7 +18,7 @@
 
 GithubCredentials = APIKeyCredentials | OAuth2Credentials
 GithubCredentialsInput = CredentialsMetaInput[
-    Literal["github"],
+    Literal[ProviderName.GITHUB],
     Literal["api_key", "oauth2"] if GITHUB_OAUTH_IS_CONFIGURED else Literal["api_key"],
 ]
 
@@ -30,10 +31,6 @@ def GithubCredentialsField(scope: str) -> GithubCredentialsInput:
         scope: The authorization scope needed for the block to work. ([list of available scopes](https://docs.github.com/en/apps/oauth-apps/building-oauth-apps/scopes-for-oauth-apps#available-scopes))
     """  # noqa
     return CredentialsField(
-        provider="github",
-        supported_credential_types=(
-            {"api_key", "oauth2"} if GITHUB_OAUTH_IS_CONFIGURED else {"api_key"}
-        ),
         required_scopes={scope},
         description="The GitHub integration can be used with OAuth, "
         "or any API key with sufficient permissions for the blocks it is used on.",

diff --git a/autogpt_platform/backend/backend/blocks/google/_auth.py b/autogpt_platform/backend/backend/blocks/google/_auth.py
@@ -3,6 +3,7 @@
 from pydantic import SecretStr
 
 from backend.data.model import CredentialsField, CredentialsMetaInput, OAuth2Credentials
+from backend.integrations.providers import ProviderName
 from backend.util.settings import Secrets
 
 # --8<-- [start:GoogleOAuthIsConfigured]
@@ -12,7 +13,9 @@
 )
 # --8<-- [end:GoogleOAuthIsConfigured]
 GoogleCredentials = OAuth2Credentials
-GoogleCredentialsInput = CredentialsMetaInput[Literal["google"], Literal["oauth2"]]
+GoogleCredentialsInput = CredentialsMetaInput[
+    Literal[ProviderName.GOOGLE], Literal["oauth2"]
+]
 
 
 def GoogleCredentialsField(scopes: list[str]) -> GoogleCredentialsInput:
@@ -23,8 +26,6 @@ def GoogleCredentialsField(scopes: list[str]) -> GoogleCredentialsInput:
         scopes: The authorization scopes needed for the block to work.
     """
     return CredentialsField(
-        provider="google",
-        supported_credential_types={"oauth2"},
         required_scopes=set(scopes),
         description="The Google integration requires OAuth2 authentication.",
     )

diff --git a/autogpt_platform/backend/backend/blocks/google_maps.py b/autogpt_platform/backend/backend/blocks/google_maps.py
@@ -10,6 +10,7 @@
     CredentialsMetaInput,
     SchemaField,
 )
+from backend.integrations.providers import ProviderName
 
 TEST_CREDENTIALS = APIKeyCredentials(
     id="01234567-89ab-cdef-0123-456789abcdef",
@@ -38,12 +39,8 @@ class Place(BaseModel):
 class GoogleMapsSearchBlock(Block):
     class Input(BlockSchema):
         credentials: CredentialsMetaInput[
-            Literal["google_maps"], Literal["api_key"]
-        ] = CredentialsField(
-            provider="google_maps",
-            supported_credential_types={"api_key"},
-            description="Google Maps API Key",
-        )
+            Literal[ProviderName.GOOGLE_MAPS], Literal["api_key"]
+        ] = CredentialsField(description="Google Maps API Key")
         query: str = SchemaField(
             description="Search query for local businesses",
             placeholder="e.g., 'restaurants in New York'",

diff --git a/autogpt_platform/backend/backend/blocks/hubspot/_auth.py b/autogpt_platform/backend/backend/blocks/hubspot/_auth.py
@@ -3,19 +3,18 @@
 from pydantic import SecretStr
 
 from backend.data.model import APIKeyCredentials, CredentialsField, CredentialsMetaInput
+from backend.integrations.providers import ProviderName
 
 HubSpotCredentials = APIKeyCredentials
 HubSpotCredentialsInput = CredentialsMetaInput[
-    Literal["hubspot"],
+    Literal[ProviderName.HUBSPOT],
     Literal["api_key"],
 ]
 
 
 def HubSpotCredentialsField() -> HubSpotCredentialsInput:
     """Creates a HubSpot credentials input on a block."""
     return CredentialsField(
-        provider="hubspot",
-        supported_credential_types={"api_key"},
         description="The HubSpot integration requires an API Key.",
     )
 

diff --git a/autogpt_platform/backend/backend/blocks/ideogram.py b/autogpt_platform/backend/backend/blocks/ideogram.py
@@ -11,6 +11,7 @@
     CredentialsMetaInput,
     SchemaField,
 )
+from backend.integrations.providers import ProviderName
 from backend.util.request import requests
 
 TEST_CREDENTIALS = APIKeyCredentials(
@@ -83,13 +84,10 @@ class UpscaleOption(str, Enum):
 
 class IdeogramModelBlock(Block):
     class Input(BlockSchema):
-
-        credentials: CredentialsMetaInput[Literal["ideogram"], Literal["api_key"]] = (
-            CredentialsField(
-                provider="ideogram",
-                supported_credential_types={"api_key"},
-                description="The Ideogram integration can be used with any API key with sufficient permissions for the blocks it is used on.",
-            )
+        credentials: CredentialsMetaInput[
+            Literal[ProviderName.IDEOGRAM], Literal["api_key"]
+        ] = CredentialsField(
+            description="The Ideogram integration can be used with any API key with sufficient permissions for the blocks it is used on.",
         )
         prompt: str = SchemaField(
             description="Text prompt for image generation",

diff --git a/autogpt_platform/backend/backend/blocks/jina/_auth.py b/autogpt_platform/backend/backend/blocks/jina/_auth.py
@@ -3,36 +3,21 @@
 from pydantic import SecretStr
 
 from backend.data.model import APIKeyCredentials, CredentialsField, CredentialsMetaInput
+from backend.integrations.providers import ProviderName
 
 JinaCredentials = APIKeyCredentials
 JinaCredentialsInput = CredentialsMetaInput[
-    Literal["jina"],
+    Literal[ProviderName.JINA],
     Literal["api_key"],
 ]
 
-TEST_CREDENTIALS = APIKeyCredentials(
-    id="01234567-89ab-cdef-0123-456789abcdef",
-    provider="jina",
-    api_key=SecretStr("mock-jina-api-key"),
-    title="Mock Jina API key",
-    expires_at=None,
-)
-TEST_CREDENTIALS_INPUT = {
-    "provider": TEST_CREDENTIALS.provider,
-    "id": TEST_CREDENTIALS.id,
-    "type": TEST_CREDENTIALS.type,
-    "title": TEST_CREDENTIALS.type,
-}
-
 
 def JinaCredentialsField() -> JinaCredentialsInput:
     """
     Creates a Jina credentials input on a block.
 
     """
     return CredentialsField(
-        provider="jina",
-        supported_credential_types={"api_key"},
         description="The Jina integration can be used with an API Key.",
     )
 

diff --git a/autogpt_platform/backend/backend/blocks/llm.py b/autogpt_platform/backend/backend/blocks/llm.py
@@ -7,6 +7,8 @@
 
 from pydantic import SecretStr
 
+from backend.integrations.providers import ProviderName
+
 if TYPE_CHECKING:
     from enum import _EnumMemberT
 
@@ -27,7 +29,13 @@
 
 logger = logging.getLogger(__name__)
 
-LLMProviderName = Literal["anthropic", "groq", "openai", "ollama", "open_router"]
+LLMProviderName = Literal[
+    ProviderName.ANTHROPIC,
+    ProviderName.GROQ,
+    ProviderName.OLLAMA,
+    ProviderName.OPENAI,
+    ProviderName.OPEN_ROUTER,
+]
 AICredentials = CredentialsMetaInput[LLMProviderName, Literal["api_key"]]
 
 TEST_CREDENTIALS = APIKeyCredentials(
@@ -48,8 +56,6 @@
 def AICredentialsField() -> AICredentials:
     return CredentialsField(
         description="API key for the LLM provider.",
-        provider=["anthropic", "groq", "openai", "ollama", "open_router"],
-        supported_credential_types={"api_key"},
         discriminator="model",
         discriminator_mapping={
             model.value: model.metadata.provider for model in LlmModel

diff --git a/autogpt_platform/backend/backend/blocks/medium.py b/autogpt_platform/backend/backend/blocks/medium.py
@@ -12,6 +12,7 @@
     SchemaField,
     SecretField,
 )
+from backend.integrations.providers import ProviderName
 from backend.util.request import requests
 
 TEST_CREDENTIALS = APIKeyCredentials(
@@ -77,12 +78,10 @@ class Input(BlockSchema):
             description="Whether to notify followers that the user has published",
             placeholder="False",
         )
-        credentials: CredentialsMetaInput[Literal["medium"], Literal["api_key"]] = (
-            CredentialsField(
-                provider="medium",
-                supported_credential_types={"api_key"},
-                description="The Medium integration can be used with any API key with sufficient permissions for the blocks it is used on.",
-            )
+        credentials: CredentialsMetaInput[
+            Literal[ProviderName.MEDIUM], Literal["api_key"]
+        ] = CredentialsField(
+            description="The Medium integration can be used with any API key with sufficient permissions for the blocks it is used on.",
         )
 
     class Output(BlockSchema):