forked from Archives/langchain
Harrison/postgres (#2691)
Co-authored-by: Ankit Jain <ankneo@users.noreply.github.com>
This commit is contained in:
parent
bb5118f4c9
commit
e0a13e9355
@ -0,0 +1,62 @@
|
||||
{
|
||||
"cells": [
|
||||
{
|
||||
"attachments": {},
|
||||
"cell_type": "markdown",
|
||||
"id": "91c6a7ef",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"# Postgres Chat Message History\n",
|
||||
"\n",
|
||||
"This notebook goes over how to use Postgres to store chat message history."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"id": "d15e3302",
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"from langchain.memory import PostgresChatMessageHistory\n",
|
||||
"\n",
|
||||
"history = PostgresChatMessageHistory(connection_string=\"postgresql://postgres:mypassword@localhost/chat_history\", session_id=\"foo\")\n",
|
||||
"\n",
|
||||
"history.add_user_message(\"hi!\")\n",
|
||||
"\n",
|
||||
"history.add_ai_message(\"whats up?\")"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"id": "64fc465e",
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"history.messages"
|
||||
]
|
||||
}
|
||||
],
|
||||
"metadata": {
|
||||
"kernelspec": {
|
||||
"display_name": "Python 3 (ipykernel)",
|
||||
"language": "python",
|
||||
"name": "python3"
|
||||
},
|
||||
"language_info": {
|
||||
"codemirror_mode": {
|
||||
"name": "ipython",
|
||||
"version": 3
|
||||
},
|
||||
"file_extension": ".py",
|
||||
"mimetype": "text/x-python",
|
||||
"name": "python",
|
||||
"nbconvert_exporter": "python",
|
||||
"pygments_lexer": "ipython3",
|
||||
"version": "3.11.2"
|
||||
}
|
||||
},
|
||||
"nbformat": 4,
|
||||
"nbformat_minor": 5
|
||||
}
|
@ -5,6 +5,7 @@ from langchain.memory.buffer import (
|
||||
from langchain.memory.buffer_window import ConversationBufferWindowMemory
|
||||
from langchain.memory.chat_message_histories.dynamodb import DynamoDBChatMessageHistory
|
||||
from langchain.memory.chat_message_histories.in_memory import ChatMessageHistory
|
||||
from langchain.memory.chat_message_histories.postgres import PostgresChatMessageHistory
|
||||
from langchain.memory.chat_message_histories.redis import RedisChatMessageHistory
|
||||
from langchain.memory.combined import CombinedMemory
|
||||
from langchain.memory.entity import (
|
||||
@ -36,4 +37,5 @@ __all__ = [
|
||||
"ConversationTokenBufferMemory",
|
||||
"RedisChatMessageHistory",
|
||||
"DynamoDBChatMessageHistory",
|
||||
"PostgresChatMessageHistory",
|
||||
]
|
||||
|
@ -1,7 +1,9 @@
|
||||
from langchain.memory.chat_message_histories.dynamodb import DynamoDBChatMessageHistory
|
||||
from langchain.memory.chat_message_histories.postgres import PostgresChatMessageHistory
|
||||
from langchain.memory.chat_message_histories.redis import RedisChatMessageHistory
|
||||
|
||||
__all__ = [
|
||||
"DynamoDBChatMessageHistory",
|
||||
"RedisChatMessageHistory",
|
||||
"PostgresChatMessageHistory",
|
||||
]
|
||||
|
86
langchain/memory/chat_message_histories/postgres.py
Normal file
86
langchain/memory/chat_message_histories/postgres.py
Normal file
@ -0,0 +1,86 @@
|
||||
import json
|
||||
import logging
|
||||
from typing import List
|
||||
|
||||
from langchain.schema import (
|
||||
AIMessage,
|
||||
BaseChatMessageHistory,
|
||||
BaseMessage,
|
||||
HumanMessage,
|
||||
_message_to_dict,
|
||||
messages_from_dict,
|
||||
)
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
DEFAULT_CONNECTION_STRING = "postgresql://postgres:mypassword@localhost/chat_history"
|
||||
|
||||
|
||||
class PostgresChatMessageHistory(BaseChatMessageHistory):
|
||||
def __init__(
|
||||
self,
|
||||
session_id: str,
|
||||
connection_string: str = DEFAULT_CONNECTION_STRING,
|
||||
table_name: str = "message_store",
|
||||
):
|
||||
import psycopg
|
||||
from psycopg.rows import dict_row
|
||||
|
||||
try:
|
||||
self.connection = psycopg.connect(connection_string)
|
||||
self.cursor = self.connection.cursor(row_factory=dict_row)
|
||||
except psycopg.OperationalError as error:
|
||||
logger.error(error)
|
||||
|
||||
self.session_id = session_id
|
||||
self.table_name = table_name
|
||||
|
||||
self._create_table_if_not_exists()
|
||||
|
||||
def _create_table_if_not_exists(self) -> None:
|
||||
create_table_query = f"""CREATE TABLE IF NOT EXISTS {self.table_name} (
|
||||
id SERIAL PRIMARY KEY,
|
||||
session_id TEXT NOT NULL,
|
||||
message JSONB NOT NULL
|
||||
);"""
|
||||
self.cursor.execute(create_table_query)
|
||||
self.connection.commit()
|
||||
|
||||
@property
|
||||
def messages(self) -> List[BaseMessage]: # type: ignore
|
||||
"""Retrieve the messages from PostgreSQL"""
|
||||
query = f"SELECT message FROM {self.table_name} WHERE session_id = %s;"
|
||||
self.cursor.execute(query, (self.session_id,))
|
||||
items = [record["message"] for record in self.cursor.fetchall()]
|
||||
messages = messages_from_dict(items)
|
||||
return messages
|
||||
|
||||
def add_user_message(self, message: str) -> None:
|
||||
self.append(HumanMessage(content=message))
|
||||
|
||||
def add_ai_message(self, message: str) -> None:
|
||||
self.append(AIMessage(content=message))
|
||||
|
||||
def append(self, message: BaseMessage) -> None:
|
||||
"""Append the message to the record in PostgreSQL"""
|
||||
from psycopg import sql
|
||||
|
||||
query = sql.SQL("INSERT INTO {} (session_id, message) VALUES (%s, %s);").format(
|
||||
sql.Identifier(self.table_name)
|
||||
)
|
||||
self.cursor.execute(
|
||||
query, (self.session_id, json.dumps(_message_to_dict(message)))
|
||||
)
|
||||
self.connection.commit()
|
||||
|
||||
def clear(self) -> None:
|
||||
"""Clear session memory from PostgreSQL"""
|
||||
query = f"DELETE FROM {self.table_name} WHERE session_id = %s;"
|
||||
self.cursor.execute(query, (self.session_id,))
|
||||
self.connection.commit()
|
||||
|
||||
def __del__(self) -> None:
|
||||
if self.cursor:
|
||||
self.cursor.close()
|
||||
if self.connection:
|
||||
self.connection.close()
|
Loading…
Reference in New Issue
Block a user