diff --git a/langchain_postgres/chat_message_histories.py b/langchain_postgres/chat_message_histories.py index 85f2aef..1f6c292 100644 --- a/langchain_postgres/chat_message_histories.py +++ b/langchain_postgres/chat_message_histories.py @@ -2,6 +2,7 @@ This client provides support for both sync and async via psycopg 3. """ + from __future__ import annotations import json @@ -370,3 +371,12 @@ async def aclear(self) -> None: async with self._aconnection.cursor() as cursor: await cursor.execute(query, {"session_id": self._session_id}) await self._aconnection.commit() + + def update_session_id(self, new_session_id: str) -> None: + """Update the session ID for the current instance and in the database""" + update_query = ( + f"UPDATE {self._table_name} SET session_id = %s WHERE session_id = %s;" + ) + with self._connection.cursor() as cursor: + cursor.execute(update_query, (new_session_id, self._session_id)) + self._connection.commit()