Archives/langchain
2
0
Fork 1
mirror of https://github.com/hwchase17/langchain synced 2024-11-04 06:00:26 +00:00
Code Issues Packages Projects Releases Wiki Activity

Add firestore_client param to FirestoreChatMessageHistory if caller already has one; also lets them specify GCP project, etc. (#8601)

Browse Source 
Existing implementation requires that you install `firebase-admin`
package, and prevents you from using an existing Firestore client
instance if available.

This adds optional `firestore_client` param to
`FirestoreChatMessageHistory`, so users can just use their existing
client/settings. If not passed, existing logic executes to initialize a
`firestore_client`.

---------

Co-authored-by: Bagatur <baskaryan@gmail.com>
This commit is contained in:
Neil Murphy 2023-08-01 15:42:13 -07:00 committed by GitHub
parent 13ccf202de
commit 31820a31e4
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 23 additions and 20 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

67
libs/langchain/langchain/memory/chat_message_histories/firestore.py
View File

@ -12,39 +12,10 @@ from langchain.schema.messages import BaseMessage, messages_from_dict, messages_
logger = logging.getLogger(__name__)
if TYPE_CHECKING:
from google.cloud.firestore import DocumentReference
from google.cloud.firestore import Client, DocumentReference
class FirestoreChatMessageHistory(BaseChatMessageHistory):
"""Chat message history backed by Google Firestore."""
def __init__(
self,
collection_name: str,
session_id: str,
user_id: str,
):
"""
Initialize a new instance of the FirestoreChatMessageHistory class.
:param collection_name: The name of the collection to use.
:param session_id: The session ID for the chat..
:param user_id: The user ID for the chat.
"""
self.collection_name = collection_name
self.session_id = session_id
self.user_id = user_id
self._document: Optional[DocumentReference] = None
self.messages: List[BaseMessage] = []
self.prepare_firestore()
def prepare_firestore(self) -> None:
"""Prepare the Firestore client.
Use this function to make sure your database is ready.
"""
def _get_firestore_client() -> Client:
try:
import firebase_admin
from firebase_admin import firestore
@ -61,7 +32,39 @@ class FirestoreChatMessageHistory(BaseChatMessageHistory):
logger.debug("Initializing Firebase app: %s", e)
firebase_admin.initialize_app()
self.firestore_client = firestore.client()
return firestore.client()
class FirestoreChatMessageHistory(BaseChatMessageHistory):
"""Chat message history backed by Google Firestore."""
def __init__(
self,
collection_name: str,
session_id: str,
user_id: str,
firestore_client: Optional[Client] = None,
):
"""
Initialize a new instance of the FirestoreChatMessageHistory class.
:param collection_name: The name of the collection to use.
:param session_id: The session ID for the chat..
:param user_id: The user ID for the chat.
"""
self.collection_name = collection_name
self.session_id = session_id
self.user_id = user_id
self._document: Optional[DocumentReference] = None
self.messages: List[BaseMessage] = []
self.firestore_client = firestore_client or _get_firestore_client()
self.prepare_firestore()
def prepare_firestore(self) -> None:
"""Prepare the Firestore client.
Use this function to make sure your database is ready.
"""
self._document = self.firestore_client.collection(
self.collection_name
).document(self.session_id)