From a69cb958501a64f19b043da6a35928f8c02eec65 Mon Sep 17 00:00:00 2001
From: Kshitij Wadhwa <kw@rockset.com>
Date: Mon, 14 Aug 2023 16:27:34 -0700
Subject: [PATCH] track langchain usage for Rockset (#9229)

Add ability to track langchain usage for Rockset. Rockset's new python
client allows setting this. To prevent old clients from failing, it
ignore if setting throws exception (we can't track old versions)

Tested locally with old and new Rockset python client

cc @baskaryan
---
 libs/langchain/langchain/document_loaders/rocksetdb.py      | 6 ++++++
 .../langchain/memory/chat_message_histories/rocksetdb.py    | 6 ++++++
 libs/langchain/langchain/vectorstores/rocksetdb.py          | 6 ++++++
 3 files changed, 18 insertions(+)

diff --git a/libs/langchain/langchain/document_loaders/rocksetdb.py b/libs/langchain/langchain/document_loaders/rocksetdb.py
index b39775af02..fd3095d23d 100644
--- a/libs/langchain/langchain/document_loaders/rocksetdb.py
+++ b/libs/langchain/langchain/document_loaders/rocksetdb.py
@@ -93,6 +93,12 @@ class RocksetLoader(BaseLoader):
         self.paginator = QueryPaginator
         self.request_model = QueryRequestSql
 
+        try:
+            self.client.set_application("langchain")
+        except AttributeError:
+            # ignore
+            pass
+
     def load(self) -> List[Document]:
         return list(self.lazy_load())
 
diff --git a/libs/langchain/langchain/memory/chat_message_histories/rocksetdb.py b/libs/langchain/langchain/memory/chat_message_histories/rocksetdb.py
index ce19d55693..9b3f35b06b 100644
--- a/libs/langchain/langchain/memory/chat_message_histories/rocksetdb.py
+++ b/libs/langchain/langchain/memory/chat_message_histories/rocksetdb.py
@@ -197,6 +197,12 @@ class RocksetChatMessageHistory(BaseChatMessageHistory):
         self.message_uuid_method = message_uuid_method
         self.sync = sync
 
+        try:
+            self.client.set_application("langchain")
+        except AttributeError:
+            # ignore
+            pass
+
         if not self._collection_exists():
             self._create_collection()
             self._wait_until_collection_created()
diff --git a/libs/langchain/langchain/vectorstores/rocksetdb.py b/libs/langchain/langchain/vectorstores/rocksetdb.py
index b71c216fcd..5b2fc234a4 100644
--- a/libs/langchain/langchain/vectorstores/rocksetdb.py
+++ b/libs/langchain/langchain/vectorstores/rocksetdb.py
@@ -84,6 +84,12 @@ class Rockset(VectorStore):
         self._embedding_key = embedding_key
         self._workspace = workspace
 
+        try:
+            self._client.set_application("langchain")
+        except AttributeError:
+            # ignore
+            pass
+
     @property
     def embeddings(self) -> Embeddings:
         return self._embeddings