nucliadb-utils 5.0.0.post877__py3-none-any.whl → 5.0.0.post899__py3-none-any.whl

nucliadb_utils/audit/audit.py

@@ -23,7 +23,6 @@ from google.protobuf.timestamp_pb2 import Timestamp
 
 from nucliadb_protos.audit_pb2 import (
     AuditField,
-    AuditKBCounter,
     AuditRequest,
     ChatContext,
 )
@@ -32,7 +31,9 @@ from nucliadb_protos.resources_pb2 import FieldID
 
 
 class AuditStorage:
-    async def report(
+    initialized: bool = False
+
+    def report_and_send(
         self,
         *,
         kbid: str,
@@ -43,15 +44,6 @@ class AuditStorage:
         rid: Optional[str] = None,
         field_metadata: Optional[List[FieldID]] = None,
         audit_fields: Optional[List[AuditField]] = None,
-        kb_counter: Optional[AuditKBCounter] = None,
-    ):
-        raise NotImplementedError
-
-    def report_resources(
-        self,
-        *,
-        kbid: str,
-        resources: int,
     ):
         raise NotImplementedError
 
@@ -61,45 +53,43 @@ class AuditStorage:
     async def finalize(self):
         pass
 
-    async def visited(self, kbid: str, uuid: str, user: str, origin: str):
-        raise NotImplementedError
-
-    async def search(
+    def visited(
         self,
         kbid: str,
+        uuid: str,
         user: str,
-        client: int,
         origin: str,
-        search: SearchRequest,
-        timeit: float,
-        resources: int,
     ):
         raise NotImplementedError
 
-    async def suggest(
+    def send(self, msg: AuditRequest):
+        raise NotImplementedError
+
+    def search(
         self,
         kbid: str,
         user: str,
         client: int,
         origin: str,
+        search: SearchRequest,
         timeit: float,
+        resources: int,
     ):
         raise NotImplementedError
 
-    async def chat(
+    def chat(
         self,
         kbid: str,
         user: str,
         client: int,
         origin: str,
-        timeit: float,
         question: str,
         rephrased_question: Optional[str],
         context: List[ChatContext],
         answer: Optional[str],
         learning_id: str,
+        rephrase_time: Optional[float] = None,
+        generative_answer_time: Optional[float] = None,
+        generative_answer_first_chunk_time: Optional[float] = None,
     ):
         raise NotImplementedError
-
-    async def delete_kb(self, kbid):
-        raise NotImplementedError

nucliadb_utils/audit/basic.py

@@ -23,7 +23,6 @@ from google.protobuf.timestamp_pb2 import Timestamp
 
 from nucliadb_protos.audit_pb2 import (
     AuditField,
-    AuditKBCounter,
     AuditRequest,
     ChatContext,
 )
@@ -35,10 +34,13 @@ from nucliadb_utils.audit.audit import AuditStorage
 
 
 class BasicAuditStorage(AuditStorage):
+    def __init__(self):
+        self.initialized = True
+
     def message_to_str(self, message: BrokerMessage) -> str:
         return f"{message.type}+{message.multiid}+{message.audit.user}+{message.kbid}+{message.uuid}+{message.audit.when.ToJsonString()}+{message.audit.origin}+{message.audit.source}"  # noqa
 
-    async def report(
+    async def report_and_send(
         self,
         *,
         kbid: str,
@@ -49,22 +51,22 @@ class BasicAuditStorage(AuditStorage):
         rid: Optional[str] = None,
         field_metadata: Optional[List[FieldID]] = None,
         audit_fields: Optional[List[AuditField]] = None,
-        kb_counter: Optional[AuditKBCounter] = None,
     ):
-        logger.debug(f"AUDIT {audit_type} {kbid} {user} {origin} {rid} {audit_fields} {kb_counter}")
+        logger.debug(f"AUDIT {audit_type} {kbid} {user} {origin} {rid} {audit_fields}")
 
-    def report_resources(
+    async def visited(
         self,
-        *,
         kbid: str,
-        resources: int,
+        uuid: str,
+        user: str,
+        origin: str,
     ):
-        logger.debug(f"REPORT RESOURCES {kbid} {resources}")
-
-    async def visited(self, kbid: str, uuid: str, user: str, origin: str):
         logger.debug(f"VISITED {kbid} {uuid} {user} {origin}")
 
-    async def search(
+    def send(self, msg: AuditRequest):
+        logger.debug(f"sending a {msg.type} queued message")
+
+    def search(
         self,
         kbid: str,
         user: str,
@@ -76,30 +78,19 @@ class BasicAuditStorage(AuditStorage):
     ):
         logger.debug(f"SEARCH {kbid} {user} {origin} ''{search}'' {timeit} {resources}")
 
-    async def suggest(
-        self,
-        kbid: str,
-        user: str,
-        client: int,
-        origin: str,
-        timeit: float,
-    ):
-        logger.debug(f"SUGGEST {kbid} {user} {origin} {timeit}")
-
-    async def chat(
+    def chat(
         self,
         kbid: str,
         user: str,
         client_type: int,
         origin: str,
-        timeit: float,
         question: str,
         rephrased_question: Optional[str],
         context: List[ChatContext],
         answer: Optional[str],
         learning_id: str,
+        rephrase_time: Optional[float] = None,
+        generative_answer_time: Optional[float] = None,
+        generative_answer_first_chunk_time: Optional[float] = None,
     ):
-        logger.debug(f"CHAT {kbid} {user} {origin} {timeit}")
-
-    async def delete_kb(self, kbid):
-        logger.debug(f"KB DELETED {kbid}")
+        logger.debug(f"CHAT {kbid} {user} {origin}")

nucliadb_utils/audit/stream.py

@@ -18,47 +18,115 @@
 # along with this program. If not, see <http://www.gnu.org/licenses/>.
 #
 import asyncio
-from datetime import datetime
-from typing import List, Optional
+import contextvars
+import time
+from datetime import datetime, timezone
+from typing import Callable, List, Optional
 
 import backoff
 import mmh3
 import nats
+from fastapi import Request
 from google.protobuf.timestamp_pb2 import Timestamp
 from opentelemetry.trace import format_trace_id, get_current_span
+from starlette.middleware.base import BaseHTTPMiddleware, RequestResponseEndpoint
+from starlette.responses import Response, StreamingResponse
+from starlette.types import ASGIApp
 
 from nucliadb_protos.audit_pb2 import (
     AuditField,
-    AuditKBCounter,
     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
 
-KB_USAGE_STREAM_AUDIT = "kb-usage.nuclia_db"
+
+class RequestContext:
+    def __init__(self):
+        self.audit_request: AuditRequest = AuditRequest()
+        self.start_time: float = time.monotonic()
+
+
+request_context_var = contextvars.ContextVar[Optional[RequestContext]]("request_context", default=None)
+
+
+def get_trace_id() -> str:
+    span = get_current_span()
+    if span is None:
+        return ""
+    return format_trace_id(span.get_span_context().trace_id)
+
+
+def get_request_context() -> Optional[RequestContext]:
+    return request_context_var.get()
+
+
+class AuditMiddleware(BaseHTTPMiddleware):
+    def __init__(self, app: ASGIApp, audit_utility_getter: Callable[[], Optional[AuditStorage]]) -> None:
+        self.audit_utility_getter = audit_utility_getter
+        super().__init__(app)
+
+    @property
+    def audit_utility(self):
+        return self.audit_utility_getter()
+
+    async def dispatch(self, request: Request, call_next: RequestResponseEndpoint) -> Response:
+        context = RequestContext()
+        token = request_context_var.set(context)
+        context.audit_request.time.FromDatetime(datetime.now(tz=timezone.utc))
+        context.audit_request.trace_id = get_trace_id()
+        response = await call_next(request)
+
+        if isinstance(response, StreamingResponse):
+            response = self.wrap_streaming_response(response, context)
+        else:
+            self.enqueue_pending(context)
+
+        request_context_var.reset(token)
+
+        return response
+
+    def enqueue_pending(self, context: RequestContext):
+        if context.audit_request.kbid:
+            # an audit request with no kbid makes no sense, we use this as an heuristic
+            # mark that no audit has been set during this request
+
+            context.audit_request.request_time = time.monotonic() - context.start_time
+            if self.audit_utility is not None:
+                self.audit_utility.send(context.audit_request)
+
+    def wrap_streaming_response(
+        self, response: StreamingResponse, context: RequestContext
+    ) -> StreamingResponse:
+        """
+        When dealing with streaming responses, AND if we depend on any state that only will be available once
+        the request is fully finished, the response we have after the dispatch call_next is not enough, as
+        there, no iteration of the streaming response has been done yet.
+
+        This is why we need to rewrap to be able to to the auditing at the _real_ request end without losing
+        any audit bits.
+        """
+        original_body_iterator = response.body_iterator
+
+        async def custom_body_iterator():
+            try:
+                async for chunk in original_body_iterator:
+                    yield chunk
+            finally:
+                self.enqueue_pending(context)
+
+        response.body_iterator = custom_body_iterator()
+        return response
 
 
 class StreamAuditStorage(AuditStorage):
-    task: Optional[asyncio.Task] = None
-    initialized: bool = False
+    task: Optional[asyncio.Task]
+    initialized: bool
     queue: asyncio.Queue
-    lock: asyncio.Lock
 
     def __init__(
         self,
@@ -74,9 +142,10 @@ class StreamAuditStorage(AuditStorage):
         self.nats_target = nats_target
         self.partitions = partitions
         self.seed = seed
-        self.lock = asyncio.Lock()
         self.queue = asyncio.Queue()
         self.service = service
+        self.task = None
+        self.initialized = False
 
     def get_partition(self, kbid: str):
         return mmh3.hash(kbid, self.seed, signed=False) % self.partitions
@@ -112,16 +181,9 @@ 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_AUDIT
-        )
-        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:
@@ -144,7 +206,7 @@ class StreamAuditStorage(AuditStorage):
                 if item_dequeued:
                     self.queue.task_done()
 
-    async def send(self, message: AuditRequest):
+    def send(self, message: AuditRequest):
         self.queue.put_nowait(message)
 
     @backoff.on_exception(backoff.expo, (Exception,), jitter=backoff.random_jitter, max_tries=4)
@@ -163,7 +225,7 @@ class StreamAuditStorage(AuditStorage):
         )
         return res.seq
 
-    async def report(
+    def report_and_send(
         self,
         *,
         kbid: str,
@@ -174,18 +236,19 @@ class StreamAuditStorage(AuditStorage):
         rid: Optional[str] = None,
         field_metadata: Optional[List[FieldID]] = None,
         audit_fields: Optional[List[AuditField]] = None,
-        kb_counter: Optional[AuditKBCounter] = None,
     ):
-        # Reports MODIFIED / DELETED / NEW events
         auditrequest = AuditRequest()
+
+        # Reports MODIFIED / DELETED / NEW events
+
+        auditrequest.trace_id = get_trace_id()
         auditrequest.kbid = kbid
         auditrequest.userid = user or ""
         auditrequest.rid = rid or ""
         auditrequest.origin = origin or ""
         auditrequest.type = audit_type
-        if when is None or when.SerializeToString() == b"":
-            auditrequest.time.FromDatetime(datetime.now())
-        else:
+        # If defined, when needs to overwrite any previously set time
+        if not (when is None or when.SerializeToString() == b""):
             auditrequest.time.CopyFrom(when)
 
         auditrequest.field_metadata.extend(field_metadata or [])
@@ -193,66 +256,28 @@ class StreamAuditStorage(AuditStorage):
         if audit_fields:
             auditrequest.fields_audit.extend(audit_fields)
 
-        if kb_counter:
-            auditrequest.kb_counter.CopyFrom(kb_counter)
-
-            self.kb_usage_utility.send_kb_usage(
-                service=Service.NUCLIA_DB,
-                account_id=None,
-                kb_id=kbid,
-                kb_source=KBSource.HOSTED,
-                storage=Storage(paragraphs=kb_counter.paragraphs, fields=kb_counter.fields),
-            )
+        self.send(auditrequest)
 
-        auditrequest.trace_id = get_trace_id()
-
-        await self.send(auditrequest)
-
-    def report_resources(
+    def visited(
         self,
-        *,
         kbid: str,
-        resources: int,
+        uuid: str,
+        user: str,
+        origin: str,
+        send: bool = False,
     ):
-        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),
-        )
+        context = get_request_context()
+        if context is None:
+            return
+        auditrequest = context.audit_request
 
-    async def visited(self, kbid: str, uuid: str, user: str, origin: str):
-        auditrequest = AuditRequest()
         auditrequest.origin = origin
         auditrequest.userid = user
         auditrequest.rid = uuid
         auditrequest.kbid = kbid
         auditrequest.type = AuditRequest.VISITED
-        auditrequest.time.FromDatetime(datetime.now())
-
-        auditrequest.trace_id = get_trace_id()
-
-        await self.send(auditrequest)
 
-    async def delete_kb(self, kbid):
-        # Search is a base64 encoded search
-        auditrequest = AuditRequest()
-        auditrequest.kbid = kbid
-        auditrequest.type = AuditRequest.KB_DELETED
-        auditrequest.time.FromDatetime(datetime.now())
-        auditrequest.trace_id = get_trace_id()
-        await self.send(auditrequest)
-
-        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),
-        )
-
-    async def search(
+    def search(
         self,
         kbid: str,
         user: str,
@@ -262,95 +287,53 @@ class StreamAuditStorage(AuditStorage):
         timeit: float,
         resources: int,
     ):
-        # Search is a base64 encoded search
-        auditrequest = AuditRequest()
+        context = get_request_context()
+        if context is None:
+            return
+
+        auditrequest = context.audit_request
+
         auditrequest.origin = origin
         auditrequest.client_type = client_type  # type: ignore
         auditrequest.userid = user
         auditrequest.kbid = kbid
         auditrequest.search.CopyFrom(search)
-        auditrequest.timeit = timeit
+        auditrequest.retrieval_time = timeit
         auditrequest.resources = resources
         auditrequest.type = AuditRequest.SEARCH
-        auditrequest.time.FromDatetime(datetime.now())
-
-        auditrequest.trace_id = get_trace_id()
-        await self.send(auditrequest)
-
-        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,
-                )
-            ],
-        )
 
-    async def suggest(
+    def chat(
         self,
         kbid: str,
         user: str,
         client_type: int,
         origin: str,
-        timeit: float,
-    ):
-        auditrequest = AuditRequest()
-        auditrequest.origin = origin
-        auditrequest.client_type = client_type  # type: ignore
-        auditrequest.userid = user
-        auditrequest.kbid = kbid
-        auditrequest.timeit = timeit
-        auditrequest.type = AuditRequest.SUGGEST
-        auditrequest.time.FromDatetime(datetime.now())
-        auditrequest.trace_id = get_trace_id()
-
-        await self.send(auditrequest)
-
-        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,
-                )
-            ],
-        )
-
-    async def chat(
-        self,
-        kbid: str,
-        user: str,
-        client_type: int,
-        origin: str,
-        timeit: float,
         question: str,
         rephrased_question: Optional[str],
         context: List[ChatContext],
         answer: Optional[str],
         learning_id: str,
+        rephrase_time: Optional[float] = None,
+        generative_answer_time: Optional[float] = None,
+        generative_answer_first_chunk_time: Optional[float] = None,
     ):
-        auditrequest = AuditRequest()
+        rcontext = get_request_context()
+        if rcontext is None:
+            return
+
+        auditrequest = rcontext.audit_request
+
         auditrequest.origin = origin
         auditrequest.client_type = client_type  # type: ignore
         auditrequest.userid = user
         auditrequest.kbid = kbid
-        auditrequest.timeit = timeit
+        if rephrase_time is not None:
+            auditrequest.rephrase_time = rephrase_time
+        if generative_answer_time is not None:
+            auditrequest.generative_answer_time = generative_answer_time
+        if generative_answer_first_chunk_time is not None:
+            auditrequest.generative_answer_first_chunk_time = generative_answer_first_chunk_time
         auditrequest.type = AuditRequest.CHAT
-        auditrequest.time.FromDatetime(datetime.now())
-        auditrequest.trace_id = get_trace_id()
         auditrequest.chat.question = question
         auditrequest.chat.context.extend(context)
         auditrequest.chat.learning_id = learning_id
@@ -358,11 +341,3 @@ class StreamAuditStorage(AuditStorage):
             auditrequest.chat.rephrased_question = rephrased_question
         if answer is not None:
             auditrequest.chat.answer = answer
-        await self.send(auditrequest)
-
-
-def get_trace_id() -> str:
-    span = get_current_span()
-    if span is None:
-        return ""
-    return format_trace_id(span.get_span_context().trace_id)

nucliadb_utils/nuclia_usage/utils/kb_usage_report.py

@@ -22,10 +22,9 @@ import logging
 from collections.abc import Iterable
 from contextlib import suppress
 from datetime import datetime, timezone
-from typing import Optional
-
-from nats.js.client import JetStreamContext
+from typing import List, Optional
 
+from nucliadb_utils.nats import NatsConnectionManager
 from nucliadb_utils.nuclia_usage.protos.kb_usage_pb2 import (
     KBSource,
     KbUsage,
@@ -40,29 +39,48 @@ logger = logging.getLogger(__name__)
 
 
 class KbUsageReportUtility:
+    task: Optional[asyncio.Task]
+    initialized: bool
     queue: asyncio.Queue
-    lock: asyncio.Lock
+    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_stream = nats_stream
+        self.nats_connection_manager = NatsConnectionManager(
+            service_name=service,
+            nats_servers=nats_servers,
+            nats_creds=nats_creds,
+        )
         self.nats_subject = nats_subject
         self.queue = asyncio.Queue(max_queue_size)
         self.task = None
+        self.initialized = False
 
     async def initialize(self):
-        if self.task is None:
-            self.task = asyncio.create_task(self.run())
+        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
 
     async def finalize(self):
-        if self.task is not None:
-            self.task.cancel()
-            with suppress(asyncio.CancelledError, asyncio.exceptions.TimeoutError):
-                await asyncio.wait_for(self.task, timeout=2)
+        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
 
     async def run(self) -> None:
         while True:
@@ -75,13 +93,15 @@ 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_stream.publish(
+        res = await self.nats_connection_manager.js.publish(
             self.nats_subject,
             message.SerializeToString(),
         )

nucliadb_utils/settings.py

@@ -200,6 +200,16 @@ class AuditSettings(BaseSettings):
 audit_settings = AuditSettings()
 
 
+class UsageSettings(BaseSettings):
+    usage_jetstream_subject: Optional[str] = "kb-usage.nuclia_db"
+    usage_jetstream_servers: List[str] = []
+    usage_jetstream_auth: Optional[str] = None
+    usage_stream: str = "kb-usage"
+
+
+usage_settings = UsageSettings()
+
+
 class NATSConsumerSettings(BaseSettings):
     # Read about message ordering:
     #   https://docs.nats.io/nats-concepts/subject_mapping#when-is-deterministic-partitioning-needed

nucliadb_utils/tests/nats.py

@@ -191,7 +191,11 @@ def start_gnatsd(gnatsd: Gnatsd):  # pragma: no cover
 
 @pytest.fixture(scope="session")
 def natsd_server():  # pragma: no cover
-    if not os.path.isfile("nats-server"):
+    # Create a persistent temporary directory
+    tmpdir = tempfile.mkdtemp()
+    nats_server_path = os.path.join(tmpdir, "nats-server")
+
+    if not os.path.isfile(nats_server_path):
         version = "v2.10.12"
         arch = platform.machine()
         if arch == "x86_64":
@@ -205,13 +209,13 @@ def natsd_server():  # pragma: no cover
 
         file = zipfile.open(f"nats-server-{version}-{system}-{arch}/nats-server")
         content = file.read()
-        with open("nats-server", "wb") as f:
+        with open(nats_server_path, "wb") as f:
             f.write(content)
-        os.chmod("nats-server", 755)
+        os.chmod(nats_server_path, 0o755)
 
     server = Gnatsd(port=4222)
     server.bin_name = "nats-server"
-    server.path = os.getcwd()
+    server.path = tmpdir
     return server
 
 

nucliadb_utils/utilities.py

@@ -40,6 +40,7 @@ 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,
@@ -48,6 +49,7 @@ 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
@@ -81,6 +83,7 @@ class Utility(str, Enum):
     LOCAL_STORAGE = "local_storage"
     NUCLIA_STORAGE = "nuclia_storage"
     MAINDB_DRIVER = "driver"
+    USAGE = "usage"
     ENDECRYPTOR = "endecryptor"
     PINECONE_SESSION = "pinecone_session"
 
@@ -322,16 +325,41 @@ def get_audit() -> Optional[AuditStorage]:
     return get_utility(Utility.AUDIT)
 
 
-async def start_audit_utility(service: str):
-    audit_utility: Optional[AuditStorage] = get_utility(Utility.AUDIT)
-    if audit_utility is not None:
+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":
-        audit_utility = BasicAuditStorage()
+        b_audit_utility: AuditStorage = BasicAuditStorage()
+        set_utility(Utility.AUDIT, b_audit_utility)
         logger.info("Configuring basic audit log")
+        return b_audit_utility
     elif audit_settings.audit_driver == "stream":
-        audit_utility = StreamAuditStorage(
+        s_audit_utility: AuditStorage = StreamAuditStorage(
             nats_creds=audit_settings.audit_jetstream_auth,
             nats_servers=audit_settings.audit_jetstream_servers,
             nats_target=cast(str, audit_settings.audit_jetstream_target),
@@ -339,11 +367,22 @@ async def start_audit_utility(service: str):
             seed=audit_settings.audit_hash_seed,
             service=service,
         )
+        set_utility(Utility.AUDIT, s_audit_utility)
         logger.info(f"Configuring stream audit log {audit_settings.audit_jetstream_target}")
+        return s_audit_utility
     else:
         raise ConfigurationError("Invalid audit driver")
-    await audit_utility.initialize()
-    set_utility(Utility.AUDIT, audit_utility)
+
+
+async def start_audit_utility(service: str):
+    audit_utility: Optional[AuditStorage] = get_utility(Utility.AUDIT)
+    if audit_utility is not None and audit_utility.initialized is True:
+        return
+
+    if audit_utility is None:
+        audit_utility = register_audit_utility(service)
+    if audit_utility.initialized is False:
+        await audit_utility.initialize()
 
 
 async def stop_audit_utility():

{nucliadb_utils-5.0.0.post877.dist-info → nucliadb_utils-5.0.0.post899.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: nucliadb_utils
-Version: 5.0.0.post877
+Version: 5.0.0.post899
 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.0.post877
-Requires-Dist: nucliadb-telemetry >=5.0.0.post877
+Requires-Dist: nucliadb-protos >=5.0.0.post899
+Requires-Dist: nucliadb-telemetry >=5.0.0.post899
 Provides-Extra: cache
 Requires-Dist: redis >=4.3.4 ; extra == 'cache'
 Requires-Dist: orjson >=3.6.7 ; extra == 'cache'

{nucliadb_utils-5.0.0.post877.dist-info → nucliadb_utils-5.0.0.post899.dist-info}/RECORD

@@ -12,19 +12,19 @@ nucliadb_utils/nats.py,sha256=zTAXECDXeCPtydk3F_6EMFDZ059kK0UYUU_tnWoxgXs,8208
 nucliadb_utils/partition.py,sha256=jBgy4Hu5Iwn4gjbPPcthSykwf-qNx-GcLAIwbzPd1d0,1157
 nucliadb_utils/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 nucliadb_utils/run.py,sha256=HpAIM8xbR7UpVC2_7xOjB4fYbUVykyPP6yHrv2RD3DI,1707
-nucliadb_utils/settings.py,sha256=AaOtQZVRqRcMnUyN1l1MpR10lANaDT2uPrbhmTyn6uk,7647
+nucliadb_utils/settings.py,sha256=QR51SX0T17-_YofMNpci-nkI77l_CUWFy4H7i8hNOHU,7911
 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=0NNoWLAR88Kb6QuHOK08ZIhrjAKsQ46iQqwN3-I9S6w,15355
+nucliadb_utils/utilities.py,sha256=jjapoJvgZ-H0uMoThVK8P6EvCTTwL0kM_EEJY-Q22EU,16747
 nucliadb_utils/aiopynecone/__init__.py,sha256=cp15ZcFnHvpcu_5-aK2A4uUyvuZVV_MJn4bIXMa20ks,835
 nucliadb_utils/aiopynecone/client.py,sha256=jlTLOdphrLZElKW9ajVzDcebN7CSAbjNyX6tnD37eqY,18721
 nucliadb_utils/aiopynecone/exceptions.py,sha256=hFhq-UEY4slqNWjObXr_LPnRf_AQ1vpcG4SF2XRFd1E,2873
 nucliadb_utils/aiopynecone/models.py,sha256=B_ihJhHZGp3ivQVUxhV49uoUnHe1PLDKxTgHNbHgSS0,2937
 nucliadb_utils/audit/__init__.py,sha256=cp15ZcFnHvpcu_5-aK2A4uUyvuZVV_MJn4bIXMa20ks,835
-nucliadb_utils/audit/audit.py,sha256=dn5ZnCVQUlCcvdjzaORghbrjk9QgVGrtkfIftq30Bp8,2819
-nucliadb_utils/audit/basic.py,sha256=NViey6mKbCXqRTLDBX2xNTcCg9I-2e4oB2xkekuhDvM,3392
-nucliadb_utils/audit/stream.py,sha256=CSNFrwCT_IeHFStx4q4ImgOQ7v0WzJr4YG6QRzE7YlQ,11952
+nucliadb_utils/audit/audit.py,sha256=kpjiTLgWtLXA5DNgJcmeIaodel11Mz0Y-JfQt-yfhis,2641
+nucliadb_utils/audit/basic.py,sha256=LwCCvUAwy8LNzYdrw3NebxqNdUumCEyRQ244dJ93JfE,3190
+nucliadb_utils/audit/stream.py,sha256=_RZf1G9_c5eaxfutq1OdMmnRaRM7EiaXH33Kb4mN9Gw,11525
 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=E1eUSFXBVNzQP9Q2rWj9y3koCO5S7iKwckny_AoLKuk,3870
+nucliadb_utils/nuclia_usage/utils/kb_usage_report.py,sha256=t4-QXUTOTeMBWFFO8hdqiIco1fgC2Jla4NIWuRSs6N4,4566
 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
@@ -62,10 +62,10 @@ nucliadb_utils/tests/fixtures.py,sha256=rmix1VGpPULBxJd_5ScgVD4Z0nRrkNOVt-uTW_Wo
 nucliadb_utils/tests/gcs.py,sha256=Ii8BCHUAAxFIzX67pKTRFRgbqv3FJ6DrPAdAx2Xod1Y,3036
 nucliadb_utils/tests/indexing.py,sha256=YW2QhkhO9Q_8A4kKWJaWSvXvyQ_AiAwY1VylcfVQFxk,1513
 nucliadb_utils/tests/local.py,sha256=c3gZJJWmvOftruJkIQIwB3q_hh3uxEhqGIAVWim1Bbk,1343
-nucliadb_utils/tests/nats.py,sha256=Tosonm9A9cusImyji80G4pgdXEHNVPaCLT5TbFK_ra0,7543
+nucliadb_utils/tests/nats.py,sha256=xqpww4jZjTKY9oPGlJdDJG67L3FIBQsa9qDHxILR8r8,7687
 nucliadb_utils/tests/s3.py,sha256=YB8QqDaBXxyhHonEHmeBbRRDmvB7sTOaKBSi8KBGokg,2330
-nucliadb_utils-5.0.0.post877.dist-info/METADATA,sha256=IeNihiMykhRSexgDnkLuVbwXnJ9Y-LMq_Z6z9Ff_D9k,2073
-nucliadb_utils-5.0.0.post877.dist-info/WHEEL,sha256=Wyh-_nZ0DJYolHNn1_hMa4lM7uDedD_RGVwbmTjyItk,91
-nucliadb_utils-5.0.0.post877.dist-info/top_level.txt,sha256=fE3vJtALTfgh7bcAWcNhcfXkNPp_eVVpbKK-2IYua3E,15
-nucliadb_utils-5.0.0.post877.dist-info/zip-safe,sha256=AbpHGcgLb-kRsJGnwFEktk7uzpZOCcBY74-YBdrKVGs,1
-nucliadb_utils-5.0.0.post877.dist-info/RECORD,,
+nucliadb_utils-5.0.0.post899.dist-info/METADATA,sha256=--jd5upWgPLSU_qpT3aCoyyMFyHhxq5vkrtrGdXGXSM,2073
+nucliadb_utils-5.0.0.post899.dist-info/WHEEL,sha256=Wyh-_nZ0DJYolHNn1_hMa4lM7uDedD_RGVwbmTjyItk,91
+nucliadb_utils-5.0.0.post899.dist-info/top_level.txt,sha256=fE3vJtALTfgh7bcAWcNhcfXkNPp_eVVpbKK-2IYua3E,15
+nucliadb_utils-5.0.0.post899.dist-info/zip-safe,sha256=AbpHGcgLb-kRsJGnwFEktk7uzpZOCcBY74-YBdrKVGs,1
+nucliadb_utils-5.0.0.post899.dist-info/RECORD,,