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 <baskaryan@gmail.com>
pull/8604/head
Neil Murphy 1 year ago committed by GitHub
parent 13ccf202de
commit 31820a31e4
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23

@ -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)

Loading…
Cancel
Save