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

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/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/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,
         )

lm_deluge/api_requests/mistral.py

@@ -7,7 +7,8 @@ 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 ..models import APIModel
@@ -35,6 +36,7 @@ class MistralRequest(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,
@@ -53,7 +55,15 @@ class MistralRequest(APIRequestBase):
             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 = {
@@ -81,8 +91,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,8 +108,7 @@ 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"]
+                    usage = Usage.from_mistral_usage(data["usage"])
                     if self.logprobs and "logprobs" in data["choices"][0]:
                         logprobs = data["choices"][0]["logprobs"]["content"]
                 except Exception:
@@ -134,6 +142,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

@@ -7,7 +7,8 @@ from tqdm.auto import tqdm
 from typing import Callable
 
 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 ..models import APIModel
@@ -35,6 +36,7 @@ class OpenAIRequest(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,
@@ -54,7 +56,14 @@ class OpenAIRequest(APIRequestBase):
             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 = {
@@ -97,8 +106,7 @@ class OpenAIRequest(APIRequestBase):
         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,8 +150,7 @@ class OpenAIRequest(APIRequestBase):
 
                     content = Message("assistant", parts)
 
-                    input_tokens = data["usage"]["prompt_tokens"]
-                    output_tokens = data["usage"]["completion_tokens"]
+                    usage = Usage.from_openai_usage(data["usage"])
                     if self.logprobs and "logprobs" in data["choices"][0]:
                         logprobs = data["choices"][0]["logprobs"]["content"]
                 except Exception:
@@ -178,6 +185,5 @@ 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,
         )

lm_deluge/client.py

@@ -8,7 +8,7 @@ from dataclasses import dataclass
 from typing import Sequence, overload, Literal, Any
 from tqdm.auto import tqdm
 
-from lm_deluge.prompt import Conversation
+from lm_deluge.prompt import Conversation, CachePattern
 from lm_deluge.tool import Tool
 
 from .tracker import StatusTracker
@@ -232,6 +232,7 @@ class LLMClient:
         dry_run: Literal[True],
         verbose: bool = ...,
         tools: list[Tool] | None = ...,
+        cache: CachePattern | None = ...,
     ) -> dict[str, int]: ...
 
     @overload
@@ -244,6 +245,7 @@ class LLMClient:
         dry_run: bool = ...,
         verbose: bool = ...,
         tools: list[Tool] | None = ...,
+        cache: CachePattern | None = ...,
     ) -> list[str | None]: ...
 
     @overload
@@ -256,6 +258,7 @@ class LLMClient:
         dry_run: bool = ...,
         verbose: bool = ...,
         tools: list[Tool] | None = ...,
+        cache: CachePattern | None = ...,
     ) -> list[APIResponse | None]: ...
 
     async def process_prompts_async(
@@ -267,6 +270,7 @@ class LLMClient:
         dry_run: bool = False,
         verbose: bool = False,
         tools: list[Tool] | None = None,
+        cache: CachePattern | None = None,
     ) -> list[APIResponse | None] | list[str | None] | dict[str, int]:
         # if prompts are not Conversations, convert them.
         # can only handle strings for now
@@ -345,6 +349,7 @@ class LLMClient:
                     use_qps=self.use_qps,
                     verbose=verbose,
                     tools=tools,
+                    cache=cache,
                 )
             )
             api_results: list[APIResponse] = await api_task
@@ -373,6 +378,7 @@ class LLMClient:
         dry_run: bool = False,
         verbose: bool = False,
         tools: list[Tool] | None = None,
+        cache: CachePattern | None = None,
     ):
         return asyncio.run(
             self.process_prompts_async(
@@ -382,6 +388,7 @@ class LLMClient:
                 dry_run=dry_run,
                 verbose=verbose,
                 tools=tools,
+                cache=cache,
             )
         )
 
@@ -570,6 +577,7 @@ async def process_api_prompts_async(
     use_qps: bool = False,
     verbose: bool = False,
     tools: list[Tool] | None = None,
+    cache: CachePattern | None = None,
 ):
     """Processes API requests in parallel, throttling to stay under rate limits."""
     # change ids to integer list
@@ -654,6 +662,7 @@ async def process_api_prompts_async(
                         all_model_names=models,
                         all_sampling_params=sampling_params,
                         tools=tools,
+                        cache=cache,
                     )
                     status_tracker.num_tasks_started += 1
                     status_tracker.num_tasks_in_progress += 1

lm_deluge/prompt.py

@@ -8,6 +8,14 @@ from typing import Literal
 from lm_deluge.models import APIModel
 from lm_deluge.image import Image
 
+CachePattern = Literal[
+    "tools_only",
+    "system_and_tools",
+    "last_user_message",
+    "last_2_user_messages",
+    "last_3_user_messages",
+]
+
 ###############################################################################
 # 1. Low-level content blocks – either text or an image                       #
 ###############################################################################
@@ -516,7 +524,9 @@ class Conversation:
         # OpenAI Responses = single “input” array, role must be user/assistant
         return {"input": [m.oa_resp() for m in self.messages if m.role != "system"]}
 
-    def to_anthropic(self) -> tuple[str | None, list[dict]]:
+    def to_anthropic(
+        self, cache_pattern: CachePattern | None = None
+    ) -> tuple[str | list[dict] | None, list[dict]]:
         system_msg = next(
             (
                 m.parts[0].text
@@ -535,8 +545,84 @@ class Conversation:
                 other.append(user_msg.anthropic())
             else:
                 other.append(m.anthropic())
+
+        # Apply cache control if specified
+        if cache_pattern is not None:
+            system_msg, other = self._apply_cache_control(
+                system_msg, other, cache_pattern
+            )
+
         return system_msg, other
 
+    def _apply_cache_control(
+        self,
+        system_msg: str | None | list[dict],
+        messages: list[dict],
+        cache_pattern: CachePattern,
+    ) -> tuple[str | list[dict] | None, list[dict]]:
+        """Apply cache control to system message and/or messages based on the pattern."""
+
+        if cache_pattern == "system_and_tools" and system_msg is not None:
+            # Convert system message to structured format with cache control
+            # This caches tools+system prefix (since system comes after tools)
+            system_msg = [
+                {
+                    "type": "text",
+                    "text": system_msg,
+                    "cache_control": {"type": "ephemeral"},
+                }
+            ]
+
+        if cache_pattern == "last_user_message":
+            # Cache the last user message
+            user_messages = [i for i, m in enumerate(messages) if m["role"] == "user"]
+            if user_messages:
+                last_user_idx = user_messages[-1]
+                self._add_cache_control_to_message(messages[last_user_idx])
+
+        elif cache_pattern == "last_2_user_messages":
+            # Cache the last 2 user messages
+            user_messages = [i for i, m in enumerate(messages) if m["role"] == "user"]
+            for idx in user_messages[-2:]:
+                self._add_cache_control_to_message(messages[idx])
+
+        elif cache_pattern == "last_3_user_messages":
+            # Cache the last 3 user messages
+            user_messages = [i for i, m in enumerate(messages) if m["role"] == "user"]
+            for idx in user_messages[-3:]:
+                self._add_cache_control_to_message(messages[idx])
+
+        return system_msg, messages
+
+    def lock_images_as_bytes(self) -> "Conversation":
+        """
+        Convert all images to bytes format to ensure they remain unchanged for caching.
+        This should be called when caching is enabled to prevent cache invalidation
+        from image reference changes.
+        """
+        for message in self.messages:
+            for part in message.parts:
+                if isinstance(part, Image):
+                    # Force conversion to bytes if not already
+                    part.data = part._bytes()
+        return self
+
+    def _add_cache_control_to_message(self, message: dict) -> None:
+        """Add cache control to a message's content."""
+        content = message["content"]
+        if isinstance(content, str):
+            # Convert string content to structured format with cache control
+            message["content"] = [
+                {
+                    "type": "text",
+                    "text": content,
+                    "cache_control": {"type": "ephemeral"},
+                }
+            ]
+        elif isinstance(content, list) and content:
+            # Add cache control to the last content block
+            content[-1]["cache_control"] = {"type": "ephemeral"}
+
     def to_gemini(self) -> tuple[str | None, list[dict]]:
         system_msg = next(
             (

lm_deluge/usage.py

@@ -0,0 +1,114 @@
+from dataclasses import dataclass
+from typing import Optional
+
+
+@dataclass
+class Usage:
+    """
+    Unified usage tracking for all API providers.
+
+    Tracks token usage including cache hits and writes for providers that support it.
+    For providers that don't support caching, cache_read and cache_write will be None.
+    """
+
+    input_tokens: int = 0
+    output_tokens: int = 0
+    cache_read_tokens: Optional[int] = None  # Tokens read from cache (Anthropic)
+    cache_write_tokens: Optional[int] = None  # Tokens written to cache (Anthropic)
+
+    @property
+    def total_input_tokens(self) -> int:
+        """Total input tokens including both fresh input, cache writes, and cache reads."""
+        result = self.input_tokens
+        if self.cache_read_tokens is not None:
+            result += self.cache_read_tokens
+        if self.cache_write_tokens is not None:
+            result += self.cache_write_tokens
+        return result
+
+    @property
+    def total_tokens(self) -> int:
+        """Total tokens processed (input + output)."""
+        return self.total_input_tokens + self.output_tokens
+
+    @property
+    def has_cache_hit(self) -> bool:
+        """Whether this request had any cache hits."""
+        return self.cache_read_tokens is not None and self.cache_read_tokens > 0
+
+    @property
+    def has_cache_write(self) -> bool:
+        """Whether this request wrote to cache."""
+        return self.cache_write_tokens is not None and self.cache_write_tokens > 0
+
+    @classmethod
+    def from_anthropic_usage(cls, usage_data: dict) -> "Usage":
+        """Create Usage from Anthropic API response usage data."""
+        return cls(
+            input_tokens=usage_data.get("input_tokens", 0),
+            output_tokens=usage_data.get("output_tokens", 0),
+            cache_read_tokens=usage_data.get("cache_read_input_tokens"),
+            cache_write_tokens=usage_data.get("cache_creation_input_tokens"),
+        )
+
+    @classmethod
+    def from_openai_usage(cls, usage_data: dict) -> "Usage":
+        """Create Usage from OpenAI API response usage data."""
+        return cls(
+            input_tokens=usage_data.get("prompt_tokens", 0),
+            output_tokens=usage_data.get("completion_tokens", 0),
+            cache_read_tokens=None,  # OpenAI doesn't support caching yet
+            cache_write_tokens=None,
+        )
+
+    @classmethod
+    def from_mistral_usage(cls, usage_data: dict) -> "Usage":
+        """Create Usage from Mistral API response usage data."""
+        return cls(
+            input_tokens=usage_data.get("prompt_tokens", 0),
+            output_tokens=usage_data.get("completion_tokens", 0),
+            cache_read_tokens=None,  # Mistral doesn't support caching
+            cache_write_tokens=None,
+        )
+
+    def to_dict(self) -> dict:
+        """Convert to dictionary for serialization."""
+        return {
+            "input_tokens": self.input_tokens,
+            "output_tokens": self.output_tokens,
+            "cache_read_tokens": self.cache_read_tokens,
+            "cache_write_tokens": self.cache_write_tokens,
+            "total_input_tokens": self.total_input_tokens,
+            "total_tokens": self.total_tokens,
+            "has_cache_hit": self.has_cache_hit,
+            "has_cache_write": self.has_cache_write,
+        }
+
+    @classmethod
+    def from_dict(cls, data: dict) -> "Usage":
+        """Create Usage from dictionary."""
+        return cls(
+            input_tokens=data.get("input_tokens", 0),
+            output_tokens=data.get("output_tokens", 0),
+            cache_read_tokens=data.get("cache_read_tokens"),
+            cache_write_tokens=data.get("cache_write_tokens"),
+        )
+
+    def __add__(self, other: "Usage") -> "Usage":
+        """Add two Usage objects together."""
+        return Usage(
+            input_tokens=self.input_tokens + other.input_tokens,
+            output_tokens=self.output_tokens + other.output_tokens,
+            cache_read_tokens=(
+                (self.cache_read_tokens or 0) + (other.cache_read_tokens or 0)
+                if self.cache_read_tokens is not None
+                or other.cache_read_tokens is not None
+                else None
+            ),
+            cache_write_tokens=(
+                (self.cache_write_tokens or 0) + (other.cache_write_tokens or 0)
+                if self.cache_write_tokens is not None
+                or other.cache_write_tokens is not None
+                else None
+            ),
+        )

{lm_deluge-0.0.11.dist-info → lm_deluge-0.0.12.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: lm_deluge
-Version: 0.0.11
+Version: 0.0.12
 Summary: Python utility for using LLM API models.
 Author-email: Benjamin Anderson <ben@trytaylor.ai>
 Requires-Python: >=3.10
@@ -47,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
 
@@ -63,7 +63,7 @@ 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
@@ -84,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:
 
@@ -123,7 +123,9 @@ 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.
 
-## Basic Tool Use
+See a full multi-turn chat example in `examples/multiturn.md`.
+
+## Tool Use
 
 Define tools from Python functions and use them with any model:
 
@@ -135,27 +137,83 @@ def get_weather(city: str) -> str:
 
 tool = Tool.from_function(get_weather)
 client = LLMClient.basic("claude-3-haiku")
-resp = client.process_prompts_sync(["What's the weather in Paris?"], tools=[tool])
-```
+resps = client.process_prompts_sync(
+    ["What's the weather in Paris?"],
+    tools=[tool]
+)
 
-## MCP Integration
+# 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)
+```
 
-Connect to MCP servers to extend your models with external tools:
+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
-mcp_tool = Tool.from_mcp("filesystem", command="npx -y @modelcontextprotocol/server-filesystem", args=["/path/to/directory"])
-client = LLMClient.basic("gpt-4o-mini", tools=[mcp_tool])
-resp = client.process_prompts_sync(["List the files in the current directory"])
+# 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()
 ```
 
-## Caching
+### 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
 
-`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.
+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 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!).
+**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.11.dist-info → lm_deluge-0.0.12.dist-info}/RECORD

@@ -1,23 +1,24 @@
 lm_deluge/__init__.py,sha256=rndOr4Rcfnpttz-onWU3vVEm-MM0WDFgz6KexKPAx0k,222
 lm_deluge/cache.py,sha256=VB1kv8rM2t5XWPR60uhszFcxLDnVKOe1oA5hYjVDjIo,4375
-lm_deluge/client.py,sha256=lGD4rqT7qHkTKddjRvKK_1bh7s8GNIzXzQ52GCZhfCg,28932
+lm_deluge/client.py,sha256=yIC86BlH31oW1umrXnuOxaCGMGCbNVvGWXdz2yc6n1g,29296
 lm_deluge/embed.py,sha256=m-X8UK4gV9KKD7Wv3yarAceMQaj7gR1JwzD_sB0MOQY,13183
 lm_deluge/errors.py,sha256=oHjt7YnxWbh-eXMScIzov4NvpJMo0-2r5J6Wh5DQ1tk,209
 lm_deluge/gemini_limits.py,sha256=V9mpS9JtXYz7AY6OuKyQp5TuIMRH1BVv9YrSNmGmHNA,1569
 lm_deluge/image.py,sha256=hFbRajqEVQbkirAfOxsTPkeq-27Zl-so4AWBFeUbpBI,7161
 lm_deluge/models.py,sha256=oYrt0x0iVfTwoHjP-l1WWennzEDGwnZczj6ds6a6-xc,45406
-lm_deluge/prompt.py,sha256=_pJYwgjL39lDzMNmae8pPIBoORm_ekSM_9qU2iGGpOc,25445
+lm_deluge/prompt.py,sha256=O46eLd6e68_7mGHkCMwLitnsxE4NUvrd1X9vtb_Y7qc,28787
 lm_deluge/rerank.py,sha256=tW1c3gQCAqaF8Ez-r-4qxYAcdKqxnLMxwHApKOUKwk4,11289
 lm_deluge/sampling_params.py,sha256=E2kewh1vz-1Qcy5xNBCzihfGgT_GcHYMfzaWb3FLiXs,739
 lm_deluge/tool.py,sha256=5nFbHchv12C1jkL8nkEh6v9WfxpC0O6rALP25z60WsI,9476
 lm_deluge/tracker.py,sha256=Dk99scN_NeDEO0gkLO5efXiZq11Ga-k6cerUHWN7IWY,1292
+lm_deluge/usage.py,sha256=oS-rmF3ZJ1RMtR7WI6BB2uVOAjJg0scvGF3zZRahWVg,4449
 lm_deluge/api_requests/__init__.py,sha256=_aSpD6CJL9g6OpLPoChXiHjl4MH_OlGcKgfZaW8cgLM,71
-lm_deluge/api_requests/anthropic.py,sha256=MMI_w9hVbevQpcqP3NVVindpTmLb2KHqjJQpIzCi5RM,7240
-lm_deluge/api_requests/base.py,sha256=w0MEOCIccxxy2c67Y2Y-QBox9rinIxQ7MLnp8953sjQ,15954
-lm_deluge/api_requests/bedrock.py,sha256=cvB85BFvL9HKTUsP9qFUCLQzJh83IQNAcLXuW6ReZK8,10520
+lm_deluge/api_requests/anthropic.py,sha256=F1bDb1Pyzwq08LSd4K4Pu6bugi7IFpUf4tZNWySgdx0,7646
+lm_deluge/api_requests/base.py,sha256=h7AdWwd9sjZnBb2ETZmzC64E9fNYGp2vrOHGXXo8W2g,16803
+lm_deluge/api_requests/bedrock.py,sha256=vf1pkVv4wBqZX7iiBHqHzDEHTlYNEG07rfQr9XE8Pr0,10832
 lm_deluge/api_requests/common.py,sha256=U0mX_wC3Tzg2-1u9nYUCTQqYzuYJqvLrICCNW_dbbJM,287
-lm_deluge/api_requests/mistral.py,sha256=lU9AOyb46uTzRjKw6Sd5iojEbBIMF432fRex7q6Xtwk,5423
-lm_deluge/api_requests/openai.py,sha256=BuMiM_2zJQXfnUjTT94JxJi3ZX5V-KQQueRG-R0SGuc,7361
+lm_deluge/api_requests/mistral.py,sha256=DvyriHbUApNOEdyiyyn-_HJnTY3tpKty18uqJ8u6n5Y,5640
+lm_deluge/api_requests/openai.py,sha256=RjSq6LyDt9g7FASjBJd1JzyD09xz6D-abNO8wiELR1M,7552
 lm_deluge/api_requests/deprecated/bedrock.py,sha256=WrcIShCoO8JCUSlFOCHxg6KQCNTZfw3TpYTvSpYk4mA,11320
 lm_deluge/api_requests/deprecated/cohere.py,sha256=KgDScD6_bWhAzOY5BHZQKSA3kurt4KGENqC4wLsGmcU,5142
 lm_deluge/api_requests/deprecated/deepseek.py,sha256=FEApI93VAWDwuaqTooIyKMgONYqRhdUmiAPBRme-IYs,4582
@@ -31,8 +32,8 @@ lm_deluge/util/json.py,sha256=dCeG9j1D17rXmQJbKJH79X0CGof4Wlqd55TDg4D6ky8,5388
 lm_deluge/util/logprobs.py,sha256=UkBZakOxWluaLqHrjARu7xnJ0uCHVfLGHJdnYlEcutk,11768
 lm_deluge/util/validation.py,sha256=hz5dDb3ebvZrZhnaWxOxbNSVMI6nmaOODBkk0htAUhs,1575
 lm_deluge/util/xml.py,sha256=Ft4zajoYBJR3HHCt2oHwGfymGLdvp_gegVmJ-Wqk4Ck,10547
-lm_deluge-0.0.11.dist-info/licenses/LICENSE,sha256=uNNXGXPCw2TC7CUs7SEBkA-Mz6QBQFWUUEWDMgEs1dU,1058
-lm_deluge-0.0.11.dist-info/METADATA,sha256=jdPdmbo_F8ecKTAHtPTg2GeyCOFmmsJ6T4-4RUleU24,9210
-lm_deluge-0.0.11.dist-info/WHEEL,sha256=zaaOINJESkSfm_4HQVc5ssNzHCPXhJm0kEUakpsEHaU,91
-lm_deluge-0.0.11.dist-info/top_level.txt,sha256=hqU-TJX93yBwpgkDtYcXyLr3t7TLSCCZ_reytJjwBaE,10
-lm_deluge-0.0.11.dist-info/RECORD,,
+lm_deluge-0.0.12.dist-info/licenses/LICENSE,sha256=uNNXGXPCw2TC7CUs7SEBkA-Mz6QBQFWUUEWDMgEs1dU,1058
+lm_deluge-0.0.12.dist-info/METADATA,sha256=lpl7mGKp096-Ccp6kgd57vuvZOGXs8ska41z_RfvNls,11663
+lm_deluge-0.0.12.dist-info/WHEEL,sha256=zaaOINJESkSfm_4HQVc5ssNzHCPXhJm0kEUakpsEHaU,91
+lm_deluge-0.0.12.dist-info/top_level.txt,sha256=hqU-TJX93yBwpgkDtYcXyLr3t7TLSCCZ_reytJjwBaE,10
+lm_deluge-0.0.12.dist-info/RECORD,,