lm-deluge 0.0.40__py3-none-any.whl → 0.0.41__py3-none-any.whl

lm_deluge/batches.py

@@ -1,21 +1,22 @@
-import os
+import asyncio
 import json
+import os
+import tempfile
 import time
-import asyncio
+from typing import Literal, Sequence
+
 import aiohttp
-import tempfile
-from lm_deluge.prompt import CachePattern, Conversation, prompts_to_conversations
-from lm_deluge.config import SamplingParams
-from lm_deluge.models import APIModel
-from typing import Sequence, Literal
-from lm_deluge.api_requests.openai import _build_oa_chat_request
-from lm_deluge.api_requests.anthropic import _build_anthropic_request
 from rich.console import Console
 from rich.live import Live
 from rich.spinner import Spinner
 from rich.table import Table
 from rich.text import Text
-from lm_deluge.models import registry
+
+from lm_deluge.api_requests.anthropic import _build_anthropic_request
+from lm_deluge.api_requests.openai import _build_oa_chat_request
+from lm_deluge.config import SamplingParams
+from lm_deluge.models import APIModel, registry
+from lm_deluge.prompt import CachePattern, Conversation, prompts_to_conversations
 from lm_deluge.request_context import RequestContext
 
 
@@ -162,6 +163,91 @@ async def _submit_anthropic_batch(file_path: str, headers: dict, model: str):
         return batch_id
 
 
+async def create_batch_files_oa(
+    model: str,
+    sampling_params: SamplingParams,
+    prompts: Sequence[str | list[dict] | Conversation],
+    batch_size: int = 50_000,
+    destination: str | None = None,  # if none provided, temp files
+):
+    MAX_BATCH_SIZE_BYTES = 200 * 1024 * 1024  # 200MB
+    MAX_BATCH_SIZE_ITEMS = batch_size
+
+    prompts = prompts_to_conversations(prompts)
+    if any(p is None for p in prompts):
+        raise ValueError("All prompts must be valid.")
+
+    model_obj = APIModel.from_registry(model)
+
+    current_batch = []
+    current_batch_size = 0
+    file_paths = []
+
+    for idx, prompt in enumerate(prompts):
+        assert isinstance(prompt, Conversation)
+        context = RequestContext(
+            task_id=idx,
+            model_name=model,
+            prompt=prompt,
+            sampling_params=sampling_params,
+        )
+        request = {
+            "custom_id": str(idx),
+            "method": "POST",
+            "url": "/v1/chat/completions",
+            "body": await _build_oa_chat_request(model_obj, context),
+        }
+
+        # Calculate size of this request
+        request_json = json.dumps(request) + "\n"
+        request_size = len(request_json.encode("utf-8"))
+
+        # Check if adding this request would exceed limits
+        would_exceed_size = current_batch_size + request_size > MAX_BATCH_SIZE_BYTES
+        would_exceed_items = len(current_batch) >= MAX_BATCH_SIZE_ITEMS
+
+        if current_batch and (would_exceed_size or would_exceed_items):
+            # Submit current batch
+            def write_batch_file():
+                with tempfile.NamedTemporaryFile(
+                    mode="w+", suffix=".jsonl", delete=False
+                ) as f:
+                    for batch_request in current_batch:
+                        json.dump(batch_request, f)
+                        f.write("\n")
+                    print("wrote", len(current_batch), "items")
+                    return f.name
+
+            file_path = await asyncio.to_thread(write_batch_file)
+            file_paths.append(file_path)
+            # Start new batch
+            current_batch = []
+            current_batch_size = 0
+            # current_batch_start_idx = idx
+
+        # Add request to current batch
+        current_batch.append(request)
+        current_batch_size += request_size
+
+    # Submit final batch if it has items
+    if current_batch:
+
+        def write_final_batch_file():
+            with tempfile.NamedTemporaryFile(
+                mode="w+", suffix=".jsonl", delete=False
+            ) as f:
+                for batch_request in current_batch:
+                    json.dump(batch_request, f)
+                    f.write("\n")
+                print("wrote", len(current_batch), "items")
+                return f.name
+
+        file_path = await asyncio.to_thread(write_final_batch_file)
+        file_paths.append(file_path)
+
+    return file_paths
+
+
 async def submit_batches_oa(
     model: str,
     sampling_params: SamplingParams,

lm_deluge/client.py

@@ -356,16 +356,16 @@ class _LLMClient(BaseModel):
         prompts = prompts_to_conversations(prompts)
         ids = list(range(len(prompts)))
         results: list[APIResponse | None] = [None for _ in range(len(prompts))]
-
-        # Create StatusTracker
-        tracker = StatusTracker(
-            max_requests_per_minute=self.max_requests_per_minute,
-            max_tokens_per_minute=self.max_tokens_per_minute,
-            max_concurrent_requests=self.max_concurrent_requests,
-            progress_style=self.progress,
-            use_progress_bar=show_progress,
-        )
-        tracker.init_progress_bar(total=len(prompts), disable=not show_progress)
+        # Use existing tracker if client has been opened; otherwise open/close automatically
+        tracker: StatusTracker
+        tracker_preopened = self._tracker is not None
+        if tracker_preopened:
+            tracker = self._tracker  # type: ignore[assignment]
+            tracker.add_to_total(len(prompts))
+        else:
+            self.open(total=len(prompts), show_progress=show_progress)
+            tracker = self._tracker  # type: ignore[assignment]
+        assert tracker is not None
 
         # Create retry queue for failed requests
         retry_queue: asyncio.Queue[RequestContext] = asyncio.Queue()
@@ -458,7 +458,8 @@ class _LLMClient(BaseModel):
             # Sleep - original logic
             await asyncio.sleep(seconds_to_sleep_each_loop + tracker.seconds_to_pause)
 
-        tracker.log_final_status()
+        if not tracker_preopened:
+            self.close()
 
         if return_completions_only:
             return [r.completion if r is not None else None for r in results]

lm_deluge/tracker.py

@@ -14,7 +14,7 @@ from rich.progress import (
     TextColumn,
 )
 from rich.text import Text
-from tqdm import tqdm
+from tqdm.auto import tqdm
 
 SECONDS_TO_PAUSE_AFTER_RATE_LIMIT_ERROR = 5
 

{lm_deluge-0.0.40.dist-info → lm_deluge-0.0.41.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: lm_deluge
-Version: 0.0.40
+Version: 0.0.41
 Summary: Python utility for using LLM API models.
 Author-email: Benjamin Anderson <ben@trytaylor.ai>
 Requires-Python: >=3.10

{lm_deluge-0.0.40.dist-info → lm_deluge-0.0.41.dist-info}/RECORD

@@ -1,9 +1,9 @@
 lm_deluge/__init__.py,sha256=mAztMuxINmh7dGbYnT8tsmw1eryQAvd0jpY8yHzd0EE,315
 lm_deluge/agent.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-lm_deluge/batches.py,sha256=vJXVnuuGkIQnXoDPODPERrvdG9X1Ov1jnXExnPe6ZAc,21772
+lm_deluge/batches.py,sha256=rQocJLyIs3Ko_nRdAE9jT__5cKWYxiIRAH_Lw3L0E1k,24653
 lm_deluge/cache.py,sha256=VB1kv8rM2t5XWPR60uhszFcxLDnVKOe1oA5hYjVDjIo,4375
 lm_deluge/cli.py,sha256=Ilww5gOw3J5v0NReq_Ra4hhxU4BCIJBl1oTGxJZKedc,12065
-lm_deluge/client.py,sha256=ddLk9p4nHMQdDKTV2klPJ9D5Xoj89WepWMHcFFo0we0,34102
+lm_deluge/client.py,sha256=TeaFqAzMKD-oFEfX-L5Qsh-o6G-SKQgPSLWCNhjmMYE,34193
 lm_deluge/config.py,sha256=H1tQyJDNHGFuwxqQNL5Z-CjWAC0luHSBA3iY_pxmACM,932
 lm_deluge/embed.py,sha256=CO-TOlC5kOTAM8lcnicoG4u4K664vCBwHF1vHa-nAGg,13382
 lm_deluge/errors.py,sha256=oHjt7YnxWbh-eXMScIzov4NvpJMo0-2r5J6Wh5DQ1tk,209
@@ -14,7 +14,7 @@ lm_deluge/prompt.py,sha256=quG4dzK2yMiRqPiF67USY5Gl2TqT3rwYkZDJTfwZFHw,37183
 lm_deluge/request_context.py,sha256=o33LSEwnK6YPhZeulUoSE_VrdKCXiCQa0tjjixK2K6M,2540
 lm_deluge/rerank.py,sha256=-NBAJdHz9OB-SWWJnHzkFmeVO4wR6lFV7Vw-SxG7aVo,11457
 lm_deluge/tool.py,sha256=_coOKB9nPNVZoseMRumRyQ8BMR7_d0IlstzMHNT69JY,15732
-lm_deluge/tracker.py,sha256=5ehC-mcLWXpI-OBm6B9W8R4wUegA9OqkRc1PxyotvnY,11526
+lm_deluge/tracker.py,sha256=EHFPsS94NmsON2u97rSE70q1t6pwCsixUmGV-kIphMs,11531
 lm_deluge/usage.py,sha256=VMEKghePFIID5JFBObqYxFpgYxnbYm_dnHy7V1-_T6M,4866
 lm_deluge/api_requests/__init__.py,sha256=AbpHGcgLb-kRsJGnwFEktk7uzpZOCcBY74-YBdrKVGs,1
 lm_deluge/api_requests/anthropic.py,sha256=J5BzYV7aYNoL6FPArB6usyS267z1BguZTRY5JLMd0So,8159
@@ -64,8 +64,8 @@ lm_deluge/util/logprobs.py,sha256=UkBZakOxWluaLqHrjARu7xnJ0uCHVfLGHJdnYlEcutk,11
 lm_deluge/util/spatial.py,sha256=BsF_UKhE-x0xBirc-bV1xSKZRTUhsOBdGqsMKme20C8,4099
 lm_deluge/util/validation.py,sha256=hz5dDb3ebvZrZhnaWxOxbNSVMI6nmaOODBkk0htAUhs,1575
 lm_deluge/util/xml.py,sha256=Ft4zajoYBJR3HHCt2oHwGfymGLdvp_gegVmJ-Wqk4Ck,10547
-lm_deluge-0.0.40.dist-info/licenses/LICENSE,sha256=uNNXGXPCw2TC7CUs7SEBkA-Mz6QBQFWUUEWDMgEs1dU,1058
-lm_deluge-0.0.40.dist-info/METADATA,sha256=QMWHzRuDQohIf6ZWjWpajRmQhAh2aJqEqAnXgVZ613s,13443
-lm_deluge-0.0.40.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-lm_deluge-0.0.40.dist-info/top_level.txt,sha256=hqU-TJX93yBwpgkDtYcXyLr3t7TLSCCZ_reytJjwBaE,10
-lm_deluge-0.0.40.dist-info/RECORD,,
+lm_deluge-0.0.41.dist-info/licenses/LICENSE,sha256=uNNXGXPCw2TC7CUs7SEBkA-Mz6QBQFWUUEWDMgEs1dU,1058
+lm_deluge-0.0.41.dist-info/METADATA,sha256=JQLcTcp79Lqe5WPz8F2JPTZNVIWi5Z0NHoXakPkvzfw,13443
+lm_deluge-0.0.41.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+lm_deluge-0.0.41.dist-info/top_level.txt,sha256=hqU-TJX93yBwpgkDtYcXyLr3t7TLSCCZ_reytJjwBaE,10
+lm_deluge-0.0.41.dist-info/RECORD,,