Merge branch 'main' into remove-noisy-log

nuclia · Jul 26, 2024 · 9188867 · 9188867
2 parents 241f6bb + fe8db55
commit 9188867
Show file tree

Hide file tree

Showing 2 changed files with 57 additions and 1 deletion.
diff --git a/nucliadb/src/nucliadb/search/api/v1/suggest.py b/nucliadb/src/nucliadb/search/api/v1/suggest.py
@@ -40,6 +40,16 @@
     SuggestOptions,
 )
 from nucliadb_utils.authentication import requires
+from nucliadb_utils.nuclia_usage.protos.kb_usage_pb2 import (
+    ClientType as ClientTypeKbUsage,
+)
+from nucliadb_utils.nuclia_usage.protos.kb_usage_pb2 import (
+    KBSource,
+    Search,
+    SearchType,
+    Service,
+)
+from nucliadb_utils.utilities import get_usage_utility
 
 
 @api.get(
@@ -125,6 +135,7 @@ async def suggest(
     debug: bool,
     highlight: bool,
 ) -> KnowledgeboxSuggestResults:
+    usage = get_usage_utility()
     # We need the nodes/shards that are connected to the KB
     # We need to query all nodes
     pb_query = suggest_query_to_pb(
@@ -155,4 +166,21 @@ async def suggest(
     if debug and queried_shards:
         search_results.shards = queried_shards
 
+    if usage is not None:
+        usage.send_kb_usage(
+            service=Service.NUCLIA_DB,  # type: ignore
+            account_id=None,
+            kb_id=kbid,
+            kb_source=KBSource.HOSTED,  # type: ignore
+            # TODO unify AuditRequest client type and Nuclia Usage client type
+            searches=[
+                Search(
+                    client=ClientTypeKbUsage.Value(x_ndb_client.name),  # type: ignore
+                    type=SearchType.SUGGEST,
+                    tokens=0,
+                    num_searches=1,
+                )
+            ],
+        )
+
     return search_results
diff --git a/nucliadb/src/nucliadb/search/search/find.py b/nucliadb/src/nucliadb/search/search/find.py
@@ -46,7 +46,16 @@
     NucliaDBClientType,
     SearchOptions,
 )
-from nucliadb_utils.utilities import get_audit
+from nucliadb_utils.nuclia_usage.protos.kb_usage_pb2 import (
+    ClientType as ClientTypeKbUsage,
+)
+from nucliadb_utils.nuclia_usage.protos.kb_usage_pb2 import (
+    KBSource,
+    Search,
+    SearchType,
+    Service,
+)
+from nucliadb_utils.utilities import get_audit, get_usage_utility
 
 logger = logging.getLogger(__name__)
 
@@ -90,6 +99,7 @@ async def _index_node_retrieval(
     metrics: RAGMetrics = RAGMetrics(),
 ) -> tuple[KnowledgeboxFindResults, bool, QueryParser]:
     audit = get_audit()
+    usage = get_usage_utility()
     start_time = time()
 
     item.min_score = min_score_from_payload(item.min_score)
@@ -159,6 +169,24 @@ async def _index_node_retrieval(
             search_time,
             len(search_results.resources),
         )
+
+    if usage is not None:
+        usage.send_kb_usage(
+            service=Service.NUCLIA_DB,  # type: ignore
+            account_id=None,
+            kb_id=kbid,
+            kb_source=KBSource.HOSTED,  # type: ignore
+            # TODO unify AuditRequest client type and Nuclia Usage client type
+            searches=[
+                Search(
+                    client=ClientTypeKbUsage.Value(x_ndb_client.name),  # type: ignore
+                    type=SearchType.SEARCH,
+                    tokens=2000,
+                    num_searches=1,
+                )
+            ],
+        )
+
     if item.debug:
         search_results.nodes = debug_nodes_info(queried_nodes)