database.py

import os
from datetime import datetime
from sqlalchemy import create_engine, Column, Integer, String, Boolean, DateTime, ForeignKey, Float, text
from sqlalchemy.ext.declarative import declarative_base
from sqlalchemy.orm import sessionmaker, relationship

DATABASE_URL = "sqlite:///vapi.db"

engine = create_engine(DATABASE_URL, connect_args={"check_same_thread": False})
SessionLocal = sessionmaker(autocommit=False, autoflush=False, bind=engine)
Base = declarative_base()

class Provider(Base):
    __tablename__ = "providers"
    id = Column(Integer, primary_key=True, index=True)
    name = Column(String, nullable=False)
    api_key_encrypted = Column(String, nullable=False)
    is_active = Column(Boolean, default=True)
    created_at = Column(DateTime, default=datetime.utcnow)

    virtual_keys = relationship("VirtualKey", back_populates="provider", cascade="all, delete-orphan")

class VirtualKey(Base):
    __tablename__ = "virtual_keys"
    id = Column(Integer, primary_key=True, index=True)
    key = Column(String, unique=True, index=True, nullable=False)
    name = Column(String, nullable=False)
    provider_id = Column(Integer, ForeignKey("providers.id"), nullable=False)
    is_active = Column(Boolean, default=True)
    daily_budget = Column(Float, nullable=True)
    monthly_budget = Column(Float, nullable=True)
    allowed_models = Column(String, nullable=True)
    created_at = Column(DateTime, default=datetime.utcnow)

    provider = relationship("Provider", back_populates="virtual_keys")
    usage_logs = relationship("UsageLog", back_populates="virtual_key", cascade="all, delete-orphan")

class UsageLog(Base):
    __tablename__ = "usage_logs"
    id = Column(Integer, primary_key=True, index=True)
    virtual_key_id = Column(Integer, ForeignKey("virtual_keys.id"), nullable=False)
    endpoint = Column(String, nullable=False)
    method = Column(String, nullable=False)
    status_code = Column(Integer, nullable=False)
    cost = Column(Float, default=0.0)
    timestamp = Column(DateTime, default=datetime.utcnow)

    virtual_key = relationship("VirtualKey", back_populates="usage_logs")

class Setting(Base):
    __tablename__ = "settings"
    id = Column(Integer, primary_key=True, index=True)
    key = Column(String, unique=True, index=True, nullable=False)
    value = Column(String, nullable=False)

def init_db():
    Base.metadata.create_all(bind=engine)
    try:
        with engine.begin() as conn:
            conn.execute(text("ALTER TABLE virtual_keys ADD COLUMN daily_budget FLOAT"))
    except Exception: pass
    try:
        with engine.begin() as conn:
            conn.execute(text("ALTER TABLE virtual_keys ADD COLUMN monthly_budget FLOAT"))
    except Exception: pass
    try:
        with engine.begin() as conn:
            conn.execute(text("ALTER TABLE virtual_keys ADD COLUMN allowed_models VARCHAR"))
    except Exception: pass
    try:
        with engine.begin() as conn:
            conn.execute(text("ALTER TABLE usage_logs ADD COLUMN cost FLOAT DEFAULT 0.0"))
    except Exception: pass

def get_db():
    db = SessionLocal()
    try:
        yield db
    finally:
        db.close()