mirror of
https://github.com/aljazceru/enclava.git
synced 2025-12-18 07:54:29 +01:00
clean commit
This commit is contained in:
296
backend/app/models/budget.py
Normal file
296
backend/app/models/budget.py
Normal file
@@ -0,0 +1,296 @@
|
||||
"""
|
||||
Budget model for managing spending limits and cost control
|
||||
"""
|
||||
|
||||
from datetime import datetime, timedelta
|
||||
from typing import Optional, Dict, Any
|
||||
from enum import Enum
|
||||
from sqlalchemy import Column, Integer, String, DateTime, Boolean, Text, JSON, ForeignKey, Float
|
||||
from sqlalchemy.orm import relationship
|
||||
from app.db.database import Base
|
||||
|
||||
|
||||
class BudgetType(str, Enum):
|
||||
"""Budget type enumeration"""
|
||||
USER = "user"
|
||||
API_KEY = "api_key"
|
||||
GLOBAL = "global"
|
||||
|
||||
|
||||
class BudgetPeriod(str, Enum):
|
||||
"""Budget period types"""
|
||||
DAILY = "daily"
|
||||
WEEKLY = "weekly"
|
||||
MONTHLY = "monthly"
|
||||
YEARLY = "yearly"
|
||||
CUSTOM = "custom"
|
||||
|
||||
|
||||
class Budget(Base):
|
||||
"""Budget model for setting and managing spending limits"""
|
||||
|
||||
__tablename__ = "budgets"
|
||||
|
||||
id = Column(Integer, primary_key=True, index=True)
|
||||
name = Column(String, nullable=False) # Human-readable name for the budget
|
||||
|
||||
# User and API Key relationships
|
||||
user_id = Column(Integer, ForeignKey("users.id"), nullable=False)
|
||||
user = relationship("User", back_populates="budgets")
|
||||
|
||||
api_key_id = Column(Integer, ForeignKey("api_keys.id"), nullable=True) # Optional: specific to an API key
|
||||
api_key = relationship("APIKey", back_populates="budgets")
|
||||
|
||||
# Usage tracking relationship
|
||||
usage_tracking = relationship("UsageTracking", back_populates="budget", cascade="all, delete-orphan")
|
||||
|
||||
# Budget limits (in cents)
|
||||
limit_cents = Column(Integer, nullable=False) # Maximum spend limit
|
||||
warning_threshold_cents = Column(Integer, nullable=True) # Warning threshold (e.g., 80% of limit)
|
||||
|
||||
# Time period settings
|
||||
period_type = Column(String, nullable=False, default="monthly") # daily, weekly, monthly, yearly, custom
|
||||
period_start = Column(DateTime, nullable=False) # Start of current period
|
||||
period_end = Column(DateTime, nullable=False) # End of current period
|
||||
|
||||
# Current usage (in cents)
|
||||
current_usage_cents = Column(Integer, default=0) # Spent in current period
|
||||
|
||||
# Budget status
|
||||
is_active = Column(Boolean, default=True)
|
||||
is_exceeded = Column(Boolean, default=False)
|
||||
is_warning_sent = Column(Boolean, default=False)
|
||||
|
||||
# Enforcement settings
|
||||
enforce_hard_limit = Column(Boolean, default=True) # Block requests when limit exceeded
|
||||
enforce_warning = Column(Boolean, default=True) # Send warnings at threshold
|
||||
|
||||
# Allowed resources (optional filters)
|
||||
allowed_models = Column(JSON, default=list) # Specific models this budget applies to
|
||||
allowed_endpoints = Column(JSON, default=list) # Specific endpoints this budget applies to
|
||||
|
||||
# Metadata
|
||||
description = Column(Text, nullable=True)
|
||||
tags = Column(JSON, default=list)
|
||||
currency = Column(String, default="USD")
|
||||
|
||||
# Auto-renewal settings
|
||||
auto_renew = Column(Boolean, default=True) # Automatically renew budget for next period
|
||||
rollover_unused = Column(Boolean, default=False) # Rollover unused budget to next period
|
||||
|
||||
# Notification settings
|
||||
notification_settings = Column(JSON, default=dict) # Email, webhook, etc.
|
||||
|
||||
# Timestamps
|
||||
created_at = Column(DateTime, default=datetime.utcnow)
|
||||
updated_at = Column(DateTime, default=datetime.utcnow, onupdate=datetime.utcnow)
|
||||
last_reset_at = Column(DateTime, nullable=True) # Last time budget was reset
|
||||
|
||||
def __repr__(self):
|
||||
return f"<Budget(id={self.id}, name='{self.name}', user_id={self.user_id}, limit=${self.limit_cents/100:.2f})>"
|
||||
|
||||
def to_dict(self):
|
||||
"""Convert budget to dictionary for API responses"""
|
||||
return {
|
||||
"id": self.id,
|
||||
"name": self.name,
|
||||
"user_id": self.user_id,
|
||||
"api_key_id": self.api_key_id,
|
||||
"limit_cents": self.limit_cents,
|
||||
"limit_dollars": self.limit_cents / 100,
|
||||
"warning_threshold_cents": self.warning_threshold_cents,
|
||||
"warning_threshold_dollars": self.warning_threshold_cents / 100 if self.warning_threshold_cents else None,
|
||||
"period_type": self.period_type,
|
||||
"period_start": self.period_start.isoformat() if self.period_start else None,
|
||||
"period_end": self.period_end.isoformat() if self.period_end else None,
|
||||
"current_usage_cents": self.current_usage_cents,
|
||||
"current_usage_dollars": self.current_usage_cents / 100,
|
||||
"remaining_cents": max(0, self.limit_cents - self.current_usage_cents),
|
||||
"remaining_dollars": max(0, (self.limit_cents - self.current_usage_cents) / 100),
|
||||
"usage_percentage": (self.current_usage_cents / self.limit_cents * 100) if self.limit_cents > 0 else 0,
|
||||
"is_active": self.is_active,
|
||||
"is_exceeded": self.is_exceeded,
|
||||
"is_warning_sent": self.is_warning_sent,
|
||||
"enforce_hard_limit": self.enforce_hard_limit,
|
||||
"enforce_warning": self.enforce_warning,
|
||||
"allowed_models": self.allowed_models,
|
||||
"allowed_endpoints": self.allowed_endpoints,
|
||||
"description": self.description,
|
||||
"tags": self.tags,
|
||||
"currency": self.currency,
|
||||
"auto_renew": self.auto_renew,
|
||||
"rollover_unused": self.rollover_unused,
|
||||
"notification_settings": self.notification_settings,
|
||||
"created_at": self.created_at.isoformat() if self.created_at else None,
|
||||
"updated_at": self.updated_at.isoformat() if self.updated_at else None,
|
||||
"last_reset_at": self.last_reset_at.isoformat() if self.last_reset_at else None
|
||||
}
|
||||
|
||||
def is_in_period(self) -> bool:
|
||||
"""Check if current time is within budget period"""
|
||||
now = datetime.utcnow()
|
||||
return self.period_start <= now <= self.period_end
|
||||
|
||||
def is_expired(self) -> bool:
|
||||
"""Check if budget period has expired"""
|
||||
return datetime.utcnow() > self.period_end
|
||||
|
||||
def can_spend(self, amount_cents: int) -> bool:
|
||||
"""Check if spending amount is within budget"""
|
||||
if not self.is_active or not self.is_in_period():
|
||||
return False
|
||||
|
||||
if not self.enforce_hard_limit:
|
||||
return True
|
||||
|
||||
return (self.current_usage_cents + amount_cents) <= self.limit_cents
|
||||
|
||||
def would_exceed_warning(self, amount_cents: int) -> bool:
|
||||
"""Check if spending amount would exceed warning threshold"""
|
||||
if not self.warning_threshold_cents:
|
||||
return False
|
||||
|
||||
return (self.current_usage_cents + amount_cents) >= self.warning_threshold_cents
|
||||
|
||||
def add_usage(self, amount_cents: int):
|
||||
"""Add usage to current budget"""
|
||||
self.current_usage_cents += amount_cents
|
||||
|
||||
# Check if budget is exceeded
|
||||
if self.current_usage_cents >= self.limit_cents:
|
||||
self.is_exceeded = True
|
||||
|
||||
# Check if warning threshold is reached
|
||||
if self.warning_threshold_cents and self.current_usage_cents >= self.warning_threshold_cents:
|
||||
if not self.is_warning_sent:
|
||||
self.is_warning_sent = True
|
||||
|
||||
self.updated_at = datetime.utcnow()
|
||||
|
||||
def reset_period(self):
|
||||
"""Reset budget for new period"""
|
||||
if self.rollover_unused and self.current_usage_cents < self.limit_cents:
|
||||
# Rollover unused budget
|
||||
unused_amount = self.limit_cents - self.current_usage_cents
|
||||
self.limit_cents += unused_amount
|
||||
|
||||
self.current_usage_cents = 0
|
||||
self.is_exceeded = False
|
||||
self.is_warning_sent = False
|
||||
self.last_reset_at = datetime.utcnow()
|
||||
|
||||
# Calculate next period
|
||||
if self.period_type == "daily":
|
||||
self.period_start = self.period_end
|
||||
self.period_end = self.period_start + timedelta(days=1)
|
||||
elif self.period_type == "weekly":
|
||||
self.period_start = self.period_end
|
||||
self.period_end = self.period_start + timedelta(weeks=1)
|
||||
elif self.period_type == "monthly":
|
||||
self.period_start = self.period_end
|
||||
# Handle month boundaries properly
|
||||
if self.period_start.month == 12:
|
||||
next_month = self.period_start.replace(year=self.period_start.year + 1, month=1)
|
||||
else:
|
||||
next_month = self.period_start.replace(month=self.period_start.month + 1)
|
||||
self.period_end = next_month
|
||||
elif self.period_type == "yearly":
|
||||
self.period_start = self.period_end
|
||||
self.period_end = self.period_start.replace(year=self.period_start.year + 1)
|
||||
|
||||
self.updated_at = datetime.utcnow()
|
||||
|
||||
def get_period_days_remaining(self) -> int:
|
||||
"""Get number of days remaining in current period"""
|
||||
if self.is_expired():
|
||||
return 0
|
||||
return (self.period_end - datetime.utcnow()).days
|
||||
|
||||
def get_daily_burn_rate(self) -> float:
|
||||
"""Get average daily spend rate in current period"""
|
||||
if not self.is_in_period():
|
||||
return 0.0
|
||||
|
||||
days_elapsed = (datetime.utcnow() - self.period_start).days
|
||||
if days_elapsed == 0:
|
||||
days_elapsed = 1 # Avoid division by zero
|
||||
|
||||
return self.current_usage_cents / days_elapsed / 100 # Return in dollars
|
||||
|
||||
def get_projected_spend(self) -> float:
|
||||
"""Get projected spend for entire period based on current burn rate"""
|
||||
daily_burn = self.get_daily_burn_rate()
|
||||
total_period_days = (self.period_end - self.period_start).days
|
||||
return daily_burn * total_period_days
|
||||
|
||||
@classmethod
|
||||
def create_monthly_budget(
|
||||
cls,
|
||||
user_id: int,
|
||||
name: str,
|
||||
limit_dollars: float,
|
||||
api_key_id: Optional[int] = None,
|
||||
warning_threshold_percentage: float = 0.8
|
||||
) -> "Budget":
|
||||
"""Create a monthly budget"""
|
||||
now = datetime.utcnow()
|
||||
# Start of current month
|
||||
period_start = now.replace(day=1, hour=0, minute=0, second=0, microsecond=0)
|
||||
# Start of next month
|
||||
if now.month == 12:
|
||||
period_end = period_start.replace(year=now.year + 1, month=1)
|
||||
else:
|
||||
period_end = period_start.replace(month=now.month + 1)
|
||||
|
||||
limit_cents = int(limit_dollars * 100)
|
||||
warning_threshold_cents = int(limit_cents * warning_threshold_percentage)
|
||||
|
||||
return cls(
|
||||
name=name,
|
||||
user_id=user_id,
|
||||
api_key_id=api_key_id,
|
||||
limit_cents=limit_cents,
|
||||
warning_threshold_cents=warning_threshold_cents,
|
||||
period_type="monthly",
|
||||
period_start=period_start,
|
||||
period_end=period_end,
|
||||
is_active=True,
|
||||
enforce_hard_limit=True,
|
||||
enforce_warning=True,
|
||||
auto_renew=True,
|
||||
notification_settings={
|
||||
"email_on_warning": True,
|
||||
"email_on_exceeded": True
|
||||
}
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def create_daily_budget(
|
||||
cls,
|
||||
user_id: int,
|
||||
name: str,
|
||||
limit_dollars: float,
|
||||
api_key_id: Optional[int] = None
|
||||
) -> "Budget":
|
||||
"""Create a daily budget"""
|
||||
now = datetime.utcnow()
|
||||
period_start = now.replace(hour=0, minute=0, second=0, microsecond=0)
|
||||
period_end = period_start + timedelta(days=1)
|
||||
|
||||
limit_cents = int(limit_dollars * 100)
|
||||
warning_threshold_cents = int(limit_cents * 0.8) # 80% warning threshold
|
||||
|
||||
return cls(
|
||||
name=name,
|
||||
user_id=user_id,
|
||||
api_key_id=api_key_id,
|
||||
limit_cents=limit_cents,
|
||||
warning_threshold_cents=warning_threshold_cents,
|
||||
period_type="daily",
|
||||
period_start=period_start,
|
||||
period_end=period_end,
|
||||
is_active=True,
|
||||
enforce_hard_limit=True,
|
||||
enforce_warning=True,
|
||||
auto_renew=True
|
||||
)
|
||||
Reference in New Issue
Block a user