feat(backend): Type for API block data response (#10763)

Moving to auto-generated frontend types caused returned blocks data to
no longer have proper typing.

### Changes 🏗️

- Add `BlockInfo` model and `get_info` function that returns it to the
`Block` class, including costs
- Move `BlockCost` and `BlockCostType` to `block.py` to prevent circular
imports

### 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] Endpoints using the new type work correctly

Co-authored-by: Abhimanyu Yadav <[email protected]>

Author: kcze

autogpt_platform/backend/backend/data/block.py

@@ -91,6 +91,45 @@ def dict(self) -> dict[str, str]:
         return {"category": self.name, "description": self.value}
 
 
+class BlockCostType(str, Enum):
+    RUN = "run"  # cost X credits per run
+    BYTE = "byte"  # cost X credits per byte
+    SECOND = "second"  # cost X credits per second
+
+
+class BlockCost(BaseModel):
+    cost_amount: int
+    cost_filter: BlockInput
+    cost_type: BlockCostType
+
+    def __init__(
+        self,
+        cost_amount: int,
+        cost_type: BlockCostType = BlockCostType.RUN,
+        cost_filter: Optional[BlockInput] = None,
+        **data: Any,
+    ) -> None:
+        super().__init__(
+            cost_amount=cost_amount,
+            cost_filter=cost_filter or {},
+            cost_type=cost_type,
+            **data,
+        )
+
+
+class BlockInfo(BaseModel):
+    id: str
+    name: str
+    inputSchema: dict[str, Any]
+    outputSchema: dict[str, Any]
+    costs: list[BlockCost]
+    description: str
+    categories: list[dict[str, str]]
+    contributors: list[dict[str, Any]]
+    staticOutput: bool
+    uiType: str
+
+
 class BlockSchema(BaseModel):
     cached_jsonschema: ClassVar[dict[str, Any]]
 
@@ -454,6 +493,24 @@ def to_dict(self):
             "uiType": self.block_type.value,
         }
 
+    def get_info(self) -> BlockInfo:
+        from backend.data.credit import get_block_cost
+
+        return BlockInfo(
+            id=self.id,
+            name=self.name,
+            inputSchema=self.input_schema.jsonschema(),
+            outputSchema=self.output_schema.jsonschema(),
+            costs=get_block_cost(self),
+            description=self.description,
+            categories=[category.dict() for category in self.categories],
+            contributors=[
+                contributor.model_dump() for contributor in self.contributors
+            ],
+            staticOutput=self.static_output,
+            uiType=self.block_type.value,
+        )
+
     async def execute(self, input_data: BlockInput, **kwargs) -> BlockOutput:
         if error := self.input_schema.validate_data(input_data):
             raise ValueError(

autogpt_platform/backend/backend/data/block_cost_config.py

@@ -29,8 +29,7 @@
 from backend.blocks.smart_decision_maker import SmartDecisionMakerBlock
 from backend.blocks.talking_head import CreateTalkingAvatarVideoBlock
 from backend.blocks.text_to_speech_block import UnrealTextToSpeechBlock
-from backend.data.block import Block
-from backend.data.cost import BlockCost, BlockCostType
+from backend.data.block import Block, BlockCost, BlockCostType
 from backend.integrations.credentials_store import (
     aiml_api_credentials,
     anthropic_credentials,

autogpt_platform/backend/backend/data/cost.py

@@ -2,7 +2,7 @@
 from abc import ABC, abstractmethod
 from collections import defaultdict
 from datetime import datetime, timezone
-from typing import Any, cast
+from typing import TYPE_CHECKING, Any, cast
 
 import stripe
 from prisma import Json
@@ -23,7 +23,6 @@
 
 from backend.data import db
 from backend.data.block_cost_config import BLOCK_COSTS
-from backend.data.cost import BlockCost
 from backend.data.model import (
     AutoTopUpConfig,
     RefundRequest,
@@ -41,6 +40,9 @@
 from backend.util.retry import func_retry
 from backend.util.settings import Settings
 
+if TYPE_CHECKING:
+    from backend.data.block import Block, BlockCost
+
 settings = Settings()
 stripe.api_key = settings.secrets.stripe_api_key
 logger = logging.getLogger(__name__)
@@ -997,10 +999,14 @@ def get_user_credit_model() -> UserCreditBase:
     return UserCredit()
 
 
-def get_block_costs() -> dict[str, list[BlockCost]]:
+def get_block_costs() -> dict[str, list["BlockCost"]]:
     return {block().id: costs for block, costs in BLOCK_COSTS.items()}
 
 
+def get_block_cost(block: "Block") -> list["BlockCost"]:
+    return BLOCK_COSTS.get(block.__class__, [])
+
+
 async def get_stripe_customer_id(user_id: str) -> str:
     user = await get_user_by_id(user_id)
 

autogpt_platform/backend/backend/data/credit.py

@@ -10,9 +10,15 @@
 
 from backend.data import execution as execution_db
 from backend.data import graph as graph_db
-from backend.data.block import Block, BlockInput, BlockOutputEntry, BlockType, get_block
+from backend.data.block import (
+    Block,
+    BlockCostType,
+    BlockInput,
+    BlockOutputEntry,
+    BlockType,
+    get_block,
+)
 from backend.data.block_cost_config import BLOCK_COSTS
-from backend.data.cost import BlockCostType
 from backend.data.db import prisma
 from backend.data.execution import (
     ExecutionStatus,

autogpt_platform/backend/backend/executor/utils.py

@@ -63,7 +63,7 @@
 
 # Cost System
 try:
-    from backend.data.cost import BlockCost, BlockCostType
+    from backend.data.block import BlockCost, BlockCostType
 except ImportError:
     from backend.data.block_cost_config import BlockCost, BlockCostType
 

autogpt_platform/backend/backend/sdk/__init__.py

@@ -8,7 +8,7 @@
 
 from pydantic import SecretStr
 
-from backend.data.cost import BlockCost, BlockCostType
+from backend.data.block import BlockCost, BlockCostType
 from backend.data.model import (
     APIKeyCredentials,
     Credentials,

autogpt_platform/backend/backend/sdk/builder.py

@@ -8,9 +8,8 @@
 import logging
 from typing import List, Type
 
-from backend.data.block import Block
+from backend.data.block import Block, BlockCost
 from backend.data.block_cost_config import BLOCK_COSTS
-from backend.data.cost import BlockCost
 from backend.sdk.registry import AutoRegistry
 
 logger = logging.getLogger(__name__)

autogpt_platform/backend/backend/sdk/cost_integration.py

@@ -7,7 +7,7 @@
 
 from pydantic import BaseModel, SecretStr
 
-from backend.data.cost import BlockCost
+from backend.data.block import BlockCost
 from backend.data.model import (
     APIKeyCredentials,
     Credentials,

autogpt_platform/backend/backend/sdk/provider.py

@@ -7,12 +7,10 @@
 import backend.data.block
 from backend.blocks import load_all_blocks
 from backend.blocks.llm import LlmModel
-from backend.data.block import Block, BlockCategory, BlockSchema
-from backend.data.credit import get_block_costs
+from backend.data.block import Block, BlockCategory, BlockInfo, BlockSchema
 from backend.integrations.providers import ProviderName
 from backend.server.v2.builder.model import (
     BlockCategoryResponse,
-    BlockData,
     BlockResponse,
     BlockType,
     CountResponse,
@@ -25,7 +23,7 @@
 logger = logging.getLogger(__name__)
 llm_models = [name.name.lower().replace("_", " ") for name in LlmModel]
 _static_counts_cache: dict | None = None
-_suggested_blocks: list[BlockData] | None = None
+_suggested_blocks: list[BlockInfo] | None = None
 
 
 def get_block_categories(category_blocks: int = 3) -> list[BlockCategoryResponse]:
@@ -53,7 +51,7 @@ def get_block_categories(category_blocks: int = 3) -> list[BlockCategoryResponse
 
             # Append if the category has less than the specified number of blocks
             if len(categories[category].blocks) < category_blocks:
-                categories[category].blocks.append(block.to_dict())
+                categories[category].blocks.append(block.get_info())
 
     # Sort categories by name
     return sorted(categories.values(), key=lambda x: x.name)
@@ -109,10 +107,8 @@ def get_blocks(
             take -= 1
             blocks.append(block)
 
-    costs = get_block_costs()
-
     return BlockResponse(
-        blocks=[{**b.to_dict(), "costs": costs.get(b.id, [])} for b in blocks],
+        blocks=[b.get_info() for b in blocks],
         pagination=Pagination(
             total_items=total,
             total_pages=(total + page_size - 1) // page_size,
@@ -174,11 +170,9 @@ def search_blocks(
             take -= 1
             blocks.append(block)
 
-    costs = get_block_costs()
-
     return SearchBlocksResponse(
         blocks=BlockResponse(
-            blocks=[{**b.to_dict(), "costs": costs.get(b.id, [])} for b in blocks],
+            blocks=[b.get_info() for b in blocks],
             pagination=Pagination(
                 total_items=total,
                 total_pages=(total + page_size - 1) // page_size,
@@ -323,7 +317,7 @@ def _get_all_providers() -> dict[ProviderName, Provider]:
     return providers
 
 
-async def get_suggested_blocks(count: int = 5) -> list[BlockData]:
+async def get_suggested_blocks(count: int = 5) -> list[BlockInfo]:
     global _suggested_blocks
 
     if _suggested_blocks is not None and len(_suggested_blocks) >= count:
@@ -351,7 +345,7 @@ async def get_suggested_blocks(count: int = 5) -> list[BlockData]:
 
     # Get the top blocks based on execution count
     # But ignore Input and Output blocks
-    blocks: list[tuple[BlockData, int]] = []
+    blocks: list[tuple[BlockInfo, int]] = []
 
     for block_type in load_all_blocks().values():
         block: Block[BlockSchema, BlockSchema] = block_type()
@@ -366,7 +360,7 @@ async def get_suggested_blocks(count: int = 5) -> list[BlockData]:
             (row["execution_count"] for row in results if row["block_id"] == block.id),
             0,
         )
-        blocks.append((block.to_dict(), execution_count))
+        blocks.append((block.get_info(), execution_count))
     # Sort blocks by execution count
     blocks.sort(key=lambda x: x[1], reverse=True)