TakwayDisplayPlatform/app/model.py

from sqlalchemy import create_engine, Column, Integer, String, CHAR
from sqlalchemy.ext.declarative import declarative_base
from sqlalchemy.orm import sessionmaker, Session
from contextlib import contextmanager
from config import Config


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

MESSAGE_LENGTH_LIMIT = 2**31-1

class User(Base):
    __tablename__ = "user"
    id = Column(CHAR(36), primary_key=True)
    name = Column(String(32))
    email = Column(String(64))
    password = Column(String(128))

class Assistant(Base):
    __tablename__ = "assistant"
    id = Column(CHAR(36), primary_key=True)
    user_id = Column(CHAR(36))
    name = Column(String(32))
    system_prompt = Column(String(512))
    messages = Column(String(MESSAGE_LENGTH_LIMIT))
    user_info = Column(String(256))
    llm_info = Column(String(256))
    tts_info = Column(String(256))
    token = Column(Integer)

Base.metadata.create_all(bind=engine)

@contextmanager
def get_db_context():
    db = SessionLocal()
    try:
        yield db
    finally:
        db.close()

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