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

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: lm_deluge
-Version: 0.0.9
+Version: 0.0.12
 Summary: Python utility for using LLM API models.
 Author-email: Benjamin Anderson <ben@trytaylor.ai>
 Requires-Python: >=3.10
@@ -21,6 +21,7 @@ Requires-Dist: bs4
 Requires-Dist: lxml
 Requires-Dist: pdf2image
 Requires-Dist: pillow
+Requires-Dist: fastmcp>=2.4
 Requires-Dist: fasttext-wheel
 Requires-Dist: fasttext-langdetect
 Dynamic: license-file
@@ -32,6 +33,8 @@ Dynamic: license-file
 - **Unified client** – Send prompts to all relevant models with a single client.
 - **Massive concurrency with throttling** – Set `max_tokens_per_minute` and `max_requests_per_minute` and let it fly. The client will process as many requests as possible while respecting rate limits and retrying failures.
 - **Spray across models/providers** – Configure a client with multiple models from any provider(s), and sampling weights. The client samples a model for each request.
+- **Tool Use** – Unified API for defining tools for all providers, and creating tools automatically from python functions.
+- **MCP Support** – Instantiate a `Tool` from a local or remote MCP server so that any LLM can use it, whether or not that provider natively supports MCP.
 - **Caching** – Save completions in a local or distributed cache to avoid repeated LLM calls to process the same input.
 - **Convenient message constructor** – No more looking up how to build an Anthropic messages list with images. Our `Conversation` and `Message` classes work great with our client or with the `openai` and `anthropic` packages.
 - **Sync and async APIs** – Use the client from sync or async code.
@@ -44,7 +47,7 @@ Dynamic: license-file
 pip install lm-deluge
 ```
 
-The package relies on environment variables for API keys. Typical variables include `OPENAI_API_KEY`, `ANTHROPIC_API_KEY`, `COHERE_API_KEY`, `META_API_KEY`, and `GOOGLE_API_KEY`. `LLMClient` will automatically load the `.env` file when imported; we recommend using that to set the environment variables.
+The package relies on environment variables for API keys. Typical variables include `OPENAI_API_KEY`, `ANTHROPIC_API_KEY`, `COHERE_API_KEY`, `META_API_KEY`, and `GOOGLE_API_KEY`. `LLMClient` will automatically load the `.env` file when imported; we recommend using that to set the environment variables. For Bedrock, you'll need to set `AWS_ACCESS_KEY_ID` and `AWS_SECRET_ACCESS_KEY`.
 
 ## Quickstart
 
@@ -60,13 +63,13 @@ print(resp[0].completion)
 
 ## Spraying Across Models
 
-To distribute your requests across models, just provide a list of more than one model to the constructor. The rate limits for the client apply to the client as a whole, not per-model, so you may want to increase them:
+To distribute your requests across models, just provide a list of more than one model to the constructor. See all available models in `models.py`. The rate limits for the client apply to the client as a whole, not per-model, so you may want to increase them:
 
 ```python
 from lm_deluge import LLMClient
 
 client = LLMClient.basic(
-    ["gpt-4o-mini", "claude-haiku-anthropic"],
+    ["gpt-4o-mini", "claude-3-haiku"],
     max_requests_per_minute=10_000
 )
 resps = client.process_prompts_sync(
@@ -81,7 +84,7 @@ API calls can be customized in a few ways.
 
 1. **Sampling Parameters.** This determines things like structured outputs, maximum completion tokens, nucleus sampling, etc. Provide a custom `SamplingParams` to the `LLMClient` to set temperature, top_p, json_mode, max_new_tokens, and/or reasoning_effort. You can pass 1 `SamplingParams` to use for all models, or a list of `SamplingParams` that's the same length as the list of models. You can also pass many of these arguments directly to `LLMClient.basic` so you don't have to construct an entire `SamplingParams` object.
 2. **Arguments to LLMClient.** This is where you set request timeout, rate limits, model name(s), model weight(s) for distributing requests across models, retries, and caching.
-3. **Arguments to process_prompts.** Per-call, you can set verbosity, whether to display progress, and whether to return just completions (rather than the full APIResponse object).
+3. **Arguments to process_prompts.** Per-call, you can set verbosity, whether to display progress, and whether to return just completions (rather than the full APIResponse object). This is also where you provide tools.
 
 Putting it all together:
 
@@ -120,11 +123,97 @@ resps = client.process_prompts_sync([prompt])
 
 This just works. Images can be local images on disk, URLs, bytes, base64 data URLs... go wild. You can use `Conversation.to_openai` or `Conversation.to_anthropic` to format your messages for the OpenAI or Anthropic clients directly.
 
-## Caching
+See a full multi-turn chat example in `examples/multiturn.md`.
 
-`lm_deluge.cache` includes LevelDB, SQLite and custom dictionary based caches.  Pass an instance via `LLMClient(..., cache=my_cache)` and previously seen prompts will not be re‑sent across different `process_prompts_[...]` calls.
+## Tool Use
 
-**IMPORTANT:** Caching does not currently work for prompts in the SAME batch. That is, if you call `process_prompts_sync` with the same prompt 100 times, there will be 0 cache hits. If you call `process_prompts_sync` a *second* time with those same 100 prompts, all 100 will be cache hits. The cache is intended to be persistent and help you save costs across many invocations, but it can't help with a single batch-inference session (yet!).
+Define tools from Python functions and use them with any model:
+
+```python
+from lm_deluge import LLMClient, Tool
+
+def get_weather(city: str) -> str:
+    return f"The weather in {city} is sunny and 72°F"
+
+tool = Tool.from_function(get_weather)
+client = LLMClient.basic("claude-3-haiku")
+resps = client.process_prompts_sync(
+    ["What's the weather in Paris?"],
+    tools=[tool]
+)
+
+# you can iterate over the tool calls in the response automatically
+for tool_call in resps[0].tool_calls:
+    print(tool_call.name, tool_call.arguments)
+```
+
+You can also automatically instantiate tools from MCP servers. Under the hood, the the constructor connects to the server, asks it what tools it has, and then creates a `Tool` from each of them, *with a built-in `call` and `acall` interface*.
+
+```python
+from lm_deluge import LLMClient, Tool
+
+# Connect to a local MCP server and get all of its tools
+filesystem_tools = Tool.from_mcp(
+    "filesystem",
+    command="npx",
+    args=["-y", "@modelcontextprotocol/server-filesystem", "/path/to/directory"]
+)
+
+# or load ALL the tools from a Claude Desktop like config
+config = {
+    "mcpServers": {
+        "exa": {
+            "url": f"https://mcp.exa.ai/mcp?exaApiKey={os.getenv('EXA_API_KEY')}"
+        },
+        "zapier": {
+            "url": f"https://mcp.zapier.com/api/mcp/s/{os.getenv('ZAPIER_MCP_SECRET')}/mcp"
+        }
+    }
+}
+all_tools = Tool.from_mcp_config(config)
+
+# let the model use the tools
+client = LLMClient.basic("gpt-4o-mini")
+resps = client.process_prompts_sync(
+    ["List the files in the current directory"],
+    tools=tools
+)
+
+# call the tools
+for tool_call in resps[0].tool_calls:
+    # this is dumb sorry will make it better
+    tool_to_call = [x for x in tools if x.name == tool_call.name][0]
+    tool_to_call.call(**tool_call.arguments) # in async code, use .acall()
+```
+
+### Prompt Caching (Anthropic)
+
+For Anthropic models, you can use prompt caching to reduce costs and latency for repeated context. This uses Anthropic's server-side prompt caching. Other providers like OpenAI and Google do this automatically, but Anthropic requires you to manually set cache-control on messages. You can do this in lm-deluge with a simple "cache" argument to `process_prompts_sync` or `process_prompts_async`:
+
+```python
+from lm_deluge import LLMClient, Conversation, Message
+
+# Create a conversation with system message
+conv = (
+    Conversation.system("You are an expert Python developer with deep knowledge of async programming.")
+    .add(Message.user("How do I use asyncio.gather?"))
+)
+
+# Use prompt caching to cache system message and tools
+client = LLMClient.basic("claude-3-5-sonnet")
+resps = client.process_prompts_sync(
+    [conv],
+    cache="system_and_tools"  # Cache system message and any tools
+)
+```
+
+Available cache patterns: `"system_and_tools"`, `"tools_only"`, `"last_user_message"`, `"last_2_user_messages"`, `"last_3_user_messages"`.
+
+## Local Caching
+
+Besides caching from model providers (which provides cache reads at a discount, but not for free) `lm_deluge.cache` includes LevelDB, SQLite and custom dictionary based caches to cache prompts locally. Pass an instance via `LLMClient(..., cache=my_cache)` and previously seen prompts will not be re‑sent across different `process_prompts_[...]` calls.
+
+**IMPORTANT:** Caching does not currently work for prompts in the SAME batch. That is, if you call `process_prompts_sync` with the same prompt 100 times, there will be 0 cache hits. If you call `process_prompts_sync` a *second* time with those same 100 prompts, all 100 will be cache hits. The local cache is intended to be persistent and help you save costs across many invocations, but it can't help with a single batch-inference session (yet!).
 
 ## Asynchronous Client
 Use this in asynchronous code, or in a Jupyter notebook. If you try to use the sync client in a Jupyter notebook, you'll have to use `nest-asyncio`, because internally the sync client uses async code. Don't do it! Just use the async client!

{lm_deluge-0.0.9 → lm_deluge-0.0.12}/README.md

@@ -5,6 +5,8 @@
 - **Unified client** – Send prompts to all relevant models with a single client.
 - **Massive concurrency with throttling** – Set `max_tokens_per_minute` and `max_requests_per_minute` and let it fly. The client will process as many requests as possible while respecting rate limits and retrying failures.
 - **Spray across models/providers** – Configure a client with multiple models from any provider(s), and sampling weights. The client samples a model for each request.
+- **Tool Use** – Unified API for defining tools for all providers, and creating tools automatically from python functions.
+- **MCP Support** – Instantiate a `Tool` from a local or remote MCP server so that any LLM can use it, whether or not that provider natively supports MCP.
 - **Caching** – Save completions in a local or distributed cache to avoid repeated LLM calls to process the same input.
 - **Convenient message constructor** – No more looking up how to build an Anthropic messages list with images. Our `Conversation` and `Message` classes work great with our client or with the `openai` and `anthropic` packages.
 - **Sync and async APIs** – Use the client from sync or async code.
@@ -17,7 +19,7 @@
 pip install lm-deluge
 ```
 
-The package relies on environment variables for API keys. Typical variables include `OPENAI_API_KEY`, `ANTHROPIC_API_KEY`, `COHERE_API_KEY`, `META_API_KEY`, and `GOOGLE_API_KEY`. `LLMClient` will automatically load the `.env` file when imported; we recommend using that to set the environment variables.
+The package relies on environment variables for API keys. Typical variables include `OPENAI_API_KEY`, `ANTHROPIC_API_KEY`, `COHERE_API_KEY`, `META_API_KEY`, and `GOOGLE_API_KEY`. `LLMClient` will automatically load the `.env` file when imported; we recommend using that to set the environment variables. For Bedrock, you'll need to set `AWS_ACCESS_KEY_ID` and `AWS_SECRET_ACCESS_KEY`.
 
 ## Quickstart
 
@@ -33,13 +35,13 @@ print(resp[0].completion)
 
 ## Spraying Across Models
 
-To distribute your requests across models, just provide a list of more than one model to the constructor. The rate limits for the client apply to the client as a whole, not per-model, so you may want to increase them:
+To distribute your requests across models, just provide a list of more than one model to the constructor. See all available models in `models.py`. The rate limits for the client apply to the client as a whole, not per-model, so you may want to increase them:
 
 ```python
 from lm_deluge import LLMClient
 
 client = LLMClient.basic(
-    ["gpt-4o-mini", "claude-haiku-anthropic"],
+    ["gpt-4o-mini", "claude-3-haiku"],
     max_requests_per_minute=10_000
 )
 resps = client.process_prompts_sync(
@@ -54,7 +56,7 @@ API calls can be customized in a few ways.
 
 1. **Sampling Parameters.** This determines things like structured outputs, maximum completion tokens, nucleus sampling, etc. Provide a custom `SamplingParams` to the `LLMClient` to set temperature, top_p, json_mode, max_new_tokens, and/or reasoning_effort. You can pass 1 `SamplingParams` to use for all models, or a list of `SamplingParams` that's the same length as the list of models. You can also pass many of these arguments directly to `LLMClient.basic` so you don't have to construct an entire `SamplingParams` object.
 2. **Arguments to LLMClient.** This is where you set request timeout, rate limits, model name(s), model weight(s) for distributing requests across models, retries, and caching.
-3. **Arguments to process_prompts.** Per-call, you can set verbosity, whether to display progress, and whether to return just completions (rather than the full APIResponse object).
+3. **Arguments to process_prompts.** Per-call, you can set verbosity, whether to display progress, and whether to return just completions (rather than the full APIResponse object). This is also where you provide tools.
 
 Putting it all together:
 
@@ -93,11 +95,97 @@ resps = client.process_prompts_sync([prompt])
 
 This just works. Images can be local images on disk, URLs, bytes, base64 data URLs... go wild. You can use `Conversation.to_openai` or `Conversation.to_anthropic` to format your messages for the OpenAI or Anthropic clients directly.
 
-## Caching
+See a full multi-turn chat example in `examples/multiturn.md`.
 
-`lm_deluge.cache` includes LevelDB, SQLite and custom dictionary based caches.  Pass an instance via `LLMClient(..., cache=my_cache)` and previously seen prompts will not be re‑sent across different `process_prompts_[...]` calls.
+## Tool Use
 
-**IMPORTANT:** Caching does not currently work for prompts in the SAME batch. That is, if you call `process_prompts_sync` with the same prompt 100 times, there will be 0 cache hits. If you call `process_prompts_sync` a *second* time with those same 100 prompts, all 100 will be cache hits. The cache is intended to be persistent and help you save costs across many invocations, but it can't help with a single batch-inference session (yet!).
+Define tools from Python functions and use them with any model:
+
+```python
+from lm_deluge import LLMClient, Tool
+
+def get_weather(city: str) -> str:
+    return f"The weather in {city} is sunny and 72°F"
+
+tool = Tool.from_function(get_weather)
+client = LLMClient.basic("claude-3-haiku")
+resps = client.process_prompts_sync(
+    ["What's the weather in Paris?"],
+    tools=[tool]
+)
+
+# you can iterate over the tool calls in the response automatically
+for tool_call in resps[0].tool_calls:
+    print(tool_call.name, tool_call.arguments)
+```
+
+You can also automatically instantiate tools from MCP servers. Under the hood, the the constructor connects to the server, asks it what tools it has, and then creates a `Tool` from each of them, *with a built-in `call` and `acall` interface*.
+
+```python
+from lm_deluge import LLMClient, Tool
+
+# Connect to a local MCP server and get all of its tools
+filesystem_tools = Tool.from_mcp(
+    "filesystem",
+    command="npx",
+    args=["-y", "@modelcontextprotocol/server-filesystem", "/path/to/directory"]
+)
+
+# or load ALL the tools from a Claude Desktop like config
+config = {
+    "mcpServers": {
+        "exa": {
+            "url": f"https://mcp.exa.ai/mcp?exaApiKey={os.getenv('EXA_API_KEY')}"
+        },
+        "zapier": {
+            "url": f"https://mcp.zapier.com/api/mcp/s/{os.getenv('ZAPIER_MCP_SECRET')}/mcp"
+        }
+    }
+}
+all_tools = Tool.from_mcp_config(config)
+
+# let the model use the tools
+client = LLMClient.basic("gpt-4o-mini")
+resps = client.process_prompts_sync(
+    ["List the files in the current directory"],
+    tools=tools
+)
+
+# call the tools
+for tool_call in resps[0].tool_calls:
+    # this is dumb sorry will make it better
+    tool_to_call = [x for x in tools if x.name == tool_call.name][0]
+    tool_to_call.call(**tool_call.arguments) # in async code, use .acall()
+```
+
+### Prompt Caching (Anthropic)
+
+For Anthropic models, you can use prompt caching to reduce costs and latency for repeated context. This uses Anthropic's server-side prompt caching. Other providers like OpenAI and Google do this automatically, but Anthropic requires you to manually set cache-control on messages. You can do this in lm-deluge with a simple "cache" argument to `process_prompts_sync` or `process_prompts_async`:
+
+```python
+from lm_deluge import LLMClient, Conversation, Message
+
+# Create a conversation with system message
+conv = (
+    Conversation.system("You are an expert Python developer with deep knowledge of async programming.")
+    .add(Message.user("How do I use asyncio.gather?"))
+)
+
+# Use prompt caching to cache system message and tools
+client = LLMClient.basic("claude-3-5-sonnet")
+resps = client.process_prompts_sync(
+    [conv],
+    cache="system_and_tools"  # Cache system message and any tools
+)
+```
+
+Available cache patterns: `"system_and_tools"`, `"tools_only"`, `"last_user_message"`, `"last_2_user_messages"`, `"last_3_user_messages"`.
+
+## Local Caching
+
+Besides caching from model providers (which provides cache reads at a discount, but not for free) `lm_deluge.cache` includes LevelDB, SQLite and custom dictionary based caches to cache prompts locally. Pass an instance via `LLMClient(..., cache=my_cache)` and previously seen prompts will not be re‑sent across different `process_prompts_[...]` calls.
+
+**IMPORTANT:** Caching does not currently work for prompts in the SAME batch. That is, if you call `process_prompts_sync` with the same prompt 100 times, there will be 0 cache hits. If you call `process_prompts_sync` a *second* time with those same 100 prompts, all 100 will be cache hits. The local cache is intended to be persistent and help you save costs across many invocations, but it can't help with a single batch-inference session (yet!).
 
 ## Asynchronous Client
 Use this in asynchronous code, or in a Jupyter notebook. If you try to use the sync client in a Jupyter notebook, you'll have to use `nest-asyncio`, because internally the sync client uses async code. Don't do it! Just use the async client!

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

@@ -3,7 +3,7 @@ requires = ["setuptools", "wheel"]
 
 [project]
 name = "lm_deluge"
-version = "0.0.9"
+version = "0.0.12"
 authors = [{ name = "Benjamin Anderson", email = "ben@trytaylor.ai" }]
 description = "Python utility for using LLM API models."
 readme = "README.md"
@@ -27,6 +27,7 @@ dependencies = [
     "lxml",
     "pdf2image",
     "pillow",
+    "fastmcp>=2.4",
     "fasttext-wheel",
     "fasttext-langdetect",
 ]

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

@@ -6,7 +6,15 @@ import warnings
 from tqdm import tqdm
 from typing import Callable
 
-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
@@ -35,6 +43,7 @@ class AnthropicRequest(APIRequestBase):
         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,
@@ -52,11 +61,16 @@ class AnthropicRequest(APIRequestBase):
             all_model_names=all_model_names,
             all_sampling_params=all_sampling_params,
             tools=tools,
+            cache=cache,
         )
         self.model = APIModel.from_registry(model_name)
         self.url = f"{self.model.api_base}/messages"
 
-        self.system_message, messages = prompt.to_anthropic()
+        # Lock images as bytes if caching is enabled
+        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",
@@ -97,15 +111,18 @@ 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]
+            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
 
     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
         status_code = http_response.status
         mimetype = http_response.headers.get("Content-Type", None)
         rate_limits = {}
@@ -143,8 +160,7 @@ class AnthropicRequest(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 = (
@@ -182,6 +198,5 @@ class AnthropicRequest(APIRequestBase):
             thinking=thinking,
             model_internal=self.model_name,
             sampling_params=self.sampling_params,
-            input_tokens=input_tokens,
-            output_tokens=output_tokens,
+            usage=usage,
         )

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

@@ -7,7 +7,8 @@ from dataclasses import dataclass
 from abc import ABC, abstractmethod
 from typing import Callable
 
-from lm_deluge.prompt import Conversation, Message
+from lm_deluge.prompt import Conversation, Message, CachePattern
+from lm_deluge.usage import Usage
 
 from ..tracker import StatusTracker
 from ..sampling_params import SamplingParams
@@ -29,9 +30,8 @@ class APIResponse:
     is_error: bool | None
     error_message: str | None
 
-    # completion information
-    input_tokens: int | None
-    output_tokens: int | None
+    # completion information - unified usage tracking
+    usage: Usage | None = None
 
     # response content - structured format
     content: Message | None = None
@@ -56,6 +56,26 @@ class APIResponse:
             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)
@@ -63,14 +83,13 @@ class APIResponse:
         self.model_external = api_model.name
         self.cost = None
         if (
-            self.input_tokens is not None
-            and self.output_tokens is not None
+            self.usage is not None
             and api_model.input_cost is not None
             and api_model.output_cost is not None
         ):
             self.cost = (
-                self.input_tokens * api_model.input_cost / 1e6
-                + self.output_tokens * api_model.output_cost / 1e6
+                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(
@@ -90,8 +109,7 @@ class APIResponse:
             "error_message": self.error_message,
             "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,
+            "usage": self.usage.to_dict() if self.usage else None,
             "finish_reason": self.finish_reason,
             "cost": self.cost,
         }
@@ -107,6 +125,10 @@ class APIResponse:
             # 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"],
@@ -115,8 +137,7 @@ class APIResponse:
             status_code=data["status_code"],
             is_error=data["is_error"],
             error_message=data["error_message"],
-            input_tokens=data["input_tokens"],
-            output_tokens=data["output_tokens"],
+            usage=usage,
             content=content,
             thinking=data.get("thinking"),
             model_external=data.get("model_external"),
@@ -168,6 +189,7 @@ class APIRequestBase(ABC):
         all_model_names: list[str] | None = None,
         all_sampling_params: list[SamplingParams] | None = None,
         tools: list | None = None,
+        cache: CachePattern | None = None,
     ):
         if all_model_names is None:
             raise ValueError("all_model_names must be provided.")
@@ -190,6 +212,7 @@ class APIRequestBase(ABC):
         self.all_model_names = all_model_names
         self.all_sampling_params = all_sampling_params
         self.tools = tools
+        self.cache: CachePattern | None = cache
         self.result = []  # list of APIResponse objects from each attempt
 
         # these should be set in the __init__ of the subclass
@@ -280,6 +303,7 @@ class APIRequestBase(ABC):
                         all_model_names=self.all_model_names,
                         all_sampling_params=self.all_sampling_params,
                         tools=self.tools,
+                        cache=self.cache,
                     )
                     # PROBLEM: new request is never put into results array, so we can't get the result.
                     self.retry_queue.put_nowait(new_request)
@@ -323,8 +347,7 @@ class APIRequestBase(ABC):
                     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)
@@ -341,8 +364,7 @@ class APIRequestBase(ABC):
                     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,
                 )
             )
             # maybe consider making True?
@@ -370,6 +392,7 @@ def create_api_request(
     all_model_names: list[str] | None = None,
     all_sampling_params: list[SamplingParams] | None = None,
     tools: list | None = None,
+    cache: CachePattern | None = None,
 ) -> APIRequestBase:
     from .common import CLASSES  # circular import so made it lazy, does this work?
 
@@ -395,5 +418,6 @@ def create_api_request(
         all_model_names=all_model_names,
         all_sampling_params=all_sampling_params,
         tools=tools,
+        cache=cache,
         **kwargs,
     )

{lm_deluge-0.0.9 → lm_deluge-0.0.12}/src/lm_deluge/api_requests/bedrock.py

@@ -12,7 +12,15 @@ 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
@@ -38,6 +46,7 @@ class BedrockRequest(APIRequestBase):
         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,
@@ -55,8 +64,13 @@ class BedrockRequest(APIRequestBase):
             all_model_names=all_model_names,
             all_sampling_params=all_sampling_params,
             tools=tools,
+            cache=cache,
         )
 
+        # 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 +101,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 = {
@@ -102,7 +116,11 @@ class BedrockRequest(APIRequestBase):
             self.request_json["system"] = self.system_message
 
         if tools:
-            self.request_json["tools"] = [tool.dump_for("anthropic") for tool in 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
 
         # Setup AWS4Auth for signing
         self.auth = AWS4Auth(
@@ -179,8 +197,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 +216,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 +226,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 +253,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 +292,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,
         )