1  from pydantic import BaseModel, Field, EmailStr, validator
  2  from typing import Optional, List, Dict, Any
  3  from datetime import datetime
  4  from enum import Enum
  5  import uuid
  6  
  7  class ApiKeyLevel(str, Enum):
  8      FREE = "free"
  9      BASIC = "basic"
 10      PREMIUM = "premium"
 11      ENTERPRISE = "enterprise"
 12  
 13  class UsageLimit(BaseModel):
 14      """Usage limits for different API levels"""
 15      daily_requests: int = Field(..., description="Maximum number of requests per day")
 16      monthly_requests: int = Field(..., description="Maximum number of requests per month")
 17      max_tokens_per_request: int = Field(..., description="Maximum number of tokens per request")
 18      max_text_length: int = Field(..., description="Maximum text length in characters")
 19      batch_processing: bool = Field(..., description="Access to batch processing")
 20      max_concurrent_requests: int = Field(..., description="Maximum number of concurrent requests")
 21      advanced_models: bool = Field(..., description="Access to advanced models")
 22  
 23  # Define User before UserInDB
 24  class User(BaseModel):
 25      """Model for users"""
 26      id: str = Field(default_factory=lambda: str(uuid.uuid4()), description="Unique user ID")
 27      username: str = Field(..., description="Username")
 28      email: EmailStr = Field(..., description="User email")
 29      hashed_password: str = Field(..., description="Hashed password")
 30      full_name: Optional[str] = Field(None, description="Full name")
 31      disabled: bool = Field(default=False, description="Whether the account is disabled")
 32      roles: List[str] = Field(default_factory=lambda: ["user"], description="User roles")
 33      created_at: datetime = Field(default_factory=datetime.utcnow, description="Creation date")
 34      subscription: ApiKeyLevel = Field(default=ApiKeyLevel.FREE, description="Subscription level")
 35      
 36      class Config:
 37          json_encoders = {
 38              datetime: lambda v: v.isoformat()
 39          }
 40  
 41  class UserCreate(BaseModel):
 42      """Model for user creation"""
 43      username: str = Field(..., description="Username", min_length=3, max_length=50)
 44      email: EmailStr = Field(..., description="User email")
 45      password: str = Field(..., description="Password", min_length=8)
 46      full_name: Optional[str] = Field(None, description="Full name")
 47      
 48      @validator('username')
 49      def username_alphanumeric(cls, v):
 50          if not v.isalnum():
 51              raise ValueError('Username must be alphanumeric')
 52          return v
 53  
 54  # Now UserInDB can inherit from User
 55  class UserInDB(User):
 56      """User model for database with hashed password."""
 57      hashed_password: str
 58      is_active: bool = True
 59      is_admin: bool = False  
 60      subscription_level: str = "free"
 61      last_login: Optional[datetime] = None
 62      stripe_customer_id: Optional[str] = None
 63      
 64  class ApiKey(BaseModel):
 65      """Model for API keys"""
 66      id: str = Field(default_factory=lambda: str(uuid.uuid4()), description="Unique API key ID")
 67      key: str = Field(..., description="API key (hashed in the database)")
 68      name: str = Field(..., description="Descriptive name for the key")
 69      user_id: str = Field(..., description="ID of the user who owns the key")
 70      level: ApiKeyLevel = Field(default=ApiKeyLevel.FREE, description="Access level")
 71      created_at: datetime = Field(default_factory=datetime.utcnow, description="Creation date")
 72      last_used_at: Optional[datetime] = Field(None, description="Last used date")
 73      expires_at: Optional[datetime] = Field(None, description="Expiration date")
 74      is_active: bool = Field(default=True, description="Whether the key is active")
 75      usage: Dict[str, int] = Field(default_factory=dict, description="Usage statistics")
 76      
 77      class Config:
 78          json_encoders = {
 79              datetime: lambda v: v.isoformat()
 80          }
 81  
 82  class TokenData(BaseModel):
 83      """Data contained in the JWT token"""
 84      sub: str = Field(..., description="User ID")
 85      name: Optional[str] = Field(None, description="User name")
 86      email: Optional[str] = Field(None, description="User email")
 87      roles: List[str] = Field(default_factory=list, description="User roles")
 88      api_level: ApiKeyLevel = Field(default=ApiKeyLevel.FREE, description="API level")
 89      exp: Optional[int] = Field(None, description="Token expiration (timestamp)")
 90  
 91  class UserResponse(BaseModel):
 92      """Model for user response (without sensitive data)"""
 93      id: str
 94      username: str
 95      email: EmailStr
 96      full_name: Optional[str]
 97      roles: List[str]
 98      created_at: datetime
 99      subscription: ApiKeyLevel
100      
101      class Config:
102          json_encoders = {
103              datetime: lambda v: v.isoformat()
104          }
105  
106  class Token(BaseModel):
107      """Model for authentication token"""
108      access_token: str
109      token_type: str = "bearer"
110      expires_at: int
111      user: UserResponse
112  
113  class ApiKeyCreate(BaseModel):
114      """Model for API key creation"""
115      name: str = Field(..., description="Descriptive name for the key", min_length=3, max_length=50)
116      level: Optional[ApiKeyLevel] = Field(None, description="Access level (admin only)")
117      expires_at: Optional[datetime] = Field(None, description="Expiration date (optional)")
118  
119  class ApiKeyResponse(BaseModel):
120      """Model for API key response"""
121      id: str
122      key: str
123      name: str
124      level: ApiKeyLevel
125      created_at: datetime
126      expires_at: Optional[datetime]
127      is_active: bool
128      
129      class Config:
130          json_encoders = {
131              datetime: lambda v: v.isoformat()
132          }
133  
134  class UsageRecord(BaseModel):
135      """Model for usage recording"""
136      user_id: str
137      api_key_id: str
138      request_path: str
139      request_method: str
140      tokens_input: int = 0
141      tokens_output: int = 0
142      processing_time: float
143      status_code: int
144      timestamp: datetime = Field(default_factory=datetime.utcnow)
145      
146      class Config:
147          json_encoders = {
148              datetime: lambda v: v.isoformat()
149          }