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

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:
@@ -12,11 +11,19 @@ except ImportError:
         "aws4auth is required for bedrock support. Install with: pip install requests-aws4auth"
     )
 
-from lm_deluge.prompt import Conversation, Message, Text, ToolCall, Thinking
+from lm_deluge.prompt import (
+    Conversation,
+    Message,
+    Text,
+    ToolCall,
+    Thinking,
+    CachePattern,
+)
+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
 
 
@@ -28,16 +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,
@@ -45,18 +54,24 @@ 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()
+
         self.model = APIModel.from_registry(model_name)
 
         # Get AWS credentials from environment
@@ -87,7 +102,7 @@ class BedrockRequest(APIRequestBase):
         self.url = f"https://bedrock-runtime.{self.region}.amazonaws.com/model/{self.model.name}/invoke"
 
         # Convert prompt to Anthropic format for bedrock
-        self.system_message, messages = prompt.to_anthropic()
+        self.system_message, messages = prompt.to_anthropic(cache_pattern=cache)
 
         # Prepare request body in Anthropic's bedrock format
         self.request_json = {
@@ -101,8 +116,35 @@ class BedrockRequest(APIRequestBase):
         if self.system_message is not None:
             self.request_json["system"] = self.system_message
 
-        if tools:
-            self.request_json["tools"] = [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
         self.auth = AWS4Auth(
@@ -179,8 +221,7 @@ class BedrockRequest(APIRequestBase):
                     is_error=True,
                     error_message="Request timed out (terminated by client).",
                     content=None,
-                    input_tokens=None,
-                    output_tokens=None,
+                    usage=None,
                 )
             )
             self.handle_error(create_new_request=False)
@@ -199,8 +240,7 @@ class BedrockRequest(APIRequestBase):
                     is_error=True,
                     error_message=f"Unexpected {type(e).__name__}: {str(e) or 'No message.'}",
                     content=None,
-                    input_tokens=None,
-                    output_tokens=None,
+                    usage=None,
                 )
             )
             self.handle_error(create_new_request=False)
@@ -210,8 +250,7 @@ class BedrockRequest(APIRequestBase):
         error_message = None
         thinking = None
         content = None
-        input_tokens = None
-        output_tokens = None
+        usage = None
         status_code = http_response.status
         mimetype = http_response.headers.get("Content-Type", None)
 
@@ -238,8 +277,7 @@ class BedrockRequest(APIRequestBase):
                         )
 
                 content = Message("assistant", parts)
-                input_tokens = data["usage"]["input_tokens"]
-                output_tokens = data["usage"]["output_tokens"]
+                usage = Usage.from_anthropic_usage(data["usage"])
             except Exception as e:
                 is_error = True
                 error_message = (
@@ -278,6 +316,5 @@ class BedrockRequest(APIRequestBase):
             model_internal=self.model_name,
             region=self.region,
             sampling_params=self.sampling_params,
-            input_tokens=input_tokens,
-            output_tokens=output_tokens,
+            usage=usage,
         )

lm_deluge/api_requests/common.py

@@ -1,10 +1,11 @@
-from .openai import OpenAIRequest
+from .openai import OpenAIRequest, OpenAIResponsesRequest
 from .anthropic import AnthropicRequest
 from .mistral import MistralRequest
 from .bedrock import BedrockRequest
 
 CLASSES = {
     "openai": OpenAIRequest,
+    "openai-responses": OpenAIResponsesRequest,
     "anthropic": AnthropicRequest,
     "mistral": MistralRequest,
     "bedrock": BedrockRequest,

lm_deluge/api_requests/mistral.py

@@ -1,15 +1,14 @@
-import asyncio
 import warnings
 from aiohttp import ClientResponse
 import json
 import os
-from tqdm.auto import tqdm
 from typing import Callable
 
 from .base import APIRequestBase, APIResponse
-from ..prompt import Conversation, Message
+from ..prompt import Conversation, Message, CachePattern
+from ..usage import Usage
 from ..tracker import StatusTracker
-from ..sampling_params import SamplingParams
+from ..config import SamplingParams
 from ..models import APIModel
 
 
@@ -23,18 +22,14 @@ class MistralRequest(APIRequestBase):
         prompt: Conversation,
         attempts_left: int,
         status_tracker: StatusTracker,
-        retry_queue: asyncio.Queue,
         results_arr: list,
         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,
+        cache: CachePattern | None = None,
     ):
         super().__init__(
             task_id=task_id,
@@ -42,18 +37,21 @@ class MistralRequest(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,
-            logprobs=logprobs,
-            top_logprobs=top_logprobs,
-            pbar=pbar,
             callback=callback,
-            debug=debug,
             all_model_names=all_model_names,
             all_sampling_params=all_sampling_params,
+            tools=tools,
+            cache=cache,
         )
+
+        # Warn if cache is specified for non-Anthropic model
+        if cache is not None:
+            warnings.warn(
+                f"Cache parameter '{cache}' is only supported for Anthropic models, ignoring for {model_name}"
+            )
         self.model = APIModel.from_registry(model_name)
         self.url = f"{self.model.api_base}/chat/completions"
         self.request_header = {
@@ -70,7 +68,7 @@ class MistralRequest(APIRequestBase):
             warnings.warn(
                 f"Ignoring reasoning_effort param for non-reasoning model: {model_name}"
             )
-        if logprobs:
+        if sampling_params.logprobs:
             warnings.warn(
                 f"Ignoring logprobs param for non-logprobs model: {model_name}"
             )
@@ -81,8 +79,7 @@ class MistralRequest(APIRequestBase):
         is_error = False
         error_message = None
         completion = None
-        input_tokens = None
-        output_tokens = None
+        usage = None
         logprobs = None
         status_code = http_response.status
         mimetype = http_response.headers.get("Content-Type", None)
@@ -99,9 +96,11 @@ class MistralRequest(APIRequestBase):
                 assert data is not None, "data is None"
                 try:
                     completion = data["choices"][0]["message"]["content"]
-                    input_tokens = data["usage"]["prompt_tokens"]
-                    output_tokens = data["usage"]["completion_tokens"]
-                    if self.logprobs and "logprobs" in data["choices"][0]:
+                    usage = Usage.from_mistral_usage(data["usage"])
+                    if (
+                        self.sampling_params.logprobs
+                        and "logprobs" in data["choices"][0]
+                    ):
                         logprobs = data["choices"][0]["logprobs"]["content"]
                 except Exception:
                     is_error = True
@@ -134,6 +133,5 @@ class MistralRequest(APIRequestBase):
             content=Message.ai(completion),
             model_internal=self.model_name,
             sampling_params=self.sampling_params,
-            input_tokens=input_tokens,
-            output_tokens=output_tokens,
+            usage=usage,
         )

lm_deluge/api_requests/openai.py

@@ -1,18 +1,56 @@
-import asyncio
 import warnings
 from aiohttp import ClientResponse
 import json
 import os
-from tqdm.auto import tqdm
 from typing import Callable
 
+from lm_deluge.tool import Tool
+
 from .base import APIRequestBase, APIResponse
-from ..prompt import Conversation, Message, Text, ToolCall, Thinking
+from ..prompt import Conversation, Message, Text, ToolCall, Thinking, CachePattern
+from ..usage import Usage
 from ..tracker import StatusTracker
-from ..sampling_params import SamplingParams
+from ..config import SamplingParams
 from ..models import APIModel
 
 
+def _build_oa_chat_request(
+    model: APIModel,
+    prompt: Conversation,
+    tools: list[Tool] | None,
+    sampling_params: SamplingParams,
+) -> dict:
+    request_json = {
+        "model": model.name,
+        "messages": prompt.to_openai(),
+        "temperature": sampling_params.temperature,
+        "top_p": sampling_params.top_p,
+    }
+    # set max_tokens or max_completion_tokens dep. on provider
+    if "cohere" in model.api_base:
+        request_json["max_tokens"] = sampling_params.max_new_tokens
+    else:
+        request_json["max_completion_tokens"] = sampling_params.max_new_tokens
+    if model.reasoning_model:
+        request_json["temperature"] = 1.0
+        request_json["top_p"] = 1.0
+        request_json["reasoning_effort"] = sampling_params.reasoning_effort
+    else:
+        if sampling_params.reasoning_effort:
+            warnings.warn(
+                f"Ignoring reasoning_effort param for non-reasoning model: {model.name}"
+            )
+    if sampling_params.logprobs:
+        request_json["logprobs"] = True
+        if sampling_params.top_logprobs is not None:
+            request_json["top_logprobs"] = sampling_params.top_logprobs
+    if sampling_params.json_mode and model.supports_json:
+        request_json["response_format"] = {"type": "json_object"}
+    if tools:
+        request_json["tools"] = [tool.dump_for("openai-completions") for tool in tools]
+    return request_json
+
+
 class OpenAIRequest(APIRequestBase):
     def __init__(
         self,
@@ -23,18 +61,14 @@ class OpenAIRequest(APIRequestBase):
         prompt: Conversation,
         attempts_left: int,
         status_tracker: StatusTracker,
-        retry_queue: asyncio.Queue,
         results_arr: list,
         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,
+        cache: CachePattern | None = None,
     ):
         super().__init__(
             task_id=task_id,
@@ -42,63 +76,37 @@ class OpenAIRequest(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,
-            logprobs=logprobs,
-            top_logprobs=top_logprobs,
-            pbar=pbar,
             callback=callback,
-            debug=debug,
             all_model_names=all_model_names,
             all_sampling_params=all_sampling_params,
             tools=tools,
+            cache=cache,
         )
+
+        # Warn if cache is specified for non-Anthropic model
+        if cache is not None:
+            warnings.warn(
+                f"Cache parameter '{cache}' is only supported for Anthropic models, ignoring for {model_name}"
+            )
         self.model = APIModel.from_registry(model_name)
         self.url = f"{self.model.api_base}/chat/completions"
         self.request_header = {
             "Authorization": f"Bearer {os.getenv(self.model.api_key_env_var)}"
         }
 
-        self.request_json = {
-            "model": self.model.name,
-            "messages": prompt.to_openai(),
-            "temperature": sampling_params.temperature,
-            "top_p": sampling_params.top_p,
-        }
-        # set max_tokens or max_completion_tokens dep. on provider
-        if "cohere" in self.model.api_base:
-            self.request_json["max_tokens"] = sampling_params.max_new_tokens
-        elif "openai" in self.model.api_base:
-            self.request_json["max_completion_tokens"] = sampling_params.max_new_tokens
-        if self.model.reasoning_model:
-            self.request_json["temperature"] = 1.0
-            self.request_json["top_p"] = 1.0
-            self.request_json["reasoning_effort"] = sampling_params.reasoning_effort
-        else:
-            if sampling_params.reasoning_effort:
-                warnings.warn(
-                    f"Ignoring reasoning_effort param for non-reasoning model: {model_name}"
-                )
-        if logprobs:
-            self.request_json["logprobs"] = True
-            if top_logprobs is not None:
-                self.request_json["top_logprobs"] = top_logprobs
-        if sampling_params.json_mode and self.model.supports_json:
-            self.request_json["response_format"] = {"type": "json_object"}
-        if tools:
-            self.request_json["tools"] = [
-                tool.dump_for("openai-completions") for tool in tools
-            ]
+        self.request_json = _build_oa_chat_request(
+            self.model, prompt, tools, sampling_params
+        )
 
     async def handle_response(self, http_response: ClientResponse) -> APIResponse:
         is_error = False
         error_message = None
         thinking = None
         content = None
-        input_tokens = None
-        output_tokens = None
+        usage = None
         logprobs = None
         status_code = http_response.status
         mimetype = http_response.headers.get("Content-Type", None)
@@ -142,9 +150,11 @@ class OpenAIRequest(APIRequestBase):
 
                     content = Message("assistant", parts)
 
-                    input_tokens = data["usage"]["prompt_tokens"]
-                    output_tokens = data["usage"]["completion_tokens"]
-                    if self.logprobs and "logprobs" in data["choices"][0]:
+                    usage = Usage.from_openai_usage(data["usage"])
+                    if (
+                        self.sampling_params.logprobs
+                        and "logprobs" in data["choices"][0]
+                    ):
                         logprobs = data["choices"][0]["logprobs"]["content"]
                 except Exception:
                     is_error = True
@@ -178,6 +188,228 @@ class OpenAIRequest(APIRequestBase):
             content=content,
             model_internal=self.model_name,
             sampling_params=self.sampling_params,
-            input_tokens=input_tokens,
-            output_tokens=output_tokens,
+            usage=usage,
+            raw_response=data,
+        )
+
+
+class OpenAIResponsesRequest(APIRequestBase):
+    def __init__(
+        self,
+        task_id: int,
+        model_name: str,
+        prompt: Conversation,
+        attempts_left: int,
+        status_tracker: StatusTracker,
+        results_arr: list,
+        request_timeout: int = 30,
+        sampling_params: SamplingParams = SamplingParams(),
+        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,
+    ):
+        super().__init__(
+            task_id=task_id,
+            model_name=model_name,
+            prompt=prompt,
+            attempts_left=attempts_left,
+            status_tracker=status_tracker,
+            results_arr=results_arr,
+            request_timeout=request_timeout,
+            sampling_params=sampling_params,
+            callback=callback,
+            all_model_names=all_model_names,
+            all_sampling_params=all_sampling_params,
+            tools=tools,
+            cache=cache,
+        )
+
+        # Store computer use parameters
+        self.computer_use = computer_use
+        self.display_width = display_width
+        self.display_height = display_height
+
+        # Validate computer use requirements
+        if computer_use and model_name != "openai-computer-use-preview":
+            raise ValueError(
+                f"Computer use is only supported with openai-computer-use-preview model, got {model_name}"
+            )
+
+        # Warn if cache is specified for non-Anthropic model
+        if cache is not None:
+            warnings.warn(
+                f"Cache parameter '{cache}' is only supported for Anthropic models, ignoring for {model_name}"
+            )
+        self.model = APIModel.from_registry(model_name)
+        self.url = f"{self.model.api_base}/responses"
+        self.request_header = {
+            "Authorization": f"Bearer {os.getenv(self.model.api_key_env_var)}"
+        }
+
+        # Convert conversation to input format for Responses API
+        openai_responses_format = prompt.to_openai_responses()
+
+        self.request_json = {
+            "model": self.model.name,
+            "input": openai_responses_format["input"],
+            "temperature": sampling_params.temperature,
+            "top_p": sampling_params.top_p,
+        }
+
+        # Add max_output_tokens for responses API
+        if sampling_params.max_new_tokens:
+            self.request_json["max_output_tokens"] = sampling_params.max_new_tokens
+
+        if self.model.reasoning_model:
+            self.request_json["temperature"] = 1.0
+            self.request_json["top_p"] = 1.0
+            self.request_json["reasoning"] = {
+                "effort": sampling_params.reasoning_effort
+            }
+        else:
+            if sampling_params.reasoning_effort:
+                warnings.warn(
+                    f"Ignoring reasoning_effort param for non-reasoning model: {model_name}"
+                )
+
+        if sampling_params.json_mode and self.model.supports_json:
+            self.request_json["text"] = {"format": {"type": "json_object"}}
+
+        # Handle tools
+        request_tools = []
+        if computer_use:
+            # Add computer use tool
+            request_tools.append(
+                {
+                    "type": "computer_use_preview",
+                    "display_width": display_width,
+                    "display_height": display_height,
+                    "environment": "browser",  # Default to browser, could be configurable
+                }
+            )
+            # Set truncation to auto as required for computer use
+            self.request_json["truncation"] = "auto"
+
+        if tools:
+            # Add regular function tools
+            request_tools.extend([tool.dump_for("openai-responses") for tool in tools])
+
+        if request_tools:
+            self.request_json["tools"] = request_tools
+
+    async def handle_response(self, http_response: ClientResponse) -> APIResponse:
+        is_error = False
+        error_message = None
+        thinking = None
+        content = None
+        usage = None
+        logprobs = None
+        status_code = http_response.status
+        mimetype = http_response.headers.get("Content-Type", None)
+        data = None
+
+        if status_code >= 200 and status_code < 300:
+            try:
+                data = await http_response.json()
+            except Exception:
+                is_error = True
+                error_message = (
+                    f"Error calling .json() on response w/ status {status_code}"
+                )
+            if not is_error:
+                assert data is not None, "data is None"
+                try:
+                    # Parse Responses API format
+                    parts = []
+
+                    # Get the output array from the response
+                    output = data.get("output", [])
+                    if not output:
+                        is_error = True
+                        error_message = "No output in response"
+                    else:
+                        # Process each output item
+                        for item in output:
+                            if item.get("type") == "message":
+                                message_content = item.get("content", [])
+                                for content_item in message_content:
+                                    if content_item.get("type") == "output_text":
+                                        parts.append(Text(content_item["text"]))
+                                    # Handle tool calls if present
+                                    elif content_item.get("type") == "tool_call":
+                                        tool_call = content_item["tool_call"]
+                                        parts.append(
+                                            ToolCall(
+                                                id=tool_call["id"],
+                                                name=tool_call["function"]["name"],
+                                                arguments=json.loads(
+                                                    tool_call["function"]["arguments"]
+                                                ),
+                                            )
+                                        )
+                            elif item.get("type") == "computer_call":
+                                # Handle computer use actions
+                                action = item.get("action", {})
+                                parts.append(
+                                    ToolCall(
+                                        id=item["call_id"],
+                                        name=f"_computer_{action.get('type', 'action')}",
+                                        arguments=action,
+                                    )
+                                )
+
+                        # Handle reasoning if present
+                        if "reasoning" in data and data["reasoning"].get("summary"):
+                            thinking = data["reasoning"]["summary"]
+                            parts.append(Thinking(thinking))
+
+                        content = Message("assistant", parts)
+
+                        # Extract usage information
+                        if "usage" in data:
+                            usage = Usage.from_openai_usage(data["usage"])
+
+                        # Extract response_id for computer use continuation
+                        # response_id = data.get("id")
+
+                except Exception as e:
+                    is_error = True
+                    error_message = f"Error parsing {self.model.name} responses API response: {str(e)}"
+
+        elif mimetype and "json" in mimetype.lower():
+            is_error = True
+            data = await http_response.json()
+            error_message = json.dumps(data)
+        else:
+            is_error = True
+            text = await http_response.text()
+            error_message = text
+
+        # Handle special kinds of errors
+        if is_error and error_message is not None:
+            if "rate limit" in error_message.lower() or status_code == 429:
+                error_message += " (Rate limit error, triggering cooldown.)"
+                self.status_tracker.rate_limit_exceeded()
+            if "context length" in error_message:
+                error_message += " (Context length exceeded, set retries to 0.)"
+                self.attempts_left = 0
+
+        return APIResponse(
+            id=self.task_id,
+            status_code=status_code,
+            is_error=is_error,
+            error_message=error_message,
+            prompt=self.prompt,
+            logprobs=logprobs,
+            thinking=thinking,
+            content=content,
+            model_internal=self.model_name,
+            sampling_params=self.sampling_params,
+            usage=usage,
+            raw_response=data,
         )