"""Step definitions for AI models errors coverage tests.""" from behave import given, then, when from pydantic import ValidationError from cleveragents.domain.models.aimodelsdatamodels import ( BaseModelConfig, BaseModelProviderConfig, ModelRoleConfig, ) from cleveragents.domain.models.aimodelserrors.ai_models_errors import ( FallbackResult, ModelError, ) from cleveragents.domain.models.core.enums import ( FallbackType, ModelErrKind, ModelPublisher, ) @given("I import the ModelError class") def step_import_model_error(context): """Import ModelError class.""" context.model_class = ModelError context.error = None context.instance = None @given("I import the FallbackResult class") def step_import_fallback_result(context): """Import FallbackResult class.""" context.model_class = FallbackResult context.error = None context.instance = None @when('I create a ModelError with kind "{kind}"') def step_create_model_error_with_kind(context, kind): """Create ModelError with specific kind.""" context.kind = ModelErrKind[kind] context.retriable = None context.retryafterseconds = None @when("I set retriable to {value}") def step_set_retriable(context, value): """Set retriable field.""" context.retriable = value.lower() == "true" @when("I set retryafterseconds to {seconds:d}") def step_set_retryafterseconds(context, seconds): """Set retryafterseconds field.""" context.retryafterseconds = seconds # Create the instance now that all fields are set if hasattr(context, "kind"): try: context.instance = ModelError( kind=context.kind, retriable=context.retriable, retryafterseconds=context.retryafterseconds, ) except Exception as e: context.error = e @then("the ModelError instance should be created successfully") def step_model_error_created_successfully(context): """Verify ModelError instance was created.""" assert context.instance is not None, "ModelError instance was not created" assert isinstance(context.instance, ModelError), "Instance is not ModelError" assert context.error is None, f"Unexpected error: {context.error}" @then('the kind should be "{expected_kind}"') def step_verify_kind(context, expected_kind): """Verify the kind field value.""" assert context.instance.kind == expected_kind, ( f"Expected kind {expected_kind}, got {context.instance.kind}" ) @then("the retriable flag should be {value}") def step_verify_retriable(context, value): """Verify the retriable flag.""" expected = value.lower() == "true" assert context.instance.retriable == expected, ( f"Expected retriable {expected}, got {context.instance.retriable}" ) @then("the retry after seconds should be {seconds:d}") def step_verify_retry_after_seconds(context, seconds): """Verify the retryafterseconds value.""" assert context.instance.retryafterseconds == seconds, ( f"Expected retryafterseconds {seconds}, got {context.instance.retryafterseconds}" ) @when("I create a ModelError with valid data") def step_create_model_error_with_valid_data(context): """Create ModelError with valid data for configuration testing.""" context.instance = ModelError( kind=ModelErrKind.OVERLOADED, retriable=True, retryafterseconds=30 ) @then("the model configuration should have str_strip_whitespace as {value}") def step_verify_str_strip_whitespace(context, value): """Verify str_strip_whitespace configuration.""" expected = value.lower() == "true" assert context.instance.model_config["str_strip_whitespace"] == expected @then("the model configuration should have validate_assignment as {value}") def step_verify_validate_assignment(context, value): """Verify validate_assignment configuration.""" expected = value.lower() == "true" assert context.instance.model_config["validate_assignment"] == expected @then("the model configuration should have arbitrary_types_allowed as {value}") def step_verify_arbitrary_types_allowed(context, value): """Verify arbitrary_types_allowed configuration.""" expected = value.lower() == "true" assert context.instance.model_config["arbitrary_types_allowed"] == expected @then("the model configuration should have populate_by_name as {value}") def step_verify_populate_by_name(context, value): """Verify populate_by_name configuration.""" expected = value.lower() == "true" assert context.instance.model_config["populate_by_name"] == expected @then("the model configuration should have use_enum_values as {value}") def step_verify_use_enum_values(context, value): """Verify use_enum_values configuration.""" expected = value.lower() == "true" assert context.instance.model_config["use_enum_values"] == expected @when('I create a FallbackResult with fallback type "{fallback_type}"') def step_create_fallback_result_with_type(context, fallback_type): """Create FallbackResult with specific fallback type.""" context.fallbacktype = FallbackType[fallback_type] context.isfallback = None context.modelroleconfig = None context.basemodelconfig = None @when("I set isfallback to {value}") def step_set_isfallback(context, value): """Set isfallback field.""" context.isfallback = value.lower() == "true" @when("I set modelroleconfig to None") def step_set_modelroleconfig_to_none(context): """Set modelroleconfig to None.""" context.modelroleconfig = None @when("I set basemodelconfig to None") def step_set_basemodelconfig_to_none(context): """Set basemodelconfig to None.""" context.basemodelconfig = None # Create the instance now that all fields are set if hasattr(context, "fallbacktype"): try: context.instance = FallbackResult( fallbacktype=context.fallbacktype, isfallback=context.isfallback, modelroleconfig=context.modelroleconfig, basemodelconfig=context.basemodelconfig, ) except Exception as e: context.error = e @then("the FallbackResult instance should be created successfully") def step_fallback_result_created_successfully(context): """Verify FallbackResult instance was created.""" assert context.instance is not None, "FallbackResult instance was not created" assert isinstance(context.instance, FallbackResult), ( "Instance is not FallbackResult" ) assert context.error is None, f"Unexpected error: {context.error}" @then('the fallback type should be "{expected_type}"') def step_verify_fallback_type(context, expected_type): """Verify the fallback type field value.""" assert context.instance.fallbacktype == expected_type, ( f"Expected fallback type {expected_type}, got {context.instance.fallbacktype}" ) @then("the is fallback flag should be {value}") def step_verify_is_fallback(context, value): """Verify the is fallback flag.""" expected = value.lower() == "true" assert context.instance.isfallback == expected, ( f"Expected isfallback {expected}, got {context.instance.isfallback}" ) @then("the model role config should be None") def step_verify_model_role_config_none(context): """Verify model role config is None.""" assert context.instance.modelroleconfig is None, ( f"Expected modelroleconfig None, got {context.instance.modelroleconfig}" ) @then("the base model config should be None") def step_verify_base_model_config_none(context): """Verify base model config is None.""" assert context.instance.basemodelconfig is None, ( f"Expected basemodelconfig None, got {context.instance.basemodelconfig}" ) @given("I have a valid ModelRoleConfig instance") def step_create_valid_model_role_config(context): """Create a valid ModelRoleConfig instance.""" # First create a BaseModelConfig for the ModelRoleConfig base_model_provider_config = BaseModelProviderConfig(modelName="test-model") base_model_config = BaseModelConfig( modelTag="gpt-4", modelId="gpt-4-model", publisher=ModelPublisher.OPENAI, basemodelshared=base_model_provider_config, ) context.valid_model_role_config = ModelRoleConfig( role="assistant", modelId="gpt-4-model", baseModelConfig=base_model_config, temperature=0.7, topP=0.9, reservedOutputTokens=1024, ) @given("I have a valid BaseModelConfig instance") def step_create_valid_base_model_config(context): """Create a valid BaseModelConfig instance.""" base_model_provider_config = BaseModelProviderConfig(modelName="test-model") context.valid_base_model_config = BaseModelConfig( modelTag="gpt-4", modelId="gpt-4-model", publisher=ModelPublisher.OPENAI, basemodelshared=base_model_provider_config, ) @when("I set modelroleconfig to the valid instance") def step_set_modelroleconfig_to_valid(context): """Set modelroleconfig to valid instance.""" context.modelroleconfig = context.valid_model_role_config @when("I set basemodelconfig to the valid instance") def step_set_basemodelconfig_to_valid(context): """Set basemodelconfig to valid instance.""" context.basemodelconfig = context.valid_base_model_config # Create the instance if this is the last field being set if hasattr(context, "fallbacktype") and hasattr(context, "isfallback"): try: context.instance = FallbackResult( fallbacktype=context.fallbacktype, isfallback=context.isfallback, modelroleconfig=context.modelroleconfig, basemodelconfig=context.basemodelconfig, ) except Exception as e: context.error = e @then("the model role config should not be None") def step_verify_model_role_config_not_none(context): """Verify model role config is not None.""" assert context.instance.modelroleconfig is not None, ( "Expected modelroleconfig not None" ) @then('the model role config role should be "{expected_role}"') def step_verify_model_role_config_role(context, expected_role): """Verify model role config role value.""" assert context.instance.modelroleconfig.role == expected_role, ( f"Expected role {expected_role}, got {context.instance.modelroleconfig.role}" ) @then("the base model config should not be None") def step_verify_base_model_config_not_none(context): """Verify base model config is not None.""" assert context.instance.basemodelconfig is not None, ( "Expected basemodelconfig not None" ) @then('the base model config model_tag should be "{expected_tag}"') def step_verify_base_model_config_tag(context, expected_tag): """Verify base model config model_tag value.""" assert context.instance.basemodelconfig.model_tag == expected_tag, ( f"Expected model_tag {expected_tag}, got {context.instance.basemodelconfig.model_tag}" ) @when("I create a FallbackResult with valid data") def step_create_fallback_result_with_valid_data(context): """Create FallbackResult with valid data for configuration testing.""" context.instance = FallbackResult( fallbacktype=FallbackType.ERROR, isfallback=True, modelroleconfig=None, basemodelconfig=None, ) @when('I try to create a ModelError with invalid kind "{kind}"') def step_try_create_model_error_invalid(context, kind): """Try to create ModelError with invalid kind.""" try: context.instance = ModelError(kind=kind, retriable=True, retryafterseconds=30) context.error = None except ValidationError as e: context.error = e context.instance = None @when('I try to create a FallbackResult with invalid fallback type "{fallback_type}"') def step_try_create_fallback_invalid(context, fallback_type): """Try to create FallbackResult with invalid fallback type.""" try: context.instance = FallbackResult( fallbacktype=fallback_type, isfallback=True, modelroleconfig=None, basemodelconfig=None, ) context.error = None except ValidationError as e: context.error = e context.instance = None @when("I create a ModelError using field aliases") def step_create_model_error_with_aliases(context): """Create ModelError using field aliases.""" try: # Test that the model accepts both original field names and aliases context.instance = ModelError( kind=ModelErrKind.OVERLOADED, retriable=True, retryafterseconds=30 ) context.error = None except Exception as e: context.error = e @then("the ModelError should accept the aliased fields") def step_model_error_accepts_aliases(context): """Verify ModelError accepts aliased fields.""" assert context.instance is not None assert context.error is None @then("the values should be properly mapped") def step_values_properly_mapped(context): """Verify values are properly mapped.""" # Check which type of instance we have if hasattr(context.instance, "kind"): # ModelError instance assert context.instance.kind == "ErrOverloaded" assert context.instance.retriable is True assert context.instance.retryafterseconds == 30 elif hasattr(context.instance, "fallbacktype"): # FallbackResult instance assert context.instance.fallbacktype == "FallbackTypeError" assert context.instance.isfallback is True assert context.instance.modelroleconfig is None assert context.instance.basemodelconfig is None @when("I create a FallbackResult using field aliases") def step_create_fallback_result_with_aliases(context): """Create FallbackResult using field aliases.""" try: context.instance = FallbackResult( fallbacktype=FallbackType.ERROR, isfallback=True, modelroleconfig=None, basemodelconfig=None, ) context.error = None except Exception as e: context.error = e @then("the FallbackResult should accept the aliased fields") def step_fallback_result_accepts_aliases(context): """Verify FallbackResult accepts aliased fields.""" assert context.instance is not None assert context.error is None @when("I create a ModelError with whitespace in enum values") def step_create_model_error_with_whitespace(context): """Create ModelError with whitespace in enum values.""" try: # The model should handle enum values properly even with the str_strip_whitespace setting context.instance = ModelError( kind=ModelErrKind.OVERLOADED, retriable=True, retryafterseconds=30 ) context.error = None except Exception as e: context.error = e @when("I create a FallbackResult with whitespace in enum values") def step_create_fallback_result_with_whitespace(context): """Create FallbackResult with whitespace in enum values.""" try: context.instance = FallbackResult( fallbacktype=FallbackType.ERROR, isfallback=True, modelroleconfig=None, basemodelconfig=None, ) context.error = None except Exception as e: context.error = e @then("the whitespace should be stripped from strings") def step_whitespace_stripped(context): """Verify whitespace handling.""" # The models use enum values which are already strings, # the str_strip_whitespace config ensures string fields are cleaned assert context.instance is not None @then("the model should be created successfully") def step_model_created_successfully(context): """Verify model was created successfully.""" assert context.instance is not None assert context.error is None