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
This commit is contained in:
Kshitij Wadhwa
2023-08-14 16:27:34 -07:00
committed by GitHub
parent 7810ea5812
commit a69cb95850
3 changed files with 18 additions and 0 deletions

View File

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

View File

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

View File

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