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

lm_deluge/computer_use/anthropic_tools.py

@@ -0,0 +1,75 @@
+from typing import Literal
+
+ToolVersion = Literal["2024-10-22", "2025-01-24", "2025-04-29"]
+ToolType = Literal["bash", "computer", "editor"]
+
+
+def model_to_version(model: str) -> ToolVersion:
+    if "opus" not in model and "sonnet" not in model:
+        raise ValueError("cannot use computer tools with incompatible model")
+    if "claude-4" in model:
+        return "2025-04-29"
+    elif "3.7" in model:
+        return "2025-01-24"
+    else:
+        return "2024-10-22"
+
+
+def get_anthropic_cu_tools(
+    model: str,
+    display_width: int,
+    display_height: int,
+    exclude_tools: list[ToolType] | None = None,
+):
+    version = model_to_version(model)
+    if version == "2024-10-22":
+        result = [
+            {
+                "name": "computer",
+                "type": "computer_20241022",
+                "display_width_px": display_width,
+                "display_height_px": display_height,
+                "display_number": None,
+            },
+            {"name": "str_replace_editor", "type": "text_editor_20250429"},
+            {"type": "bash_20250124", "name": "bash"},
+        ]
+    elif version == "2025-01-24":
+        result = [
+            {
+                "name": "computer",
+                "type": "computer_20250124",
+                "display_width_px": display_width,
+                "display_height_px": display_height,
+                "display_number": None,
+            },
+            {"name": "str_replace_editor", "type": "text_editor_20250124"},
+            {"type": "bash_20250124", "name": "bash"},
+        ]
+    elif version == "2025-04-29":
+        result = [
+            {
+                "name": "computer",
+                "type": "computer_20250124",
+                "display_width_px": display_width,
+                "display_height_px": display_height,
+                "display_number": None,
+            },
+            {"name": "str_replace_based_edit_tool", "type": "text_editor_20250429"},
+            {
+                "name": "bash",
+                "type": "bash_20250124",
+            },
+        ]
+    else:
+        raise ValueError("invalid tool version")
+
+    if exclude_tools is None:
+        return result
+    if "bash" in exclude_tools:
+        result = [x for x in result if x["name"] != "bash"]
+    if "editor" in exclude_tools:
+        result = [x for x in result if "edit" not in x["name"]]
+    if "computer" in exclude_tools:
+        result = [x for x in result if "computer" not in x["name"]]
+    return result

lm_deluge/{sampling_params.py → config.py}

@@ -1,14 +1,15 @@
-from dataclasses import dataclass
+from pydantic import BaseModel
 from typing import Literal
 
 
-@dataclass
-class SamplingParams:
+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
+    logprobs: bool = False
+    top_logprobs: int | None = None
 
     def to_vllm(self):
         try:
@@ -23,3 +24,9 @@ class SamplingParams:
             top_p=self.top_p,
             max_tokens=self.max_new_tokens,
         )
+
+
+class ComputerUseParams(BaseModel):
+    enabled: bool = False
+    display_width: int = 1024
+    display_height: int = 768

lm_deluge/embed.py

@@ -1,12 +1,14 @@
 ### specific utility for cohere rerank api
-import os
-import numpy as np
-import aiohttp
-from tqdm.auto import tqdm
 import asyncio
+import os
 import time
-from typing import Any
 from dataclasses import dataclass
+from typing import Any
+
+import aiohttp
+import numpy as np
+from tqdm.auto import tqdm
+
 from .tracker import StatusTracker
 
 registry = {
@@ -56,7 +58,6 @@ class EmbeddingRequest:
         texts: list[str],
         attempts_left: int,
         status_tracker: StatusTracker,
-        retry_queue: asyncio.Queue,
         request_timeout: int,
         pbar: tqdm | None = None,
         **kwargs,  # openai or cohere specific params
@@ -66,7 +67,6 @@ class EmbeddingRequest:
         self.texts = texts
         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 = []
@@ -89,7 +89,8 @@ class EmbeddingRequest:
         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.")
@@ -243,7 +244,11 @@ async def embed_parallel_async(
 
     # initialize trackers
     retry_queue = asyncio.Queue()
-    status_tracker = StatusTracker()
+    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
@@ -262,7 +267,8 @@ async def embed_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():
+            assert status_tracker.retry_queue
+            if not status_tracker.retry_queue.empty():
                 next_request = retry_queue.get_nowait()
                 print(f"Retrying request {next_request.task_id}.")
             elif prompts_not_finished:
@@ -285,7 +291,7 @@ async def embed_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/models.py

@@ -178,6 +178,21 @@ registry = {
     #                ░███
     #                █████
     #               ░░░░░
+    "openai-computer-use-preview": {
+        "id": "openai-computer-use-preview",
+        "name": "computer-use-preview",
+        "api_base": "https://api.openai.com/v1",
+        "api_key_env_var": "OPENAI_API_KEY",
+        "supports_json": True,
+        "supports_logprobs": False,
+        "supports_responses": True,
+        "api_spec": "openai-responses",
+        "input_cost": 2.0,
+        "output_cost": 8.0,
+        "requests_per_minute": 20,
+        "tokens_per_minute": 100_000,
+        "reasoning_model": False,
+    },
     "o3": {
         "id": "o3",
         "name": "o3-2025-04-16",
@@ -185,6 +200,7 @@ registry = {
         "api_key_env_var": "OPENAI_API_KEY",
         "supports_json": False,
         "supports_logprobs": True,
+        "supports_responses": True,
         "api_spec": "openai",
         "input_cost": 10.0,
         "output_cost": 40.0,
@@ -199,6 +215,7 @@ registry = {
         "api_key_env_var": "OPENAI_API_KEY",
         "supports_json": False,
         "supports_logprobs": True,
+        "supports_responses": True,
         "api_spec": "openai",
         "input_cost": 1.1,
         "output_cost": 4.4,
@@ -213,6 +230,7 @@ registry = {
         "api_key_env_var": "OPENAI_API_KEY",
         "supports_json": True,
         "supports_logprobs": True,
+        "supports_responses": True,
         "api_spec": "openai",
         "input_cost": 2.0,
         "output_cost": 8.0,
@@ -227,6 +245,7 @@ registry = {
         "api_key_env_var": "OPENAI_API_KEY",
         "supports_json": True,
         "supports_logprobs": True,
+        "supports_responses": True,
         "api_spec": "openai",
         "input_cost": 0.4,
         "output_cost": 1.6,
@@ -241,6 +260,7 @@ registry = {
         "api_key_env_var": "OPENAI_API_KEY",
         "supports_json": True,
         "supports_logprobs": True,
+        "supports_responses": True,
         "api_spec": "openai",
         "input_cost": 0.1,
         "output_cost": 0.4,
@@ -255,6 +275,7 @@ registry = {
         "api_key_env_var": "OPENAI_API_KEY",
         "supports_json": False,
         "supports_logprobs": True,
+        "supports_responses": True,
         "api_spec": "openai",
         "input_cost": 75.0,
         "output_cost": 150.0,
@@ -269,6 +290,7 @@ registry = {
         "api_key_env_var": "OPENAI_API_KEY",
         "supports_json": False,
         "supports_logprobs": True,
+        "supports_responses": True,
         "api_spec": "openai",
         "input_cost": 1.1,
         "output_cost": 4.4,
@@ -283,6 +305,7 @@ registry = {
         "api_key_env_var": "OPENAI_API_KEY",
         "supports_json": False,
         "supports_logprobs": True,
+        "supports_responses": True,
         "api_spec": "openai",
         "input_cost": 15.0,
         "output_cost": 60.0,
@@ -297,6 +320,7 @@ registry = {
         "api_key_env_var": "OPENAI_API_KEY",
         "supports_json": False,
         "supports_logprobs": True,
+        "supports_responses": True,
         "api_spec": "openai",
         "input_cost": 15.0,
         "output_cost": 60.0,
@@ -311,6 +335,7 @@ registry = {
         "api_key_env_var": "OPENAI_API_KEY",
         "supports_json": False,
         "supports_logprobs": True,
+        "supports_responses": True,
         "api_spec": "openai",
         "input_cost": 3.0,
         "output_cost": 15.0,
@@ -325,6 +350,7 @@ registry = {
         "api_key_env_var": "OPENAI_API_KEY",
         "supports_json": True,
         "supports_logprobs": True,
+        "supports_responses": True,
         "api_spec": "openai",
         "input_cost": 5.0,
         "output_cost": 15.0,
@@ -338,6 +364,7 @@ registry = {
         "api_key_env_var": "OPENAI_API_KEY",
         "supports_json": True,
         "supports_logprobs": True,
+        "supports_responses": True,
         "api_spec": "openai",
         "input_cost": 0.15,
         "output_cost": 0.6,
@@ -351,6 +378,7 @@ registry = {
         "api_key_env_var": "OPENAI_API_KEY",
         "supports_json": True,
         "supports_logprobs": True,
+        "supports_responses": True,
         "api_spec": "openai",
         "input_cost": 0.0,
         "output_cost": 0.0,
@@ -364,6 +392,7 @@ registry = {
         "api_key_env_var": "OPENAI_API_KEY",
         "supports_json": True,
         "supports_logprobs": True,
+        "supports_responses": True,
         "api_spec": "openai",
         "input_cost": 0.5,
         "output_cost": 1.5,
@@ -377,6 +406,7 @@ registry = {
         "api_key_env_var": "OPENAI_API_KEY",
         "supports_json": True,
         "supports_logprobs": True,
+        "supports_responses": True,
         "api_spec": "openai",
         "input_cost": 10.0,
         "output_cost": 30.0,
@@ -390,6 +420,7 @@ registry = {
         "api_key_env_var": "OPENAI_API_KEY",
         "supports_json": False,
         "supports_logprobs": False,
+        "supports_responses": True,
         "api_spec": "openai",
         "input_cost": 30.0,
         "output_cost": 60.0,
@@ -403,6 +434,7 @@ registry = {
         "api_key_env_var": "OPENAI_API_KEY",
         "supports_json": False,
         "supports_logprobs": False,
+        "supports_responses": True,
         "api_spec": "openai",
         "input_cost": 60.0,
         "output_cost": 120.0,
@@ -1093,6 +1125,7 @@ class APIModel:
     output_cost: float | None = 0  # $ per million output tokens
     supports_json: bool = False
     supports_logprobs: bool = False
+    supports_responses: bool = False
     reasoning_model: bool = False
     regions: list[str] | dict[str, int] = field(default_factory=list)
     tokens_per_minute: int | None = None

lm_deluge/prompt.py

@@ -4,10 +4,18 @@ import tiktoken
 import xxhash
 from dataclasses import dataclass, field
 from pathlib import Path
-from typing import Literal
+from typing import Literal, Sequence
 from lm_deluge.models import APIModel
 from lm_deluge.image import Image
 
+CachePattern = Literal[
+    "tools_only",
+    "system_and_tools",
+    "last_user_message",
+    "last_2_user_messages",
+    "last_3_user_messages",
+]
+
 ###############################################################################
 # 1. Low-level content blocks – either text or an image                       #
 ###############################################################################
@@ -91,24 +99,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
@@ -420,6 +462,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:
@@ -514,9 +564,41 @@ 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())
 
-    def to_anthropic(self) -> tuple[str | None, list[dict]]:
+        return {"input": input_items}
+
+    def to_anthropic(
+        self, cache_pattern: CachePattern | None = None
+    ) -> tuple[str | list[dict] | None, list[dict]]:
         system_msg = next(
             (
                 m.parts[0].text
@@ -535,8 +617,84 @@ class Conversation:
                 other.append(user_msg.anthropic())
             else:
                 other.append(m.anthropic())
+
+        # Apply cache control if specified
+        if cache_pattern is not None:
+            system_msg, other = self._apply_cache_control(
+                system_msg, other, cache_pattern
+            )
+
         return system_msg, other
 
+    def _apply_cache_control(
+        self,
+        system_msg: str | None | list[dict],
+        messages: list[dict],
+        cache_pattern: CachePattern,
+    ) -> tuple[str | list[dict] | None, list[dict]]:
+        """Apply cache control to system message and/or messages based on the pattern."""
+
+        if cache_pattern == "system_and_tools" and system_msg is not None:
+            # Convert system message to structured format with cache control
+            # This caches tools+system prefix (since system comes after tools)
+            system_msg = [
+                {
+                    "type": "text",
+                    "text": system_msg,
+                    "cache_control": {"type": "ephemeral"},
+                }
+            ]
+
+        if cache_pattern == "last_user_message":
+            # Cache the last user message
+            user_messages = [i for i, m in enumerate(messages) if m["role"] == "user"]
+            if user_messages:
+                last_user_idx = user_messages[-1]
+                self._add_cache_control_to_message(messages[last_user_idx])
+
+        elif cache_pattern == "last_2_user_messages":
+            # Cache the last 2 user messages
+            user_messages = [i for i, m in enumerate(messages) if m["role"] == "user"]
+            for idx in user_messages[-2:]:
+                self._add_cache_control_to_message(messages[idx])
+
+        elif cache_pattern == "last_3_user_messages":
+            # Cache the last 3 user messages
+            user_messages = [i for i, m in enumerate(messages) if m["role"] == "user"]
+            for idx in user_messages[-3:]:
+                self._add_cache_control_to_message(messages[idx])
+
+        return system_msg, messages
+
+    def lock_images_as_bytes(self) -> "Conversation":
+        """
+        Convert all images to bytes format to ensure they remain unchanged for caching.
+        This should be called when caching is enabled to prevent cache invalidation
+        from image reference changes.
+        """
+        for message in self.messages:
+            for part in message.parts:
+                if isinstance(part, Image):
+                    # Force conversion to bytes if not already
+                    part.data = part._bytes()
+        return self
+
+    def _add_cache_control_to_message(self, message: dict) -> None:
+        """Add cache control to a message's content."""
+        content = message["content"]
+        if isinstance(content, str):
+            # Convert string content to structured format with cache control
+            message["content"] = [
+                {
+                    "type": "text",
+                    "text": content,
+                    "cache_control": {"type": "ephemeral"},
+                }
+            ]
+        elif isinstance(content, list) and content:
+            # Add cache control to the last content block
+            content[-1]["cache_control"] = {"type": "ephemeral"}
+
     def to_gemini(self) -> tuple[str | None, list[dict]]:
         system_msg = next(
             (
@@ -664,6 +822,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)