From fde8454872fd71ce5459d4a844463e086af5e1ed Mon Sep 17 00:00:00 2001
From: atroyn <anton.troynikov@gmail.com>
Date: Sat, 4 Nov 2023 14:10:18 -0700
Subject: [PATCH] Telemetry

---
 chromadb/api/segment.py              |  4 ++++
 chromadb/telemetry/product/events.py | 14 ++++++++++++++
 2 files changed, 18 insertions(+)

diff --git a/chromadb/api/segment.py b/chromadb/api/segment.py
index d9ee403a100..adaf94632f3 100644
--- a/chromadb/api/segment.py
+++ b/chromadb/api/segment.py
@@ -364,6 +364,7 @@ def _add(
                 add_amount=len(ids),
                 with_metadata=len(ids) if metadatas is not None else 0,
                 with_documents=len(ids) if documents is not None else 0,
+                with_uri=len(ids) if uris is not None else 0,
             )
         )
         return True
@@ -405,6 +406,7 @@ def _update(
                 with_embeddings=len(embeddings) if embeddings else 0,
                 with_metadata=len(metadatas) if metadatas else 0,
                 with_documents=len(documents) if documents else 0,
+                with_uri=len(uris) if uris else 0,
             )
         )
 
@@ -514,6 +516,7 @@ def _get(
                 limit=limit if limit else 0,
                 include_metadata=ids_amount if "metadatas" in include else 0,
                 include_documents=ids_amount if "documents" in include else 0,
+                include_uris=ids_amount if "uris" in include else 0,
             )
         )
 
@@ -706,6 +709,7 @@ def _query(
                 with_document_filter=query_amount if where_document is not None else 0,
                 include_metadatas=query_amount if "metadatas" in include else 0,
                 include_documents=query_amount if "documents" in include else 0,
+                include_uris=query_amount if "uris" in include else 0,
                 include_distances=query_amount if "distances" in include else 0,
             )
         )
diff --git a/chromadb/telemetry/product/events.py b/chromadb/telemetry/product/events.py
index e5f6bc688c1..5b5d4163964 100644
--- a/chromadb/telemetry/product/events.py
+++ b/chromadb/telemetry/product/events.py
@@ -32,6 +32,7 @@ class CollectionAddEvent(ProductTelemetryEvent):
     add_amount: int
     with_documents: int
     with_metadata: int
+    with_uri: int
 
     def __init__(
         self,
@@ -39,6 +40,7 @@ def __init__(
         add_amount: int,
         with_documents: int,
         with_metadata: int,
+        with_uri: int,
         batch_size: int = 1,
     ):
         super().__init__()
@@ -46,6 +48,7 @@ def __init__(
         self.add_amount = add_amount
         self.with_documents = with_documents
         self.with_metadata = with_metadata
+        self.with_uri = with_uri
         self.batch_size = batch_size
 
     @property
@@ -62,6 +65,7 @@ def batch(self, other: "ProductTelemetryEvent") -> "CollectionAddEvent":
             add_amount=total_amount,
             with_documents=self.with_documents + other.with_documents,
             with_metadata=self.with_metadata + other.with_metadata,
+            with_uri=self.with_uri + other.with_uri,
             batch_size=self.batch_size + other.batch_size,
         )
 
@@ -72,6 +76,7 @@ class CollectionUpdateEvent(ProductTelemetryEvent):
     with_embeddings: int
     with_metadata: int
     with_documents: int
+    with_uri: int
 
     def __init__(
         self,
@@ -80,6 +85,7 @@ def __init__(
         with_embeddings: int,
         with_metadata: int,
         with_documents: int,
+        with_uri: int,
     ):
         super().__init__()
         self.collection_uuid = collection_uuid
@@ -87,6 +93,7 @@ def __init__(
         self.with_embeddings = with_embeddings
         self.with_metadata = with_metadata
         self.with_documents = with_documents
+        self.with_uri = with_uri
 
 
 class CollectionQueryEvent(ProductTelemetryEvent):
@@ -99,6 +106,7 @@ class CollectionQueryEvent(ProductTelemetryEvent):
     n_results: int
     include_metadatas: int
     include_documents: int
+    include_uris: int
     include_distances: int
 
     def __init__(
@@ -110,6 +118,7 @@ def __init__(
         n_results: int,
         include_metadatas: int,
         include_documents: int,
+        include_uris: int,
         include_distances: int,
         batch_size: int = 1,
     ):
@@ -121,6 +130,7 @@ def __init__(
         self.n_results = n_results
         self.include_metadatas = include_metadatas
         self.include_documents = include_documents
+        self.include_uris = include_uris
         self.include_distances = include_distances
         self.batch_size = batch_size
 
@@ -141,6 +151,7 @@ def batch(self, other: "ProductTelemetryEvent") -> "CollectionQueryEvent":
             n_results=self.n_results + other.n_results,
             include_metadatas=self.include_metadatas + other.include_metadatas,
             include_documents=self.include_documents + other.include_documents,
+            include_uris=self.include_uris + other.include_uris,
             include_distances=self.include_distances + other.include_distances,
             batch_size=self.batch_size + other.batch_size,
         )
@@ -152,6 +163,7 @@ class CollectionGetEvent(ProductTelemetryEvent):
     limit: int
     include_metadata: int
     include_documents: int
+    include_uris: int
 
     def __init__(
         self,
@@ -160,6 +172,7 @@ def __init__(
         limit: int,
         include_metadata: int,
         include_documents: int,
+        include_uris: int,
     ):
         super().__init__()
         self.collection_uuid = collection_uuid
@@ -167,6 +180,7 @@ def __init__(
         self.limit = limit
         self.include_metadata = include_metadata
         self.include_documents = include_documents
+        self.include_uris = include_uris
 
 
 class CollectionDeleteEvent(ProductTelemetryEvent):