quivr/backend/core/models/chat.py
Mamadou DICKO 6e777327aa
feat(chatMessages): add brain_id and prompt_id columns (#912)
* feat: add prompt_id and brain_id to chat history)

* feat: add prompt_id and brain_id to chat routes
2023-08-10 10:25:08 +02:00

42 lines
1.1 KiB
Python

from dataclasses import asdict, dataclass
from typing import Optional
from uuid import UUID
@dataclass
class Chat:
chat_id: str
user_id: str
creation_time: str
chat_name: str
def __init__(self, chat_dict: dict):
self.chat_id = chat_dict.get("chat_id", "")
self.user_id = chat_dict.get("user_id", "")
self.creation_time = chat_dict.get("creation_time", "")
self.chat_name = chat_dict.get("chat_name", "")
@dataclass
class ChatHistory:
chat_id: str
message_id: str
user_message: str
assistant: str
message_time: str
prompt_id: Optional[UUID]
brain_id: Optional[UUID]
def __init__(self, chat_dict: dict):
self.chat_id = chat_dict.get("chat_id", "")
self.message_id = chat_dict.get("message_id", "")
self.user_message = chat_dict.get("user_message", "")
self.assistant = chat_dict.get("assistant", "")
self.message_time = chat_dict.get("message_time", "")
self.prompt_id = chat_dict.get("prompt_id")
self.brain_id = chat_dict.get("brain_id")
def to_dict(self):
return asdict(self)