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