File size: 5,526 Bytes
86e2833 | 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 131 132 133 134 135 136 137 138 139 140 141 142 143 | from sqlalchemy import insert, MetaData, Table
from sqlalchemy.orm import sessionmaker
from sqlalchemy import create_engine, select
from typing import Any
import uuid
from sqlalchemy import Column, Integer, String, ForeignKey, DateTime
from sqlalchemy.ext.declarative import declarative_base
import uuid
from datetime import datetime
Base = declarative_base()
class Origin(Base):
__tablename__ = "origin"
id_origin = Column(
String, primary_key=True, index=True, default=lambda: str(uuid.uuid4())
)
origin = Column(String, nullable=True)
class Status(Base):
__tablename__ = "status"
id_status = Column(
String, primary_key=True, index=True, default=lambda: str(uuid.uuid4())
)
status = Column(String, nullable=True)
class Prompt(Base):
__tablename__ = "prompt"
id_prompt = Column(
String, primary_key=True, index=True, default=lambda: str(uuid.uuid4())
)
session_id = Column(String, nullable=False)
id_origin = Column(String, ForeignKey("origin.id_origin"), nullable=True)
prompt = Column(String, nullable=True)
response = Column(String, nullable=True)
timestamp = Column(DateTime, nullable=True, default=datetime.now)
class Log(Base):
__tablename__ = "log"
id_log = Column(
String, primary_key=True, index=True, default=lambda: str(uuid.uuid4())
)
timestamp = Column(String, nullable=False)
id_prompt = Column(String, ForeignKey("prompt.id_prompt"), nullable=False)
id_status = Column(String, ForeignKey("status.id_status"), nullable=False)
id_origin = Column(String, ForeignKey("origin.id_origin"), nullable=False)
class Database:
def __init__(self, db_path: str):
self.engine = create_engine(f"sqlite:///{db_path}")
self.Session = sessionmaker(bind=self.engine)
self.session = self.Session()
self.metadata = MetaData()
self.metadata.reflect(bind=self.engine) # Load all tables
def insert(self, table: str, data: dict[str, Any]) -> bool:
"""
Inserts a row into the specified table.
:param table: str.
:param data: Dictionary of column names and values.
:return: True if insertion is successful, otherwise False.
clés pour enregistrer dans la table prompt:
- session_id
- origin
- prompt
- response
clés pour enregistrer dans la table log :
- timestamp
- prompt
- status
"""
if not data:
erreur = "Error: No data provided for insertion."
return erreur, False
table_obj = self.metadata.tables.get(table)
if table_obj is None:
erreur = f"Error: Table '{table}' does not exist."
return erreur, False
try:
with self.Session() as session:
if table == "prompt":
statement = select(Origin).where(Origin.origin == data["origin"])
id_origin = session.scalar(statement=statement)
if not id_origin:
insert_origin = insert(Origin).values(
id_origin=str(uuid.uuid4()), origin=data["origin"]
)
session.execute(statement=insert_origin)
session.commit()
id_origin = session.scalar(statement=statement).id_origin
insert_prompt = insert(Prompt).values(
id_prompt=str(uuid.uuid4()),
session_id=data["session_id"],
id_origin=id_origin,
prompt=data["prompt"],
response=data["response"],
)
session.execute(insert_prompt)
session.commit()
return id_prompt, True
if table == "log":
statement = select(Status).where(Status.status == data["status"])
id_status = session.scalar(statement=statement)
if not id_status:
insert_status = insert(Status).values(
id_status=str(uuid.uuid4()), status=data["status"]
)
session.execute(insert_status)
session.commit()
id_status = session.scalar(statement).id_status
statement = select(Prompt).where(Prompt.prompt == data["prompt"])
prompt = session.scalar(statement)
if not prompt:
raise "Le prompt demandé n'existe pas."
id_prompt = prompt.id_prompt
id_origin = prompt.id_origin
insert_log = insert(Log).values(
id_log=str(uuid.uuid4()),
timestamp=data["timestamp"],
id_prompt=id_prompt,
id_status=id_status,
id_origin=id_origin,
)
session.execute(insert_log)
session.commit()
return id_status, True
except Exception as e:
erreur = f"Error inserting into table '{table}': {e}"
return erreur, False
|