airbyte-cdk 6.33.2.dev1__py3-none-any.whl → 6.33.4__py3-none-any.whl

airbyte_cdk/sources/declarative/models/declarative_component_schema.py

@@ -3,7 +3,6 @@
 
 from __future__ import annotations
 
-from datetime import datetime, timedelta
 from enum import Enum
 from typing import Any, Dict, List, Literal, Optional, Union
 
@@ -652,8 +651,11 @@ class Rate(BaseModel):
         description="The maximum number of calls allowed within the interval.",
         title="Limit",
     )
-    interval: timedelta = Field(
-        ..., description="The time interval for the rate limit.", title="Interval"
+    interval: str = Field(
+        ...,
+        description="The time interval for the rate limit.",
+        examples=["PT1H", "P1D"],
+        title="Interval",
     )
 
 
@@ -927,15 +929,6 @@ class CustomDecoder(BaseModel):
     parameters: Optional[Dict[str, Any]] = Field(None, alias="$parameters")
 
 
-class GzipJsonDecoder(BaseModel):
-    class Config:
-        extra = Extra.allow
-
-    type: Literal["GzipJsonDecoder"]
-    encoding: Optional[str] = "utf-8"
-    parameters: Optional[Dict[str, Any]] = Field(None, alias="$parameters")
-
-
 class MinMaxDatetime(BaseModel):
     type: Literal["MinMaxDatetime"]
     datetime: str = Field(
@@ -1314,18 +1307,8 @@ class LegacySessionTokenAuthenticator(BaseModel):
     parameters: Optional[Dict[str, Any]] = Field(None, alias="$parameters")
 
 
-class JsonParser(BaseModel):
-    type: Literal["JsonParser"]
-    encoding: Optional[str] = "utf-8"
-
-
-class JsonLineParser(BaseModel):
-    type: Literal["JsonLineParser"]
-    encoding: Optional[str] = "utf-8"
-
-
-class CsvParser(BaseModel):
-    type: Literal["CsvParser"]
+class CsvDecoder(BaseModel):
+    type: Literal["CsvDecoder"]
     encoding: Optional[str] = "utf-8"
     delimiter: Optional[str] = ","
 
@@ -1629,12 +1612,7 @@ class FixedWindowCallRatePolicy(BaseModel):
         extra = Extra.allow
 
     type: Literal["FixedWindowCallRatePolicy"]
-    next_reset_ts: datetime = Field(
-        ...,
-        description="The timestamp when the rate limit will reset.",
-        title="Next Reset Timestamp",
-    )
-    period: timedelta = Field(
+    period: str = Field(
         ..., description="The time interval for the rate limit window.", title="Period"
     )
     call_limit: int = Field(
@@ -1774,9 +1752,9 @@ class RecordSelector(BaseModel):
     parameters: Optional[Dict[str, Any]] = Field(None, alias="$parameters")
 
 
-class GzipParser(BaseModel):
-    type: Literal["GzipParser"]
-    inner_parser: Union[JsonLineParser, CsvParser, JsonParser]
+class GzipDecoder(BaseModel):
+    type: Literal["GzipDecoder"]
+    decoder: Union[CsvDecoder, GzipDecoder, JsonDecoder, JsonlDecoder]
 
 
 class Spec(BaseModel):
@@ -1809,29 +1787,6 @@ class CompositeErrorHandler(BaseModel):
     parameters: Optional[Dict[str, Any]] = Field(None, alias="$parameters")
 
 
-class APIBudget(BaseModel):
-    class Config:
-        extra = Extra.allow
-
-    type: Literal["APIBudget"]
-    policies: List[
-        Union[
-            FixedWindowCallRatePolicy,
-            MovingWindowCallRatePolicy,
-            UnlimitedCallRatePolicy,
-        ]
-    ] = Field(
-        ...,
-        description="List of call rate policies that define how many calls are allowed.",
-        title="Policies",
-    )
-    maximum_attempts_to_acquire: Optional[int] = Field(
-        100000,
-        description="The maximum number of attempts to acquire a call before giving up.",
-        title="Maximum Attempts to Acquire",
-    )
-
-
 class HTTPAPIBudget(BaseModel):
     class Config:
         extra = Extra.allow
@@ -1863,11 +1818,6 @@ class HTTPAPIBudget(BaseModel):
         description="List of HTTP status codes that indicate a rate limit has been hit.",
         title="Status Codes for Rate Limit Hit",
     )
-    maximum_attempts_to_acquire: Optional[int] = Field(
-        100000,
-        description="The maximum number of attempts to acquire a call before giving up.",
-        title="Maximum Attempts to Acquire",
-    )
 
 
 class ZipfileDecoder(BaseModel):
@@ -1875,18 +1825,13 @@ class ZipfileDecoder(BaseModel):
         extra = Extra.allow
 
     type: Literal["ZipfileDecoder"]
-    parser: Union[GzipParser, JsonParser, JsonLineParser, CsvParser] = Field(
+    decoder: Union[CsvDecoder, GzipDecoder, JsonDecoder, JsonlDecoder] = Field(
         ...,
         description="Parser to parse the decompressed data from the zipfile(s).",
         title="Parser",
     )
 
 
-class CompositeRawDecoder(BaseModel):
-    type: Literal["CompositeRawDecoder"]
-    parser: Union[GzipParser, JsonParser, JsonLineParser, CsvParser]
-
-
 class DeclarativeSource1(BaseModel):
     class Config:
         extra = Extra.forbid
@@ -1903,11 +1848,7 @@ class DeclarativeSource1(BaseModel):
     definitions: Optional[Dict[str, Any]] = None
     spec: Optional[Spec] = None
     concurrency_level: Optional[ConcurrencyLevel] = None
-    api_budget: Optional[Union[APIBudget, HTTPAPIBudget]] = Field(
-        None,
-        description="Defines how many requests can be made to the API in a given time frame. This field accepts either a generic APIBudget or an HTTP-specific configuration (HTTPAPIBudget) to be applied across all streams.",
-        title="API Budget",
-    )
+    api_budget: Optional[HTTPAPIBudget] = None
     metadata: Optional[Dict[str, Any]] = Field(
         None,
         description="For internal Airbyte use only - DO NOT modify manually. Used by consumers of declarative manifests for storing related metadata.",
@@ -1934,11 +1875,7 @@ class DeclarativeSource2(BaseModel):
     definitions: Optional[Dict[str, Any]] = None
     spec: Optional[Spec] = None
     concurrency_level: Optional[ConcurrencyLevel] = None
-    api_budget: Optional[Union[APIBudget, HTTPAPIBudget]] = Field(
-        None,
-        description="Defines how many requests can be made to the API in a given time frame. This field accepts either a generic APIBudget or an HTTP-specific configuration (HTTPAPIBudget) to be applied across all streams.",
-        title="API Budget",
-    )
+    api_budget: Optional[HTTPAPIBudget] = None
     metadata: Optional[Dict[str, Any]] = Field(
         None,
         description="For internal Airbyte use only - DO NOT modify manually. Used by consumers of declarative manifests for storing related metadata.",
@@ -2098,7 +2035,7 @@ class SessionTokenAuthenticator(BaseModel):
         description="Authentication method to use for requests sent to the API, specifying how to inject the session token.",
         title="Data Request Authentication",
     )
-    decoder: Optional[Union[JsonDecoder, XmlDecoder, CompositeRawDecoder]] = Field(
+    decoder: Optional[Union[JsonDecoder, XmlDecoder]] = Field(
         None, description="Component used to decode the response.", title="Decoder"
     )
     parameters: Optional[Dict[str, Any]] = Field(None, alias="$parameters")
@@ -2298,12 +2235,12 @@ class SimpleRetriever(BaseModel):
     decoder: Optional[
         Union[
             CustomDecoder,
+            CsvDecoder,
+            GzipDecoder,
             JsonDecoder,
             JsonlDecoder,
             IterableDecoder,
             XmlDecoder,
-            GzipJsonDecoder,
-            CompositeRawDecoder,
             ZipfileDecoder,
         ]
     ] = Field(
@@ -2376,12 +2313,12 @@ class AsyncRetriever(BaseModel):
     decoder: Optional[
         Union[
             CustomDecoder,
+            CsvDecoder,
+            GzipDecoder,
             JsonDecoder,
             JsonlDecoder,
             IterableDecoder,
             XmlDecoder,
-            GzipJsonDecoder,
-            CompositeRawDecoder,
             ZipfileDecoder,
         ]
     ] = Field(
@@ -2392,12 +2329,12 @@ class AsyncRetriever(BaseModel):
     download_decoder: Optional[
         Union[
             CustomDecoder,
+            CsvDecoder,
+            GzipDecoder,
             JsonDecoder,
             JsonlDecoder,
             IterableDecoder,
             XmlDecoder,
-            GzipJsonDecoder,
-            CompositeRawDecoder,
             ZipfileDecoder,
         ]
     ] = Field(
@@ -2442,6 +2379,7 @@ class DynamicDeclarativeStream(BaseModel):
 
 
 ComplexFieldType.update_forward_refs()
+GzipDecoder.update_forward_refs()
 CompositeErrorHandler.update_forward_refs()
 DeclarativeSource1.update_forward_refs()
 DeclarativeSource2.update_forward_refs()

airbyte_cdk/sources/declarative/parsers/model_to_component_factory.py

@@ -60,10 +60,8 @@ from airbyte_cdk.sources.declarative.datetime import MinMaxDatetime
 from airbyte_cdk.sources.declarative.declarative_stream import DeclarativeStream
 from airbyte_cdk.sources.declarative.decoders import (
     Decoder,
-    GzipJsonDecoder,
     IterableDecoder,
     JsonDecoder,
-    JsonlDecoder,
     PaginationDecoderDecorator,
     XmlDecoder,
     ZipfileDecoder,
@@ -103,8 +101,8 @@ from airbyte_cdk.sources.declarative.migrations.legacy_to_per_partition_state_mi
     LegacyToPerPartitionStateMigration,
 )
 from airbyte_cdk.sources.declarative.models import (
-    Clamping,
     CustomStateMigration,
+    GzipDecoder,
 )
 from airbyte_cdk.sources.declarative.models.declarative_component_schema import (
     AddedFieldDefinition as AddedFieldDefinitionModel,
@@ -112,9 +110,6 @@ from airbyte_cdk.sources.declarative.models.declarative_component_schema import
 from airbyte_cdk.sources.declarative.models.declarative_component_schema import (
     AddFields as AddFieldsModel,
 )
-from airbyte_cdk.sources.declarative.models.declarative_component_schema import (
-    APIBudget as APIBudgetModel,
-)
 from airbyte_cdk.sources.declarative.models.declarative_component_schema import (
     ApiKeyAuthenticator as ApiKeyAuthenticatorModel,
 )
@@ -145,9 +140,6 @@ from airbyte_cdk.sources.declarative.models.declarative_component_schema import
 from airbyte_cdk.sources.declarative.models.declarative_component_schema import (
     CompositeErrorHandler as CompositeErrorHandlerModel,
 )
-from airbyte_cdk.sources.declarative.models.declarative_component_schema import (
-    CompositeRawDecoder as CompositeRawDecoderModel,
-)
 from airbyte_cdk.sources.declarative.models.declarative_component_schema import (
     ConcurrencyLevel as ConcurrencyLevelModel,
 )
@@ -158,7 +150,7 @@ from airbyte_cdk.sources.declarative.models.declarative_component_schema import
     ConstantBackoffStrategy as ConstantBackoffStrategyModel,
 )
 from airbyte_cdk.sources.declarative.models.declarative_component_schema import (
-    CsvParser as CsvParserModel,
+    CsvDecoder as CsvDecoderModel,
 )
 from airbyte_cdk.sources.declarative.models.declarative_component_schema import (
     CursorPagination as CursorPaginationModel,
@@ -236,10 +228,7 @@ from airbyte_cdk.sources.declarative.models.declarative_component_schema import
     FlattenFields as FlattenFieldsModel,
 )
 from airbyte_cdk.sources.declarative.models.declarative_component_schema import (
-    GzipJsonDecoder as GzipJsonDecoderModel,
-)
-from airbyte_cdk.sources.declarative.models.declarative_component_schema import (
-    GzipParser as GzipParserModel,
+    GzipDecoder as GzipDecoderModel,
 )
 from airbyte_cdk.sources.declarative.models.declarative_component_schema import (
     HTTPAPIBudget as HTTPAPIBudgetModel,
@@ -271,12 +260,6 @@ from airbyte_cdk.sources.declarative.models.declarative_component_schema import
 from airbyte_cdk.sources.declarative.models.declarative_component_schema import (
     JsonlDecoder as JsonlDecoderModel,
 )
-from airbyte_cdk.sources.declarative.models.declarative_component_schema import (
-    JsonLineParser as JsonLineParserModel,
-)
-from airbyte_cdk.sources.declarative.models.declarative_component_schema import (
-    JsonParser as JsonParserModel,
-)
 from airbyte_cdk.sources.declarative.models.declarative_component_schema import (
     JwtAuthenticator as JwtAuthenticatorModel,
 )
@@ -562,9 +545,9 @@ class ModelToComponentFactory:
             CheckStreamModel: self.create_check_stream,
             CheckDynamicStreamModel: self.create_check_dynamic_stream,
             CompositeErrorHandlerModel: self.create_composite_error_handler,
-            CompositeRawDecoderModel: self.create_composite_raw_decoder,
             ConcurrencyLevelModel: self.create_concurrency_level,
             ConstantBackoffStrategyModel: self.create_constant_backoff_strategy,
+            CsvDecoderModel: self.create_csv_decoder,
             CursorPaginationModel: self.create_cursor_pagination,
             CustomAuthenticatorModel: self.create_custom_component,
             CustomBackoffStrategyModel: self.create_custom_component,
@@ -594,10 +577,7 @@ class ModelToComponentFactory:
             InlineSchemaLoaderModel: self.create_inline_schema_loader,
             JsonDecoderModel: self.create_json_decoder,
             JsonlDecoderModel: self.create_jsonl_decoder,
-            JsonLineParserModel: self.create_json_line_parser,
-            JsonParserModel: self.create_json_parser,
-            GzipJsonDecoderModel: self.create_gzipjson_decoder,
-            GzipParserModel: self.create_gzip_parser,
+            GzipDecoderModel: self.create_gzip_decoder,
             KeysToLowerModel: self.create_keys_to_lower_transformation,
             KeysToSnakeCaseModel: self.create_keys_to_snake_transformation,
             KeysReplaceModel: self.create_keys_replace_transformation,
@@ -638,7 +618,6 @@ class ModelToComponentFactory:
             StreamConfigModel: self.create_stream_config,
             ComponentMappingDefinitionModel: self.create_components_mapping_definition,
             ZipfileDecoderModel: self.create_zipfile_decoder,
-            APIBudgetModel: self.create_api_budget,
             HTTPAPIBudgetModel: self.create_http_api_budget,
             FixedWindowCallRatePolicyModel: self.create_fixed_window_call_rate_policy,
             MovingWindowCallRatePolicyModel: self.create_moving_window_call_rate_policy,
@@ -972,6 +951,17 @@ class ModelToComponentFactory:
             parameters={},
         )
 
+    @staticmethod
+    def apply_stream_state_migrations(
+        stream_state_migrations: List[Any] | None, stream_state: MutableMapping[str, Any]
+    ) -> MutableMapping[str, Any]:
+        if stream_state_migrations:
+            for state_migration in stream_state_migrations:
+                if state_migration.should_migrate(stream_state):
+                    # The state variable is expected to be mutable but the migrate method returns an immutable mapping.
+                    stream_state = dict(state_migration.migrate(stream_state))
+        return stream_state
+
     def create_concurrent_cursor_from_datetime_based_cursor(
         self,
         model_type: Type[BaseModel],
@@ -981,6 +971,7 @@ class ModelToComponentFactory:
         config: Config,
         message_repository: Optional[MessageRepository] = None,
         runtime_lookback_window: Optional[datetime.timedelta] = None,
+        stream_state_migrations: Optional[List[Any]] = None,
         **kwargs: Any,
     ) -> ConcurrentCursor:
         # Per-partition incremental streams can dynamically create child cursors which will pass their current
@@ -991,6 +982,7 @@ class ModelToComponentFactory:
             if "stream_state" not in kwargs
             else kwargs["stream_state"]
         )
+        stream_state = self.apply_stream_state_migrations(stream_state_migrations, stream_state)
 
         component_type = component_definition.get("type")
         if component_definition.get("type") != model_type.__name__:
@@ -1226,6 +1218,7 @@ class ModelToComponentFactory:
         config: Config,
         stream_state: MutableMapping[str, Any],
         partition_router: PartitionRouter,
+        stream_state_migrations: Optional[List[Any]] = None,
         **kwargs: Any,
     ) -> ConcurrentPerPartitionCursor:
         component_type = component_definition.get("type")
@@ -1274,8 +1267,10 @@ class ModelToComponentFactory:
                 stream_namespace=stream_namespace,
                 config=config,
                 message_repository=NoopMessageRepository(),
+                stream_state_migrations=stream_state_migrations,
             )
         )
+        stream_state = self.apply_stream_state_migrations(stream_state_migrations, stream_state)
 
         # Return the concurrent cursor and state converter
         return ConcurrentPerPartitionCursor(
@@ -1784,6 +1779,7 @@ class ModelToComponentFactory:
                     stream_name=model.name or "",
                     stream_namespace=None,
                     config=config or {},
+                    stream_state_migrations=model.state_migrations,
                 )
             return (
                 self._create_component_from_model(model=model.incremental_sync, config=config)
@@ -2092,25 +2088,26 @@ class ModelToComponentFactory:
         )
 
     @staticmethod
-    def create_json_decoder(model: JsonDecoderModel, config: Config, **kwargs: Any) -> JsonDecoder:
+    def create_json_decoder(model: JsonDecoderModel, config: Config, **kwargs: Any) -> Decoder:
         return JsonDecoder(parameters={})
 
     @staticmethod
-    def create_json_parser(model: JsonParserModel, config: Config, **kwargs: Any) -> JsonParser:
-        encoding = model.encoding if model.encoding else "utf-8"
-        return JsonParser(encoding=encoding)
+    def create_csv_decoder(model: CsvDecoderModel, config: Config, **kwargs: Any) -> Decoder:
+        return CompositeRawDecoder(
+            parser=ModelToComponentFactory._get_parser(model, config), stream_response=True
+        )
 
     @staticmethod
-    def create_jsonl_decoder(
-        model: JsonlDecoderModel, config: Config, **kwargs: Any
-    ) -> JsonlDecoder:
-        return JsonlDecoder(parameters={})
+    def create_jsonl_decoder(model: JsonlDecoderModel, config: Config, **kwargs: Any) -> Decoder:
+        return CompositeRawDecoder(
+            parser=ModelToComponentFactory._get_parser(model, config), stream_response=True
+        )
 
     @staticmethod
-    def create_json_line_parser(
-        model: JsonLineParserModel, config: Config, **kwargs: Any
-    ) -> JsonLineParser:
-        return JsonLineParser(encoding=model.encoding)
+    def create_gzip_decoder(model: GzipDecoderModel, config: Config, **kwargs: Any) -> Decoder:
+        return CompositeRawDecoder(
+            parser=ModelToComponentFactory._get_parser(model, config), stream_response=True
+        )
 
     @staticmethod
     def create_iterable_decoder(
@@ -2122,33 +2119,30 @@ class ModelToComponentFactory:
     def create_xml_decoder(model: XmlDecoderModel, config: Config, **kwargs: Any) -> XmlDecoder:
         return XmlDecoder(parameters={})
 
-    @staticmethod
-    def create_gzipjson_decoder(
-        model: GzipJsonDecoderModel, config: Config, **kwargs: Any
-    ) -> GzipJsonDecoder:
-        return GzipJsonDecoder(parameters={}, encoding=model.encoding)
-
     def create_zipfile_decoder(
         self, model: ZipfileDecoderModel, config: Config, **kwargs: Any
     ) -> ZipfileDecoder:
-        parser = self._create_component_from_model(model=model.parser, config=config)
-        return ZipfileDecoder(parser=parser)
-
-    def create_gzip_parser(
-        self, model: GzipParserModel, config: Config, **kwargs: Any
-    ) -> GzipParser:
-        inner_parser = self._create_component_from_model(model=model.inner_parser, config=config)
-        return GzipParser(inner_parser=inner_parser)
+        return ZipfileDecoder(parser=ModelToComponentFactory._get_parser(model.decoder, config))
 
     @staticmethod
-    def create_csv_parser(model: CsvParserModel, config: Config, **kwargs: Any) -> CsvParser:
-        return CsvParser(encoding=model.encoding, delimiter=model.delimiter)
+    def _get_parser(model: BaseModel, config: Config) -> Parser:
+        if isinstance(model, JsonDecoderModel):
+            # Note that the logic is a bit different from the JsonDecoder as there is some legacy that is maintained to return {} on error cases
+            return JsonParser()
+        elif isinstance(model, JsonlDecoderModel):
+            return JsonLineParser()
+        elif isinstance(model, CsvDecoderModel):
+            return CsvParser(encoding=model.encoding, delimiter=model.delimiter)
+        elif isinstance(model, GzipDecoderModel):
+            return GzipParser(
+                inner_parser=ModelToComponentFactory._get_parser(model.decoder, config)
+            )
+        elif isinstance(
+            model, (CustomDecoderModel, IterableDecoderModel, XmlDecoderModel, ZipfileDecoderModel)
+        ):
+            raise ValueError(f"Decoder type {model} does not have parser associated to it")
 
-    def create_composite_raw_decoder(
-        self, model: CompositeRawDecoderModel, config: Config, **kwargs: Any
-    ) -> CompositeRawDecoder:
-        parser = self._create_component_from_model(model=model.parser, config=config)
-        return CompositeRawDecoder(parser=parser)
+        raise ValueError(f"Unknown decoder type {model}")
 
     @staticmethod
     def create_json_file_schema_loader(
@@ -2421,6 +2415,8 @@ class ModelToComponentFactory:
             if model.record_filter
             else None
         )
+
+        transform_before_filtering = False
         if client_side_incremental_sync:
             record_filter = ClientSideIncrementalRecordFilterDecorator(
                 config=config,
@@ -2430,6 +2426,8 @@ class ModelToComponentFactory:
                 else None,
                 **client_side_incremental_sync,
             )
+            transform_before_filtering = True
+
         schema_normalization = (
             TypeTransformer(SCHEMA_TRANSFORMER_TYPE_MAPPING[model.schema_normalization])
             if isinstance(model.schema_normalization, SchemaNormalizationModel)
@@ -2444,6 +2442,7 @@ class ModelToComponentFactory:
             transformations=transformations or [],
             schema_normalization=schema_normalization,
             parameters=model.parameters or {},
+            transform_before_filtering=transform_before_filtering,
         )
 
     @staticmethod
@@ -2965,17 +2964,6 @@ class ModelToComponentFactory:
         else:
             return False
 
-    def create_api_budget(self, model: APIBudgetModel, config: Config, **kwargs: Any) -> APIBudget:
-        policies = [
-            self._create_component_from_model(model=policy, config=config)
-            for policy in model.policies
-        ]
-
-        return APIBudget(
-            policies=policies,
-            maximum_attempts_to_acquire=model.maximum_attempts_to_acquire or 100000,
-        )
-
     def create_http_api_budget(
         self, model: HTTPAPIBudgetModel, config: Config, **kwargs: Any
     ) -> HttpAPIBudget:
@@ -2986,10 +2974,9 @@ class ModelToComponentFactory:
 
         return HttpAPIBudget(
             policies=policies,
-            maximum_attempts_to_acquire=model.maximum_attempts_to_acquire or 100000,
             ratelimit_reset_header=model.ratelimit_reset_header or "ratelimit-reset",
             ratelimit_remaining_header=model.ratelimit_remaining_header or "ratelimit-remaining",
-            status_codes_for_ratelimit_hit=model.status_codes_for_ratelimit_hit or (429,),
+            status_codes_for_ratelimit_hit=model.status_codes_for_ratelimit_hit or [429],
         )
 
     def create_fixed_window_call_rate_policy(
@@ -2999,9 +2986,12 @@ class ModelToComponentFactory:
             self._create_component_from_model(model=matcher, config=config)
             for matcher in model.matchers
         ]
+
+        # Set the initial reset timestamp to 10 days from now.
+        # This value will be updated by the first request.
         return FixedWindowCallRatePolicy(
-            next_reset_ts=model.next_reset_ts,
-            period=model.period,
+            next_reset_ts=datetime.datetime.now() + datetime.timedelta(days=10),
+            period=parse_duration(model.period),
             call_limit=model.call_limit,
             matchers=matchers,
         )
@@ -3036,7 +3026,7 @@ class ModelToComponentFactory:
     def create_rate(self, model: RateModel, config: Config, **kwargs: Any) -> Rate:
         return Rate(
             limit=model.limit,
-            interval=model.interval,
+            interval=parse_duration(model.interval),
         )
 
     def create_http_request_matcher(
@@ -3051,16 +3041,6 @@ class ModelToComponentFactory:
         )
 
     def set_api_budget(self, component_definition: ComponentDefinition, config: Config) -> None:
-        model_str = component_definition.get("type")
-        if model_str == "APIBudget":
-            # Annotate model_type as a type that is a subclass of BaseModel
-            model_type: Union[Type[APIBudgetModel], Type[HTTPAPIBudgetModel]] = APIBudgetModel
-        elif model_str == "HTTPAPIBudget":
-            model_type = HTTPAPIBudgetModel
-        else:
-            raise ValueError(f"Unknown API Budget type: {model_str}")
-
-        # create_component expects a type[BaseModel] and returns an instance of that model.
         self._api_budget = self.create_component(
-            model_type=model_type, component_definition=component_definition, config=config
+            model_type=HTTPAPIBudgetModel, component_definition=component_definition, config=config
         )