Spaces:
Sleeping
Sleeping
File size: 5,009 Bytes
c19c7bf | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 | """
Database models for billing and usage tracking.
"""
from sqlalchemy import Column, String, Integer, Float, DateTime, Boolean, ForeignKey, Text
from sqlalchemy.orm import relationship
from datetime import datetime
from typing import Optional
from app.db.database import Base
class Tenant(Base):
"""Tenant/organization model."""
__tablename__ = "tenants"
id = Column(String, primary_key=True, index=True)
name = Column(String, nullable=False)
created_at = Column(DateTime, default=datetime.utcnow, nullable=False)
updated_at = Column(DateTime, default=datetime.utcnow, onupdate=datetime.utcnow, nullable=False)
# Relationships
plan = relationship("TenantPlan", back_populates="tenant", uselist=False)
usage_events = relationship("UsageEvent", back_populates="tenant")
daily_usage = relationship("UsageDaily", back_populates="tenant")
monthly_usage = relationship("UsageMonthly", back_populates="tenant")
class TenantPlan(Base):
"""Tenant subscription plan."""
__tablename__ = "tenant_plans"
id = Column(Integer, primary_key=True, autoincrement=True)
tenant_id = Column(String, ForeignKey("tenants.id"), unique=True, nullable=False, index=True)
plan_name = Column(String, nullable=False, index=True) # "starter", "growth", "pro"
monthly_chat_limit = Column(Integer, nullable=False) # -1 for unlimited
created_at = Column(DateTime, default=datetime.utcnow, nullable=False)
updated_at = Column(DateTime, default=datetime.utcnow, onupdate=datetime.utcnow, nullable=False)
# Relationships
tenant = relationship("Tenant", back_populates="plan")
class UsageEvent(Base):
"""Individual usage event (each /chat request)."""
__tablename__ = "usage_events"
id = Column(Integer, primary_key=True, autoincrement=True)
request_id = Column(String, unique=True, nullable=False, index=True)
tenant_id = Column(String, ForeignKey("tenants.id"), nullable=False, index=True)
user_id = Column(String, nullable=False, index=True)
kb_id = Column(String, nullable=False)
# LLM details
provider = Column(String, nullable=False) # "gemini" or "openai"
model = Column(String, nullable=False)
# Token usage
prompt_tokens = Column(Integer, nullable=False, default=0)
completion_tokens = Column(Integer, nullable=False, default=0)
total_tokens = Column(Integer, nullable=False, default=0)
# Cost tracking
estimated_cost_usd = Column(Float, nullable=False, default=0.0)
# Timestamp
request_timestamp = Column(DateTime, default=datetime.utcnow, nullable=False, index=True)
# Relationships
tenant = relationship("Tenant", back_populates="usage_events")
class UsageDaily(Base):
"""Daily aggregated usage per tenant."""
__tablename__ = "usage_daily"
id = Column(Integer, primary_key=True, autoincrement=True)
tenant_id = Column(String, ForeignKey("tenants.id"), nullable=False, index=True)
date = Column(DateTime, nullable=False, index=True)
# Aggregated metrics
total_requests = Column(Integer, nullable=False, default=0)
total_tokens = Column(Integer, nullable=False, default=0)
total_cost_usd = Column(Float, nullable=False, default=0.0)
# Provider breakdown
gemini_requests = Column(Integer, nullable=False, default=0)
openai_requests = Column(Integer, nullable=False, default=0)
created_at = Column(DateTime, default=datetime.utcnow, nullable=False)
updated_at = Column(DateTime, default=datetime.utcnow, onupdate=datetime.utcnow, nullable=False)
# Unique constraint: one record per tenant per day
__table_args__ = (
{"sqlite_autoincrement": True},
)
# Relationships
tenant = relationship("Tenant", back_populates="daily_usage")
class UsageMonthly(Base):
"""Monthly aggregated usage per tenant."""
__tablename__ = "usage_monthly"
id = Column(Integer, primary_key=True, autoincrement=True)
tenant_id = Column(String, ForeignKey("tenants.id"), nullable=False, index=True)
year = Column(Integer, nullable=False, index=True)
month = Column(Integer, nullable=False, index=True) # 1-12
# Aggregated metrics
total_requests = Column(Integer, nullable=False, default=0)
total_tokens = Column(Integer, nullable=False, default=0)
total_cost_usd = Column(Float, nullable=False, default=0.0)
# Provider breakdown
gemini_requests = Column(Integer, nullable=False, default=0)
openai_requests = Column(Integer, nullable=False, default=0)
created_at = Column(DateTime, default=datetime.utcnow, nullable=False)
updated_at = Column(DateTime, default=datetime.utcnow, onupdate=datetime.utcnow, nullable=False)
# Unique constraint: one record per tenant per month
__table_args__ = (
{"sqlite_autoincrement": True},
)
# Relationships
tenant = relationship("Tenant", back_populates="monthly_usage")
|