nv-ingest-api 2025.4.17.dev20250417__py3-none-any.whl → 2025.4.18.dev20250418__py3-none-any.whl

nv_ingest_api/util/message_brokers/simple_message_broker/simple_client.py

@@ -11,10 +11,10 @@ import socket
 import json
 import time
 import logging
-from typing import Optional
+from typing import Optional, Tuple, Union
 
 from nv_ingest_api.internal.schemas.message_brokers.response_schema import ResponseSchema
-from nv_ingest_api.util.service_clients.client_base import MessageBrokerClientBase
+from nv_ingest_api.util.service_clients.client_base import MessageBrokerClientBase, FetchMode
 
 logger = logging.getLogger(__name__)
 
@@ -80,7 +80,11 @@ class SimpleClient(MessageBrokerClientBase):
         return self
 
     def submit_message(
-        self, queue_name: str, message: str, timeout: Optional[float] = None, for_nv_ingest: bool = False
+        self,
+        queue_name: str,
+        message: str,
+        timeout: Optional[Tuple[int, Union[float]]] = (100, None),
+        for_nv_ingest: bool = False,
     ) -> ResponseSchema:
         """
         Submit a message to the specified queue.
@@ -103,7 +107,12 @@ class SimpleClient(MessageBrokerClientBase):
         """
         return self._handle_push(queue_name, message, timeout, for_nv_ingest)
 
-    def fetch_message(self, queue_name: str, timeout: Optional[float] = None) -> ResponseSchema:
+    def fetch_message(
+        self,
+        queue_name: str,
+        timeout: Optional[Tuple[int, Union[float]]] = (100, None),
+        override_fetch_mode: FetchMode = None,
+    ) -> ResponseSchema:
         """
         Fetch a message from the specified queue.
 
@@ -119,6 +128,9 @@ class SimpleClient(MessageBrokerClientBase):
         ResponseSchema
             The response containing the fetched message.
         """
+        if isinstance(timeout, int):
+            timeout = (timeout, None)
+
         return self._handle_pop(queue_name, timeout)
 
     def ping(self) -> ResponseSchema:
@@ -151,7 +163,7 @@ class SimpleClient(MessageBrokerClientBase):
         return self._execute_simple_command(command)
 
     def _handle_push(
-        self, queue_name: str, message: str, timeout: Optional[float], for_nv_ingest: bool
+        self, queue_name: str, message: str, timeout: Optional[Tuple[int, Union[float, None]]], for_nv_ingest: bool
     ) -> ResponseSchema:
         """
         Push a message to the queue with optional timeout.
@@ -183,6 +195,7 @@ class SimpleClient(MessageBrokerClientBase):
         else:
             command = {"command": "PUSH", "queue_name": queue_name, "message": message}
 
+        timeout = int(timeout[0])
         if timeout is not None:
             command["timeout"] = timeout
 
@@ -237,7 +250,7 @@ class SimpleClient(MessageBrokerClientBase):
 
             time.sleep(0.5)  # Backoff delay before retry
 
-    def _handle_pop(self, queue_name: str, timeout: Optional[float]) -> ResponseSchema:
+    def _handle_pop(self, queue_name: str, timeout: Optional[Tuple[int, Union[float, None]]]) -> ResponseSchema:
         """
         Pop a message from the queue with optional timeout.
 
@@ -258,6 +271,9 @@ class SimpleClient(MessageBrokerClientBase):
             return ResponseSchema(response_code=1, response_reason="Invalid queue name.")
 
         command = {"command": "POP", "queue_name": queue_name}
+
+        timeout = int(timeout[0])
+
         if timeout is not None:
             command["timeout"] = timeout
 

nv_ingest_api/util/service_clients/client_base.py

@@ -4,6 +4,16 @@
 
 from abc import ABC
 from abc import abstractmethod
+from enum import Enum, auto
+from typing import Tuple
+
+from nv_ingest_api.internal.schemas.message_brokers.response_schema import ResponseSchema
+
+
+class FetchMode(Enum):
+    DESTRUCTIVE = auto()  # Read and delete immediately (current behavior)
+    NON_DESTRUCTIVE = auto()  # Read without deleting (requires TTL on Redis data)
+    CACHE_BEFORE_DELETE = auto()  # Read, write to local cache, then delete from Redis
 
 
 class MessageBrokerClientBase(ABC):
@@ -49,24 +59,28 @@ class MessageBrokerClientBase(ABC):
         """
 
     @abstractmethod
-    def fetch_message(self, job_index: str, timeout: float = 0) -> str:
+    def fetch_message(
+        self, job_index: str, timeout: Tuple[int, float] = (100, None), override_fetch_mode: FetchMode = None
+    ) -> ResponseSchema:
         """
         Fetches a message from the specified queue with retries on failure.
 
         Parameters:
             job_index (str): The index of the job to fetch the message for.
             timeout (float): The timeout in seconds for blocking until a message is available.
+            override_fetch_mode: Optional; overrides the default fetch mode.
 
         Returns:
             The fetched message, or None if no message could be fetched.
         """
 
     @abstractmethod
-    def submit_message(self, channel_name: str, message: str) -> str:
+    def submit_message(self, channel_name: str, message: str, for_nv_ingest=False) -> ResponseSchema:
         """
         Submits a message to a specified queue with retries on failure.
 
         Parameters:
             channel_name (str): The name of the queue to submit the message to.
             message (str): The message to submit.
+            for_nv_ingest (bool): Whether the message is for NV Ingest.
         """