lm-deluge 0.0.8__tar.gz → 0.0.9__tar.gz

{lm_deluge-0.0.8/src/lm_deluge.egg-info → lm_deluge-0.0.9}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: lm_deluge
-Version: 0.0.8
+Version: 0.0.9
 Summary: Python utility for using LLM API models.
 Author-email: Benjamin Anderson <ben@trytaylor.ai>
 Requires-Python: >=3.10

{lm_deluge-0.0.8 → lm_deluge-0.0.9}/pyproject.toml

@@ -3,7 +3,7 @@ requires = ["setuptools", "wheel"]
 
 [project]
 name = "lm_deluge"
-version = "0.0.8"
+version = "0.0.9"
 authors = [{ name = "Benjamin Anderson", email = "ben@trytaylor.ai" }]
 description = "Python utility for using LLM API models."
 readme = "README.md"

{lm_deluge-0.0.8 → lm_deluge-0.0.9}/src/lm_deluge/api_requests/anthropic.py

@@ -6,7 +6,7 @@ import warnings
 from tqdm import tqdm
 from typing import Callable
 
-from lm_deluge.prompt import Conversation
+from lm_deluge.prompt import Conversation, Message, Text, ToolCall, Thinking
 from .base import APIRequestBase, APIResponse
 
 from ..tracker import StatusTracker
@@ -34,6 +34,7 @@ class AnthropicRequest(APIRequestBase):
         # for retries
         all_model_names: list[str] | None = None,
         all_sampling_params: list[SamplingParams] | None = None,
+        tools: list | None = None,
     ):
         super().__init__(
             task_id=task_id,
@@ -50,6 +51,7 @@ class AnthropicRequest(APIRequestBase):
             debug=debug,
             all_model_names=all_model_names,
             all_sampling_params=all_sampling_params,
+            tools=tools,
         )
         self.model = APIModel.from_registry(model_name)
         self.url = f"{self.model.api_base}/messages"
@@ -94,12 +96,14 @@ class AnthropicRequest(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]
 
     async def handle_response(self, http_response: ClientResponse) -> APIResponse:
         is_error = False
         error_message = None
         thinking = None
-        completion = None
+        content = None
         input_tokens = None
         output_tokens = None
         status_code = http_response.status
@@ -119,14 +123,26 @@ class AnthropicRequest(APIRequestBase):
         if status_code >= 200 and status_code < 300:
             try:
                 data = await http_response.json()
-                content = data["content"]  # [0]["text"]
-                for item in content:
+                response_content = data["content"]
+
+                # Parse response into Message with parts
+                parts = []
+                for item in response_content:
                     if item["type"] == "text":
-                        completion = item["text"]
+                        parts.append(Text(item["text"]))
                     elif item["type"] == "thinking":
                         thinking = item["thinking"]
+                        parts.append(Thinking(item["thinking"]))
                     elif item["type"] == "tool_use":
-                        continue  # TODO: implement and report tool use
+                        parts.append(
+                            ToolCall(
+                                id=item["id"],
+                                name=item["name"],
+                                arguments=item["input"],
+                            )
+                        )
+
+                content = Message("assistant", parts)
                 input_tokens = data["usage"]["input_tokens"]
                 output_tokens = data["usage"]["output_tokens"]
             except Exception as e:
@@ -162,7 +178,7 @@ class AnthropicRequest(APIRequestBase):
             is_error=is_error,
             error_message=error_message,
             prompt=self.prompt,
-            completion=completion,
+            content=content,
             thinking=thinking,
             model_internal=self.model_name,
             sampling_params=self.sampling_params,

{lm_deluge-0.0.8 → lm_deluge-0.0.9}/src/lm_deluge/api_requests/base.py

@@ -7,7 +7,7 @@ from dataclasses import dataclass
 from abc import ABC, abstractmethod
 from typing import Callable
 
-from lm_deluge.prompt import Conversation
+from lm_deluge.prompt import Conversation, Message
 
 from ..tracker import StatusTracker
 from ..sampling_params import SamplingParams
@@ -30,10 +30,12 @@ class APIResponse:
     error_message: str | None
 
     # completion information
-    completion: str | None
     input_tokens: int | None
     output_tokens: int | 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
@@ -47,6 +49,13 @@ class APIResponse:
     # 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
+
     def __post_init__(self):
         # calculate cost & get external model name
         self.id = int(self.id)
@@ -63,7 +72,7 @@ class APIResponse:
                 self.input_tokens * api_model.input_cost / 1e6
                 + self.output_tokens * api_model.output_cost / 1e6
             )
-        elif self.completion is not None:
+        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}."
             )
@@ -79,7 +88,8 @@ class APIResponse:
             "status_code": self.status_code,
             "is_error": self.is_error,
             "error_message": self.error_message,
-            "completion": self.completion,
+            "completion": self.completion,  # computed property
+            "content": self.content.to_log() if self.content else None,
             "input_tokens": self.input_tokens,
             "output_tokens": self.output_tokens,
             "finish_reason": self.finish_reason,
@@ -88,11 +98,18 @@ class APIResponse:
 
     @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"])
+
         return cls(
             id=data.get("id", random.randint(0, 1_000_000_000)),
             model_internal=data["model_internal"],
-            model_external=data["model_external"],
-            region=data["region"],
             prompt=Conversation.from_log(data["prompt"]),
             sampling_params=SamplingParams(**data["sampling_params"]),
             status_code=data["status_code"],
@@ -100,9 +117,14 @@ class APIResponse:
             error_message=data["error_message"],
             input_tokens=data["input_tokens"],
             output_tokens=data["output_tokens"],
-            completion=data["completion"],
-            finish_reason=data["finish_reason"],
-            cost=data["cost"],
+            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):
@@ -145,6 +167,7 @@ class APIRequestBase(ABC):
         debug: bool = False,
         all_model_names: list[str] | None = None,
         all_sampling_params: list[SamplingParams] | None = None,
+        tools: list | None = None,
     ):
         if all_model_names is None:
             raise ValueError("all_model_names must be provided.")
@@ -166,6 +189,7 @@ class APIRequestBase(ABC):
         self.debug = debug
         self.all_model_names = all_model_names
         self.all_sampling_params = all_sampling_params
+        self.tools = tools
         self.result = []  # list of APIResponse objects from each attempt
 
         # these should be set in the __init__ of the subclass
@@ -255,6 +279,7 @@ class APIRequestBase(ABC):
                         callback=self.callback,
                         all_model_names=self.all_model_names,
                         all_sampling_params=self.all_sampling_params,
+                        tools=self.tools,
                     )
                     # PROBLEM: new request is never put into results array, so we can't get the result.
                     self.retry_queue.put_nowait(new_request)
@@ -297,7 +322,7 @@ class APIRequestBase(ABC):
                     status_code=None,
                     is_error=True,
                     error_message="Request timed out (terminated by client).",
-                    completion=None,
+                    content=None,
                     input_tokens=None,
                     output_tokens=None,
                 )
@@ -315,7 +340,7 @@ class APIRequestBase(ABC):
                     status_code=None,
                     is_error=True,
                     error_message=f"Unexpected {type(e).__name__}: {str(e) or 'No message.'}",
-                    completion=None,
+                    content=None,
                     input_tokens=None,
                     output_tokens=None,
                 )
@@ -344,6 +369,7 @@ def create_api_request(
     callback: Callable | None = None,
     all_model_names: list[str] | None = None,
     all_sampling_params: list[SamplingParams] | None = None,
+    tools: list | None = None,
 ) -> APIRequestBase:
     from .common import CLASSES  # circular import so made it lazy, does this work?
 
@@ -368,5 +394,6 @@ def create_api_request(
         callback=callback,
         all_model_names=all_model_names,
         all_sampling_params=all_sampling_params,
+        tools=tools,
         **kwargs,
     )

lm_deluge-0.0.9/src/lm_deluge/api_requests/bedrock.py

@@ -0,0 +1,283 @@
+import asyncio
+import json
+import os
+from aiohttp import ClientResponse
+from tqdm import tqdm
+from typing import Callable
+
+try:
+    from requests_aws4auth import AWS4Auth
+except ImportError:
+    raise ImportError(
+        "aws4auth is required for bedrock support. Install with: pip install requests-aws4auth"
+    )
+
+from lm_deluge.prompt import Conversation, Message, Text, ToolCall, Thinking
+from .base import APIRequestBase, APIResponse
+
+from ..tracker import StatusTracker
+from ..sampling_params import SamplingParams
+from ..models import APIModel
+
+
+class BedrockRequest(APIRequestBase):
+    def __init__(
+        self,
+        task_id: int,
+        model_name: str,
+        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,
+    ):
+        super().__init__(
+            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,
+            debug=debug,
+            all_model_names=all_model_names,
+            all_sampling_params=all_sampling_params,
+            tools=tools,
+        )
+
+        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")
+
+        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."
+            )
+
+        # 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()
+
+        # 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:
+            self.request_json["tools"] = [tool.dump_for("anthropic") for tool in tools]
+
+        # Setup AWS4Auth for signing
+        self.auth = AWS4Auth(
+            self.access_key,
+            self.secret_key,
+            self.region,
+            self.service,
+            session_token=self.session_token,
+        )
+
+        # Setup basic headers (AWS4Auth will add the Authorization header)
+        self.request_header = {
+            "Content-Type": "application/json",
+        }
+
+    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,
+                    headers=signed_headers,
+                    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)
+
+        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,
+                    input_tokens=None,
+                    output_tokens=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,
+                    input_tokens=None,
+                    output_tokens=None,
+                )
+            )
+            self.handle_error(create_new_request=False)
+
+    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
+        status_code = http_response.status
+        mimetype = http_response.headers.get("Content-Type", None)
+
+        if status_code >= 200 and status_code < 300:
+            try:
+                data = await http_response.json()
+                response_content = data["content"]
+
+                # Parse response into Message with parts
+                parts = []
+                for item in response_content:
+                    if item["type"] == "text":
+                        parts.append(Text(item["text"]))
+                    elif item["type"] == "thinking":
+                        thinking = item["thinking"]
+                        parts.append(Thinking(item["thinking"]))
+                    elif item["type"] == "tool_use":
+                        parts.append(
+                            ToolCall(
+                                id=item["id"],
+                                name=item["name"],
+                                arguments=item["input"],
+                            )
+                        )
+
+                content = Message("assistant", parts)
+                input_tokens = data["usage"]["input_tokens"]
+                output_tokens = data["usage"]["output_tokens"]
+            except Exception as e:
+                is_error = True
+                error_message = (
+                    f"Error calling .json() on response w/ status {status_code}: {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 "throttling" 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 or "too long" 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,
+            content=content,
+            thinking=thinking,
+            model_internal=self.model_name,
+            region=self.region,
+            sampling_params=self.sampling_params,
+            input_tokens=input_tokens,
+            output_tokens=output_tokens,
+        )

{lm_deluge-0.0.8 → lm_deluge-0.0.9}/src/lm_deluge/api_requests/common.py

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

{lm_deluge-0.0.8 → lm_deluge-0.0.9}/src/lm_deluge/api_requests/mistral.py

@@ -7,7 +7,7 @@ from tqdm.auto import tqdm
 from typing import Callable
 
 from .base import APIRequestBase, APIResponse
-from ..prompt import Conversation
+from ..prompt import Conversation, Message
 from ..tracker import StatusTracker
 from ..sampling_params import SamplingParams
 from ..models import APIModel
@@ -130,7 +130,7 @@ class MistralRequest(APIRequestBase):
             error_message=error_message,
             prompt=self.prompt,
             logprobs=logprobs,
-            completion=completion,
+            content=Message.ai(completion),
             model_internal=self.model_name,
             sampling_params=self.sampling_params,
             input_tokens=input_tokens,

{lm_deluge-0.0.8 → lm_deluge-0.0.9}/src/lm_deluge/api_requests/openai.py

@@ -7,7 +7,7 @@ from tqdm.auto import tqdm
 from typing import Callable
 
 from .base import APIRequestBase, APIResponse
-from ..prompt import Conversation
+from ..prompt import Conversation, Message, Text, ToolCall, Thinking
 from ..tracker import StatusTracker
 from ..sampling_params import SamplingParams
 from ..models import APIModel
@@ -34,6 +34,7 @@ class OpenAIRequest(APIRequestBase):
         debug: bool = False,
         all_model_names: list[str] | None = None,
         all_sampling_params: list[SamplingParams] | None = None,
+        tools: list | None = None,
     ):
         super().__init__(
             task_id=task_id,
@@ -52,6 +53,7 @@ class OpenAIRequest(APIRequestBase):
             debug=debug,
             all_model_names=all_model_names,
             all_sampling_params=all_sampling_params,
+            tools=tools,
         )
         self.model = APIModel.from_registry(model_name)
         self.url = f"{self.model.api_base}/chat/completions"
@@ -85,12 +87,16 @@ class OpenAIRequest(APIRequestBase):
                 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
+            ]
 
     async def handle_response(self, http_response: ClientResponse) -> APIResponse:
         is_error = False
         error_message = None
         thinking = None
-        completion = None
+        content = None
         input_tokens = None
         output_tokens = None
         logprobs = None
@@ -108,9 +114,34 @@ class OpenAIRequest(APIRequestBase):
             if not is_error:
                 assert data is not None, "data is None"
                 try:
-                    completion = data["choices"][0]["message"]["content"]
-                    if "reasoning_content" in data["choices"][0]["message"]:
-                        thinking = data["choices"][0]["message"]["reasoning_content"]
+                    # Parse response into Message with parts
+                    parts = []
+                    message = data["choices"][0]["message"]
+
+                    # Add text content if present
+                    if message.get("content"):
+                        parts.append(Text(message["content"]))
+
+                    # Add thinking content if present (reasoning models)
+                    if "reasoning_content" in message:
+                        thinking = message["reasoning_content"]
+                        parts.append(Thinking(thinking))
+
+                    # Add tool calls if present
+                    if "tool_calls" in message:
+                        for tool_call in message["tool_calls"]:
+                            parts.append(
+                                ToolCall(
+                                    id=tool_call["id"],
+                                    name=tool_call["function"]["name"],
+                                    arguments=json.loads(
+                                        tool_call["function"]["arguments"]
+                                    ),
+                                )
+                            )
+
+                    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]:
@@ -144,7 +175,7 @@ class OpenAIRequest(APIRequestBase):
             prompt=self.prompt,
             logprobs=logprobs,
             thinking=thinking,
-            completion=completion,
+            content=content,
             model_internal=self.model_name,
             sampling_params=self.sampling_params,
             input_tokens=input_tokens,