From 31820a31e498fbb7ed911cdc6f6e9e0fd551829c Mon Sep 17 00:00:00 2001 From: Neil Murphy Date: Tue, 1 Aug 2023 15:42:13 -0700 Subject: [PATCH] Add firestore_client param to FirestoreChatMessageHistory if caller already has one; also lets them specify GCP project, etc. (#8601) 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 --- .../chat_message_histories/firestore.py | 43 ++++++++++--------- 1 file changed, 23 insertions(+), 20 deletions(-) diff --git a/libs/langchain/langchain/memory/chat_message_histories/firestore.py b/libs/langchain/langchain/memory/chat_message_histories/firestore.py index 9e20a9ba965..e1f2435b27f 100644 --- a/libs/langchain/langchain/memory/chat_message_histories/firestore.py +++ b/libs/langchain/langchain/memory/chat_message_histories/firestore.py @@ -12,7 +12,27 @@ 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 + + +def _get_firestore_client() -> Client: + try: + import firebase_admin + from firebase_admin import firestore + except ImportError: + raise ImportError( + "Could not import firebase-admin python package. " + "Please install it with `pip install firebase-admin`." + ) + + # For multiple instances, only initialize the app once. + try: + firebase_admin.get_app() + except ValueError as e: + logger.debug("Initializing Firebase app: %s", e) + firebase_admin.initialize_app() + + return firestore.client() class FirestoreChatMessageHistory(BaseChatMessageHistory): @@ -23,6 +43,7 @@ class FirestoreChatMessageHistory(BaseChatMessageHistory): collection_name: str, session_id: str, user_id: str, + firestore_client: Optional[Client] = None, ): """ Initialize a new instance of the FirestoreChatMessageHistory class. @@ -34,10 +55,9 @@ class FirestoreChatMessageHistory(BaseChatMessageHistory): 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: @@ -45,23 +65,6 @@ class FirestoreChatMessageHistory(BaseChatMessageHistory): Use this function to make sure your database is ready. """ - try: - import firebase_admin - from firebase_admin import firestore - except ImportError: - raise ImportError( - "Could not import firebase-admin python package. " - "Please install it with `pip install firebase-admin`." - ) - - # For multiple instances, only initialize the app once. - try: - firebase_admin.get_app() - except ValueError as e: - logger.debug("Initializing Firebase app: %s", e) - firebase_admin.initialize_app() - - self.firestore_client = firestore.client() self._document = self.firestore_client.collection( self.collection_name ).document(self.session_id)