|
|
import logging |
|
|
from typing import List, TYPE_CHECKING, Optional |
|
|
from datetime import datetime |
|
|
import pytz |
|
|
|
|
|
from langchain_core.chat_history import BaseChatMessageHistory |
|
|
from langchain_core.messages import ( |
|
|
BaseMessage, |
|
|
message_to_dict, |
|
|
messages_from_dict, |
|
|
) |
|
|
from langchain_core.utils import get_from_env |
|
|
|
|
|
if TYPE_CHECKING: |
|
|
from supabase import Client |
|
|
|
|
|
logger = logging.getLogger(__name__) |
|
|
|
|
|
class SupabaseChatMessageHistory(BaseChatMessageHistory): |
|
|
"""Chat message history stored in a Supabase project database.""" |
|
|
|
|
|
def __init__( |
|
|
self, |
|
|
session_id: str, |
|
|
table_name: str = "message_store", |
|
|
session_name: str = "session", |
|
|
client: Optional['Client'] = None, |
|
|
supabase_url: Optional[str] = None, |
|
|
supabase_key: Optional[str] = None, |
|
|
): |
|
|
try: |
|
|
from supabase import create_client |
|
|
except ImportError: |
|
|
raise ImportError( |
|
|
"Could not import supabase python package. " |
|
|
"Please install it with `pip install supabase`." |
|
|
) |
|
|
|
|
|
|
|
|
if not session_id: |
|
|
raise ValueError("Please ensure that the session_id parameter is provided") |
|
|
|
|
|
self.client = client |
|
|
if client is None: |
|
|
supabase_url = get_from_env("url", "SUPABASE_URL", supabase_url) |
|
|
supabase_key = get_from_env("key", "SUPABASE_KEY", supabase_key) |
|
|
|
|
|
self.client = create_client( |
|
|
supabase_url=supabase_url, |
|
|
supabase_key=supabase_key |
|
|
) |
|
|
|
|
|
self.session_id = session_id |
|
|
self.table_name = table_name |
|
|
self.session_name = session_name |
|
|
|
|
|
@property |
|
|
def messages(self) -> List[BaseMessage]: |
|
|
"""Retrieve the messages from the Supabase project database""" |
|
|
response = self.client.table(self.table_name) \ |
|
|
.select("id", "query_id", "message", "error_log") \ |
|
|
.eq(f"{self.session_name}_id", self.session_id) \ |
|
|
.order('created_at', desc=False) \ |
|
|
.execute() |
|
|
|
|
|
failed_messages = [record for record in response.data if record["message"]["data"]["content"] == "" or record["error_log"] is not None] |
|
|
|
|
|
failed_ids = [] |
|
|
for failed_message in failed_messages: |
|
|
failed_ids.extend([failed_message["id"], failed_message["query_id"]]) |
|
|
|
|
|
items = [record["message"] for record in response.data if record["id"] not in failed_ids] |
|
|
messages = messages_from_dict(items) |
|
|
|
|
|
return messages |
|
|
|
|
|
def add_message(self, message: BaseMessage, query_id: Optional[str] = None) -> None: |
|
|
"""Append the message to the record in the Supabase project database""" |
|
|
response = self.client.table(self.table_name).insert( |
|
|
{ |
|
|
f"{self.session_name}_id": self.session_id, |
|
|
"message": message_to_dict(message), |
|
|
"query_id": query_id, |
|
|
} |
|
|
).execute() |
|
|
|
|
|
return response.data[0]["id"] |
|
|
|
|
|
def update_message( |
|
|
self, |
|
|
message_id:str, |
|
|
message: Optional[BaseMessage] = None, |
|
|
error_log: Optional[dict] = None |
|
|
) -> None: |
|
|
"""Append the message to the record in the Supabase project database""" |
|
|
|
|
|
updated_dict = { |
|
|
"updated_at": datetime.now(pytz.utc).isoformat() |
|
|
} |
|
|
|
|
|
if message is not None: |
|
|
updated_dict["message"] = message_to_dict(message) |
|
|
|
|
|
if error_log is not None: |
|
|
updated_dict["error_log"] = error_log |
|
|
|
|
|
|
|
|
self.client.table(self.table_name).update(updated_dict) \ |
|
|
.eq('id', message_id) \ |
|
|
.execute() |
|
|
|
|
|
|
|
|
def clear(self) -> None: |
|
|
"""Clear session memory from the Supabase project database""" |
|
|
self.client.table(self.table_name) \ |
|
|
.delete() \ |
|
|
.eq(f"{self.session_name}_id", self.session_id) \ |
|
|
.execute() |
|
|
|
