65 lines
1.8 KiB
Python
65 lines
1.8 KiB
Python
from typing import TYPE_CHECKING, Dict, List
|
|
|
|
from chainlit.context import context
|
|
|
|
if TYPE_CHECKING:
|
|
from chainlit.message import Message
|
|
|
|
chat_contexts: Dict[str, List["Message"]] = {}
|
|
|
|
|
|
class ChatContext:
|
|
def get(self) -> List["Message"]:
|
|
if not context.session:
|
|
return []
|
|
|
|
if context.session.id not in chat_contexts:
|
|
# Create a new chat context
|
|
chat_contexts[context.session.id] = []
|
|
|
|
return chat_contexts[context.session.id].copy()
|
|
|
|
def add(self, message: "Message"):
|
|
if not context.session:
|
|
return
|
|
|
|
if context.session.id not in chat_contexts:
|
|
chat_contexts[context.session.id] = []
|
|
|
|
if message not in chat_contexts[context.session.id]:
|
|
chat_contexts[context.session.id].append(message)
|
|
|
|
return message
|
|
|
|
def remove(self, message: "Message") -> bool:
|
|
if not context.session:
|
|
return False
|
|
|
|
if context.session.id not in chat_contexts:
|
|
return False
|
|
|
|
if message in chat_contexts[context.session.id]:
|
|
chat_contexts[context.session.id].remove(message)
|
|
return True
|
|
|
|
return False
|
|
|
|
def clear(self) -> None:
|
|
if context.session and context.session.id in chat_contexts:
|
|
chat_contexts[context.session.id] = []
|
|
|
|
def to_openai(self):
|
|
messages = []
|
|
for message in self.get():
|
|
if message.type == "assistant_message":
|
|
messages.append({"role": "assistant", "content": message.content})
|
|
elif message.type == "user_message":
|
|
messages.append({"role": "user", "content": message.content})
|
|
else:
|
|
messages.append({"role": "system", "content": message.content})
|
|
|
|
return messages
|
|
|
|
|
|
chat_context = ChatContext()
|