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

lm_deluge/api_requests/bedrock.py

@@ -1,8 +1,8 @@
 import asyncio
 import json
 import os
+
 from aiohttp import ClientResponse
-from typing import Callable
 
 try:
     from requests_aws4auth import AWS4Auth
@@ -12,186 +12,178 @@ except ImportError:
     )
 
 from lm_deluge.prompt import (
+    CachePattern,
     Conversation,
     Message,
     Text,
-    ToolCall,
     Thinking,
-    CachePattern,
+    ToolCall,
 )
+from lm_deluge.request_context import RequestContext
+from lm_deluge.tool import MCPServer, Tool
 from lm_deluge.usage import Usage
-from .base import APIRequestBase, APIResponse
 
-from ..tracker import StatusTracker
 from ..config import SamplingParams
 from ..models import APIModel
+from .base import APIRequestBase, APIResponse
 
 
-class BedrockRequest(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 support
-        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,
+# according to bedrock docs the header is "anthropic_beta" vs. "anthropic-beta"
+# for anthropic. i don't know if this is a typo or the worst ever UX
+def _add_beta(headers: dict, beta: str):
+    if "anthropic_beta" in headers and headers["anthropic_beta"]:
+        if beta not in headers["anthropic_beta"]:
+            headers["anthropic_beta"] += f",{beta}"
+    else:
+        headers["anthropic_beta"] = beta
+
+
+def _build_anthropic_bedrock_request(
+    model: APIModel,
+    prompt: Conversation,
+    tools: list[Tool | dict | MCPServer] | None,
+    sampling_params: SamplingParams,
+    cache_pattern: CachePattern | None = None,
+):
+    system_message, messages = prompt.to_anthropic(cache_pattern=cache_pattern)
+
+    # handle AWS auth
+    access_key = os.getenv("AWS_ACCESS_KEY_ID")
+    secret_key = os.getenv("AWS_SECRET_ACCESS_KEY")
+    session_token = os.getenv("AWS_SESSION_TOKEN")
+
+    if not access_key or not secret_key:
+        raise ValueError(
+            "AWS credentials not found. Please set AWS_ACCESS_KEY_ID and AWS_SECRET_ACCESS_KEY environment variables."
         )
 
-        self.computer_use = computer_use
-        self.display_width = display_width
-        self.display_height = display_height
+    # Determine region - use us-west-2 for cross-region inference models
+    if model.name.startswith("us.anthropic."):
+        # Cross-region inference profiles should use us-west-2
+        region = "us-west-2"
+    else:
+        raise ValueError("only cross-region inference for bedrock")
+        # # Direct model IDs can use default region
+        # region = getattr(model, "region", "us-east-1")
+        # if hasattr(model, "regions") and model.regions:
+        #     if isinstance(model.regions, list):
+        #         region = model.regions[0]
+        #     elif isinstance(model.regions, dict):
+        #         region = list(model.regions.keys())[0]
+
+    # Construct the endpoint URL
+    service = "bedrock"  # Service name for signing is 'bedrock' even though endpoint is bedrock-runtime
+    url = f"https://bedrock-runtime.{region}.amazonaws.com/model/{model.name}/invoke"
+
+    # Prepare headers
+    auth = AWS4Auth(
+        access_key,
+        secret_key,
+        region,
+        service,
+        session_token=session_token,
+    )
 
-        # Lock images as bytes if caching is enabled
-        if cache is not None:
-            prompt.lock_images_as_bytes()
+    # Setup basic headers (AWS4Auth will add the Authorization header)
+    request_header = {
+        "Content-Type": "application/json",
+    }
+
+    # Prepare request body in Anthropic's bedrock format
+    request_json = {
+        "anthropic_version": "bedrock-2023-05-31",
+        "max_tokens": sampling_params.max_new_tokens,
+        "temperature": sampling_params.temperature,
+        "top_p": sampling_params.top_p,
+        "messages": messages,
+    }
+
+    if system_message is not None:
+        request_json["system"] = system_message
+
+    if tools:
+        mcp_servers = []
+        tool_definitions = []
+        for tool in tools:
+            if isinstance(tool, Tool):
+                tool_definitions.append(tool.dump_for("anthropic"))
+            elif isinstance(tool, dict):
+                tool_definitions.append(tool)
+                # add betas if needed
+                if tool["type"] in [
+                    "computer_20241022",
+                    "text_editor_20241022",
+                    "bash_20241022",
+                ]:
+                    _add_beta(request_header, "computer-use-2024-10-22")
+                elif tool["type"] == "computer_20250124":
+                    _add_beta(request_header, "computer-use-2025-01-24")
+                elif tool["type"] == "code_execution_20250522":
+                    _add_beta(request_header, "code-execution-2025-05-22")
+            elif isinstance(tool, MCPServer):
+                raise ValueError("bedrock doesn't support MCP connector right now")
+                # _add_beta(request_header, "mcp-client-2025-04-04")
+                # mcp_servers.append(tool.for_anthropic())
+
+        # 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
+        if len(mcp_servers) > 0:
+            request_json["mcp_servers"] = mcp_servers
+
+    return request_json, request_header, auth, url
 
-        self.model = APIModel.from_registry(model_name)
 
-        # Get AWS credentials from environment
-        self.access_key = os.getenv("AWS_ACCESS_KEY_ID")
-        self.secret_key = os.getenv("AWS_SECRET_ACCESS_KEY")
-        self.session_token = os.getenv("AWS_SESSION_TOKEN")
+class BedrockRequest(APIRequestBase):
+    def __init__(self, context: RequestContext):
+        super().__init__(context=context)
+
+        self.model = APIModel.from_registry(self.context.model_name)
+        self.url = f"{self.model.api_base}/messages"
 
-        if not self.access_key or not self.secret_key:
-            raise ValueError(
-                "AWS credentials not found. Please set AWS_ACCESS_KEY_ID and AWS_SECRET_ACCESS_KEY environment variables."
+        # Lock images as bytes if caching is enabled
+        if self.context.cache is not None:
+            self.context.prompt.lock_images_as_bytes()
+
+        self.request_json, self.request_header, self.auth, self.url = (
+            _build_anthropic_bedrock_request(
+                self.model,
+                context.prompt,
+                context.tools,
+                context.sampling_params,
+                context.cache,
             )
+        )
 
-        # Determine region - use us-west-2 for cross-region inference models
-        if self.model.name.startswith("us.anthropic."):
-            # Cross-region inference profiles should use us-west-2
-            self.region = "us-west-2"
-        else:
-            # Direct model IDs can use default region
-            self.region = getattr(self.model, "region", "us-east-1")
-            if hasattr(self.model, "regions") and self.model.regions:
-                if isinstance(self.model.regions, list):
-                    self.region = self.model.regions[0]
-                elif isinstance(self.model.regions, dict):
-                    self.region = list(self.model.regions.keys())[0]
-
-        # Construct the endpoint URL
-        self.service = "bedrock"  # Service name for signing is 'bedrock' even though endpoint is bedrock-runtime
-        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(cache_pattern=cache)
-
-        # Prepare request body in Anthropic's bedrock format
-        self.request_json = {
-            "anthropic_version": "bedrock-2023-05-31",
-            "max_tokens": sampling_params.max_new_tokens,
-            "temperature": sampling_params.temperature,
-            "top_p": sampling_params.top_p,
-            "messages": messages,
-        }
-
-        if self.system_message is not None:
-            self.request_json["system"] = self.system_message
-
-        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)
+    async def execute_once(self) -> APIResponse:
+        """Override execute_once to handle AWS4Auth signing."""
+        import aiohttp
 
-                # 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
-                )
+        assert self.context.status_tracker
 
-            # Add user-provided tools
-            if tools:
-                tool_definitions.extend([tool.dump_for("anthropic") for tool in tools])
+        self.context.status_tracker.total_requests += 1
+        timeout = aiohttp.ClientTimeout(total=self.context.request_timeout)
 
-            # 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"}
+        # Prepare the request data
+        payload = json.dumps(self.request_json, separators=(",", ":")).encode("utf-8")
 
-            self.request_json["tools"] = tool_definitions
+        # Create a fake requests.PreparedRequest object for AWS4Auth to sign
+        import requests
 
-        # Setup AWS4Auth for signing
-        self.auth = AWS4Auth(
-            self.access_key,
-            self.secret_key,
-            self.region,
-            self.service,
-            session_token=self.session_token,
+        fake_request = requests.Request(
+            method="POST",
+            url=self.url,
+            data=payload,
+            headers=self.request_header.copy(),
         )
 
-        # Setup basic headers (AWS4Auth will add the Authorization header)
-        self.request_header = {
-            "Content-Type": "application/json",
-        }
+        prepared_request = fake_request.prepare()
+        signed_request = self.auth(prepared_request)
+        signed_headers = dict(signed_request.headers)
 
-    async def call_api(self):
-        """Override call_api to handle AWS4Auth signing."""
         try:
-            import aiohttp
-
-            self.status_tracker.total_requests += 1
-            timeout = aiohttp.ClientTimeout(total=self.request_timeout)
-
-            # Prepare the request data
-            payload = json.dumps(self.request_json, separators=(",", ":")).encode(
-                "utf-8"
-            )
-
-            # Create a fake requests.PreparedRequest object for AWS4Auth to sign
-            import requests
-
-            fake_request = requests.Request(
-                method="POST",
-                url=self.url,
-                data=payload,
-                headers=self.request_header.copy(),
-            )
-
-            # Prepare the request so AWS4Auth can sign it properly
-            prepared_request = fake_request.prepare()
-
-            # Let AWS4Auth sign the prepared request
-            signed_request = self.auth(prepared_request)
-
-            # Extract the signed headers
-            signed_headers = dict(signed_request.headers)
-
             async with aiohttp.ClientSession(timeout=timeout) as session:
                 async with session.post(
                     url=self.url,
@@ -199,51 +191,36 @@ class BedrockRequest(APIRequestBase):
                     data=payload,
                 ) as http_response:
                     response: APIResponse = await self.handle_response(http_response)
-
-            self.result.append(response)
-            if response.is_error:
-                self.handle_error(
-                    create_new_request=response.retry_with_different_model or False,
-                    give_up_if_no_other_models=response.give_up_if_no_other_models
-                    or False,
-                )
-            else:
-                self.handle_success(response)
+            return response
 
         except asyncio.TimeoutError:
-            self.result.append(
-                APIResponse(
-                    id=self.task_id,
-                    model_internal=self.model_name,
-                    prompt=self.prompt,
-                    sampling_params=self.sampling_params,
-                    status_code=None,
-                    is_error=True,
-                    error_message="Request timed out (terminated by client).",
-                    content=None,
-                    usage=None,
-                )
+            return APIResponse(
+                id=self.context.task_id,
+                model_internal=self.context.model_name,
+                prompt=self.context.prompt,
+                sampling_params=self.context.sampling_params,
+                status_code=None,
+                is_error=True,
+                error_message="Request timed out (terminated by client).",
+                content=None,
+                usage=None,
             )
-            self.handle_error(create_new_request=False)
 
         except Exception as e:
             from ..errors import raise_if_modal_exception
 
             raise_if_modal_exception(e)
-            self.result.append(
-                APIResponse(
-                    id=self.task_id,
-                    model_internal=self.model_name,
-                    prompt=self.prompt,
-                    sampling_params=self.sampling_params,
-                    status_code=None,
-                    is_error=True,
-                    error_message=f"Unexpected {type(e).__name__}: {str(e) or 'No message.'}",
-                    content=None,
-                    usage=None,
-                )
+            return APIResponse(
+                id=self.context.task_id,
+                model_internal=self.context.model_name,
+                prompt=self.context.prompt,
+                sampling_params=self.context.sampling_params,
+                status_code=None,
+                is_error=True,
+                error_message=f"Unexpected {type(e).__name__}: {str(e) or 'No message.'}",
+                content=None,
+                usage=None,
             )
-            self.handle_error(create_new_request=False)
 
     async def handle_response(self, http_response: ClientResponse) -> APIResponse:
         is_error = False
@@ -253,6 +230,7 @@ class BedrockRequest(APIRequestBase):
         usage = None
         status_code = http_response.status
         mimetype = http_response.headers.get("Content-Type", None)
+        assert self.context.status_tracker
 
         if status_code >= 200 and status_code < 300:
             try:
@@ -300,21 +278,21 @@ class BedrockRequest(APIRequestBase):
                 or status_code == 429
             ):
                 error_message += " (Rate limit error, triggering cooldown.)"
-                self.status_tracker.rate_limit_exceeded()
+                self.context.status_tracker.rate_limit_exceeded()
             if "context length" in error_message or "too long" in error_message:
                 error_message += " (Context length exceeded, set retries to 0.)"
-                self.attempts_left = 0
+                self.context.attempts_left = 0
 
         return APIResponse(
-            id=self.task_id,
+            id=self.context.task_id,
             status_code=status_code,
             is_error=is_error,
             error_message=error_message,
-            prompt=self.prompt,
+            prompt=self.context.prompt,
             content=content,
             thinking=thinking,
-            model_internal=self.model_name,
+            model_internal=self.context.model_name,
             region=self.region,
-            sampling_params=self.sampling_params,
+            sampling_params=self.context.sampling_params,
             usage=usage,
         )

lm_deluge/api_requests/gemini.py

@@ -1,16 +1,15 @@
 import json
 import os
 import warnings
-from typing import Callable
 
 from aiohttp import ClientResponse
 
+from lm_deluge.request_context import RequestContext
 from lm_deluge.tool import Tool
 
 from ..config import SamplingParams
 from ..models import APIModel
-from ..prompt import CachePattern, Conversation, Message, Text, Thinking, ToolCall
-from ..tracker import StatusTracker
+from ..prompt import Conversation, Message, Text, Thinking, ToolCall
 from ..usage import Usage
 from .base import APIRequestBase, APIResponse
 
@@ -66,45 +65,16 @@ def _build_gemini_request(
 
 
 class GeminiRequest(APIRequestBase):
-    def __init__(
-        self,
-        task_id: int,
-        model_name: str,  # must correspond to registry
-        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,
-    ):
-        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,
-        )
+    def __init__(self, context: RequestContext):
+        super().__init__(context=context)
 
         # Warn if cache is specified for Gemini model
-        if cache is not None:
+        if self.context.cache is not None:
             warnings.warn(
-                f"Cache parameter '{cache}' is not supported for Gemini models, ignoring for {model_name}"
+                f"Cache parameter '{self.context.cache}' is not supported for Gemini models, ignoring for {self.context.model_name}"
             )
 
-        self.model = APIModel.from_registry(model_name)
+        self.model = APIModel.from_registry(self.context.model_name)
         # Gemini API endpoint format: https://generativelanguage.googleapis.com/v1beta/models/{model}:generateContent
         self.url = f"{self.model.api_base}/models/{self.model.name}:generateContent"
         self.request_header = {
@@ -120,7 +90,10 @@ class GeminiRequest(APIRequestBase):
         self.url += f"?key={api_key}"
 
         self.request_json = _build_gemini_request(
-            self.model, prompt, tools, sampling_params
+            self.model,
+            self.context.prompt,
+            self.context.tools,
+            self.context.sampling_params,
         )
 
     async def handle_response(self, http_response: ClientResponse) -> APIResponse:
@@ -132,6 +105,7 @@ class GeminiRequest(APIRequestBase):
         status_code = http_response.status
         mimetype = http_response.headers.get("Content-Type", None)
         data = None
+        assert self.context.status_tracker
 
         if status_code >= 200 and status_code < 300:
             try:
@@ -199,24 +173,24 @@ class GeminiRequest(APIRequestBase):
         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()
+                self.context.status_tracker.rate_limit_exceeded()
             if (
                 "context length" in error_message.lower()
                 or "token limit" in error_message.lower()
             ):
                 error_message += " (Context length exceeded, set retries to 0.)"
-                self.attempts_left = 0
+                self.context.attempts_left = 0
 
         return APIResponse(
-            id=self.task_id,
+            id=self.context.task_id,
             status_code=status_code,
             is_error=is_error,
             error_message=error_message,
-            prompt=self.prompt,
+            prompt=self.context.prompt,
             content=content,
             thinking=thinking,
-            model_internal=self.model_name,
-            sampling_params=self.sampling_params,
+            model_internal=self.context.model_name,
+            sampling_params=self.context.sampling_params,
             usage=usage,
             raw_response=data,
         )