schemas/schema.py

from pydantic import BaseModel, Field, ConfigDict
from typing import List

class SMSRequest(BaseModel):
    messages: List[str]

class EmbeddingResponse(BaseModel):
    dimensions: int
    embeddings: List[List[float]]

    model_config = ConfigDict(
        json_schema_extra={
            "example": {
                "dimensions": 768,
                "embeddings": [
                    [0.1, 0.2, 0.3, 0.4, 0.5, 0.6],  # Example values for first message
                    [0.7, 0.8, 0.9, 1.0, 1.1, 1.2]   # Example values for second message
                ]
            }
        }
    )

class SimilarityRequest(BaseModel):
    message1: str
    message2: str

class SimilarityResponse(BaseModel):
    similarity_score: float
    
class PredictionRequest(BaseModel):
    message: str

class PredictionResponse(BaseModel):
    label: str
    probability: float