From a14a93b1084449bfb63f7da1d0d78528e27ff5aa Mon Sep 17 00:00:00 2001
From: Matthew Tran <0e4ef622@gmail.com>
Date: Tue, 28 Jan 2025 09:03:25 -0600
Subject: [PATCH 1/2] Apply extra_kwargs in batches wait_for_results

---
 ntropy_sdk/batches.py | 18 ++++++++++++++----
 1 file changed, 14 insertions(+), 4 deletions(-)

diff --git a/ntropy_sdk/batches.py b/ntropy_sdk/batches.py
index ac73ccd..5f0fb46 100644
--- a/ntropy_sdk/batches.py
+++ b/ntropy_sdk/batches.py
@@ -174,11 +174,12 @@ def _wait(
         poll_interval: int,
         timeout: int,
         stop_fn: Callable[[Batch], bool],
+        extra_kwargs: "ExtraKwargs",
     ) -> Batch:
         start_time = time.monotonic()
         batch = None
         while time.monotonic() - start_time < timeout:
-            batch = self.get(id=id)
+            batch = self.get(id=id, **extra_kwargs)
             if stop_fn(batch):
                 break
             time.sleep(poll_interval)
@@ -191,6 +192,7 @@ def _wait_with_progress(
         poll_interval: int,
         timeout: int,
         stop_fn: Callable[[Batch], bool],
+        extra_kwargs: "ExtraKwargs",
     ) -> Batch:
         from tqdm.auto import tqdm
 
@@ -199,7 +201,7 @@ def _wait_with_progress(
         total_set = False
         with tqdm() as p:
             while time.monotonic() - start_time < timeout:
-                batch = self.get(id=id)
+                batch = self.get(id=id, **extra_kwargs)
                 if not total_set:
                     p.total = batch.total
                 p.desc = batch.status
@@ -230,11 +232,19 @@ def stop_fn(b: Batch):
 
         if with_progress:
             batch = self._wait_with_progress(
-                id=id, poll_interval=poll_interval, timeout=timeout, stop_fn=stop_fn
+                id=id,
+                poll_interval=poll_interval,
+                timeout=timeout,
+                stop_fn=stop_fn,
+                extra_kwargs=extra_kwargs,
             )
         else:
             batch = self._wait(
-                id=id, poll_interval=poll_interval, timeout=timeout, stop_fn=stop_fn
+                id=id,
+                poll_interval=poll_interval,
+                timeout=timeout,
+                stop_fn=stop_fn,
+                extra_kwargs=extra_kwargs,
             )
 
         if batch and batch.status not in finish_statuses:

From f07b9701dc9ef69d752c1c15c99ea41957e3a17e Mon Sep 17 00:00:00 2001
From: Matthew Tran <0e4ef622@gmail.com>
Date: Tue, 28 Jan 2025 13:06:44 -0600
Subject: [PATCH 2/2] =?UTF-8?q?Bump=20version:=205.1.1=20=E2=86=92=205.1.2?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 ntropy_sdk/__init__.py | 2 +-
 setup.cfg              | 2 +-
 setup.py               | 2 +-
 3 files changed, 3 insertions(+), 3 deletions(-)

diff --git a/ntropy_sdk/__init__.py b/ntropy_sdk/__init__.py
index 1fc25b3..acb6186 100644
--- a/ntropy_sdk/__init__.py
+++ b/ntropy_sdk/__init__.py
@@ -1,4 +1,4 @@
-__version__ = "5.1.1"
+__version__ = "5.1.2"
 
 from typing import TYPE_CHECKING, Optional
 
diff --git a/setup.cfg b/setup.cfg
index c539062..c05cd1d 100644
--- a/setup.cfg
+++ b/setup.cfg
@@ -1,5 +1,5 @@
 [bumpversion]
-current_version = 5.1.1
+current_version = 5.1.2
 commit = True
 tag = True
 parse = (?P<major>\d+)\.(?P<minor>\d+)\.(?P<patch>\d+)(rc(?P<pre>\d+))?
diff --git a/setup.py b/setup.py
index 3eb5db3..bc83dad 100644
--- a/setup.py
+++ b/setup.py
@@ -55,6 +55,6 @@
     test_suite="tests",
     tests_require=test_requirements,
     url="https://github.com/ntropy-network/ntropy-sdk",
-    version="5.1.1",
+    version="5.1.2",
     zip_safe=False,
 )