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 <122007096+Abhi1992002@users.noreply.github.com>
This commit is contained in:
Krzysztof Czerwinski
2025-09-06 13:21:48 +09:00
committed by GitHub
parent 46e0f6cc45
commit cfc975d39b
13 changed files with 169 additions and 71 deletions

View File

@@ -91,6 +91,45 @@ class BlockCategory(Enum):
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 @@ class Block(ABC, Generic[BlockSchemaInputType, BlockSchemaOutputType]):
"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(

View File

@@ -29,8 +29,7 @@ from backend.blocks.replicate.replicate_block import ReplicateModelBlock
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,

View File

@@ -1,32 +0,0 @@
from enum import Enum
from typing import Any, Optional
from pydantic import BaseModel
from backend.data.block import BlockInput
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,
)

View File

@@ -2,7 +2,7 @@ import logging
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 pydantic import BaseModel
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.models import Pagination
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)

View File

@@ -10,9 +10,15 @@ from pydantic import BaseModel, JsonValue, ValidationError
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,

View File

@@ -63,7 +63,7 @@ except ImportError:
# 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

View File

@@ -8,7 +8,7 @@ from typing import Callable, List, Optional, Type
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,

View File

@@ -8,9 +8,8 @@ BLOCK_COSTS configuration used by the execution system.
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__)

View File

@@ -7,7 +7,7 @@ from typing import Any, Callable, List, Optional, Set, Type
from pydantic import BaseModel, SecretStr
from backend.data.cost import BlockCost
from backend.data.block import BlockCost
from backend.data.model import (
APIKeyCredentials,
Credentials,

View File

@@ -7,12 +7,10 @@ import prisma
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 @@ from backend.util.models import Pagination
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)

View File

@@ -1,9 +1,10 @@
from typing import Any, Literal
from typing import Literal
from pydantic import BaseModel
import backend.server.v2.library.model as library_model
import backend.server.v2.store.model as store_model
from backend.data.block import BlockInfo
from backend.integrations.providers import ProviderName
from backend.util.models import Pagination
@@ -16,29 +17,27 @@ FilterType = Literal[
BlockType = Literal["all", "input", "action", "output"]
BlockData = dict[str, Any]
# Suggestions
class SuggestionsResponse(BaseModel):
otto_suggestions: list[str]
recent_searches: list[str]
providers: list[ProviderName]
top_blocks: list[BlockData]
top_blocks: list[BlockInfo]
# All blocks
class BlockCategoryResponse(BaseModel):
name: str
total_blocks: int
blocks: list[BlockData]
blocks: list[BlockInfo]
model_config = {"use_enum_values": False} # <== use enum names like "AI"
# Input/Action/Output and see all for block categories
class BlockResponse(BaseModel):
blocks: list[BlockData]
blocks: list[BlockInfo]
pagination: Pagination
@@ -71,7 +70,7 @@ class SearchBlocksResponse(BaseModel):
class SearchResponse(BaseModel):
items: list[BlockData | library_model.LibraryAgent | store_model.StoreAgent]
items: list[BlockInfo | library_model.LibraryAgent | store_model.StoreAgent]
total_items: dict[FilterType, int]
page: int
more_pages: bool

View File

@@ -370,7 +370,7 @@ class TestProviderBuilder:
def test_provider_builder_with_base_cost(self):
"""Test building a provider with base costs."""
from backend.data.cost import BlockCostType
from backend.data.block import BlockCostType
provider = (
ProviderBuilder("cost_test")
@@ -411,7 +411,7 @@ class TestProviderBuilder:
def test_provider_builder_complete_example(self):
"""Test building a complete provider with all features."""
from backend.data.cost import BlockCostType
from backend.data.block import BlockCostType
class TestOAuth(BaseOAuthHandler):
PROVIDER_NAME = ProviderName.GITHUB

View File

@@ -4759,7 +4759,7 @@
"name": { "type": "string", "title": "Name" },
"total_blocks": { "type": "integer", "title": "Total Blocks" },
"blocks": {
"items": { "additionalProperties": true, "type": "object" },
"items": { "$ref": "#/components/schemas/BlockInfo" },
"type": "array",
"title": "Blocks"
}
@@ -4768,10 +4768,80 @@
"required": ["name", "total_blocks", "blocks"],
"title": "BlockCategoryResponse"
},
"BlockCost": {
"properties": {
"cost_amount": { "type": "integer", "title": "Cost Amount" },
"cost_filter": {
"additionalProperties": true,
"type": "object",
"title": "Cost Filter"
},
"cost_type": { "$ref": "#/components/schemas/BlockCostType" }
},
"type": "object",
"required": ["cost_amount", "cost_filter", "cost_type"],
"title": "BlockCost"
},
"BlockCostType": {
"type": "string",
"enum": ["run", "byte", "second"],
"title": "BlockCostType"
},
"BlockInfo": {
"properties": {
"id": { "type": "string", "title": "Id" },
"name": { "type": "string", "title": "Name" },
"inputSchema": {
"additionalProperties": true,
"type": "object",
"title": "Inputschema"
},
"outputSchema": {
"additionalProperties": true,
"type": "object",
"title": "Outputschema"
},
"costs": {
"items": { "$ref": "#/components/schemas/BlockCost" },
"type": "array",
"title": "Costs"
},
"description": { "type": "string", "title": "Description" },
"categories": {
"items": {
"additionalProperties": { "type": "string" },
"type": "object"
},
"type": "array",
"title": "Categories"
},
"contributors": {
"items": { "additionalProperties": true, "type": "object" },
"type": "array",
"title": "Contributors"
},
"staticOutput": { "type": "boolean", "title": "Staticoutput" },
"uiType": { "type": "string", "title": "Uitype" }
},
"type": "object",
"required": [
"id",
"name",
"inputSchema",
"outputSchema",
"costs",
"description",
"categories",
"contributors",
"staticOutput",
"uiType"
],
"title": "BlockInfo"
},
"BlockResponse": {
"properties": {
"blocks": {
"items": { "additionalProperties": true, "type": "object" },
"items": { "$ref": "#/components/schemas/BlockInfo" },
"type": "array",
"title": "Blocks"
},
@@ -7039,7 +7109,7 @@
"items": {
"items": {
"anyOf": [
{ "additionalProperties": true, "type": "object" },
{ "$ref": "#/components/schemas/BlockInfo" },
{ "$ref": "#/components/schemas/LibraryAgent" },
{ "$ref": "#/components/schemas/StoreAgent" }
]
@@ -7528,7 +7598,7 @@
"title": "Providers"
},
"top_blocks": {
"items": { "additionalProperties": true, "type": "object" },
"items": { "$ref": "#/components/schemas/BlockInfo" },
"type": "array",
"title": "Top Blocks"
}