lm-deluge 0.0.12__py3-none-any.whl → 0.0.14__py3-none-any.whl

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
-from lm_deluge.models import APIModel
+from typing import Literal, Sequence
+
+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",
@@ -99,24 +102,58 @@ class ToolCall:
 @dataclass(slots=True)
 class ToolResult:
     tool_call_id: str  # references the ToolCall.id
-    result: str  # tool execution result
+    result: (
+        str | dict | list[dict]
+    )  # tool execution result - can be string or list for images
     type: str = field(init=False, default="tool_result")
 
     @property
     def fingerprint(self) -> str:
-        return xxhash.xxh64(f"{self.tool_call_id}:{self.result}".encode()).hexdigest()
+        result_str = (
+            json.dumps(self.result, sort_keys=True)
+            if isinstance(self.result, list) or isinstance(self.result, dict)
+            else str(self.result)
+        )
+        return xxhash.xxh64(f"{self.tool_call_id}:{result_str}".encode()).hexdigest()
 
     # ── provider-specific emission ────────────────────────────────────────────
     def oa_chat(
         self,
     ) -> dict:  # OpenAI Chat Completions - tool results are separate messages
-        return {"tool_call_id": self.tool_call_id, "content": self.result}
+        content = (
+            json.dumps(self.result) if isinstance(self.result, list) else self.result
+        )
+        return {"tool_call_id": self.tool_call_id, "content": content}
 
     def oa_resp(self) -> dict:  # OpenAI Responses
+        # Check if this is a computer use output (special case)
+        if isinstance(self.result, dict) and self.result.get("_computer_use_output"):
+            # This is a computer use output, emit it properly
+            output_data = self.result.copy()
+            output_data.pop("_computer_use_output")  # Remove marker
+
+            result = {
+                "type": "computer_call_output",
+                "call_id": self.tool_call_id,
+                "output": output_data.get("output", {}),
+            }
+
+            # Add acknowledged safety checks if present
+            if "acknowledged_safety_checks" in output_data:
+                result["acknowledged_safety_checks"] = output_data[
+                    "acknowledged_safety_checks"
+                ]
+
+            return result
+
+        # Regular function result
+        result = (
+            json.dumps(self.result) if isinstance(self.result, list) else self.result
+        )
         return {
             "type": "function_result",
             "call_id": self.tool_call_id,
-            "result": self.result,
+            "result": result,
         }
 
     def anthropic(self) -> dict:  # Anthropic Messages
@@ -169,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
 
 
 ###############################################################################
@@ -212,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."""
@@ -228,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(
                     {
@@ -262,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"])
@@ -306,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))
@@ -328,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
@@ -369,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)
@@ -428,6 +506,14 @@ class Message:
 
     def oa_resp(self) -> dict:
         content = [p.oa_resp() for p in self.parts]
+        # For OpenAI Responses API, handle tool results specially
+        if self.role == "tool" or (
+            self.role == "user" and any(isinstance(p, ToolResult) for p in self.parts)
+        ):
+            # Tool results are returned directly, not wrapped in a message
+            # This handles computer_call_output when stored as ToolResult
+            if len(self.parts) == 1 and isinstance(self.parts[0], ToolResult):
+                return self.parts[0].oa_resp()
         return {"role": self.role, "content": content}
 
     def anthropic(self) -> dict:
@@ -469,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
@@ -522,7 +614,37 @@ class Conversation:
 
     def to_openai_responses(self) -> dict:
         # OpenAI Responses = single “input” array, role must be user/assistant
-        return {"input": [m.oa_resp() for m in self.messages if m.role != "system"]}
+        input_items = []
+
+        for m in self.messages:
+            if m.role == "system":
+                continue
+            elif m.role == "assistant":
+                # For assistant messages, extract computer calls as separate items
+                text_parts = []
+                for p in m.parts:
+                    if isinstance(p, ToolCall) and p.name.startswith("_computer_"):
+                        # Computer calls become separate items in the input array
+                        action_type = p.name.replace("_computer_", "")
+                        input_items.append(
+                            {
+                                "type": "computer_call",
+                                "call_id": p.id,
+                                "action": {"type": action_type, **p.arguments},
+                            }
+                        )
+                    elif isinstance(p, Text):
+                        text_parts.append({"type": "output_text", "text": p.text})
+                    # TODO: Handle other part types as needed
+
+                # Add message if it has text content
+                if text_parts:
+                    input_items.append({"role": m.role, "content": text_parts})
+            else:
+                # User and tool messages use normal format
+                input_items.append(m.oa_resp())
+
+        return {"input": input_items}
 
     def to_anthropic(
         self, cache_pattern: CachePattern | None = None
@@ -605,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:
@@ -693,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(
                         {
@@ -723,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":
@@ -732,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"])
@@ -750,6 +883,14 @@ class Conversation:
         return cls(msgs)
 
 
+def prompts_to_conversations(prompts: Sequence[str | list[dict] | Conversation]):
+    if any(isinstance(x, list) for x in prompts):
+        raise ValueError("can't convert list[dict] to conversation yet")
+    return [  # type: ignore
+        Conversation.user(p) if isinstance(p, str) else p for p in prompts
+    ]
+
+
 ###############################################################################
 # --------------------------------------------------------------------------- #
 # Basic usage examples                                                        #

lm_deluge/rerank.py

@@ -1,10 +1,12 @@
 ### specific utility for cohere rerank api
-import os
-import aiohttp
-from tqdm.auto import tqdm
 import asyncio
+import os
 import time
 from dataclasses import dataclass
+
+import aiohttp
+from tqdm.auto import tqdm
+
 from .tracker import StatusTracker
 
 registry = [
@@ -25,7 +27,6 @@ class RerankingRequest:
         top_k: int,
         attempts_left: int,
         status_tracker: StatusTracker,
-        retry_queue: asyncio.Queue,
         request_timeout: int,
         pbar: tqdm | None = None,
     ):
@@ -36,7 +37,6 @@ class RerankingRequest:
         self.top_k = top_k
         self.attempts_left = attempts_left
         self.status_tracker = status_tracker
-        self.retry_queue = retry_queue
         self.request_timeout = request_timeout
         self.pbar = pbar
         self.result = []
@@ -63,7 +63,8 @@ class RerankingRequest:
         print(error_to_print)
         if self.attempts_left > 0:
             self.attempts_left -= 1
-            self.retry_queue.put_nowait(self)
+            assert self.status_tracker.retry_queue
+            self.status_tracker.retry_queue.put_nowait(self)
             return
         else:
             print(f"Task {self.task_id} out of tries.")
@@ -203,8 +204,12 @@ async def rerank_parallel_async(
     seconds_to_sleep_each_loop = 0.003  # so concurrent tasks can run
 
     # initialize trackers
-    retry_queue = asyncio.Queue()
-    status_tracker = StatusTracker()
+    # retry_queue = asyncio.Queue()
+    status_tracker = StatusTracker(
+        max_tokens_per_minute=10_000_000,
+        max_requests_per_minute=max_requests_per_minute,
+        max_concurrent_requests=1_000,
+    )
     next_request = None  # variable to hold the next request to call
 
     # initialize available capacity counts
@@ -222,8 +227,10 @@ async def rerank_parallel_async(
     while True:
         # get next request (if one is not already waiting for capacity)
         if next_request is None:
-            if not retry_queue.empty():
-                next_request = retry_queue.get_nowait()
+            assert status_tracker.retry_queue
+
+            if not status_tracker.retry_queue.empty():
+                next_request = status_tracker.retry_queue.get_nowait()
                 print(f"Retrying request {next_request.task_id}.")
             elif prompts_not_finished:
                 try:
@@ -237,7 +244,6 @@ async def rerank_parallel_async(
                         top_k=top_k,
                         attempts_left=max_attempts,
                         status_tracker=status_tracker,
-                        retry_queue=retry_queue,
                         request_timeout=request_timeout,
                         pbar=progress_bar,
                     )
@@ -246,7 +252,7 @@ async def rerank_parallel_async(
 
                 except StopIteration:
                     prompts_not_finished = False
-                    print("API requests finished, only retries remain.")
+                    # print("API requests finished, only retries remain.")
 
         # update available capacity
         current_time = time.time()

lm_deluge/tool.py

@@ -4,7 +4,7 @@ import asyncio
 
 from fastmcp import Client  # pip install fastmcp >= 2.0
 from mcp.types import Tool as MCPTool
-from pydantic import BaseModel, Field
+from pydantic import BaseModel, Field, field_validator
 
 
 async def _load_all_mcp_tools(client: Client) -> list["Tool"]:
@@ -46,6 +46,16 @@ class Tool(BaseModel):
     # if desired, can provide a callable to run the tool
     run: Callable | None = None
 
+    @field_validator("name")
+    @classmethod
+    def validate_name(cls, v: str) -> str:
+        if v.startswith("_computer_"):
+            raise ValueError(
+                f"Tool name '{v}' uses reserved prefix '_computer_'. "
+                "This prefix is reserved for computer use actions."
+            )
+        return v
+
     def _is_async(self) -> bool:
         return inspect.iscoroutinefunction(self.run)
 

lm_deluge/tracker.py

@@ -1,21 +1,111 @@
+import asyncio
 import time
-from dataclasses import dataclass
+from dataclasses import dataclass, field
+
+from rich.console import Console, Group
+from rich.live import Live
+from rich.progress import (
+    BarColumn,
+    MofNCompleteColumn,
+    Progress,
+    SpinnerColumn,
+    TextColumn,
+)
+from rich.text import Text
+from tqdm import tqdm
+
+SECONDS_TO_PAUSE_AFTER_RATE_LIMIT_ERROR = 5
 
 
 @dataclass
 class StatusTracker:
+    max_requests_per_minute: int
+    max_tokens_per_minute: int
+    max_concurrent_requests: int
     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
+    total_requests: int = 0
+    retry_queue: asyncio.Queue = field(default_factory=asyncio.Queue)
+
+    # Progress bar configuration
+    use_progress_bar: bool = True
+    progress_bar_total: int | None = None
+    progress_bar_disable: bool = False
+    _pbar: tqdm | None = None
+
+    # Rich display configuration
+    use_rich: bool = True
+    _rich_console: Console | None = None
+    _rich_live: object | None = None
+    _rich_progress: object | None = None
+    _rich_task_id: object | None = None
+    _rich_display_task: asyncio.Task | None = None
+    _rich_stop_event: asyncio.Event | None = None
+
+    def __post_init__(self):
+        self.available_request_capacity = self.max_requests_per_minute
+        self.available_token_capacity = self.max_tokens_per_minute
+        self.last_update_time = time.time() - 1
+        self.last_pbar_update_time = time.time() - 1
+        self.limiting_factor = None
 
     @property
     def time_since_rate_limit_error(self):
         return time.time() - self.time_of_last_rate_limit_error
 
+    @property
+    def seconds_to_pause(self):
+        return max(
+            0,
+            SECONDS_TO_PAUSE_AFTER_RATE_LIMIT_ERROR - self.time_since_rate_limit_error,
+        )
+
+    def set_limiting_factor(self, factor):
+        self.limiting_factor = factor
+
+    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 = (
+            self.num_tasks_in_progress < self.max_concurrent_requests
+        )
+        if request_available and tokens_available and concurrent_request_available:
+            self.available_request_capacity -= 1
+            self.available_token_capacity -= num_tokens
+            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:
+            # update reason why
+            if not request_available:
+                self.set_limiting_factor("Requests")
+            elif not concurrent_request_available:
+                self.set_limiting_factor("Concurrent Requests")
+            elif not tokens_available:
+                self.set_limiting_factor("Tokens")
+
+    def update_capacity(self):
+        current_time = time.time()
+        seconds_since_update = current_time - self.last_update_time
+        self.available_request_capacity = min(
+            self.available_request_capacity
+            + self.max_requests_per_minute * seconds_since_update / 60.0,
+            self.max_requests_per_minute,
+        )
+        self.available_token_capacity = min(
+            self.available_token_capacity
+            + self.max_tokens_per_minute * seconds_since_update / 60.0,
+            self.max_tokens_per_minute,
+        )
+        self.last_update_time = current_time
+
     def start_task(self, task_id):
         self.num_tasks_started += 1
         self.num_tasks_in_progress += 1
@@ -27,12 +117,16 @@ class StatusTracker:
     def task_succeeded(self, task_id):
         self.num_tasks_in_progress -= 1
         self.num_tasks_succeeded += 1
+        self.increment_pbar()
 
     def task_failed(self, task_id):
         self.num_tasks_in_progress -= 1
         self.num_tasks_failed += 1
 
     def log_final_status(self):
+        # Close progress bar before printing final status
+        self.close_progress_bar()
+
         if self.num_tasks_failed > 0:
             print(
                 f"{self.num_tasks_failed} / {self.num_tasks_started} requests failed."
@@ -41,3 +135,121 @@ class StatusTracker:
             print(
                 f"{self.num_rate_limit_errors} rate limit errors received. Consider running at a lower rate."
             )
+
+    @property
+    def pbar(self) -> tqdm | None:
+        """Backward compatibility property to access progress bar."""
+        return self._pbar
+
+    def init_progress_bar(self, total: int | None = None, disable: bool | None = None):
+        """Initialize progress bar if enabled."""
+        if not self.use_progress_bar:
+            return
+
+        if self.use_rich:
+            self._init_rich_display(total, disable)
+        else:
+            # Use provided values or fall back to instance defaults
+            pbar_total = total if total is not None else self.progress_bar_total
+            pbar_disable = disable if disable is not None else self.progress_bar_disable
+            self._pbar = tqdm(total=pbar_total, disable=pbar_disable)
+        self.update_pbar()
+
+    def close_progress_bar(self):
+        """Close progress bar if it exists."""
+        if self.use_rich and self._rich_stop_event:
+            self._close_rich_display()
+        elif self._pbar is not None:
+            self._pbar.close()
+            self._pbar = None
+
+    def _init_rich_display(self, total: int | None = None, disable: bool | None = None):
+        """Initialize Rich display components."""
+        if disable:
+            return
+
+        pbar_total = total if total is not None else self.progress_bar_total
+        if pbar_total is None:
+            pbar_total = 100  # Default fallback
+
+        self._rich_console = Console()
+        self._rich_stop_event = asyncio.Event()
+
+        # Start the display updater task
+        self._rich_display_task = asyncio.create_task(
+            self._rich_display_updater(pbar_total)
+        )
+
+    async def _rich_display_updater(self, total: int):
+        """Update Rich display independently."""
+        if not self._rich_console or self._rich_stop_event is None:
+            return
+
+        # Create progress bar without console so we can use it in Live
+        progress = Progress(
+            SpinnerColumn(),
+            TextColumn("Processing requests..."),
+            BarColumn(),
+            MofNCompleteColumn(),
+        )
+        main_task = progress.add_task("requests", total=total)
+
+        # Use Live to combine progress + text
+
+        with Live(console=self._rich_console, refresh_per_second=10) as live:
+            while not self._rich_stop_event.is_set():
+                completed = self.num_tasks_succeeded
+                progress.update(main_task, completed=completed)
+
+                # Create capacity info text
+                tokens_info = f"TPM Capacity: {self.available_token_capacity / 1000:.1f}k/{self.max_tokens_per_minute / 1000:.1f}k"
+                reqs_info = f"RPM Capacity: {int(self.available_request_capacity)}/{self.max_requests_per_minute}"
+                in_progress = f"In Progress: {int(self.num_tasks_in_progress)}"
+                capacity_text = Text(f"{in_progress} • {tokens_info} • {reqs_info}")
+
+                # Group progress bar and text
+                display = Group(progress, capacity_text)
+                live.update(display)
+
+                await asyncio.sleep(0.1)
+
+    def _close_rich_display(self):
+        """Clean up Rich display."""
+        if self._rich_stop_event:
+            self._rich_stop_event.set()
+        if self._rich_display_task and not self._rich_display_task.done():
+            self._rich_display_task.cancel()
+
+        self._rich_console = None
+        self._rich_live = None
+        self._rich_display_task = None
+        self._rich_stop_event = None
+
+    def update_pbar(self, n: int = 0):
+        """Update progress bar status and optionally increment.
+
+        Args:
+            n: Number of items to increment (0 means just update postfix)
+        """
+        current_time = time.time()
+        if self._pbar and (current_time - self.last_pbar_update_time > 1):
+            self.last_pbar_update_time = current_time
+            self._pbar.set_postfix(
+                {
+                    "Token Capacity": f"{self.available_token_capacity / 1_000:.1f}k",
+                    "Req. Capacity": f"{int(self.available_request_capacity)}",
+                    "Reqs. in Progress": self.num_tasks_in_progress,
+                    "Limiting Factor": self.limiting_factor,
+                }
+            )
+
+        if n > 0 and self._pbar:
+            self._pbar.update(n)
+
+    def increment_pbar(self):
+        """Increment progress bar by 1."""
+        if self.use_rich:
+            # Rich display is updated automatically by the display updater
+            pass
+        elif self._pbar:
+            self._pbar.update(1)

lm_deluge/util/json.py

@@ -1,5 +1,6 @@
-import re
 import json
+import re
+
 import json5
 
 
@@ -166,3 +167,19 @@ def load_json(
             pass
 
     raise ValueError(f"Invalid JSON string: {json_string}")
+
+
+def try_load_json(
+    json_string: str | None,
+    allow_json5: bool = True,
+    allow_partial: bool = False,
+    allow_healing: bool = True,
+):
+    """
+    Like the above, except it returns None instead of raising an error.
+    """
+    try:
+        return load_json(json_string, allow_json5, allow_partial, allow_healing)
+    except Exception as e:
+        print(f"Failed to load json: {e}. Returning None.")
+        return None