lm-deluge 0.0.13__py3-none-any.whl → 0.0.15__py3-none-any.whl

lm_deluge/client.py

@@ -6,6 +6,7 @@ import yaml
 from pydantic import BaseModel
 from pydantic.functional_validators import model_validator
 
+from lm_deluge.api_requests.openai import stream_chat
 from lm_deluge.batches import (
     submit_batches_anthropic,
     submit_batches_oa,
@@ -34,6 +35,12 @@ class LLMClient(BaseModel):
     """
 
     model_names: list[str] = ["gpt-4.1-mini"]
+
+    def __init__(self, model_name: str | list[str] | None = None, **kwargs):
+        if model_name is not None:
+            kwargs["model_names"] = model_name
+        super().__init__(**kwargs)
+
     max_requests_per_minute: int = 1_000
     max_tokens_per_minute: int = 100_000
     max_concurrent_requests: int = 225
@@ -81,7 +88,7 @@ class LLMClient(BaseModel):
     @model_validator(mode="before")
     @classmethod
     def fix_lists(cls, data) -> "LLMClient":
-        if isinstance(data["model_names"], str):
+        if isinstance(data.get("model_names"), str):
             data["model_names"] = [data["model_names"]]
         if "sampling_params" not in data or len(data.get("sampling_params", [])) == 0:
             data["sampling_params"] = [
@@ -162,6 +169,11 @@ class LLMClient(BaseModel):
         kwargs["model_names"] = model
         return cls(**kwargs)
 
+    def _select_model(self):
+        assert isinstance(self.model_weights, list)
+        model_idx = np.random.choice(range(len(self.models)), p=self.model_weights)
+        return self.models[model_idx], self.sampling_params[model_idx]
+
     @overload
     async def process_prompts_async(
         self,
@@ -249,41 +261,6 @@ class LLMClient(BaseModel):
             if len(cache_hit_ids) > 0:
                 tracker.update_pbar(len(cache_hit_ids))
 
-            # api_task = asyncio.create_task(
-            #     process_api_prompts_async(
-            #         ids,
-            #         prompts,  # type: ignore -- fix later for dry running conversations
-            #         self.models,
-            #         self.model_weights,  # type: ignore
-            #         self.sampling_params,  # type: ignore
-            #         max_attempts=self.max_attempts,
-            #         max_concurrent_requests=self.max_concurrent_requests,
-            #         request_timeout=self.request_timeout,
-            #         status_tracker=tracker,
-            #         tools=tools,
-            #         cache=cache,
-            #         computer_use=computer_use,
-            #         display_width=display_width,
-            #         display_height=display_height,
-            #         use_responses_api=use_responses_api,
-            #     )
-            # )
-            # async def process_api_prompts_async(
-
-            #     models: str | list[str],
-            #     model_weights: list[float],
-            #     sampling_params: list[SamplingParams],
-            #     max_attempts: int = 5,
-            #     max_concurrent_requests: int = 1_000,
-            #     request_timeout: int = 30,
-            #     status_tracker: StatusTracker | None = None,
-            #     tools: list[Tool] | None = None,
-            #     cache: CachePattern | None = None,
-            #     computer_use: bool = False,
-            #     display_width: int = 1024,
-            #     display_height: int = 768,
-            #     use_responses_api: bool = False,
-            # ):
             if isinstance(ids, np.ndarray):
                 ids = ids.tolist()  # pyright: ignore
 
@@ -296,28 +273,28 @@ class LLMClient(BaseModel):
             assert tracker.retry_queue, "retry queue not initialized"
             while True:
                 # get next request (if one is not already waiting for capacity)
+                retry_request = False
                 if next_request is None:
                     if not tracker.retry_queue.empty():
                         next_request = tracker.retry_queue.get_nowait()
+                        retry_request = True
                         print(f"Retrying request {next_request.task_id}.")
                     elif prompts_not_finished:
                         try:
                             # get new request
                             id, prompt = next(prompts_iter)
                             # select model
-                            assert isinstance(self.model_weights, list)
-                            model_idx = np.random.choice(
-                                range(len(self.models)), p=self.model_weights
-                            )
+                            model, sampling_params = self._select_model()
+
                             next_request = create_api_request(
                                 task_id=id,
-                                model_name=self.models[model_idx],
+                                model_name=model,
                                 prompt=prompt,  # type: ignore
                                 request_timeout=self.request_timeout,
                                 attempts_left=self.max_attempts,
                                 status_tracker=tracker,
                                 results_arr=requests,
-                                sampling_params=self.sampling_params[model_idx],
+                                sampling_params=sampling_params,
                                 all_model_names=self.models,
                                 all_sampling_params=self.sampling_params,
                                 tools=tools,
@@ -339,10 +316,9 @@ class LLMClient(BaseModel):
                 # if enough capacity available, call API
                 if next_request:
                     next_request_tokens = next_request.num_tokens
-                    if tracker.check_capacity(next_request_tokens):
+                    if tracker.check_capacity(next_request_tokens, retry=retry_request):
                         tracker.set_limiting_factor(None)
-                        next_request.attempts_left -= 1
-                        # call API
+                        # call API (attempts_left will be decremented in handle_error if it fails)
                         asyncio.create_task(next_request.call_api())
                         next_request = None  # reset next_request to empty
                 # update pbar status
@@ -360,9 +336,10 @@ class LLMClient(BaseModel):
                     await asyncio.sleep(tracker.seconds_to_pause)
                     print(f"Pausing {tracker.seconds_to_pause}s to cool down.")
 
-                # after finishing, log final status
-                tracker.log_final_status()
-                # deduplicate results by id
+            # after finishing, log final status
+            tracker.log_final_status()
+
+            # deduplicate results by id
             api_results = deduplicate_responses(requests)
             for res in api_results:
                 results[res.id] = res
@@ -399,6 +376,17 @@ class LLMClient(BaseModel):
             )
         )
 
+    async def stream(self, prompt: str | Conversation, tools: list[Tool] | None = None):
+        model, sampling_params = self._select_model()
+        if isinstance(prompt, str):
+            prompt = Conversation.user(prompt)
+        async for item in stream_chat(model, prompt, sampling_params, tools, None):
+            if isinstance(item, str):
+                print(item, end="", flush=True)
+            else:
+                # final item
+                return item
+
     async def submit_batch_job(
         self,
         prompts: Sequence[str | list[dict] | Conversation],

lm_deluge/config.py

@@ -1,13 +1,14 @@
-from pydantic import BaseModel
 from typing import Literal
 
+from pydantic import BaseModel
+
 
 class SamplingParams(BaseModel):
     temperature: float = 0.0
     top_p: float = 1.0
     json_mode: bool = False
     max_new_tokens: int = 512
-    reasoning_effort: Literal["low", "medium", "high", None] = None
+    reasoning_effort: Literal["low", "medium", "high", "none", None] = None
     logprobs: bool = False
     top_logprobs: int | None = None
 

lm_deluge/file.py

@@ -0,0 +1,154 @@
+import os
+import io
+import requests
+import base64
+import mimetypes
+import xxhash
+from dataclasses import dataclass, field
+from pathlib import Path
+
+
+@dataclass(slots=True)
+class File:
+    # raw bytes, pathlike, http url, base64 data url, or file_id
+    data: bytes | io.BytesIO | Path | str
+    media_type: str | None = None  # inferred if None
+    filename: str | None = None  # optional filename for uploads
+    file_id: str | None = None  # for OpenAI file uploads or Anthropic file API
+    type: str = field(init=False, default="file")
+
+    # helpers -----------------------------------------------------------------
+    def _bytes(self) -> bytes:
+        if isinstance(self.data, bytes):
+            return self.data
+        elif isinstance(self.data, io.BytesIO):
+            return self.data.getvalue()
+        elif isinstance(self.data, str) and self.data.startswith("http"):
+            res = requests.get(self.data)
+            res.raise_for_status()
+            return res.content
+        elif isinstance(self.data, str) and os.path.exists(self.data):
+            with open(self.data, "rb") as f:
+                return f.read()
+        elif isinstance(self.data, Path) and self.data.exists():
+            return Path(self.data).read_bytes()
+        elif isinstance(self.data, str) and self.data.startswith("data:"):
+            header, encoded = self.data.split(",", 1)
+            return base64.b64decode(encoded)
+        else:
+            raise ValueError("unreadable file format")
+
+    def _mime(self) -> str:
+        if self.media_type:
+            return self.media_type
+        if isinstance(self.data, (Path, str)):
+            # For URL or path, try to guess from the string
+            path_str = str(self.data)
+            guess = mimetypes.guess_type(path_str)[0]
+            if guess:
+                return guess
+        return "application/pdf"  # default to PDF
+
+    def _filename(self) -> str:
+        if self.filename:
+            return self.filename
+        if isinstance(self.data, (Path, str)):
+            path_str = str(self.data)
+            if path_str.startswith("http"):
+                # Extract filename from URL
+                return path_str.split("/")[-1].split("?")[0] or "document.pdf"
+            else:
+                # Extract from local path
+                return os.path.basename(path_str) or "document.pdf"
+        return "document.pdf"
+
+    def _base64(self, include_header: bool = True) -> str:
+        encoded = base64.b64encode(self._bytes()).decode("utf-8")
+        if not include_header:
+            return encoded
+        return f"data:{self._mime()};base64,{encoded}"
+
+    @property
+    def fingerprint(self) -> str:
+        # Hash the file contents for fingerprinting
+        file_bytes = self._bytes()
+        return xxhash.xxh64(file_bytes).hexdigest()
+
+    @property
+    def size(self) -> int:
+        """Return file size in bytes."""
+        return len(self._bytes())
+
+    # ── provider-specific emission ────────────────────────────────────────────
+    def oa_chat(self) -> dict:
+        """For OpenAI Chat Completions - file content as base64 or file_id."""
+        if self.file_id:
+            return {
+                "type": "file",
+                "file": {
+                    "file_id": self.file_id,
+                },
+            }
+        else:
+            return {
+                "type": "file",
+                "file": {
+                    "filename": self._filename(),
+                    "file_data": self._base64(),
+                },
+            }
+
+    def oa_resp(self) -> dict:
+        """For OpenAI Responses API - file content as base64 or file_id."""
+        if self.file_id:
+            return {
+                "type": "input_file",
+                "file_id": self.file_id,
+            }
+        else:
+            return {
+                "type": "input_file",
+                "filename": self._filename(),
+                "file_data": self._base64(),
+            }
+
+    def anthropic(self) -> dict:
+        """For Anthropic Messages API - file content as base64 or file_id."""
+        if self.file_id:
+            return {
+                "type": "document",
+                "source": {
+                    "type": "file",
+                    "file_id": self.file_id,
+                },
+            }
+        else:
+            b64 = base64.b64encode(self._bytes()).decode()
+            return {
+                "type": "document",
+                "source": {
+                    "type": "base64",
+                    "media_type": self._mime(),
+                    "data": b64,
+                },
+            }
+
+    def anthropic_file_upload(self) -> tuple[str, bytes, str]:
+        """For Anthropic Files API - return tuple for file upload."""
+        filename = self._filename()
+        content = self._bytes()
+        media_type = self._mime()
+        return filename, content, media_type
+
+    def gemini(self) -> dict:
+        """For Gemini API - files are provided as inline data."""
+        return {
+            "inlineData": {
+                "mimeType": self._mime(),
+                "data": self._base64(include_header=False),
+            }
+        }
+
+    def mistral(self) -> dict:
+        """For Mistral API - not yet supported."""
+        raise NotImplementedError("File support for Mistral is not yet implemented")

lm_deluge/models.py

@@ -167,6 +167,63 @@ registry = {
         "tokens_per_minute": 100_000,
         "reasoning_model": True,
     },
+    # Native Gemini API versions with file support
+    "gemini-2.0-flash-gemini": {
+        "id": "gemini-2.0-flash-gemini",
+        "name": "gemini-2.0-flash",
+        "api_base": "https://generativelanguage.googleapis.com/v1beta",
+        "api_key_env_var": "GEMINI_API_KEY",
+        "supports_json": True,
+        "supports_logprobs": False,
+        "api_spec": "gemini",
+        "input_cost": 0.1,
+        "output_cost": 0.4,
+        "requests_per_minute": 20,
+        "tokens_per_minute": 100_000,
+        "reasoning_model": False,
+    },
+    "gemini-2.0-flash-lite-gemini": {
+        "id": "gemini-2.0-flash-lite-gemini",
+        "name": "gemini-2.0-flash-lite",
+        "api_base": "https://generativelanguage.googleapis.com/v1beta",
+        "api_key_env_var": "GEMINI_API_KEY",
+        "supports_json": True,
+        "supports_logprobs": False,
+        "api_spec": "gemini",
+        "input_cost": 0.1,
+        "output_cost": 0.4,
+        "requests_per_minute": 20,
+        "tokens_per_minute": 100_000,
+        "reasoning_model": False,
+    },
+    "gemini-2.5-pro-gemini": {
+        "id": "gemini-2.5-pro-gemini",
+        "name": "gemini-2.5-pro-preview-05-06",
+        "api_base": "https://generativelanguage.googleapis.com/v1beta",
+        "api_key_env_var": "GEMINI_API_KEY",
+        "supports_json": True,
+        "supports_logprobs": False,
+        "api_spec": "gemini",
+        "input_cost": 0.1,
+        "output_cost": 0.4,
+        "requests_per_minute": 20,
+        "tokens_per_minute": 100_000,
+        "reasoning_model": True,
+    },
+    "gemini-2.5-flash-gemini": {
+        "id": "gemini-2.5-flash-gemini",
+        "name": "gemini-2.5-flash-preview-05-20",
+        "api_base": "https://generativelanguage.googleapis.com/v1beta",
+        "api_key_env_var": "GEMINI_API_KEY",
+        "supports_json": True,
+        "supports_logprobs": False,
+        "api_spec": "gemini",
+        "input_cost": 0.1,
+        "output_cost": 0.4,
+        "requests_per_minute": 20,
+        "tokens_per_minute": 100_000,
+        "reasoning_model": True,
+    },
     #     ███████                                    █████████   █████
     #   ███░░░░░███                                 ███░░░░░███ ░░███
     #  ███     ░░███ ████████   ██████  ████████   ░███    ░███  ░███

lm_deluge/prompt.py

@@ -1,12 +1,15 @@
 import io
 import json
-import tiktoken
-import xxhash
 from dataclasses import dataclass, field
 from pathlib import Path
 from typing import Literal, Sequence
-from lm_deluge.models import APIModel
+
+import tiktoken
+import xxhash
+
+from lm_deluge.file import File
 from lm_deluge.image import Image
+from lm_deluge.models import APIModel
 
 CachePattern = Literal[
     "tools_only",
@@ -203,7 +206,7 @@ class Thinking:
         return {"type": "text", "text": f"[Thinking: {self.content}]"}
 
 
-Part = Text | Image | ToolCall | ToolResult | Thinking
+Part = Text | Image | File | ToolCall | ToolResult | Thinking
 
 
 ###############################################################################
@@ -246,6 +249,11 @@ class Message:
         """Get all image parts with proper typing."""
         return [part for part in self.parts if part.type == "image"]  # type: ignore
 
+    @property
+    def files(self) -> list[File]:
+        """Get all file parts with proper typing."""
+        return [part for part in self.parts if part.type == "file"]  # type: ignore
+
     @property
     def thinking_parts(self) -> list["Thinking"]:
         """Get all thinking parts with proper typing."""
@@ -262,6 +270,9 @@ class Message:
             elif isinstance(p, Image):  # Image – redact the bytes, keep a hint
                 w, h = p.size
                 content_blocks.append({"type": "image", "tag": f"<Image ({w}×{h})>"})
+            elif isinstance(p, File):  # File – redact the bytes, keep a hint
+                size = p.size
+                content_blocks.append({"type": "file", "tag": f"<File ({size} bytes)>"})
             elif isinstance(p, ToolCall):
                 content_blocks.append(
                     {
@@ -296,6 +307,9 @@ class Message:
             elif p["type"] == "image":
                 # We only stored a placeholder tag, so keep that placeholder.
                 parts.append(Image(p["tag"], detail="low"))
+            elif p["type"] == "file":
+                # We only stored a placeholder tag, so keep that placeholder.
+                parts.append(File(p["tag"]))
             elif p["type"] == "tool_call":
                 parts.append(
                     ToolCall(id=p["id"], name=p["name"], arguments=p["arguments"])
@@ -340,6 +354,20 @@ class Message:
         self.parts.append(img)
         return self
 
+    def add_file(
+        self,
+        data: bytes | str | Path | io.BytesIO,
+        *,
+        media_type: str | None = None,
+        filename: str | None = None,
+    ) -> "Message":
+        """
+        Append a file block and return self for chaining.
+        """
+        file = File(data, media_type=media_type, filename=filename)
+        self.parts.append(file)
+        return self
+
     def add_tool_call(self, id: str, name: str, arguments: dict) -> "Message":
         """Append a tool call block and return self for chaining."""
         self.parts.append(ToolCall(id=id, name=name, arguments=arguments))
@@ -362,12 +390,15 @@ class Message:
         text: str | None = None,
         *,
         image: str | bytes | Path | io.BytesIO | None = None,
+        file: str | bytes | Path | io.BytesIO | None = None,
     ) -> "Message":
         res = cls("user", [])
         if text is not None:
             res.add_text(text)
         if image is not None:
             res.add_image(image)
+        if file is not None:
+            res.add_file(file)
         return res
 
     @classmethod
@@ -403,6 +434,19 @@ class Message:
                     part_list.append(Text(item["text"]))
                 elif item["type"] == "image_url":
                     part_list.append(Image(data=item["image_url"]["url"]))
+                elif item["type"] == "file":
+                    file_data = item["file"]
+                    if "file_id" in file_data:
+                        # Handle file ID reference (not implemented yet)
+                        part_list.append(File(data=file_data["file_id"]))
+                    elif "file_data" in file_data:
+                        # Handle base64 file data
+                        part_list.append(
+                            File(
+                                data=file_data["file_data"],
+                                filename=file_data.get("filename"),
+                            )
+                        )
             parts = part_list
 
         # Handle tool calls (assistant messages)
@@ -511,11 +555,17 @@ class Conversation:
 
     @classmethod
     def user(
-        cls, text: str, *, image: bytes | str | Path | None = None
+        cls,
+        text: str,
+        *,
+        image: bytes | str | Path | None = None,
+        file: bytes | str | Path | None = None,
     ) -> "Conversation":
-        msg = (
-            Message.user(text) if image is None else Message.user(text).add_image(image)
-        )
+        msg = Message.user(text)
+        if image is not None:
+            msg.add_image(image)
+        if file is not None:
+            msg.add_file(file)
         return cls([msg])
 
     @classmethod
@@ -677,6 +727,9 @@ class Conversation:
                 if isinstance(part, Image):
                     # Force conversion to bytes if not already
                     part.data = part._bytes()
+                elif isinstance(part, File):
+                    # Force conversion to bytes if not already
+                    part.data = part._bytes()
         return self
 
     def _add_cache_control_to_message(self, message: dict) -> None:
@@ -765,6 +818,11 @@ class Conversation:
                     content_blocks.append(
                         {"type": "image", "tag": f"<Image ({w}×{h})>"}
                     )
+                elif isinstance(p, File):  # File – redact the bytes, keep a hint
+                    size = p.size
+                    content_blocks.append(
+                        {"type": "file", "tag": f"<File ({size} bytes)>"}
+                    )
                 elif isinstance(p, ToolCall):
                     content_blocks.append(
                         {
@@ -795,7 +853,7 @@ class Conversation:
 
         for m in payload.get("messages", []):
             role: Role = m["role"]  # 'system' | 'user' | 'assistant'
-            parts: list[Text | Image | ToolCall | ToolResult | Thinking] = []
+            parts: list[Part] = []
 
             for p in m["content"]:
                 if p["type"] == "text":
@@ -804,6 +862,9 @@ class Conversation:
                     # We only stored a placeholder tag, so keep that placeholder.
                     # You could raise instead if real image bytes are required.
                     parts.append(Image(p["tag"], detail="low"))
+                elif p["type"] == "file":
+                    # We only stored a placeholder tag, so keep that placeholder.
+                    parts.append(File(p["tag"]))
                 elif p["type"] == "tool_call":
                     parts.append(
                         ToolCall(id=p["id"], name=p["name"], arguments=p["arguments"])

lm_deluge/tracker.py

@@ -67,7 +67,7 @@ class StatusTracker:
     def set_limiting_factor(self, factor):
         self.limiting_factor = factor
 
-    def check_capacity(self, num_tokens: int):
+    def check_capacity(self, num_tokens: int, retry: bool = False):
         request_available = self.available_request_capacity >= 1
         tokens_available = self.available_token_capacity >= num_tokens
         concurrent_request_available = (
@@ -76,8 +76,10 @@ class StatusTracker:
         if request_available and tokens_available and concurrent_request_available:
             self.available_request_capacity -= 1
             self.available_token_capacity -= num_tokens
-            self.num_tasks_started += 1
-            self.num_tasks_in_progress += 1
+            if not retry:
+                # Only count new tasks, not retries
+                self.num_tasks_started += 1
+                self.num_tasks_in_progress += 1
             self.set_limiting_factor(None)
             return True
         else:

lm_deluge/usage.py

@@ -71,6 +71,16 @@ class Usage:
             cache_write_tokens=None,
         )
 
+    @classmethod
+    def from_gemini_usage(cls, usage_data: dict) -> "Usage":
+        """Create Usage from Gemini API response usage data."""
+        return cls(
+            input_tokens=usage_data.get("promptTokenCount", 0),
+            output_tokens=usage_data.get("candidatesTokenCount", 0),
+            cache_read_tokens=None,  # Gemini doesn't support caching yet
+            cache_write_tokens=None,
+        )
+
     def to_dict(self) -> dict:
         """Convert to dictionary for serialization."""
         return {

{lm_deluge-0.0.13.dist-info → lm_deluge-0.0.15.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: lm_deluge
-Version: 0.0.13
+Version: 0.0.15
 Summary: Python utility for using LLM API models.
 Author-email: Benjamin Anderson <ben@trytaylor.ai>
 Requires-Python: >=3.10
@@ -30,6 +30,7 @@ Dynamic: license-file
 `lm-deluge` is a lightweight helper library for maxing out your rate limits with LLM providers. It provides the following:
 
 - **Unified client** – Send prompts to all relevant models with a single client.
+- **Files and Images** - Include images easily for multimodal models, and PDF files for models that support them (OpenAI and Anthropic).
 - **Massive concurrency with throttling** – Set `max_tokens_per_minute` and `max_requests_per_minute` and let it fly. The client will process as many requests as possible while respecting rate limits and retrying failures.
 - **Spray across models/providers** – Configure a client with multiple models from any provider(s), and sampling weights. The client samples a model for each request.
 - **Tool Use** – Unified API for defining tools for all providers, and creating tools automatically from python functions.
@@ -41,6 +42,8 @@ Dynamic: license-file
 
 **STREAMING IS NOT IN SCOPE.** There are plenty of packages that let you stream chat completions across providers. The sole purpose of this package is to do very fast batch inference using APIs. Sorry!
 
+**Update 06/02/2025:** I lied, it supports (very basic) streaming now via client.stream(...). It will print tokens as they arrive, then return an APIResponse at the end. More sophisticated streaming may or may not be implemented later, don't count on it.
+
 ## Installation
 
 ```bash