forked from Archives/langchain
Add connection string auth to cosmos (#3867)
Adds a connection string option for the cosmos memory, in case AAD auth is not enabled on the cosmos instance.
This commit is contained in:
parent
bc7e4d5cd4
commit
c38cafd6c2
@ -28,26 +28,34 @@ class CosmosDBChatMessageHistory(BaseChatMessageHistory):
|
||||
cosmos_endpoint: str,
|
||||
cosmos_database: str,
|
||||
cosmos_container: str,
|
||||
credential: Any,
|
||||
session_id: str,
|
||||
user_id: str,
|
||||
credential: Any = None,
|
||||
connection_string: Optional[str] = None,
|
||||
ttl: Optional[int] = None,
|
||||
):
|
||||
"""
|
||||
Initializes a new instance of the CosmosDBChatMessageHistory class.
|
||||
|
||||
Make sure to call prepare_cosmos or use the context manager to make
|
||||
sure your database is ready.
|
||||
|
||||
Either a credential or a connection string must be provided.
|
||||
|
||||
:param cosmos_endpoint: The connection endpoint for the Azure Cosmos DB account.
|
||||
:param cosmos_database: The name of the database to use.
|
||||
:param cosmos_container: The name of the container to use.
|
||||
:param credential: The credential to use to authenticate to Azure Cosmos DB.
|
||||
:param session_id: The session ID to use, can be overwritten while loading.
|
||||
:param user_id: The user ID to use, can be overwritten while loading.
|
||||
:param credential: The credential to use to authenticate to Azure Cosmos DB.
|
||||
:param connection_string: The connection string to use to authenticate.
|
||||
:param ttl: The time to live (in seconds) to use for documents in the container.
|
||||
"""
|
||||
self.cosmos_endpoint = cosmos_endpoint
|
||||
self.cosmos_database = cosmos_database
|
||||
self.cosmos_container = cosmos_container
|
||||
self.credential = credential
|
||||
self.conn_string = connection_string
|
||||
self.session_id = session_id
|
||||
self.user_id = user_id
|
||||
self.ttl = ttl
|
||||
@ -70,9 +78,16 @@ class CosmosDBChatMessageHistory(BaseChatMessageHistory):
|
||||
raise ImportError(
|
||||
"You must install the azure-cosmos package to use the CosmosDBChatMessageHistory." # noqa: E501
|
||||
) from exc
|
||||
if self.credential:
|
||||
self._client = CosmosClient(
|
||||
url=self.cosmos_endpoint, credential=self.credential
|
||||
)
|
||||
elif self.conn_string:
|
||||
self._client = CosmosClient.from_connection_string(
|
||||
conn_str=self.conn_string
|
||||
)
|
||||
else:
|
||||
raise ValueError("Either a connection string or a credential must be set.")
|
||||
database = self._client.create_database_if_not_exists(self.cosmos_database)
|
||||
self._container = database.create_container_if_not_exists(
|
||||
self.cosmos_container,
|
||||
|
Loading…
Reference in New Issue
Block a user