nucliadb-utils 5.0.1.post994__py3-none-any.whl → 5.0.1.post1009__py3-none-any.whl

nucliadb_utils/audit/audit.py

@@ -93,3 +93,24 @@ class AuditStorage:
         generative_answer_first_chunk_time: Optional[float] = None,
     ):
         raise NotImplementedError
+
+    def report_fields_and_paragraphs(self, kbid: str, paragraphs: int, fields: int):
+        raise NotImplementedError
+
+    def report_resources(
+        self,
+        *,
+        kbid: str,
+        resources: int,
+    ):
+        raise NotImplementedError
+
+    def delete_kb(self, kbid: str):
+        raise NotImplementedError
+
+    def suggest(
+        self,
+        kbid: str,
+        client_type: int,
+    ):
+        raise NotImplementedError

nucliadb_utils/audit/basic.py

@@ -94,3 +94,24 @@ class BasicAuditStorage(AuditStorage):
         generative_answer_first_chunk_time: Optional[float] = None,
     ):
         logger.debug(f"CHAT {kbid} {user} {origin}")
+
+    def report_fields_and_paragraphs(self, kbid: str, paragraphs: int, fields: int):
+        logger.debug(f"FIELDS & PARAGRAPHS {kbid} {paragraphs} {fields}")
+
+    def report_resources(
+        self,
+        *,
+        kbid: str,
+        resources: int,
+    ):
+        logger.debug(f"RESOURCES {kbid} {resources}")
+
+    def delete_kb(self, kbid: str):
+        logger.debug(f"DELETE_KB {kbid}")
+
+    def suggest(
+        self,
+        kbid: str,
+        client_type: int,
+    ):
+        logger.debug(f"SUGGEST {kbid} {client_type}")

nucliadb_utils/audit/stream.py

@@ -37,12 +37,24 @@ from nucliadb_protos.audit_pb2 import (
     AuditField,
     AuditRequest,
     ChatContext,
+    ClientType,
 )
 from nucliadb_protos.nodereader_pb2 import SearchRequest
 from nucliadb_protos.resources_pb2 import FieldID
 from nucliadb_utils import logger
 from nucliadb_utils.audit.audit import AuditStorage
 from nucliadb_utils.nats import get_traced_jetstream
+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,
+    Storage,
+)
+from nucliadb_utils.nuclia_usage.utils.kb_usage_report import KbUsageReportUtility
 
 
 class RequestContext:
@@ -123,6 +135,9 @@ class AuditMiddleware(BaseHTTPMiddleware):
         return response
 
 
+KB_USAGE_STREAM_SUBJECT = "kb-usage.nuclia_db"
+
+
 class StreamAuditStorage(AuditStorage):
     task: Optional[asyncio.Task]
     initialized: bool
@@ -181,9 +196,16 @@ class StreamAuditStorage(AuditStorage):
         self.js = get_traced_jetstream(self.nc, self.service)
         self.task = asyncio.create_task(self.run())
 
+        self.kb_usage_utility = KbUsageReportUtility(
+            nats_stream=self.js, nats_subject=KB_USAGE_STREAM_SUBJECT
+        )
+        await self.kb_usage_utility.initialize()
+
         self.initialized = True
 
     async def finalize(self):
+        await self.kb_usage_utility.finalize()
+
         if self.task is not None:
             self.task.cancel()
         if self.nc:
@@ -258,6 +280,29 @@ class StreamAuditStorage(AuditStorage):
 
         self.send(auditrequest)
 
+    def report_fields_and_paragraphs(self, kbid: str, paragraphs: int, fields: int):
+        self.kb_usage_utility.send_kb_usage(
+            service=Service.NUCLIA_DB,
+            account_id=None,
+            kb_id=kbid,
+            kb_source=KBSource.HOSTED,
+            storage=Storage(paragraphs=paragraphs, fields=fields),
+        )
+
+    def report_resources(
+        self,
+        *,
+        kbid: str,
+        resources: int,
+    ):
+        self.kb_usage_utility.send_kb_usage(
+            service=Service.NUCLIA_DB,
+            account_id=None,
+            kb_id=kbid,
+            kb_source=KBSource.HOSTED,
+            storage=Storage(resources=resources),
+        )
+
     def visited(
         self,
         kbid: str,
@@ -277,6 +322,15 @@ class StreamAuditStorage(AuditStorage):
         auditrequest.kbid = kbid
         auditrequest.type = AuditRequest.VISITED
 
+    def delete_kb(self, kbid: str):
+        self.kb_usage_utility.send_kb_usage(
+            service=Service.NUCLIA_DB,
+            account_id=None,
+            kb_id=kbid,
+            kb_source=KBSource.HOSTED,
+            storage=Storage(paragraphs=0, fields=0, resources=0),
+        )
+
     def search(
         self,
         kbid: str,
@@ -302,6 +356,43 @@ class StreamAuditStorage(AuditStorage):
         auditrequest.resources = resources
         auditrequest.type = AuditRequest.SEARCH
 
+        self.kb_usage_utility.send_kb_usage(
+            service=Service.NUCLIA_DB,
+            account_id=None,
+            kb_id=kbid,
+            kb_source=KBSource.HOSTED,
+            # TODO unify AuditRequest client type and Nuclia Usage client type
+            searches=[
+                Search(
+                    client=ClientTypeKbUsage.Value(ClientType.Name(client_type)),  # type: ignore
+                    type=SearchType.SEARCH,
+                    tokens=2000,
+                    num_searches=1,
+                )
+            ],
+        )
+
+    def suggest(
+        self,
+        kbid: str,
+        client_type: int,
+    ):
+        self.kb_usage_utility.send_kb_usage(
+            service=Service.NUCLIA_DB,
+            account_id=None,
+            kb_id=kbid,
+            kb_source=KBSource.HOSTED,
+            # TODO unify AuditRequest client type and Nuclia Usage client type
+            searches=[
+                Search(
+                    client=ClientTypeKbUsage.Value(ClientType.Name(client_type)),  # type: ignore
+                    type=SearchType.SUGGEST,
+                    tokens=0,
+                    num_searches=1,
+                )
+            ],
+        )
+
     def chat(
         self,
         kbid: str,

nucliadb_utils/nuclia_usage/utils/kb_usage_report.py

@@ -22,9 +22,10 @@ import logging
 from collections.abc import Iterable
 from contextlib import suppress
 from datetime import datetime, timezone
-from typing import List, Optional
+from typing import Optional
+
+from nats.js.client import JetStreamContext
 
-from nucliadb_utils.nats import NatsConnectionManager
 from nucliadb_utils.nuclia_usage.protos.kb_usage_pb2 import (
     KBSource,
     KbUsage,
@@ -39,48 +40,28 @@ logger = logging.getLogger(__name__)
 
 
 class KbUsageReportUtility:
-    task: Optional[asyncio.Task]
-    initialized: bool
     queue: asyncio.Queue
-    service: str
 
     def __init__(
         self,
+        nats_stream: JetStreamContext,
         nats_subject: str,
-        nats_servers: List[str],
-        nats_creds: Optional[str] = None,
         max_queue_size: int = 100,
-        service: str = "",
     ):
-        self.nats_connection_manager = NatsConnectionManager(
-            service_name=service,
-            nats_servers=nats_servers,
-            nats_creds=nats_creds,
-        )
+        self.nats_stream = nats_stream
         self.nats_subject = nats_subject
         self.queue = asyncio.Queue(max_queue_size)
         self.task = None
-        self.initialized = False
 
     async def initialize(self):
-        if not self.initialized and self.nats_connection_manager._nats_servers:
-            await self.nats_connection_manager.initialize()
-
-            if self.task is None:
-                self.task = asyncio.create_task(self.run())
-
-            self.initialized = True
+        if self.task is None:
+            self.task = asyncio.create_task(self.run())
 
     async def finalize(self):
-        if self.initialized:
-            if self.task is not None:
-                self.task.cancel()
-                with suppress(asyncio.CancelledError, asyncio.exceptions.TimeoutError):
-                    await asyncio.wait_for(self.task, timeout=2)
-
-            await self.nats_connection_manager.finalize()
-
-            self.initialized = False
+        if self.task is not None:
+            self.task.cancel()
+            with suppress(asyncio.CancelledError, asyncio.exceptions.TimeoutError):
+                await asyncio.wait_for(self.task, timeout=2)
 
     async def run(self) -> None:
         while True:
@@ -93,15 +74,13 @@ class KbUsageReportUtility:
                 self.queue.task_done()
 
     def send(self, message: KbUsage):
-        if not self.initialized:
-            return
         try:
             self.queue.put_nowait(message)
         except asyncio.QueueFull:
             logger.warning("KbUsage utility queue is full, dropping message")
 
     async def _send(self, message: KbUsage) -> int:
-        res = await self.nats_connection_manager.js.publish(
+        res = await self.nats_stream.publish(
             self.nats_subject,
             message.SerializeToString(),
         )

nucliadb_utils/utilities.py

@@ -40,7 +40,6 @@ from nucliadb_utils.encryption.settings import settings as encryption_settings
 from nucliadb_utils.exceptions import ConfigurationError
 from nucliadb_utils.indexing import IndexingUtility
 from nucliadb_utils.nats import NatsConnectionManager
-from nucliadb_utils.nuclia_usage.utils.kb_usage_report import KbUsageReportUtility
 from nucliadb_utils.partition import PartitionUtility
 from nucliadb_utils.settings import (
     FileBackendConfig,
@@ -49,7 +48,6 @@ from nucliadb_utils.settings import (
     nuclia_settings,
     storage_settings,
     transaction_settings,
-    usage_settings,
 )
 from nucliadb_utils.storages.settings import settings as extended_storage_settings
 from nucliadb_utils.store import MAIN
@@ -325,33 +323,6 @@ def get_audit() -> Optional[AuditStorage]:
     return get_utility(Utility.AUDIT)
 
 
-def get_usage_utility() -> Optional[KbUsageReportUtility]:
-    return get_utility(Utility.USAGE)
-
-
-async def start_usage_utility(service: str):
-    usage_utility: Optional[KbUsageReportUtility] = get_utility(Utility.USAGE)
-    if usage_utility is not None:
-        return
-
-    usage_utility = KbUsageReportUtility(
-        nats_subject=cast(str, usage_settings.usage_jetstream_subject),
-        nats_servers=usage_settings.usage_jetstream_servers,
-        nats_creds=usage_settings.usage_jetstream_auth,
-        service=service,
-    )
-    logger.info(f"Configuring usage report utility {usage_settings.usage_jetstream_subject}")
-    await usage_utility.initialize()
-    set_utility(Utility.USAGE, usage_utility)
-
-
-async def stop_usage_utility():
-    usage_utility = get_usage_utility()
-    if usage_utility:
-        await usage_utility.finalize()
-        clean_utility(Utility.USAGE)
-
-
 def register_audit_utility(service: str) -> AuditStorage:
     if audit_settings.audit_driver == "basic":
         b_audit_utility: AuditStorage = BasicAuditStorage()

{nucliadb_utils-5.0.1.post994.dist-info → nucliadb_utils-5.0.1.post1009.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: nucliadb_utils
-Version: 5.0.1.post994
+Version: 5.0.1.post1009
 Home-page: https://nuclia.com
 License: BSD
 Classifier: Development Status :: 4 - Beta
@@ -23,8 +23,8 @@ Requires-Dist: PyNaCl
 Requires-Dist: pyjwt >=2.4.0
 Requires-Dist: memorylru >=1.1.2
 Requires-Dist: mrflagly
-Requires-Dist: nucliadb-protos >=5.0.1.post994
-Requires-Dist: nucliadb-telemetry >=5.0.1.post994
+Requires-Dist: nucliadb-protos >=5.0.1.post1009
+Requires-Dist: nucliadb-telemetry >=5.0.1.post1009
 Provides-Extra: cache
 Requires-Dist: redis >=4.3.4 ; extra == 'cache'
 Requires-Dist: orjson >=3.6.7 ; extra == 'cache'

{nucliadb_utils-5.0.1.post994.dist-info → nucliadb_utils-5.0.1.post1009.dist-info}/RECORD

@@ -16,15 +16,15 @@ nucliadb_utils/settings.py,sha256=QR51SX0T17-_YofMNpci-nkI77l_CUWFy4H7i8hNOHU,79
 nucliadb_utils/signals.py,sha256=JRNv2y9zLtBjOANBf7krGfDGfOc9qcoXZ6N1nKWS2FE,2674
 nucliadb_utils/store.py,sha256=kQ35HemE0v4_Qg6xVqNIJi8vSFAYQtwI3rDtMsNy62Y,890
 nucliadb_utils/transaction.py,sha256=mwcI3aIHAvU5KOGqd_Uz_d1XQzXhk_-NWY8NqU1lfb0,7307
-nucliadb_utils/utilities.py,sha256=trAez4j2TYYT5n1qpimyHOEAnsQqDqD2qCFgVSDi184,16747
+nucliadb_utils/utilities.py,sha256=idajCm_4Sojh7b3HTkP0fTfG2Mb6PIB9xtMmcfB7Nl0,15758
 nucliadb_utils/aiopynecone/__init__.py,sha256=cp15ZcFnHvpcu_5-aK2A4uUyvuZVV_MJn4bIXMa20ks,835
 nucliadb_utils/aiopynecone/client.py,sha256=pl3WWGpjwoc-qzgDChuIE2LU7Mopu1aKL9iPHkr8QD0,20069
 nucliadb_utils/aiopynecone/exceptions.py,sha256=EEE0XoGs1zIB5yOJ_fy6yoG4uIb4cWIawYdJeNe4eDo,3012
 nucliadb_utils/aiopynecone/models.py,sha256=tUugQ-ACQAyT-lhsfLwKAOgb6ilLchLwZfnRV7xwFck,3047
 nucliadb_utils/audit/__init__.py,sha256=cp15ZcFnHvpcu_5-aK2A4uUyvuZVV_MJn4bIXMa20ks,835
-nucliadb_utils/audit/audit.py,sha256=kpjiTLgWtLXA5DNgJcmeIaodel11Mz0Y-JfQt-yfhis,2641
-nucliadb_utils/audit/basic.py,sha256=IXrsdaXaXHVRLkBWPmq4HIO4SxWnri5VFns_h3_8v7k,3184
-nucliadb_utils/audit/stream.py,sha256=_RZf1G9_c5eaxfutq1OdMmnRaRM7EiaXH33Kb4mN9Gw,11525
+nucliadb_utils/audit/audit.py,sha256=ictuErJy3X36S7s_f_sdNQQrQ4ackAYCT6M2-4nyenE,3086
+nucliadb_utils/audit/basic.py,sha256=-Yztp0I745J6dz_zHPKdXv60OH1m9_d69NlpbhsGFZI,3717
+nucliadb_utils/audit/stream.py,sha256=yPWe9iCJmTKLtB1iktSOKR8HDB0I3KF7vsWOwloljGM,14299
 nucliadb_utils/cache/__init__.py,sha256=itSI7dtTwFP55YMX4iK7JzdMHS5CQVUiB1XzQu4UBh8,833
 nucliadb_utils/cache/exceptions.py,sha256=Zu-O_-0-yctOEgoDGI92gPzWfBMRrpiAyESA62ld6MA,975
 nucliadb_utils/cache/nats.py,sha256=-AjCfkFgKVdJUlGR0hT9JDSNkPVFg4S6w9eW-ZIcXPM,7037
@@ -43,7 +43,7 @@ nucliadb_utils/nuclia_usage/protos/kb_usage_pb2.pyi,sha256=xhyc3jJBh0KZuWcgmIbwS
 nucliadb_utils/nuclia_usage/protos/kb_usage_pb2_grpc.py,sha256=dhop8WwjplPfORYPYb9HtcS9gHMzqXPJQGqXYRjV-6M,1008
 nucliadb_utils/nuclia_usage/protos/kb_usage_pb2_grpc.pyi,sha256=6RIsZ2934iodEckflpBStgLKEkFhKfNmZ72UKg2Bwb4,911
 nucliadb_utils/nuclia_usage/utils/__init__.py,sha256=cp15ZcFnHvpcu_5-aK2A4uUyvuZVV_MJn4bIXMa20ks,835
-nucliadb_utils/nuclia_usage/utils/kb_usage_report.py,sha256=t4-QXUTOTeMBWFFO8hdqiIco1fgC2Jla4NIWuRSs6N4,4566
+nucliadb_utils/nuclia_usage/utils/kb_usage_report.py,sha256=P1fKvrcYdSPNJ1ycd0ArSXOqhUcvz1BaQhxAWHo0kcc,3847
 nucliadb_utils/storages/__init__.py,sha256=5Qc8AUWiJv9_JbGCBpAn88AIJhwDlm0OPQpg2ZdRL4U,872
 nucliadb_utils/storages/azure.py,sha256=egMDwLNIGSQyVevuySt2AswzFdNAcih05BbRg3-p8IU,16015
 nucliadb_utils/storages/exceptions.py,sha256=mm_wX4YRtp7u7enkk_4pMSlX5AQQuFbq4xLmupVDt3Y,2502
@@ -64,8 +64,8 @@ nucliadb_utils/tests/indexing.py,sha256=YW2QhkhO9Q_8A4kKWJaWSvXvyQ_AiAwY1VylcfVQ
 nucliadb_utils/tests/local.py,sha256=7nuP8EFUAiA8ZH50R1iPV9EUXBySQxOanVm3Zht_e0g,1835
 nucliadb_utils/tests/nats.py,sha256=xqpww4jZjTKY9oPGlJdDJG67L3FIBQsa9qDHxILR8r8,7687
 nucliadb_utils/tests/s3.py,sha256=IdMxK_cNdSHLvO1u8BwsKFzD87Hk1MVPDZ57zx6h-rA,3656
-nucliadb_utils-5.0.1.post994.dist-info/METADATA,sha256=N4xswJv9a86AvuhIlJfSmCkc4Z2nHXX4dvvgAFZ4Vm0,2073
-nucliadb_utils-5.0.1.post994.dist-info/WHEEL,sha256=Rp8gFpivVLXx-k3U95ozHnQw8yDcPxmhOpn_Gx8d5nc,91
-nucliadb_utils-5.0.1.post994.dist-info/top_level.txt,sha256=fE3vJtALTfgh7bcAWcNhcfXkNPp_eVVpbKK-2IYua3E,15
-nucliadb_utils-5.0.1.post994.dist-info/zip-safe,sha256=AbpHGcgLb-kRsJGnwFEktk7uzpZOCcBY74-YBdrKVGs,1
-nucliadb_utils-5.0.1.post994.dist-info/RECORD,,
+nucliadb_utils-5.0.1.post1009.dist-info/METADATA,sha256=0FahfbwL2O5bWC097JPQ8UqGCD-Z42pyX-ieCiPT0Cw,2076
+nucliadb_utils-5.0.1.post1009.dist-info/WHEEL,sha256=R0nc6qTxuoLk7ShA2_Y-UWkN8ZdfDBG2B6Eqpz2WXbs,91
+nucliadb_utils-5.0.1.post1009.dist-info/top_level.txt,sha256=fE3vJtALTfgh7bcAWcNhcfXkNPp_eVVpbKK-2IYua3E,15
+nucliadb_utils-5.0.1.post1009.dist-info/zip-safe,sha256=AbpHGcgLb-kRsJGnwFEktk7uzpZOCcBY74-YBdrKVGs,1
+nucliadb_utils-5.0.1.post1009.dist-info/RECORD,,

{nucliadb_utils-5.0.1.post994.dist-info → nucliadb_utils-5.0.1.post1009.dist-info}/WHEEL

@@ -1,5 +1,5 @@
 Wheel-Version: 1.0
-Generator: setuptools (72.0.0)
+Generator: setuptools (72.1.0)
 Root-Is-Purelib: true
 Tag: py3-none-any