lm-deluge 0.0.3__tar.gz → 0.0.4__tar.gz

{lm_deluge-0.0.3 → lm_deluge-0.0.4}/PKG-INFO

@@ -1,9 +1,9 @@
 Metadata-Version: 2.4
 Name: lm_deluge
-Version: 0.0.3
+Version: 0.0.4
 Summary: Python utility for using LLM API models.
 Author-email: Benjamin Anderson <ben@trytaylor.ai>
-Requires-Python: >=3.9
+Requires-Python: >=3.10
 Description-Content-Type: text/markdown
 Requires-Dist: python-dotenv
 Requires-Dist: json5

{lm_deluge-0.0.3 → lm_deluge-0.0.4}/pyproject.toml

@@ -3,11 +3,11 @@ requires = ["setuptools", "wheel"]
 
 [project]
 name = "lm_deluge"
-version = "0.0.3"
+version = "0.0.4"
 authors = [{ name = "Benjamin Anderson", email = "ben@trytaylor.ai" }]
 description = "Python utility for using LLM API models."
 readme = "README.md"
-requires-python = ">=3.9"
+requires-python = ">=3.10"
 keywords = []
 license = { text = "" }
 classifiers = []

{lm_deluge-0.0.3 → lm_deluge-0.0.4}/src/lm_deluge/api_requests/anthropic.py

@@ -3,9 +3,8 @@ from aiohttp import ClientResponse
 import json
 import os
 import warnings
-import time
 from tqdm import tqdm
-from typing import Optional, Callable
+from typing import Callable
 
 from lm_deluge.prompt import Conversation
 from .base import APIRequestBase, APIResponse
@@ -29,8 +28,8 @@ class AnthropicRequest(APIRequestBase):
         results_arr: list,
         request_timeout: int = 30,
         sampling_params: SamplingParams = SamplingParams(),
-        pbar: Optional[tqdm] = None,
-        callback: Optional[Callable] = None,
+        pbar: tqdm | None = None,
+        callback: Callable | None = None,
         debug: bool = False,
         # for retries
         all_model_names: list[str] | None = None,
@@ -96,8 +95,6 @@ class AnthropicRequest(APIRequestBase):
         if self.system_message is not None:
             self.request_json["system"] = self.system_message
 
-        # print("request data:", self.request_json)
-
     async def handle_response(self, http_response: ClientResponse) -> APIResponse:
         is_error = False
         error_message = None
@@ -156,8 +153,7 @@ class AnthropicRequest(APIRequestBase):
                 or "overloaded" in error_message.lower()
             ):
                 error_message += " (Rate limit error, triggering cooldown.)"
-                self.status_tracker.time_of_last_rate_limit_error = time.time()
-                self.status_tracker.num_rate_limit_errors += 1
+                self.status_tracker.rate_limit_exceeded()
             if "context length" in error_message:
                 error_message += " (Context length exceeded, set retries to 0.)"
                 self.attempts_left = 0

{lm_deluge-0.0.3 → lm_deluge-0.0.4}/src/lm_deluge/api_requests/base.py

@@ -5,7 +5,7 @@ import random
 from tqdm import tqdm
 from dataclasses import dataclass
 from abc import ABC, abstractmethod
-from typing import Optional, Callable
+from typing import Callable
 
 from lm_deluge.prompt import Conversation
 
@@ -26,25 +26,25 @@ class APIResponse:
 
     # http response information
     status_code: int | None
-    is_error: Optional[bool]
-    error_message: Optional[str]
+    is_error: bool | None
+    error_message: str | None
 
     # completion information
-    completion: Optional[str]
-    input_tokens: Optional[int]
-    output_tokens: Optional[int]
+    completion: str | None
+    input_tokens: int | None
+    output_tokens: int | None
 
     # optional or calculated automatically
-    thinking: Optional[str] = None  # if model shows thinking tokens
-    model_external: Optional[str] = None  # the model tag used by the API
-    region: Optional[str] = None
-    logprobs: Optional[list] = None
-    finish_reason: Optional[str] = None  # make required later
-    cost: Optional[float] = None  # calculated automatically
+    thinking: str | None = None  # if model shows thinking tokens
+    model_external: str | None = None  # the model tag used by the API
+    region: str | None = None
+    logprobs: list | None = None
+    finish_reason: str | None = None  # make required later
+    cost: float | None = None  # calculated automatically
     # set to true if is_error and should be retried with a different model
-    retry_with_different_model: Optional[bool] = False
+    retry_with_different_model: bool | None = False
     # set to true if should NOT retry with the same model (unrecoverable error)
-    give_up_if_no_other_models: Optional[bool] = False
+    give_up_if_no_other_models: bool | None = False
 
     def __post_init__(self):
         # calculate cost & get external model name
@@ -138,9 +138,9 @@ class APIRequestBase(ABC):
         request_timeout: int = 30,
         sampling_params: SamplingParams = SamplingParams(),
         logprobs: bool = False,
-        top_logprobs: Optional[int] = None,
-        pbar: Optional[tqdm] = None,
-        callback: Optional[Callable] = None,
+        top_logprobs: int | None = None,
+        pbar: tqdm | None = None,
+        callback: Callable | None = None,
         debug: bool = False,
         all_model_names: list[str] | None = None,
         all_sampling_params: list[SamplingParams] | None = None,
@@ -185,8 +185,7 @@ class APIRequestBase(ABC):
     def handle_success(self, data):
         self.call_callback()
         self.increment_pbar()
-        self.status_tracker.num_tasks_in_progress -= 1
-        self.status_tracker.num_tasks_succeeded += 1
+        self.status_tracker.task_succeeded(self.task_id)
 
     def handle_error(self, create_new_request=False, give_up_if_no_other_models=False):
         """
@@ -215,8 +214,7 @@ class APIRequestBase(ABC):
                         print(
                             f"No other models to try for task {self.task_id}. Giving up."
                         )
-                        self.status_tracker.num_tasks_in_progress -= 1
-                        self.status_tracker.num_tasks_failed += 1
+                        self.status_tracker.task_failed(self.task_id)
                     else:
                         print(
                             f"No other models to try for task {self.task_id}. Retrying with same model."
@@ -263,8 +261,7 @@ class APIRequestBase(ABC):
                     self.results_arr.append(new_request)
         else:
             print(f"Task {self.task_id} out of tries.")
-            self.status_tracker.num_tasks_in_progress -= 1
-            self.status_tracker.num_tasks_failed += 1
+            self.status_tracker.task_failed(self.task_id)
 
     async def call_api(self):
         try:
@@ -308,7 +305,6 @@ class APIRequestBase(ABC):
 
         except Exception as e:
             raise_if_modal_exception(e)
-            # print(f"Unexpected error {type(e).__name__}: {str(e) or 'No message.'}")
             self.result.append(
                 APIResponse(
                     id=self.task_id,
@@ -342,9 +338,9 @@ def create_api_request(
     request_timeout: int = 30,
     sampling_params: SamplingParams = SamplingParams(),
     logprobs: bool = False,
-    top_logprobs: Optional[int] = None,
-    pbar: Optional[tqdm] = None,
-    callback: Optional[Callable] = None,
+    top_logprobs: int | None = None,
+    pbar: tqdm | None = None,
+    callback: Callable | None = None,
     all_model_names: list[str] | None = None,
     all_sampling_params: list[SamplingParams] | None = None,
 ) -> APIRequestBase:

{lm_deluge-0.0.3 → lm_deluge-0.0.4}/src/lm_deluge/api_requests/cohere.py

@@ -4,9 +4,8 @@ import asyncio
 from aiohttp import ClientResponse
 import json
 import os
-import time
 from tqdm import tqdm
-from typing import Optional, Callable
+from typing import Callable
 from lm_deluge.prompt import Conversation
 from .base import APIRequestBase, APIResponse
 
@@ -29,8 +28,8 @@ class CohereRequest(APIRequestBase):
         retry_queue: asyncio.Queue,
         request_timeout: int = 30,
         sampling_params: SamplingParams = SamplingParams(),
-        pbar: Optional[tqdm] = None,
-        callback: Optional[Callable] = None,
+        pbar: tqdm | None = None,
+        callback: Callable | None = None,
         debug: bool = False,
         all_model_names: list[str] | None = None,
         all_sampling_params: list[SamplingParams] | None = None,
@@ -118,8 +117,7 @@ class CohereRequest(APIRequestBase):
                 or "overloaded" in error_message.lower()
             ):
                 error_message += " (Rate limit error, triggering cooldown.)"
-                self.status_tracker.time_of_last_rate_limit_error = time.time()
-                self.status_tracker.num_rate_limit_errors += 1
+                self.status_tracker.rate_limit_exceeded()
             if "context length" in error_message:
                 error_message += " (Context length exceeded, set retries to 0.)"
                 self.attempts_left = 0

{lm_deluge-0.0.3 → lm_deluge-0.0.4}/src/lm_deluge/api_requests/deprecated/bedrock.py

@@ -55,8 +55,8 @@
 #         retry_queue: asyncio.Queue,
 #         request_timeout: int = 30,
 #         sampling_params: SamplingParams = SamplingParams(),
-#         pbar: Optional[tqdm] = None,
-#         callback: Optional[Callable] = None,
+#         pbar: tqdm | None = None,
+#         callback: Callable | None = None,
 #         debug: bool = False,
 #         all_model_names: list[str] | None = None,
 #         all_sampling_params: list[SamplingParams] | None = None,
@@ -175,8 +175,8 @@
 #         results_arr: list,
 #         request_timeout: int = 30,
 #         sampling_params: SamplingParams = SamplingParams(),
-#         pbar: Optional[tqdm] = None,
-#         callback: Optional[Callable] = None,
+#         pbar: tqdm | None = None,
+#         callback: Callable | None = None,
 #         debug: bool = False,
 #         all_model_names: list[str] | None = None,
 #         all_sampling_params: list[SamplingParams] | None = None,

{lm_deluge-0.0.3 → lm_deluge-0.0.4}/src/lm_deluge/api_requests/deprecated/deepseek.py

@@ -25,8 +25,8 @@
 #         results_arr: list,
 #         request_timeout: int = 30,
 #         sampling_params: SamplingParams = SamplingParams(),
-#         pbar: Optional[tqdm] = None,
-#         callback: Optional[Callable] = None,
+#         pbar: tqdm | None = None,
+#         callback: Callable | None = None,
 #         debug: bool = False,
 #         all_model_names: list[str] = None,
 #         all_sampling_params: list[SamplingParams] = None,

{lm_deluge-0.0.3 → lm_deluge-0.0.4}/src/lm_deluge/api_requests/deprecated/mistral.py

@@ -27,8 +27,8 @@
 #         results_arr: list,
 #         request_timeout: int = 30,
 #         sampling_params: SamplingParams = SamplingParams(),
-#         pbar: Optional[tqdm] = None,
-#         callback: Optional[Callable] = None,
+#         pbar: tqdm | None = None,
+#         callback: Callable | None = None,
 #         debug: bool = False,
 #         all_model_names: list[str] = None,
 #         all_sampling_params: list[SamplingParams] = None,

{lm_deluge-0.0.3 → lm_deluge-0.0.4}/src/lm_deluge/api_requests/openai.py

@@ -3,9 +3,8 @@ import warnings
 from aiohttp import ClientResponse
 import json
 import os
-import time
 from tqdm.auto import tqdm
-from typing import Optional, Callable
+from typing import Callable
 
 from .base import APIRequestBase, APIResponse
 from ..prompt import Conversation
@@ -29,9 +28,9 @@ class OpenAIRequest(APIRequestBase):
         request_timeout: int = 30,
         sampling_params: SamplingParams = SamplingParams(),
         logprobs: bool = False,
-        top_logprobs: Optional[int] = None,
-        pbar: Optional[tqdm] = None,
-        callback: Optional[Callable] = None,
+        top_logprobs: int | None = None,
+        pbar: tqdm | None = None,
+        callback: Callable | None = None,
         debug: bool = False,
         all_model_names: list[str] | None = None,
         all_sampling_params: list[SamplingParams] | None = None,
@@ -124,8 +123,7 @@ class OpenAIRequest(APIRequestBase):
         if is_error and error_message is not None:
             if "rate limit" in error_message.lower() or status_code == 429:
                 error_message += " (Rate limit error, triggering cooldown.)"
-                self.status_tracker.time_of_last_rate_limit_error = time.time()
-                self.status_tracker.num_rate_limit_errors += 1
+                self.status_tracker.rate_limit_exceeded()
             if "context length" in error_message:
                 error_message += " (Context length exceeded, set retries to 0.)"
                 self.attempts_left = 0

{lm_deluge-0.0.3 → lm_deluge-0.0.4}/src/lm_deluge/api_requests/vertex.py

@@ -5,7 +5,7 @@ import json
 import os
 import time
 from tqdm import tqdm
-from typing import Optional, Callable
+from typing import Callable
 
 from lm_deluge.prompt import Conversation
 from .base import APIRequestBase, APIResponse
@@ -57,8 +57,8 @@ class VertexAnthropicRequest(APIRequestBase):
         results_arr: list,
         request_timeout: int = 30,
         sampling_params: SamplingParams = SamplingParams(),
-        pbar: Optional[tqdm] = None,
-        callback: Optional[Callable] = None,
+        pbar: tqdm | None = None,
+        callback: Callable | None = None,
         debug: bool = False,
     ):
         super().__init__(
@@ -141,8 +141,7 @@ class VertexAnthropicRequest(APIRequestBase):
                 or status_code == 429
             ):
                 error_message += " (Rate limit error, triggering cooldown.)"
-                self.status_tracker.time_of_last_rate_limit_error = time.time()
-                self.status_tracker.num_rate_limit_errors += 1
+                self.status_tracker.rate_limit_exceeded()
             if "context length" in error_message:
                 error_message += " (Context length exceeded, set retries to 0.)"
                 self.attempts_left = 0
@@ -185,8 +184,8 @@ class GeminiRequest(APIRequestBase):
         results_arr: list,
         request_timeout: int = 30,
         sampling_params: SamplingParams = SamplingParams(),
-        pbar: Optional[tqdm] = None,
-        callback: Optional[Callable] = None,
+        pbar: tqdm | None = None,
+        callback: Callable | None = None,
         debug: bool = False,
         all_model_names: list[str] | None = None,
         all_sampling_params: list[SamplingParams] | None = None,
@@ -302,16 +301,14 @@ class GeminiRequest(APIRequestBase):
                         error_message = "Finish reason SAFETY."
                         retry_with_different_model = True
                     else:
-                        print("Actual structure of response:")
-                        print(data)
+                        print("Actual structure of response:", data)
                         is_error = True
                         error_message = "No content in response."
             except Exception as e:
                 is_error = True
                 error_message = f"Error calling .json() on response w/ status {status_code}: {e.__class__} {e}"
                 if isinstance(e, KeyError):
-                    print("Actual structure of response:")
-                    print(data)
+                    print("Actual structure of response:", data)
         elif "json" in (mimetype or "").lower():
             is_error = True
             data = await http_response.json()
@@ -332,8 +329,7 @@ class GeminiRequest(APIRequestBase):
                 status_code == 429
             ):
                 error_message += " (Rate limit error, triggering cooldown & retrying with different model.)"
-                self.status_tracker.time_of_last_rate_limit_error = time.time()
-                self.status_tracker.num_rate_limit_errors += 1
+                self.status_tracker.rate_limit_exceeded()
                 retry_with_different_model = (
                     True  # if possible, retry with a different model
                 )

{lm_deluge-0.0.3 → lm_deluge-0.0.4}/src/lm_deluge/client.py

@@ -5,7 +5,7 @@ import numpy as np
 import time
 import yaml
 from dataclasses import dataclass
-from typing import Sequence, overload, Literal, Optional, Union, Any
+from typing import Sequence, overload, Literal, Any
 from tqdm.auto import tqdm
 
 from lm_deluge.prompt import Conversation
@@ -31,11 +31,11 @@ class ClientConfig:
     max_concurrent_requests: int
     max_attempts: int
     request_timeout: int
-    sampling_params: Union[SamplingParams, list[SamplingParams]]
-    model_weights: Union[list[float], Literal["uniform", "rate_limit"]]
+    sampling_params: SamplingParams | list[SamplingParams]
+    model_weights: list[float] | Literal["uniform", "rate_limit"]
     logprobs: bool = False
-    top_logprobs: Optional[int] = None
-    cache: Optional[Any] = None
+    top_logprobs: int | None = None
+    cache: Any = None
 
     @classmethod
     def from_dict(cls, config_dict: dict):
@@ -82,23 +82,21 @@ class LLMClient:
     Handles models, sampling params for each model, model weights, rate limits, etc.
     """
 
-    pass
-
     def __init__(
         self,
         model_names: list[str],
         max_requests_per_minute: int,
         max_tokens_per_minute: int,
         max_concurrent_requests: int,
-        sampling_params: Union[SamplingParams, list[SamplingParams]] = SamplingParams(),
-        model_weights: Union[list[float], Literal["uniform", "rate_limit"]] = "uniform",
+        sampling_params: SamplingParams | list[SamplingParams] = SamplingParams(),
+        model_weights: list[float] | Literal["uniform", "rate_limit"] = "uniform",
         max_attempts: int = 5,
         request_timeout: int = 30,
         logprobs: bool = False,
-        top_logprobs: Optional[int] = None,
+        top_logprobs: int | None = None,
         use_qps: bool = False,
         debug: bool = False,
-        cache: Optional[Any] = None,
+        cache: Any = None,
     ):
         self.models = model_names
         if isinstance(sampling_params, SamplingParams):
@@ -154,7 +152,7 @@ class LLMClient:
         self.cache = cache
 
     @classmethod
-    def from_config(cls, config: ClientConfig, cache: Optional[Any] = None):
+    def from_config(cls, config: ClientConfig, cache: Any = None):
         return cls(
             model_names=config.model_names,
             max_requests_per_minute=config.max_requests_per_minute,
@@ -168,25 +166,25 @@ class LLMClient:
         )
 
     @classmethod
-    def from_yaml(cls, file_path: str, cache: Optional[Any] = None):
+    def from_yaml(cls, file_path: str, cache: Any = None):
         return cls.from_config(ClientConfig.from_yaml(file_path), cache=cache)
 
     @classmethod
     def basic(
         cls,
-        model: Union[str, list[str]],
+        model: str | list[str],
         max_requests_per_minute: int = 5_000,
         max_tokens_per_minute: int = 1_000_000,
         max_concurrent_requests: int = 1_000,
         temperature: float = 0.75,
         max_new_tokens: int = 1000,
         reasoning_effort: Literal[None, "low", "medium", "high"] = None,
-        model_weights: Union[list[float], Literal["uniform", "rate_limit"]] = "uniform",
+        model_weights: list[float] | Literal["uniform", "rate_limit"] = "uniform",
         logprobs: bool = False,
-        top_logprobs: Optional[int] = None,
+        top_logprobs: int | None = None,
         max_attempts: int = 5,
         request_timeout: int = 30,
-        cache: Optional[Any] = None,
+        cache: Any = None,
     ):
         model_names = model if isinstance(model, list) else [model]
         return cls(
@@ -222,8 +220,6 @@ class LLMClient:
             top_logprobs=self.top_logprobs,
         )
 
-    from typing import Union, Literal
-
     @overload
     async def process_prompts_async(
         self,
@@ -485,9 +481,9 @@ class LLMClient:
 
 
 def api_prompts_dry_run(
-    ids: Union[np.ndarray, list[int]],
+    ids: np.ndarray | list[int],
     prompts: list[Conversation],
-    models: Union[str, list[str]],
+    models: str | list[str],
     model_weights: list[float],
     sampling_params: list[SamplingParams],
     max_tokens_per_minute: int = 500_000,
@@ -543,19 +539,19 @@ def api_prompts_dry_run(
 
 
 async def process_api_prompts_async(
-    ids: Union[np.ndarray, list[int]],
+    ids: np.ndarray | list[int],
     prompts: list[Conversation],
-    models: Union[str, list[str]],
+    models: str | list[str],
     model_weights: list[float],
     sampling_params: list[SamplingParams],
     logprobs: bool,
-    top_logprobs: Optional[int],
+    top_logprobs: int | None,
     max_attempts: int = 5,
     max_tokens_per_minute: int = 500_000,
     max_requests_per_minute: int = 1_000,
     max_concurrent_requests: int = 1_000,
     request_timeout: int = 30,
-    progress_bar: Optional[tqdm] = None,
+    progress_bar: tqdm | None = None,
     use_qps: bool = False,
     verbose: bool = False,
 ):
@@ -712,28 +708,17 @@ async def process_api_prompts_async(
         await asyncio.sleep(seconds_to_sleep_each_loop)
 
         # if a rate limit error was hit recently, pause to cool down
-        seconds_since_rate_limit_error = (
-            time.time() - status_tracker.time_of_last_rate_limit_error
+        remaining_seconds_to_pause = max(
+            0,
+            seconds_to_pause_after_rate_limit_error
+            - status_tracker.time_since_rate_limit_error,
         )
-        if seconds_since_rate_limit_error < seconds_to_pause_after_rate_limit_error:
-            remaining_seconds_to_pause = (
-                seconds_to_pause_after_rate_limit_error - seconds_since_rate_limit_error
-            )
+        if remaining_seconds_to_pause > 0:
             await asyncio.sleep(remaining_seconds_to_pause)
-            # ^e.g., if pause is 15 seconds and final limit was hit 5 seconds ago
-            print(
-                f"Pausing to cool down until {time.ctime(status_tracker.time_of_last_rate_limit_error + seconds_to_pause_after_rate_limit_error)}"
-            )
+            print(f"Pausing {remaining_seconds_to_pause}s to cool down.")
 
     # after finishing, log final status
-    if status_tracker.num_tasks_failed > 0:
-        print(
-            f"{status_tracker.num_tasks_failed} / {status_tracker.num_tasks_started} requests failed."
-        )
-    if status_tracker.num_rate_limit_errors > 0:
-        print(
-            f"{status_tracker.num_rate_limit_errors} rate limit errors received. Consider running at a lower rate."
-        )
+    status_tracker.log_final_status()
     if verbose:
         print(
             f"After processing, got {len(results)} results for {len(ids)} inputs. Removing duplicates."

{lm_deluge-0.0.3 → lm_deluge-0.0.4}/src/lm_deluge/embed.py

@@ -5,7 +5,7 @@ import aiohttp
 from tqdm.auto import tqdm
 import asyncio
 import time
-from typing import Any, Optional
+from typing import Any
 from dataclasses import dataclass
 from .tracker import StatusTracker
 
@@ -58,7 +58,7 @@ class EmbeddingRequest:
         status_tracker: StatusTracker,
         retry_queue: asyncio.Queue,
         request_timeout: int,
-        pbar: Optional[tqdm] = None,
+        pbar: tqdm | None = None,
         **kwargs,  # openai or cohere specific params
     ):
         self.task_id = task_id
@@ -78,8 +78,7 @@ class EmbeddingRequest:
 
     def handle_success(self):
         self.increment_pbar()
-        self.status_tracker.num_tasks_in_progress -= 1
-        self.status_tracker.num_tasks_succeeded += 1
+        self.status_tracker.task_succeeded(self.task_id)
 
     def handle_error(self):
         last_result: EmbeddingResponse = self.result[-1]
@@ -94,8 +93,7 @@ class EmbeddingRequest:
             return
         else:
             print(f"Task {self.task_id} out of tries.")
-            self.status_tracker.num_tasks_in_progress -= 1
-            self.status_tracker.num_tasks_failed += 1
+            self.status_tracker.task_failed(self.task_id)
 
     async def handle_response(self, response: aiohttp.ClientResponse):
         try:
@@ -217,7 +215,7 @@ class EmbeddingResponse:
     id: int
     status_code: int | None
     is_error: bool
-    error_message: Optional[str]
+    error_message: str | None
     texts: list[str]
     embeddings: list[list[float]]
 
@@ -282,8 +280,7 @@ async def embed_parallel_async(
                         pbar=pbar,
                         **kwargs,
                     )
-                    status_tracker.num_tasks_started += 1
-                    status_tracker.num_tasks_in_progress += 1
+                    status_tracker.start_task(batch_id)
                     results.append(next_request)
 
                 except StopIteration:
@@ -333,29 +330,17 @@ async def embed_parallel_async(
         await asyncio.sleep(seconds_to_sleep_each_loop)
 
         # if a rate limit error was hit recently, pause to cool down
-        seconds_since_rate_limit_error = (
-            time.time() - status_tracker.time_of_last_rate_limit_error
+        remaining_seconds_to_pause = max(
+            0,
+            seconds_to_pause_after_rate_limit_error
+            - status_tracker.time_since_rate_limit_error,
         )
-        if seconds_since_rate_limit_error < seconds_to_pause_after_rate_limit_error:
-            remaining_seconds_to_pause = (
-                seconds_to_pause_after_rate_limit_error - seconds_since_rate_limit_error
-            )
+        if remaining_seconds_to_pause > 0:
             await asyncio.sleep(remaining_seconds_to_pause)
-            # ^e.g., if pause is 15 seconds and final limit was hit 5 seconds ago
-            print(
-                f"Pausing to cool down until {time.ctime(status_tracker.time_of_last_rate_limit_error + seconds_to_pause_after_rate_limit_error)}"
-            )
+            print(f"Pausing {remaining_seconds_to_pause}s to cool down.")
 
     # after finishing, log final status
-    if status_tracker.num_tasks_failed > 0:
-        print(
-            f"{status_tracker.num_tasks_failed} / {status_tracker.num_tasks_started} requests failed."
-        )
-    if status_tracker.num_rate_limit_errors > 0:
-        print(
-            f"{status_tracker.num_rate_limit_errors} rate limit errors received. Consider running at a lower rate."
-        )
-
+    status_tracker.log_final_status()
     print(
         f"After processing, got {len(results)} results for {len(ids)} inputs. Removing duplicates."
     )

{lm_deluge-0.0.3 → lm_deluge-0.0.4}/src/lm_deluge/llm_tools/extract.py

@@ -3,7 +3,7 @@ import json
 from ..prompt import Conversation
 import asyncio
 from ..client import LLMClient
-from typing import Optional, Any
+from typing import Any
 from ..util.json import load_json
 
 try:
@@ -16,8 +16,8 @@ async def extract_async(
     inputs: list[str | Any],
     schema: Any,
     client: LLMClient,
-    document_name: Optional[str] = None,
-    object_name: Optional[str] = None,
+    document_name: str | None = None,
+    object_name: str | None = None,
     show_progress: bool = True,
     return_prompts: bool = False,
 ):
@@ -93,8 +93,8 @@ def extract(
     inputs: list[str | Any],
     schema: Any,
     client: LLMClient,
-    document_name: Optional[str] = None,
-    object_name: Optional[str] = None,
+    document_name: str | None = None,
+    object_name: str | None = None,
     show_progress: bool = True,
     return_prompts: bool = False,
 ):

{lm_deluge-0.0.3 → lm_deluge-0.0.4}/src/lm_deluge/models.py

@@ -1,6 +1,5 @@
 import random
 from dataclasses import dataclass, field
-from typing import Optional
 from .gemini_limits import gemini_1_5_pro_limits, gemini_flash_limits
 
 registry = {
@@ -928,15 +927,15 @@ class APIModel:
     api_base: str
     api_key_env_var: str
     api_spec: str
-    input_cost: Optional[float] = 0  # $ per million input tokens
-    output_cost: Optional[float] = 0  # $ per million output tokens
+    input_cost: float | None = 0  # $ per million input tokens
+    output_cost: float | None = 0  # $ per million output tokens
     supports_json: bool = False
     supports_logprobs: bool = False
     reasoning_model: bool = False
     regions: list[str] | dict[str, int] = field(default_factory=list)
     tokens_per_minute: int | None = None
     requests_per_minute: int | None = None
-    gpus: Optional[list[str]] = None
+    gpus: list[str] | None = None
 
     @classmethod
     def from_registry(cls, name: str):
@@ -950,7 +949,7 @@ class APIModel:
             regions = self.regions
             weights = [1] * len(regions)
         elif isinstance(self.regions, dict):
-            regions = self.regions.keys()
+            regions = list(self.regions.keys())
             weights = self.regions.values()
         else:
             raise ValueError("no regions to sample")

{lm_deluge-0.0.3 → lm_deluge-0.0.4}/src/lm_deluge/rerank.py

@@ -4,7 +4,6 @@ import aiohttp
 from tqdm.auto import tqdm
 import asyncio
 import time
-from typing import Optional
 from dataclasses import dataclass
 from .tracker import StatusTracker
 
@@ -28,7 +27,7 @@ class RerankingRequest:
         status_tracker: StatusTracker,
         retry_queue: asyncio.Queue,
         request_timeout: int,
-        pbar: Optional[tqdm] = None,
+        pbar: tqdm | None = None,
     ):
         self.task_id = task_id
         self.model_name = model_name
@@ -48,8 +47,7 @@ class RerankingRequest:
 
     def handle_success(self):
         self.increment_pbar()
-        self.status_tracker.num_tasks_in_progress -= 1
-        self.status_tracker.num_tasks_succeeded += 1
+        self.status_tracker.task_succeeded(self.task_id)
 
     def handle_error(self):
         """
@@ -69,8 +67,7 @@ class RerankingRequest:
             return
         else:
             print(f"Task {self.task_id} out of tries.")
-            self.status_tracker.num_tasks_in_progress -= 1
-            self.status_tracker.num_tasks_failed += 1
+            self.status_tracker.task_failed(self.task_id)
 
     async def handle_response(self, response: aiohttp.ClientResponse):
         try:
@@ -127,8 +124,9 @@ class RerankingRequest:
         try:
             self.status_tracker.total_requests += 1
             async with aiohttp.ClientSession() as session:
+                timeout = aiohttp.ClientTimeout(total=self.request_timeout)
                 async with session.post(
-                    url, headers=headers, json=data, timeout=self.request_timeout
+                    url, headers=headers, json=data, timeout=timeout
                 ) as response:
                     # print("got response!!")
                     response_obj: RerankingResponse = await self.handle_response(
@@ -176,7 +174,7 @@ class RerankingResponse:
     id: int
     status_code: int | None
     is_error: bool
-    error_message: Optional[str]
+    error_message: str | None
     query: str
     documents: list[str]
     top_k_indices: list[int]
@@ -196,7 +194,7 @@ async def rerank_parallel_async(
     max_requests_per_minute: int = 4_000,
     max_concurrent_requests: int = 500,
     request_timeout: int = 10,
-    progress_bar: Optional[tqdm] = None,
+    progress_bar: tqdm | None = None,
 ):
     """Processes rerank requests in parallel, throttling to stay under rate limits."""
     ids = range(len(queries))
@@ -243,8 +241,7 @@ async def rerank_parallel_async(
                         request_timeout=request_timeout,
                         pbar=progress_bar,
                     )
-                    status_tracker.num_tasks_started += 1
-                    status_tracker.num_tasks_in_progress += 1
+                    status_tracker.start_task(req_id)
                     results.append(next_request)
 
                 except StopIteration:
@@ -294,28 +291,17 @@ async def rerank_parallel_async(
         await asyncio.sleep(seconds_to_sleep_each_loop)
 
         # if a rate limit error was hit recently, pause to cool down
-        seconds_since_rate_limit_error = (
-            time.time() - status_tracker.time_of_last_rate_limit_error
+        remaining_seconds_to_pause = max(
+            0,
+            seconds_to_pause_after_rate_limit_error
+            - status_tracker.time_since_rate_limit_error,
         )
-        if seconds_since_rate_limit_error < seconds_to_pause_after_rate_limit_error:
-            remaining_seconds_to_pause = (
-                seconds_to_pause_after_rate_limit_error - seconds_since_rate_limit_error
-            )
+        if remaining_seconds_to_pause > 0:
             await asyncio.sleep(remaining_seconds_to_pause)
-            # ^e.g., if pause is 15 seconds and final limit was hit 5 seconds ago
-            print(
-                f"Pausing to cool down until {time.ctime(status_tracker.time_of_last_rate_limit_error + seconds_to_pause_after_rate_limit_error)}"
-            )
+            print(f"Pausing {remaining_seconds_to_pause}s to cool down.")
 
     # after finishing, log final status
-    if status_tracker.num_tasks_failed > 0:
-        print(
-            f"{status_tracker.num_tasks_failed} / {status_tracker.num_tasks_started} requests failed."
-        )
-    if status_tracker.num_rate_limit_errors > 0:
-        print(
-            f"{status_tracker.num_rate_limit_errors} rate limit errors received. Consider running at a lower rate."
-        )
+    status_tracker.log_final_status()
 
     print(
         f"After processing, got {len(results)} results for {len(ids)} inputs. Removing duplicates."

lm_deluge-0.0.4/src/lm_deluge/tracker.py

@@ -0,0 +1,43 @@
+import time
+from dataclasses import dataclass
+
+
+@dataclass
+class StatusTracker:
+    num_tasks_started: int = 0
+    num_tasks_in_progress: int = 0
+    num_tasks_succeeded: int = 0
+    num_tasks_failed: int = 0
+    num_rate_limit_errors: int = 0
+    time_of_last_rate_limit_error: int | float = 0
+    total_requests = 0
+
+    @property
+    def time_since_rate_limit_error(self):
+        return time.time() - self.time_of_last_rate_limit_error
+
+    def start_task(self, task_id):
+        self.num_tasks_started += 1
+        self.num_tasks_in_progress += 1
+
+    def rate_limit_exceeded(self):
+        self.time_of_last_rate_limit_error = time.time()
+        self.num_rate_limit_errors += 1
+
+    def task_succeeded(self, task_id):
+        self.num_tasks_in_progress -= 1
+        self.num_tasks_succeeded += 1
+
+    def task_failed(self, task_id):
+        self.num_tasks_in_progress -= 1
+        self.num_tasks_failed += 1
+
+    def log_final_status(self):
+        if self.num_tasks_failed > 0:
+            print(
+                f"{self.num_tasks_failed} / {self.num_tasks_started} requests failed."
+            )
+        if self.num_rate_limit_errors > 0:
+            print(
+                f"{self.num_rate_limit_errors} rate limit errors received. Consider running at a lower rate."
+            )

{lm_deluge-0.0.3 → lm_deluge-0.0.4}/src/lm_deluge/util/logprobs.py

@@ -1,6 +1,6 @@
 import re
 import numpy as np
-from typing import TypedDict, Optional, Callable
+from typing import TypedDict, Callable
 
 
 class TopLogprob(TypedDict):
@@ -403,7 +403,7 @@ def extract_prob(
     normalize_top_logprobs: bool = True,  # if using top_logprobs, normalize by all the present tokens so they add up to 1
     use_complement: bool = False,  # if True, assume there's 2 choices, and return 1 - p if the top token doesn't match
     token_index: int = 0,  # get from the first token of the completion by default
-    token_match_fn: Optional[Callable[[str, str], bool]] = is_match,
+    token_match_fn: Callable[[str, str], bool] | None = is_match,
 ):
     """
     Extract the probability of the token from the logprobs object of a single

{lm_deluge-0.0.3 → lm_deluge-0.0.4}/src/lm_deluge.egg-info/PKG-INFO

@@ -1,9 +1,9 @@
 Metadata-Version: 2.4
 Name: lm_deluge
-Version: 0.0.3
+Version: 0.0.4
 Summary: Python utility for using LLM API models.
 Author-email: Benjamin Anderson <ben@trytaylor.ai>
-Requires-Python: >=3.9
+Requires-Python: >=3.10
 Description-Content-Type: text/markdown
 Requires-Dist: python-dotenv
 Requires-Dist: json5

lm_deluge-0.0.3/src/lm_deluge/tracker.py

@@ -1,12 +0,0 @@
-from dataclasses import dataclass
-
-
-@dataclass
-class StatusTracker:
-    num_tasks_started: int = 0
-    num_tasks_in_progress: int = 0
-    num_tasks_succeeded: int = 0
-    num_tasks_failed: int = 0
-    num_rate_limit_errors: int = 0
-    time_of_last_rate_limit_error: int | float = 0
-    total_requests = 0