From fe8db55b1c82dcccecdab17d21a223747533edee Mon Sep 17 00:00:00 2001
From: drf7 <83355241+drf7@users.noreply.github.com>
Date: Fri, 26 Jul 2024 18:11:23 +0200
Subject: [PATCH] fix not receiving searches and suggestions to kb_usage
 (#2349)

---
 .../src/nucliadb/search/api/v1/suggest.py     | 28 +++++++++++++++++
 nucliadb/src/nucliadb/search/search/find.py   | 30 ++++++++++++++++++-
 2 files changed, 57 insertions(+), 1 deletion(-)

diff --git a/nucliadb/src/nucliadb/search/api/v1/suggest.py b/nucliadb/src/nucliadb/search/api/v1/suggest.py
index 5f8648a0a7..c11af7fada 100644
--- 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
index 62384262b8..2f4b4a0d5c 100644
--- 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)