Spaces:
Runtime error
Runtime error
File size: 1,891 Bytes
bedb8e2 bf45c7d bedb8e2 bf45c7d 27b075e bedb8e2 27b075e bedb8e2 38d63de bedb8e2 bf45c7d 38d63de |
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 |
from __future__ import annotations
from datetime import datetime
from pathlib import Path
from peewee import (
AutoField,
CharField,
DateTimeField,
ForeignKeyField,
Model,
SqliteDatabase,
TextField,
)
_DB_PATH = Path(__file__).resolve().parent.parent / "chat.db"
_db = SqliteDatabase(_DB_PATH)
class BaseModel(Model):
class Meta:
database = _db
class User(BaseModel):
id = AutoField()
username = CharField(unique=True)
class Conversation(BaseModel):
id = AutoField()
user = ForeignKeyField(User, backref="conversations")
session_name = CharField()
started_at = DateTimeField(default=datetime.utcnow)
class Meta:
indexes = ((("user", "session_name"), True),)
class Message(BaseModel):
id = AutoField()
conversation = ForeignKeyField(Conversation, backref="messages")
role = CharField()
content = TextField()
created_at = DateTimeField(default=datetime.utcnow)
__all__ = [
"_db",
"User",
"Conversation",
"Message",
"reset_history",
]
def init_db() -> None:
"""Initialise the database and create tables if they do not exist."""
if _db.is_closed():
_db.connect()
_db.create_tables([User, Conversation, Message])
def reset_history(username: str, session_name: str) -> int:
"""Delete all messages for the given user and session."""
init_db()
try:
user = User.get(User.username == username)
conv = Conversation.get(
Conversation.user == user, Conversation.session_name == session_name
)
except (User.DoesNotExist, Conversation.DoesNotExist):
return 0
deleted = Message.delete().where(Message.conversation == conv).execute()
conv.delete_instance()
if not Conversation.select().where(Conversation.user == user).exists():
user.delete_instance()
return deleted
|