lm-deluge 0.0.90__py3-none-any.whl → 0.0.91__py3-none-any.whl

lm_deluge/__init__.py

@@ -1,7 +1,6 @@
 from .client import APIResponse, LLMClient, SamplingParams
-from .file import File
-from .prompt import Conversation, Message
-from .tool import Tool
+from .prompt import Conversation, Message, File
+from .tool import Tool, MCPServer
 
 # dotenv.load_dotenv() - don't do this, fucks with other packages
 
@@ -12,5 +11,6 @@ __all__ = [
     "Conversation",
     "Message",
     "Tool",
+    "MCPServer",
     "File",
 ]

lm_deluge/api_requests/anthropic.py

@@ -10,7 +10,7 @@ from lm_deluge.prompt import (
     Thinking,
     ToolCall,
 )
-from lm_deluge.request_context import RequestContext
+from lm_deluge.api_requests.context import RequestContext
 from lm_deluge.tool import MCPServer, Tool
 from lm_deluge.usage import Usage
 from lm_deluge.util.schema import (
@@ -103,7 +103,9 @@ def _build_anthropic_request(
             if "top_p" in request_json:
                 request_json["top_p"] = max(request_json["top_p"], 0.95)
             request_json["temperature"] = 1.0
-            request_json["max_tokens"] += budget
+            max_tokens = request_json["max_tokens"]
+            assert isinstance(max_tokens, int)
+            request_json["max_tokens"] = max_tokens + budget
         else:
             request_json["thinking"] = {"type": "disabled"}
             if "kimi" in model.id and "thinking" in model.id:

lm_deluge/api_requests/base.py

@@ -10,7 +10,7 @@ from aiohttp import ClientResponse
 
 from ..errors import raise_if_modal_exception
 from ..models.openai import OPENAI_MODELS
-from ..request_context import RequestContext
+from ..api_requests.context import RequestContext
 from .response import APIResponse
 
 

lm_deluge/api_requests/bedrock.py

@@ -20,7 +20,7 @@ from lm_deluge.prompt import (
     Thinking,
     ToolCall,
 )
-from lm_deluge.request_context import RequestContext
+from lm_deluge.api_requests.context import RequestContext
 from lm_deluge.tool import MCPServer, Tool
 from lm_deluge.usage import Usage
 
@@ -263,6 +263,11 @@ class BedrockRequest(APIRequestBase):
         # Create a fake requests.PreparedRequest object for AWS4Auth to sign
         import requests
 
+        assert self.url is not None, "URL must be set after build_request"
+        assert (
+            self.request_header is not None
+        ), "Headers must be set after build_request"
+
         fake_request = requests.Request(
             method="POST",
             url=self.url,

lm_deluge/{request_context.py → api_requests/context.py}

@@ -2,9 +2,9 @@ from dataclasses import dataclass, field
 from functools import cached_property
 from typing import Any, Callable, Sequence, TYPE_CHECKING
 
-from .config import SamplingParams
-from .prompt import CachePattern, Conversation
-from .tracker import StatusTracker
+from ..config import SamplingParams
+from ..prompt import CachePattern, Conversation
+from ..tracker import StatusTracker
 
 if TYPE_CHECKING:
     from pydantic import BaseModel
@@ -83,4 +83,4 @@ class RequestContext:
         # Update with any overrides
         current_values.update(overrides)
 
-        return RequestContext(**current_values)
+        return RequestContext(**current_values)  # type: ignore[arg-type]

lm_deluge/api_requests/gemini.py

@@ -1,9 +1,10 @@
 import json
 import os
+from typing import Any
 
 from aiohttp import ClientResponse
 
-from lm_deluge.request_context import RequestContext
+from lm_deluge.api_requests.context import RequestContext
 from lm_deluge.tool import Tool
 from lm_deluge.warnings import maybe_warn
 
@@ -37,13 +38,14 @@ async def _build_gemini_request(
                             part_type="function call",
                         )
 
-    request_json = {
+    generation_config: dict[str, Any] = {
+        "temperature": sampling_params.temperature,
+        "topP": sampling_params.top_p,
+        "maxOutputTokens": sampling_params.max_new_tokens,
+    }
+    request_json: dict[str, Any] = {
         "contents": messages,
-        "generationConfig": {
-            "temperature": sampling_params.temperature,
-            "topP": sampling_params.top_p,
-            "maxOutputTokens": sampling_params.max_new_tokens,
-        },
+        "generationConfig": generation_config,
     }
 
     # Add system instruction if present
@@ -83,7 +85,7 @@ async def _build_gemini_request(
                 }
             effort = level_map[effort_key]
         thinking_config = {"thinkingLevel": effort}
-        request_json["generationConfig"]["thinkingConfig"] = thinking_config
+        generation_config["thinkingConfig"] = thinking_config
 
     elif model.reasoning_model:
         if (
@@ -126,7 +128,7 @@ async def _build_gemini_request(
             # no thoughts head empty
             thinking_config = {"includeThoughts": False, "thinkingBudget": 0}
 
-        request_json["generationConfig"]["thinkingConfig"] = thinking_config
+        generation_config["thinkingConfig"] = thinking_config
 
     else:
         if sampling_params.reasoning_effort:
@@ -171,14 +173,14 @@ async def _build_gemini_request(
 
     # Handle JSON mode
     if sampling_params.json_mode and model.supports_json:
-        request_json["generationConfig"]["responseMimeType"] = "application/json"
+        generation_config["responseMimeType"] = "application/json"
 
     # Handle media_resolution for Gemini 3 (requires v1alpha)
     if sampling_params.media_resolution is not None:
         is_gemini_3 = "gemini-3" in model.name.lower()
         if is_gemini_3:
             # Add global media resolution to generationConfig
-            request_json["generationConfig"]["mediaResolution"] = {
+            generation_config["mediaResolution"] = {
                 "level": sampling_params.media_resolution
             }
         else:

lm_deluge/api_requests/mistral.py

@@ -7,7 +7,7 @@ from lm_deluge.warnings import maybe_warn
 
 from ..models import APIModel
 from ..prompt import Message
-from ..request_context import RequestContext
+from ..api_requests.context import RequestContext
 from ..usage import Usage
 from .base import APIRequestBase, APIResponse
 

lm_deluge/api_requests/openai.py

@@ -7,7 +7,7 @@ from typing import Sequence
 import aiohttp
 from aiohttp import ClientResponse
 
-from lm_deluge.request_context import RequestContext
+from lm_deluge.api_requests.context import RequestContext
 from lm_deluge.tool import MCPServer, Tool
 from lm_deluge.util.schema import (
     prepare_output_schema,
@@ -75,7 +75,9 @@ async def _build_oa_chat_request(
             request_json["service_tier"] = context.service_tier
     # if tinker, for now hack to mush into 1 string
     if "tinker" in model.name:
-        request_json["messages"] = _message_contents_to_string(request_json["messages"])
+        messages = request_json["messages"]
+        assert isinstance(messages, list)
+        request_json["messages"] = _message_contents_to_string(messages)
 
     # set max_tokens or max_completion_tokens dep. on provider
     if "cohere" in model.api_base:

lm_deluge/batches.py

@@ -3,7 +3,7 @@ import json
 import os
 import tempfile
 import time
-from typing import Literal, Sequence, cast
+from typing import Any, Literal, Sequence, cast
 
 import aiohttp
 from rich.console import Console
@@ -22,7 +22,7 @@ from lm_deluge.prompt import (
     Prompt,
     prompts_to_conversations,
 )
-from lm_deluge.request_context import RequestContext
+from lm_deluge.api_requests.context import RequestContext
 
 
 def _create_batch_status_display(
@@ -480,7 +480,7 @@ async def _wait_for_anthropic_batch_completion_async(
 
     # Event to signal when to stop the display updater
     stop_display_event = asyncio.Event()
-    current_status = {"status": "processing", "counts": None}
+    current_status: dict[str, Any] = {"status": "processing", "counts": None}
 
     async def display_updater():
         """Update display independently of polling."""
@@ -632,7 +632,7 @@ async def _wait_for_openai_batch_completion_async(
 
     # Event to signal when to stop the display updater
     stop_display_event = asyncio.Event()
-    current_status = {"status": "pending", "counts": None}
+    current_status: dict[str, Any] = {"status": "pending", "counts": None}
 
     async def display_updater():
         """Update display independently of polling."""

lm_deluge/cache.py

@@ -8,7 +8,7 @@ from .api_requests.base import APIResponse
 try:
     import plyvel  # type: ignore
 except ImportError:
-    plyvel = None
+    plyvel: Any = None
     print("Warning: plyvel not installed, cannot use LevelDB.")