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

lm_deluge/__init__.py

@@ -1,7 +1,17 @@
 from .client import LLMClient, SamplingParams, APIResponse
 from .prompt import Conversation, Message
+from .tool import Tool
+from .file import File
 import dotenv
 
 dotenv.load_dotenv()
 
-__all__ = ["LLMClient", "SamplingParams", "APIResponse", "Conversation", "Message"]
+__all__ = [
+    "LLMClient",
+    "SamplingParams",
+    "APIResponse",
+    "Conversation",
+    "Message",
+    "Tool",
+    "File",
+]

lm_deluge/api_requests/anthropic.py

@@ -1,9 +1,6 @@
-import asyncio
 from aiohttp import ClientResponse
 import json
 import os
-import warnings
-from tqdm import tqdm
 from typing import Callable
 
 from lm_deluge.prompt import (
@@ -14,12 +11,84 @@ from lm_deluge.prompt import (
     Thinking,
     CachePattern,
 )
+from lm_deluge.tool import Tool
 from lm_deluge.usage import Usage
 from .base import APIRequestBase, APIResponse
 
 from ..tracker import StatusTracker
-from ..sampling_params import SamplingParams
+from ..config import SamplingParams
 from ..models import APIModel
+from ..computer_use.anthropic_tools import get_anthropic_cu_tools
+
+
+def _build_anthropic_request(
+    model: APIModel,
+    prompt: Conversation,
+    tools: list[Tool] | None,
+    sampling_params: SamplingParams,
+    cache_pattern: CachePattern | None = None,
+    computer_use: bool = False,
+    display_width: int = 1024,
+    display_height: int = 768,
+):
+    system_message, messages = prompt.to_anthropic(cache_pattern=cache_pattern)
+    request_header = {
+        "x-api-key": os.getenv(model.api_key_env_var),
+        "anthropic-version": "2023-06-01",
+        "content-type": "application/json",
+    }
+
+    # Add beta header for Computer Use
+    if computer_use:
+        request_header["anthropic-beta"] = "computer-use-2025-01-24"
+
+    request_json = {
+        "model": model.name,
+        "messages": messages,
+        "temperature": sampling_params.temperature,
+        "top_p": sampling_params.top_p,
+        "max_tokens": sampling_params.max_new_tokens,
+    }
+
+    # handle thinking
+    if model.reasoning_model and sampling_params.reasoning_effort:
+        # translate reasoning effort of low, medium, high to budget tokens
+        budget = {"low": 1024, "medium": 4096, "high": 16384}.get(
+            sampling_params.reasoning_effort
+        )
+        request_json["thinking"] = {
+            "type": "enabled",
+            "budget_tokens": budget,
+        }
+        request_json.pop("top_p")
+        request_json["temperature"] = 1.0
+        request_json["max_tokens"] += budget
+    else:
+        request_json["thinking"] = {"type": "disabled"}
+        if sampling_params.reasoning_effort:
+            print("ignoring reasoning_effort for non-reasoning model")
+    if system_message is not None:
+        request_json["system"] = system_message
+    if tools or computer_use:
+        tool_definitions = []
+        if tools:
+            tool_definitions.extend([tool.dump_for("anthropic") for tool in tools])
+        # Add Computer Use tools
+        if computer_use:
+            cu_tools = get_anthropic_cu_tools(
+                model=model.id,
+                display_width=display_width,  # todo: set from ComputerUseParams
+                display_height=display_height,
+            )
+            tool_definitions.extend(cu_tools)
+
+        # Add cache control to last tool if tools_only caching is specified
+        if cache_pattern == "tools_only" and tool_definitions:
+            tool_definitions[-1]["cache_control"] = {"type": "ephemeral"}
+
+        request_json["tools"] = tool_definitions
+
+    return request_json, request_header
 
 
 class AnthropicRequest(APIRequestBase):
@@ -32,18 +101,19 @@ class AnthropicRequest(APIRequestBase):
         prompt: Conversation,
         attempts_left: int,
         status_tracker: StatusTracker,
-        retry_queue: asyncio.Queue,
         results_arr: list,
         request_timeout: int = 30,
         sampling_params: SamplingParams = SamplingParams(),
-        pbar: tqdm | None = None,
         callback: Callable | None = None,
-        debug: bool = False,
         # for retries
         all_model_names: list[str] | None = None,
         all_sampling_params: list[SamplingParams] | None = None,
         tools: list | None = None,
         cache: CachePattern | None = None,
+        # Computer Use support
+        computer_use: bool = False,
+        display_width: int = 1024,
+        display_height: int = 768,
     ):
         super().__init__(
             task_id=task_id,
@@ -51,18 +121,18 @@ class AnthropicRequest(APIRequestBase):
             prompt=prompt,
             attempts_left=attempts_left,
             status_tracker=status_tracker,
-            retry_queue=retry_queue,
             results_arr=results_arr,
             request_timeout=request_timeout,
             sampling_params=sampling_params,
-            pbar=pbar,
             callback=callback,
-            debug=debug,
             all_model_names=all_model_names,
             all_sampling_params=all_sampling_params,
             tools=tools,
             cache=cache,
         )
+        self.computer_use = computer_use
+        self.display_width = display_width
+        self.display_height = display_height
         self.model = APIModel.from_registry(model_name)
         self.url = f"{self.model.api_base}/messages"
 
@@ -70,52 +140,16 @@ class AnthropicRequest(APIRequestBase):
         if cache is not None:
             prompt.lock_images_as_bytes()
 
-        self.system_message, messages = prompt.to_anthropic(cache_pattern=cache)
-        self.request_header = {
-            "x-api-key": os.getenv(self.model.api_key_env_var),
-            "anthropic-version": "2023-06-01",
-            "content-type": "application/json",
-        }
-
-        self.request_json = {
-            "model": self.model.name,
-            "messages": messages,
-            "temperature": self.sampling_params.temperature,
-            "top_p": self.sampling_params.top_p,
-            "max_tokens": self.sampling_params.max_new_tokens,
-        }
-        # handle thinking
-        if self.model.reasoning_model:
-            if sampling_params.reasoning_effort:
-                # translate reasoning effort of low, medium, high to budget tokens
-                budget = {"low": 1024, "medium": 4096, "high": 16384}.get(
-                    sampling_params.reasoning_effort
-                )
-                self.request_json["thinking"] = {
-                    "type": "enabled",
-                    "budget_tokens": budget,
-                }
-                self.request_json.pop("top_p")
-                self.request_json["temperature"] = 1.0
-                self.request_json["max_tokens"] += (
-                    budget  # assume max tokens is max completion tokens
-                )
-            else:
-                # no thinking
-                self.request_json["thinking"] = {"type": "disabled"}
-        else:
-            if sampling_params.reasoning_effort:
-                warnings.warn(
-                    f"Ignoring reasoning_effort param for non-reasoning model: {model_name}"
-                )
-        if self.system_message is not None:
-            self.request_json["system"] = self.system_message
-        if tools:
-            tool_definitions = [tool.dump_for("anthropic") for tool in tools]
-            # Add cache control to last tool if tools_only caching is specified
-            if cache == "tools_only" and tool_definitions:
-                tool_definitions[-1]["cache_control"] = {"type": "ephemeral"}
-            self.request_json["tools"] = tool_definitions
+        self.request_json, self.request_header = _build_anthropic_request(
+            self.model,
+            prompt,
+            tools,
+            sampling_params,
+            cache,
+            computer_use,
+            display_width,
+            display_height,
+        )
 
     async def handle_response(self, http_response: ClientResponse) -> APIResponse:
         is_error = False
@@ -135,8 +169,6 @@ class AnthropicRequest(APIRequestBase):
             "anthropic-ratelimit-tokens-reset",
         ]:
             rate_limits[header] = http_response.headers.get(header, None)
-        if self.debug:
-            print(f"Rate limits: {rate_limits}")
         if status_code >= 200 and status_code < 300:
             try:
                 data = await http_response.json()

lm_deluge/api_requests/base.py

@@ -1,160 +1,19 @@
-import aiohttp
 import asyncio
-import json
 import random
-from tqdm import tqdm
-from dataclasses import dataclass
+import traceback
 from abc import ABC, abstractmethod
 from typing import Callable
 
-from lm_deluge.prompt import Conversation, Message, CachePattern
-from lm_deluge.usage import Usage
-
-from ..tracker import StatusTracker
-from ..sampling_params import SamplingParams
-from ..models import APIModel
-from ..errors import raise_if_modal_exception
+import aiohttp
 from aiohttp import ClientResponse
 
+from lm_deluge.prompt import CachePattern, Conversation
 
-@dataclass
-class APIResponse:
-    # request information
-    id: int  # should be unique to the request within a given prompt-processing call
-    model_internal: str  # our internal model tag
-    prompt: Conversation
-    sampling_params: SamplingParams
-
-    # http response information
-    status_code: int | None
-    is_error: bool | None
-    error_message: str | None
-
-    # completion information - unified usage tracking
-    usage: Usage | None = None
-
-    # response content - structured format
-    content: Message | None = None
-
-    # optional or calculated automatically
-    thinking: str | None = None  # if model shows thinking tokens
-    model_external: str | None = None  # the model tag used by the API
-    region: str | None = None
-    logprobs: list | None = None
-    finish_reason: str | None = None  # make required later
-    cost: float | None = None  # calculated automatically
-    cache_hit: bool = False  # manually set if true
-    # set to true if is_error and should be retried with a different model
-    retry_with_different_model: bool | None = False
-    # set to true if should NOT retry with the same model (unrecoverable error)
-    give_up_if_no_other_models: bool | None = False
-
-    @property
-    def completion(self) -> str | None:
-        """Backward compatibility: extract text from content Message."""
-        if self.content is not None:
-            return self.content.completion
-        return None
-
-    @property
-    def input_tokens(self) -> int | None:
-        """Get input tokens from usage object."""
-        return self.usage.input_tokens if self.usage else None
-
-    @property
-    def output_tokens(self) -> int | None:
-        """Get output tokens from usage object."""
-        return self.usage.output_tokens if self.usage else None
-
-    @property
-    def cache_read_tokens(self) -> int | None:
-        """Get cache read tokens from usage object."""
-        return self.usage.cache_read_tokens if self.usage else None
-
-    @property
-    def cache_write_tokens(self) -> int | None:
-        """Get cache write tokens from usage object."""
-        return self.usage.cache_write_tokens if self.usage else None
-
-    def __post_init__(self):
-        # calculate cost & get external model name
-        self.id = int(self.id)
-        api_model = APIModel.from_registry(self.model_internal)
-        self.model_external = api_model.name
-        self.cost = None
-        if (
-            self.usage is not None
-            and api_model.input_cost is not None
-            and api_model.output_cost is not None
-        ):
-            self.cost = (
-                self.usage.input_tokens * api_model.input_cost / 1e6
-                + self.usage.output_tokens * api_model.output_cost / 1e6
-            )
-        elif self.content is not None and self.completion is not None:
-            print(
-                f"Warning: Completion provided without token counts for model {self.model_internal}."
-            )
-
-    def to_dict(self):
-        return {
-            "id": self.id,
-            "model_internal": self.model_internal,
-            "model_external": self.model_external,
-            "region": self.region,
-            "prompt": self.prompt.to_log(),  # destroys image if present
-            "sampling_params": self.sampling_params.__dict__,
-            "status_code": self.status_code,
-            "is_error": self.is_error,
-            "error_message": self.error_message,
-            "completion": self.completion,  # computed property
-            "content": self.content.to_log() if self.content else None,
-            "usage": self.usage.to_dict() if self.usage else None,
-            "finish_reason": self.finish_reason,
-            "cost": self.cost,
-        }
-
-    @classmethod
-    def from_dict(cls, data: dict):
-        # Handle backward compatibility for content/completion
-        content = None
-        if "content" in data and data["content"] is not None:
-            # Reconstruct message from log format
-            content = Message.from_log(data["content"])
-        elif "completion" in data and data["completion"] is not None:
-            # Backward compatibility: create a Message with just text
-            content = Message.ai(data["completion"])
-
-        usage = None
-        if "usage" in data and data["usage"] is not None:
-            usage = Usage.from_dict(data["usage"])
-
-        return cls(
-            id=data.get("id", random.randint(0, 1_000_000_000)),
-            model_internal=data["model_internal"],
-            prompt=Conversation.from_log(data["prompt"]),
-            sampling_params=SamplingParams(**data["sampling_params"]),
-            status_code=data["status_code"],
-            is_error=data["is_error"],
-            error_message=data["error_message"],
-            usage=usage,
-            content=content,
-            thinking=data.get("thinking"),
-            model_external=data.get("model_external"),
-            region=data.get("region"),
-            logprobs=data.get("logprobs"),
-            finish_reason=data.get("finish_reason"),
-            cost=data.get("cost"),
-            cache_hit=data.get("cache_hit", False),
-        )
-
-    def write_to_file(self, filename):
-        """
-        Writes the APIResponse as a line to a file.
-        If file exists, appends to it.
-        """
-        with open(filename, "a") as f:
-            f.write(json.dumps(self.to_dict()) + "\n")
+from ..config import SamplingParams
+from ..errors import raise_if_modal_exception
+from ..models import APIModel
+from ..tracker import StatusTracker
+from .response import APIResponse
 
 
 class APIRequestBase(ABC):
@@ -176,16 +35,11 @@ class APIRequestBase(ABC):
         prompt: Conversation,
         attempts_left: int,
         status_tracker: StatusTracker,
-        retry_queue: asyncio.Queue,
         # needed in order to retry with a different model and not throw the output away
         results_arr: list["APIRequestBase"],
         request_timeout: int = 30,
         sampling_params: SamplingParams = SamplingParams(),
-        logprobs: bool = False,
-        top_logprobs: int | None = None,
-        pbar: tqdm | None = None,
         callback: Callable | None = None,
-        debug: bool = False,
         all_model_names: list[str] | None = None,
         all_sampling_params: list[SamplingParams] | None = None,
         tools: list | None = None,
@@ -199,16 +53,11 @@ class APIRequestBase(ABC):
         self.prompt = prompt
         self.attempts_left = attempts_left
         self.status_tracker = status_tracker
-        self.retry_queue = retry_queue
         self.request_timeout = request_timeout
         self.sampling_params = sampling_params
-        self.logprobs = logprobs  # len(completion) logprobs
-        self.top_logprobs = top_logprobs
-        self.pbar = pbar
         self.callback = callback
         self.num_tokens = prompt.count_tokens(sampling_params.max_new_tokens)
         self.results_arr = results_arr
-        self.debug = debug
         self.all_model_names = all_model_names
         self.all_sampling_params = all_sampling_params
         self.tools = tools
@@ -222,8 +71,7 @@ class APIRequestBase(ABC):
         self.region = None
 
     def increment_pbar(self):
-        if self.pbar is not None:
-            self.pbar.update(1)
+        self.status_tracker.increment_pbar()
 
     def call_callback(self):
         if self.callback is not None:
@@ -232,7 +80,6 @@ class APIRequestBase(ABC):
 
     def handle_success(self, data):
         self.call_callback()
-        self.increment_pbar()
         self.status_tracker.task_succeeded(self.task_id)
 
     def handle_error(self, create_new_request=False, give_up_if_no_other_models=False):
@@ -253,7 +100,8 @@ class APIRequestBase(ABC):
         if self.attempts_left > 0:
             self.attempts_left -= 1
             if not create_new_request:
-                self.retry_queue.put_nowait(self)
+                assert self.status_tracker.retry_queue
+                self.status_tracker.retry_queue.put_nowait(self)
                 return
             else:
                 # make sure we have another model to send it to besides the current one
@@ -267,7 +115,8 @@ class APIRequestBase(ABC):
                         print(
                             f"No other models to try for task {self.task_id}. Retrying with same model."
                         )
-                        self.retry_queue.put_nowait(self)
+                        assert self.status_tracker.retry_queue
+                        self.status_tracker.retry_queue.put_nowait(self)
                 else:
                     # two things to change: model_name and sampling_params
                     new_model_name = self.model_name
@@ -292,21 +141,21 @@ class APIRequestBase(ABC):
                         prompt=self.prompt,
                         attempts_left=self.attempts_left,
                         status_tracker=self.status_tracker,
-                        retry_queue=self.retry_queue,
                         results_arr=self.results_arr,
                         request_timeout=self.request_timeout,
                         sampling_params=new_sampling_params,
-                        logprobs=self.logprobs,
-                        top_logprobs=self.top_logprobs,
-                        pbar=self.pbar,
                         callback=self.callback,
                         all_model_names=self.all_model_names,
                         all_sampling_params=self.all_sampling_params,
                         tools=self.tools,
                         cache=self.cache,
+                        computer_use=getattr(self, "computer_use", False),
+                        display_width=getattr(self, "display_width", 1024),
+                        display_height=getattr(self, "display_height", 768),
                     )
                     # PROBLEM: new request is never put into results array, so we can't get the result.
-                    self.retry_queue.put_nowait(new_request)
+                    assert self.status_tracker.retry_queue
+                    self.status_tracker.retry_queue.put_nowait(self)
                     # SOLUTION: just need to make sure it's deduplicated by task_id later.
                     self.results_arr.append(new_request)
         else:
@@ -354,6 +203,8 @@ class APIRequestBase(ABC):
 
         except Exception as e:
             raise_if_modal_exception(e)
+            tb = traceback.format_exc()
+            print(tb)
             self.result.append(
                 APIResponse(
                     id=self.task_id,
@@ -381,39 +232,52 @@ def create_api_request(
     prompt: Conversation,
     attempts_left: int,
     status_tracker: StatusTracker,
-    retry_queue: asyncio.Queue,
     results_arr: list["APIRequestBase"],
     request_timeout: int = 30,
     sampling_params: SamplingParams = SamplingParams(),
-    logprobs: bool = False,
-    top_logprobs: int | None = None,
-    pbar: tqdm | None = None,
     callback: Callable | None = None,
     all_model_names: list[str] | None = None,
     all_sampling_params: list[SamplingParams] | None = None,
     tools: list | None = None,
     cache: CachePattern | None = None,
+    computer_use: bool = False,
+    display_width: int = 1024,
+    display_height: int = 768,
+    use_responses_api: bool = False,
 ) -> APIRequestBase:
     from .common import CLASSES  # circular import so made it lazy, does this work?
 
     model_obj = APIModel.from_registry(model_name)
-    request_class = CLASSES.get(model_obj.api_spec, None)
+
+    # Choose API spec based on use_responses_api flag and model support
+    api_spec = model_obj.api_spec
+    if use_responses_api and model_obj.supports_responses and api_spec == "openai":
+        api_spec = "openai-responses"
+
+    request_class = CLASSES.get(api_spec, None)
     if request_class is None:
-        raise ValueError(f"Unsupported API spec: {model_obj.api_spec}")
-    kwargs = (
-        {} if not logprobs else {"logprobs": logprobs, "top_logprobs": top_logprobs}
-    )
+        raise ValueError(f"Unsupported API spec: {api_spec}")
+    kwargs = {}
+    # Add computer_use to kwargs if the request class supports it
+    model_obj = APIModel.from_registry(model_name)
+    if computer_use and api_spec in ["anthropic", "bedrock", "openai-responses"]:
+        kwargs.update(
+            {
+                "computer_use": computer_use,
+                "display_width": display_width,
+                "display_height": display_height,
+            }
+        )
+
     return request_class(
         task_id=task_id,
         model_name=model_name,
         prompt=prompt,
         attempts_left=attempts_left,
         status_tracker=status_tracker,
-        retry_queue=retry_queue,
         results_arr=results_arr,
         request_timeout=request_timeout,
         sampling_params=sampling_params,
-        pbar=pbar,
         callback=callback,
         all_model_names=all_model_names,
         all_sampling_params=all_sampling_params,
@@ -421,3 +285,22 @@ def create_api_request(
         cache=cache,
         **kwargs,
     )
+
+
+def deduplicate_responses(results: list[APIRequestBase]) -> list[APIResponse]:
+    deduplicated = {}
+    for request in results:
+        if request.task_id not in deduplicated:
+            deduplicated[request.task_id] = request.result[-1]
+        else:
+            current_response: APIResponse = deduplicated[request.task_id]
+            # only replace if the current request has no completion and the new one does
+            if (
+                request.result[-1].completion is not None
+                and current_response.completion is None
+            ):
+                deduplicated[request.task_id] = request.result[-1]
+
+    output = [deduplicated[request.task_id] for request in results]
+
+    return output

lm_deluge/api_requests/bedrock.py

@@ -2,7 +2,6 @@ import asyncio
 import json
 import os
 from aiohttp import ClientResponse
-from tqdm import tqdm
 from typing import Callable
 
 try:
@@ -24,7 +23,7 @@ from lm_deluge.usage import Usage
 from .base import APIRequestBase, APIResponse
 
 from ..tracker import StatusTracker
-from ..sampling_params import SamplingParams
+from ..config import SamplingParams
 from ..models import APIModel
 
 
@@ -36,17 +35,18 @@ class BedrockRequest(APIRequestBase):
         prompt: Conversation,
         attempts_left: int,
         status_tracker: StatusTracker,
-        retry_queue: asyncio.Queue,
         results_arr: list,
         request_timeout: int = 30,
         sampling_params: SamplingParams = SamplingParams(),
-        pbar: tqdm | None = None,
         callback: Callable | None = None,
-        debug: bool = False,
         all_model_names: list[str] | None = None,
         all_sampling_params: list[SamplingParams] | None = None,
         tools: list | None = None,
         cache: CachePattern | None = None,
+        # Computer Use support
+        computer_use: bool = False,
+        display_width: int = 1024,
+        display_height: int = 768,
     ):
         super().__init__(
             task_id=task_id,
@@ -54,19 +54,20 @@ class BedrockRequest(APIRequestBase):
             prompt=prompt,
             attempts_left=attempts_left,
             status_tracker=status_tracker,
-            retry_queue=retry_queue,
             results_arr=results_arr,
             request_timeout=request_timeout,
             sampling_params=sampling_params,
-            pbar=pbar,
             callback=callback,
-            debug=debug,
             all_model_names=all_model_names,
             all_sampling_params=all_sampling_params,
             tools=tools,
             cache=cache,
         )
 
+        self.computer_use = computer_use
+        self.display_width = display_width
+        self.display_height = display_height
+
         # Lock images as bytes if caching is enabled
         if cache is not None:
             prompt.lock_images_as_bytes()
@@ -115,11 +116,34 @@ class BedrockRequest(APIRequestBase):
         if self.system_message is not None:
             self.request_json["system"] = self.system_message
 
-        if tools:
-            tool_definitions = [tool.dump_for("anthropic") for tool in tools]
+        if tools or self.computer_use:
+            tool_definitions = []
+
+            # Add Computer Use tools at the beginning if enabled
+            if self.computer_use:
+                from ..computer_use.anthropic_tools import get_anthropic_cu_tools
+
+                cu_tools = get_anthropic_cu_tools(
+                    model=self.model.id,
+                    display_width=self.display_width,
+                    display_height=self.display_height,
+                )
+                tool_definitions.extend(cu_tools)
+
+                # Add computer use display parameters to the request
+                self.request_json["computer_use_display_width_px"] = self.display_width
+                self.request_json["computer_use_display_height_px"] = (
+                    self.display_height
+                )
+
+            # Add user-provided tools
+            if tools:
+                tool_definitions.extend([tool.dump_for("anthropic") for tool in tools])
+
             # Add cache control to last tool if tools_only caching is specified
             if cache == "tools_only" and tool_definitions:
                 tool_definitions[-1]["cache_control"] = {"type": "ephemeral"}
+
             self.request_json["tools"] = tool_definitions
 
         # Setup AWS4Auth for signing