cz-base / app /utils /chat.py
chengzhiedu's picture
feat(ui): opening message in chat
c2ab585
import os
from openai import OpenAI
import gradio as gr
from typing import Tuple
from supabase import create_client
from datetime import datetime
class ChatManager:
def __init__(self, client: OpenAI, assistant_id: str, supabase_url: str = None, supabase_key: str = None):
self.client = client
self.assistant_id = assistant_id
"""Initialize Supabase client"""
self.supabase = create_client(
supabase_url or os.environ.get("SUPABASE_URL"),
supabase_key or os.environ.get("SUPABASE_KEY")
)
def get_current_chat_id(self, session_state):
"""Get current chat ID from session state"""
return session_state.get('current_chat_id')
def set_current_chat_id(self, session_state, chat_id):
"""Set current chat ID in session state"""
session_state['current_chat_id'] = chat_id
def create_new_chat(self, session_state):
"""Create a new chat using OpenAI thread ID as chat_id"""
thread = self.client.beta.threads.create()
self.set_current_chat_id(session_state, thread.id)
return thread.id
def get_or_create_thread(self, session_state):
"""Get or create a thread for the chat"""
current_chat_id = self.get_current_chat_id(session_state)
if not current_chat_id:
current_chat_id = self.create_new_chat(session_state)
return current_chat_id
def reset_chat_on_load(self, username: str, session_state, init_chat=[]) -> Tuple[gr.update, list]:
"""Reset chat ID on page load and update chat selector"""
self.create_new_chat(session_state)
# Update chat selector with available chats for this user
chats = self.list_user_chats(username)
return gr.update(choices=[(c["preview"], c["chat_id"]) for c in chats], value=self.get_current_chat_id(session_state)), init_chat
def load_chat_history(self, user_id="default_user"):
"""Load chat history for a user"""
chat_id = self.get_latest_chat_id(user_id)
if chat_id:
return self.load_chat(user_id, chat_id)
return []
def list_user_chats(self, user_id="default_user"):
"""List all chats for a user"""
response = self.supabase.table("chats") \
.select("chat_id,last_updated,messages") \
.eq("user_id", user_id) \
.execute()
chats = [{
"chat_id": chat["chat_id"],
"last_updated": chat["last_updated"],
"preview": chat["messages"][0]["content"] if chat["messages"] else "Empty chat"
} for chat in response.data]
return sorted(chats, key=lambda x: x["last_updated"], reverse=True)
def switch_chat(self, chat_id, user_id="default_user", session_state=None):
"""Switch to a different chat"""
if session_state:
self.set_current_chat_id(session_state, chat_id)
messages = self.load_chat(user_id, chat_id)
return messages
def save_chat(self, user_id, chat_id, messages, current_lesson_plan, app_name):
"""Save chat history to Supabase"""
chat_data = {
"user_id": user_id,
"chat_id": chat_id,
"last_updated": datetime.now().isoformat(),
"messages": messages,
"current_lesson_plan": current_lesson_plan,
"app_name": app_name
}
# Check if chat exists
existing_chat = self.supabase.table("chats") \
.select("*") \
.eq("user_id", user_id) \
.eq("chat_id", chat_id) \
.execute()
if existing_chat.data:
# Update existing chat
self.supabase.table("chats") \
.update(chat_data) \
.eq("user_id", user_id) \
.eq("chat_id", chat_id) \
.execute()
else:
# Insert new chat
self.supabase.table("chats") \
.insert(chat_data) \
.execute()
def load_chat(self, user_id, chat_id):
"""Load specific chat history from Supabase"""
response = self.supabase.table("chats") \
.select("messages") \
.eq("user_id", user_id) \
.eq("chat_id", chat_id) \
.execute()
return response.data[0]["messages"] if response.data else []
def get_latest_chat_id(self, user_id):
"""Get the most recent chat ID for a user"""
chats = self.list_user_chats(user_id)
return chats[0]["chat_id"] if chats else None