lm-deluge 0.0.39__tar.gz → 0.0.41__tar.gz

{lm_deluge-0.0.39/src/lm_deluge.egg-info → lm_deluge-0.0.41}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: lm_deluge
-Version: 0.0.39
+Version: 0.0.41
 Summary: Python utility for using LLM API models.
 Author-email: Benjamin Anderson <ben@trytaylor.ai>
 Requires-Python: >=3.10
@@ -111,14 +111,17 @@ await client.process_prompts_async(
 
 ### Queueing individual prompts
 
-You can queue prompts one at a time and track progress explicitly:
+You can queue prompts one at a time and track progress explicitly. Iterate over
+results as they finish with `as_completed` (or gather them all at once with
+`wait_for_all`):
 
 ```python
 client = LLMClient("gpt-4.1-mini", progress="tqdm")
 client.open()
-task_id = client.start_nowait("hello there")
+client.start_nowait("hello there")
 # ... queue more tasks ...
-results = await client.wait_for_all()
+async for task_id, result in client.as_completed():
+    print(task_id, result.completion)
 client.close()
 ```
 

{lm_deluge-0.0.39 → lm_deluge-0.0.41}/README.md

@@ -84,14 +84,17 @@ await client.process_prompts_async(
 
 ### Queueing individual prompts
 
-You can queue prompts one at a time and track progress explicitly:
+You can queue prompts one at a time and track progress explicitly. Iterate over
+results as they finish with `as_completed` (or gather them all at once with
+`wait_for_all`):
 
 ```python
 client = LLMClient("gpt-4.1-mini", progress="tqdm")
 client.open()
-task_id = client.start_nowait("hello there")
+client.start_nowait("hello there")
 # ... queue more tasks ...
-results = await client.wait_for_all()
+async for task_id, result in client.as_completed():
+    print(task_id, result.completion)
 client.close()
 ```
 

{lm_deluge-0.0.39 → lm_deluge-0.0.41}/pyproject.toml

@@ -3,7 +3,7 @@ requires = ["setuptools", "wheel"]
 
 [project]
 name = "lm_deluge"
-version = "0.0.39"
+version = "0.0.41"
 authors = [{ name = "Benjamin Anderson", email = "ben@trytaylor.ai" }]
 description = "Python utility for using LLM API models."
 readme = "README.md"

{lm_deluge-0.0.39 → lm_deluge-0.0.41}/src/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-0.0.39 → lm_deluge-0.0.41}/src/lm_deluge/client.py

@@ -1,6 +1,6 @@
 import asyncio
 import random
-from typing import Any, Callable, Literal, Self, Sequence, overload
+from typing import Any, AsyncGenerator, Callable, Literal, Self, Sequence, overload
 
 import numpy as np
 import yaml
@@ -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]
@@ -557,6 +558,42 @@ class _LLMClient(BaseModel):
             task_ids = list(self._tasks.keys())
         return [await self.wait_for(tid) for tid in task_ids]
 
+    async def as_completed(
+        self, task_ids: Sequence[int] | None = None
+    ) -> AsyncGenerator[tuple[int, APIResponse | None], None]:
+        """Yield ``(task_id, result)`` pairs as tasks complete.
+
+        Args:
+            task_ids: Optional sequence of task IDs to wait on. If ``None``,
+                all queued tasks are watched.
+
+        Yields:
+            Tuples of task ID and ``APIResponse`` as each task finishes.
+        """
+
+        if task_ids is None:
+            tasks_map: dict[asyncio.Task, int] = {
+                task: tid for tid, task in self._tasks.items()
+            }
+        else:
+            tasks_map = {
+                self._tasks[tid]: tid for tid in task_ids if tid in self._tasks
+            }
+
+        # Yield any tasks that have already completed
+        for task in list(tasks_map.keys()):
+            if task.done():
+                tid = tasks_map.pop(task)
+                yield tid, self._results.get(tid, await task)
+
+        while tasks_map:
+            done, _ = await asyncio.wait(
+                set(tasks_map.keys()), return_when=asyncio.FIRST_COMPLETED
+            )
+            for task in done:
+                tid = tasks_map.pop(task)
+                yield tid, self._results.get(tid, await task)
+
     async def stream(
         self,
         prompt: str | Conversation,

{lm_deluge-0.0.39 → lm_deluge-0.0.41}/src/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.39 → lm_deluge-0.0.41/src/lm_deluge.egg-info}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: lm_deluge
-Version: 0.0.39
+Version: 0.0.41
 Summary: Python utility for using LLM API models.
 Author-email: Benjamin Anderson <ben@trytaylor.ai>
 Requires-Python: >=3.10
@@ -111,14 +111,17 @@ await client.process_prompts_async(
 
 ### Queueing individual prompts
 
-You can queue prompts one at a time and track progress explicitly:
+You can queue prompts one at a time and track progress explicitly. Iterate over
+results as they finish with `as_completed` (or gather them all at once with
+`wait_for_all`):
 
 ```python
 client = LLMClient("gpt-4.1-mini", progress="tqdm")
 client.open()
-task_id = client.start_nowait("hello there")
+client.start_nowait("hello there")
 # ... queue more tasks ...
-results = await client.wait_for_all()
+async for task_id, result in client.as_completed():
+    print(task_id, result.completion)
 client.close()
 ```