src/functions_db.py

#https://docs.sqlalchemy.org/en/20/orm/quickstart.html

from typing import List
from sqlalchemy import ForeignKey, String, Date, DateTime, JSON,  create_engine
from sqlalchemy.orm import DeclarativeBase, Mapped, mapped_column, relationship, Session
import datetime
import socket
import sentry_sdk
from sentry_sdk.integrations.serverless import serverless_function
from dotenv import load_dotenv
import os

load_dotenv()

sentry_dsn = os.getenv("SENTRY_DSN")

sentry_sdk.init(
    dsn=sentry_dsn,
    send_default_pii=True,
    traces_sample_rate=1.0,
    _experiments={
        "continuous_profiling_auto_start": True,
    },
)

######################
##Classes of the SQL##
######################

class Base(DeclarativeBase):
    pass

class User(Base):
    __tablename__ = 'user_account'

    id: Mapped[int] = mapped_column(primary_key=True)
    name: Mapped[str] = mapped_column(String(30), unique=True)

    activity: Mapped[List['Activity']] = relationship(
        "Activity", 
        back_populates="user", 
        cascade="all, delete, delete-orphan"
    )
    llm_history: Mapped[List['LLM']] = relationship("LLM", back_populates="user")

    @serverless_function
    def add_user(self, session, user):
        if not session.query(User).filter_by(name=user).first():
            user = User(name=user)
            session.add(user)
            session.commit()
            print("User added")
        else:
            print("User already exists")

    def __repr__(self):
        return f"<User(name={self.name})>"

class Activity(Base):
    __tablename__ = 'activity'

    id: Mapped[int] = mapped_column(primary_key=True)
    date: Mapped[Date] = mapped_column(Date)
    time: Mapped[DateTime] = mapped_column(DateTime)
    ip: Mapped[str] = mapped_column(String(30))
    type_of_activity: Mapped[str] = mapped_column(String(30))

    user_id: Mapped[int] = mapped_column(ForeignKey('user_account.id'))
    user: Mapped['User'] = relationship("User", back_populates="activity")

    @serverless_function
    def add_activity(self, session, user, date, time, ip, type_of_activity='default'):
        user = session.query(User).filter_by(name=user).first()
        new_activity = Activity(date=date, time=time, ip=ip, user=user, type_of_activity=type_of_activity)
        session.add(new_activity)
        session.commit()
        print("Activity added")

    def __repr__(self):
        return f"<Activity(id={self.id}, date={self.date}, time={self.time}, ip={self.ip}, type_of_activity={self.type_of_activity})>"

class LLM(Base):
    __tablename__ = 'llm_history'

    id: Mapped[int] = mapped_column(primary_key=True)
    input_sentence: Mapped[str] = mapped_column(String(255))
    output_label: Mapped[str] = mapped_column(String(30))
    model_url: Mapped[str] = mapped_column(String(255))
    
    type_of_activity: Mapped[int] = mapped_column(ForeignKey('activity.id'))
    user_id: Mapped[int] = mapped_column(ForeignKey('user_account.id'))
    user: Mapped['User'] = relationship("User", back_populates="llm_history")

    @serverless_function
    def add_llm(self, session, user_name, input_sentence, output_label, model_url, type_of_activity='default'):
        user = session.query(User).filter_by(name=user_name).first()
        new_llm = LLM(input_sentence=input_sentence, output_label=output_label, model_url=model_url, user=user, type_of_activity=type_of_activity)
        session.add(new_llm)
        session.commit()
        print("LLM added")

    def __repr__(self):
        return f"<LLM(id={self.id}, input_sentence={self.input_sentence}, output_label={self.output_label}, model_url={self.model_url})>"
    
class ChatBot(Base):
    __tablename__ = 'chatbot_history'
    id: Mapped[int] = mapped_column(primary_key=True)
    session_id: Mapped[str] = mapped_column(String(50))
    user_input: Mapped[str] = mapped_column(String(255))
    model_output: Mapped[str] = mapped_column(String(255))

    @serverless_function
    def add_gpt_history(self, session, session_id, user_input, model_output):
        new_gpt_chat = ChatBot(session_id=session_id, user_input=user_input, model_output=model_output)
        session.add(new_gpt_chat)
        session.commit()
        print("GPT history added")

    def __repr__(self):
        return f"<ChatBot(id={self.id}, session_id={self.session_id}, user_input={self.user_input}, model_output={self.model_output})"


######################
# Simple functions ###
######################

@serverless_function
def create_db_and_tables():
    '''
    format for a relative path is sqlite:///relative/path/to/file.db
    for an absolute path is sqlite:////absolute/path/to/file.db.
    '''
    engine = create_engine("sqlite:///databases/main.db", echo=True)
    Base.metadata.create_all(engine)

    with Session(engine) as session:
        user = User()
        activity = Activity()
        llm = LLM()
        chat_history = ChatBot()

        user.add_user(
            session, 
            "admin"
        )
        activity.add_activity(
            session, 
            "admin", 
            datetime.date.today(), 
            datetime.datetime.now(), 
            socket.gethostbyname(socket.gethostname()),
            'default'
        )

@serverless_function
def connect_to_db(address: str = "sqlite:///databases/main.db"):
    engine = create_engine(address, echo=True)
    Base.metadata.create_all(engine)

    session = Session(engine)
    user = User()
    activity = Activity()
    llm = LLM()
    chat_history = ChatBot()
    return session, user, activity, llm, chat_history

if __name__ == "__main__":
    create_db_and_tables()