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
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
3 changed files with 18 additions and 0 deletions

View File

@ -93,6 +93,12 @@ class RocksetLoader(BaseLoader):
self.paginator = QueryPaginator self.paginator = QueryPaginator
self.request_model = QueryRequestSql self.request_model = QueryRequestSql
try:
self.client.set_application("langchain")
except AttributeError:
# ignore
pass
def load(self) -> List[Document]: def load(self) -> List[Document]:
return list(self.lazy_load()) return list(self.lazy_load())

View File

@ -197,6 +197,12 @@ class RocksetChatMessageHistory(BaseChatMessageHistory):
self.message_uuid_method = message_uuid_method self.message_uuid_method = message_uuid_method
self.sync = sync self.sync = sync
try:
self.client.set_application("langchain")
except AttributeError:
# ignore
pass
if not self._collection_exists(): if not self._collection_exists():
self._create_collection() self._create_collection()
self._wait_until_collection_created() self._wait_until_collection_created()

View File

@ -84,6 +84,12 @@ class Rockset(VectorStore):
self._embedding_key = embedding_key self._embedding_key = embedding_key
self._workspace = workspace self._workspace = workspace
try:
self._client.set_application("langchain")
except AttributeError:
# ignore
pass
@property @property
def embeddings(self) -> Embeddings: def embeddings(self) -> Embeddings:
return self._embeddings return self._embeddings