api/schemas.py

"""API Schemas for Predict Request and Response."""
from enum import Enum

from pydantic import BaseModel, ConfigDict, ValidationError


class ProgrammingLanguage(str, Enum):
    """Programming languages supported for prediction."""

    JAVA = "java"
    PYTHON = "python"
    PHARO = "pharo"


class ModelType(str, Enum):
    """Model types for prediction."""

    SETFIT = "setfit"
    RANDOM_FOREST = "random_forest"
    TRANSFORMER = "transformer"


class PredictRequest(BaseModel):
    """Schema for Predict Request."""

    text: str
    language: ProgrammingLanguage
    model_type: ModelType

    model_config = ConfigDict(
        json_schema_extra={
            "example": {
                "text": "This method calculates the average score.",
                "language": "python",
                "model_type": "transformer",
            }
        }
    )


class PredictResponse(BaseModel):
    """Schema for Predict Response."""

    label: str
    score: float


""" Demonstration of object instantiation, printing,
and validation error handling with dummy use cases"""
if __name__ == "__main__":
    print("\n--- 1. Object Instantiation & Printing ---")

    valid_data = {
        "text": "This method calculates the average score.",
        "language": "java",
        "model_type": "setfit",
    }

    # Instantiate the object
    request = PredictRequest(**valid_data)

    # Print object as dictionary (.model_dump() is Pydantic V2 syntax)
    print(f"Valid Request Object: {request.model_dump()}")

    print("\n--- 2. Handling Invalid Data ---")

    try:
        print("Attempting to create request with language='c++'...")
        # This should fail because 'c++' is not in ProgrammingLanguage Enum
        invalid_request = PredictRequest(
            text="std::cout << 'Hello';", language="c++", model_type="setfit"
        )
    except ValidationError as e:
        print("SUCCESS: Validation Error Caught!")
        print(e.json())