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

lm_deluge/{computer_use/anthropic_tools.py → built_in_tools/anthropic.py}

@@ -11,14 +11,16 @@ def model_to_version(model: str) -> ToolVersion:
         return "2025-04-29"
     elif "3.7" in model:
         return "2025-01-24"
-    else:
+    elif "3.6" in model:
         return "2024-10-22"
+    else:
+        raise ValueError("unsupported model for anthropic CUA")
 
 
 def get_anthropic_cu_tools(
     model: str,
-    display_width: int,
-    display_height: int,
+    display_width: int = 1024,
+    display_height: int = 768,
     exclude_tools: list[ToolType] | None = None,
 ):
     version = model_to_version(model)
@@ -31,8 +33,8 @@ def get_anthropic_cu_tools(
                 "display_height_px": display_height,
                 "display_number": None,
             },
-            {"name": "str_replace_editor", "type": "text_editor_20250429"},
-            {"type": "bash_20250124", "name": "bash"},
+            {"name": "str_replace_editor", "type": "text_editor_20241022"},
+            {"name": "bash", "type": "bash_20241022"},
         ]
     elif version == "2025-01-24":
         result = [
@@ -73,3 +75,52 @@ def get_anthropic_cu_tools(
     if "computer" in exclude_tools:
         result = [x for x in result if "computer" not in x["name"]]
     return result
+
+
+def bash_tool(model: str = "claude-4-sonnet"):
+    # Claude Sonnet 3.5 requires the computer-use-2024-10-22 beta header when using the bash tool.
+    # The bash tool is generally available in Claude 4 and Sonnet 3.7.
+    if "claude-4" in model:
+        return {"type": "text_editor_20250429", "name": "str_replace_based_edit_tool"}
+    elif "3.7" in model:
+        return {"type": "text_editor_20250124", "name": "str_replace_editor"}
+    else:
+        return {"type": "text_editor_20241022", "name": "str_replace_editor"}
+
+
+def text_editor_tool(model: str = "claude-4-sonnet"):
+    if "claude-4" in model:
+        return {"type": "bash_20250124", "name": "bash"}
+    elif "3.7" in model:
+        return {"type": "bash_20250124", "name": "bash"}
+    else:
+        return {"type": "bash_20241022", "name": "bash"}
+
+
+def web_search_tool(max_uses: int = 5):
+    res = {
+        "type": "web_search_20250305",
+        "name": "web_search",
+        # Optional: Limit the number of searches per request
+        "max_uses": 5,
+        # You can use either allowed_domains or blocked_domains, but not both in the same request.
+        # Optional: Only include results from these domains
+        # "allowed_domains": ["example.com", "trusteddomain.org"],
+        #  Optional: Never include results from these domains
+        # "blocked_domains": ["untrustedsource.com"],
+        # Optional: Localize search results
+        # "user_location": {
+        #   "type": "approximate",
+        #   "city": "San Francisco",
+        #   "region": "California",
+        #   "country": "US",
+        #   "timezone": "America/Los_Angeles"
+        # }
+    }
+    return res
+
+
+def code_execution_tool():
+    # The code execution tool is currently in beta.
+    # This feature requires the beta header: "anthropic-beta": "code-execution-2025-05-22"
+    return {"type": "code_execution_20250522", "name": "code_execution"}

lm_deluge/built_in_tools/openai.py

@@ -0,0 +1,28 @@
+def image_generation_openai():
+    # TODO: handle result properly
+    return {"type": "image_generation"}
+
+
+def code_interpreter_openai(container: dict | None = None):
+    if container is None:
+        container = {"type": "auto"}
+    return {"type": "code_interpreter", "container": container}
+
+
+def local_shell_openai():
+    return {"type": "local_shell"}
+
+
+def web_search_openai():
+    return {"type": "web_search_preview"}
+
+
+def computer_use_openai(
+    display_width: int = 1024, display_height: int = 768, environment: str = "browser"
+):
+    return {
+        "type": "computer_use_preview",
+        "display_width": display_width,
+        "display_height": display_height,
+        "environment": environment,
+    }

lm_deluge/client.py

@@ -1,4 +1,5 @@
 import asyncio
+import random
 from typing import Any, Literal, Self, Sequence, overload
 
 import numpy as np
@@ -13,12 +14,12 @@ from lm_deluge.batches import (
     wait_for_batch_completion_async,
 )
 from lm_deluge.prompt import CachePattern, Conversation, prompts_to_conversations
-from lm_deluge.tool import Tool
+from lm_deluge.tool import MCPServer, Tool
 
-from .api_requests import create_api_request
-from .api_requests.base import APIRequestBase, APIResponse, deduplicate_responses
+from .api_requests.base import APIResponse
 from .config import SamplingParams
-from .models import registry
+from .models import APIModel, registry
+from .request_context import RequestContext
 from .tracker import StatusTracker
 
 # from .cache import LevelDBCache, SqliteCache
@@ -135,9 +136,7 @@ class LLMClient(BaseModel):
                     print(
                         "WARNING: using top_logprobs can result in very large outputs. consider limiting max_new_tokens."
                     )
-            if not all(
-                registry[model].get("supports_logprobs") for model in self.models
-            ):
+            if not all(registry[model].supports_logprobs for model in self.models):
                 raise ValueError(
                     "logprobs can only be enabled if all models support it."
                 )
@@ -174,6 +173,110 @@ class LLMClient(BaseModel):
         model_idx = np.random.choice(range(len(self.models)), p=self.model_weights)
         return self.models[model_idx], self.sampling_params[model_idx]
 
+    def _select_different_model(self, current_model: str):
+        """Select a model different from the provided one."""
+        other_models = [m for m in self.models if m != current_model]
+        if not other_models:
+            # No other models available, return current
+            return current_model, self.sampling_params[self.models.index(current_model)]
+
+        # Get weights for other models
+        other_indices = [self.models.index(m) for m in other_models]
+        weights = [self.model_weights[idx] for idx in other_indices]
+        weights = [w / sum(weights) for w in weights]  # type: ignore
+
+        model_idx = np.random.choice(range(len(other_models)), p=weights)
+        chosen_model = other_models[model_idx]
+        chosen_sp = self.sampling_params[self.models.index(chosen_model)]
+        return chosen_model, chosen_sp
+
+    async def _wait_for_capacity(self, num_tokens: int, tracker: StatusTracker):
+        while True:
+            if tracker.check_capacity(num_tokens):
+                tracker.set_limiting_factor(None)
+                return
+
+            if tracker.seconds_to_pause > 0:
+                await asyncio.sleep(tracker.seconds_to_pause)
+            else:
+                await asyncio.sleep(random.random())
+
+    async def _execute_request(self, context: RequestContext) -> APIResponse:
+        """Create and send a single API request using the provided context."""
+        model_obj = APIModel.from_registry(context.model_name)
+        request = model_obj.make_request(context)
+        response = await request.execute_once()
+        return response
+
+    async def process_single_request(
+        self, context: RequestContext, retry_queue: asyncio.Queue | None = None
+    ) -> APIResponse:
+        """Handle caching and single HTTP call for a request. Failed requests go to retry queue."""
+        # Check cache first
+        if self.cache:
+            cached = self.cache.get(context.prompt)
+            if cached:
+                cached.local_cache_hit = True
+                if context.status_tracker:
+                    context.status_tracker.task_succeeded(context.task_id)
+                return cached
+
+        # Execute single request
+        assert context.status_tracker
+        context.status_tracker.update_pbar()
+        response = await self._execute_request(context)
+
+        # Handle successful response
+        if not response.is_error:
+            context.status_tracker.task_succeeded(context.task_id)
+            # Cache successful responses immediately
+            if self.cache and response.completion:
+                self.cache.put(context.prompt, response)
+            # Call callback if provided
+            context.maybe_callback(response, context.status_tracker)
+            return response
+
+        # Handle error response - add to retry queue if available
+        if retry_queue and context.attempts_left > 1:
+            # Decide whether to retry with a different model
+            if response.retry_with_different_model and len(self.models) > 1:
+                # Switch to different model for retry
+                new_model, new_sp = self._select_different_model(context.model_name)
+                retry_context = context.copy(
+                    model_name=new_model,
+                    sampling_params=new_sp,
+                    attempts_left=context.attempts_left - 1,
+                )
+            else:
+                # Retry with same model
+                retry_context = context.copy(attempts_left=context.attempts_left - 1)
+
+            # Print error message for debugging
+            error_msg = (
+                f"Error task {context.task_id}. Model: {response.model_internal}"
+            )
+            if response.status_code:
+                error_msg += f" Code: {response.status_code},"
+            error_msg += f" Message: {response.error_message}. Retrying..."
+            print(error_msg)
+
+            # Add to retry queue for later processing
+            await retry_queue.put(retry_context)
+            return response  # Return the error response for now
+
+        # No retries left or no retry queue - final failure
+        context.status_tracker.task_failed(context.task_id)
+        context.maybe_callback(response, context.status_tracker)
+
+        # Print final error message
+        error_msg = f"Error task {context.task_id}. Model: {response.model_internal}"
+        if response.status_code:
+            error_msg += f" Code: {response.status_code},"
+        error_msg += f" Message: {response.error_message}. Giving up."
+        print(error_msg)
+
+        return response
+
     @overload
     async def process_prompts_async(
         self,
@@ -181,11 +284,8 @@ class LLMClient(BaseModel):
         *,
         return_completions_only: Literal[True],
         show_progress: bool = ...,
-        tools: list[Tool] | None = ...,
+        tools: list[Tool | dict | MCPServer] | None = ...,
         cache: CachePattern | None = ...,
-        computer_use: bool = ...,
-        display_width: int = ...,
-        display_height: int = ...,
         use_responses_api: bool = ...,
     ) -> list[str | None]: ...
 
@@ -196,11 +296,8 @@ class LLMClient(BaseModel):
         *,
         return_completions_only: Literal[False] = ...,
         show_progress: bool = ...,
-        tools: list[Tool] | None = ...,
+        tools: list[Tool | dict | MCPServer] | None = ...,
         cache: CachePattern | None = ...,
-        computer_use: bool = ...,
-        display_width: int = ...,
-        display_height: int = ...,
         use_responses_api: bool = ...,
     ) -> list[APIResponse | None]: ...
 
@@ -210,147 +307,117 @@ class LLMClient(BaseModel):
         *,
         return_completions_only: bool = False,
         show_progress: bool = True,
-        tools: list[Tool] | None = None,
+        tools: list[Tool | dict | MCPServer] | None = None,
         cache: CachePattern | None = None,
-        computer_use: bool = False,
-        display_width: int = 1024,
-        display_height: int = 768,
         use_responses_api: bool = False,
     ) -> list[APIResponse | None] | list[str | None] | dict[str, int]:
-        # if prompts are not Conversations, convert them.
+        # Convert prompts to Conversations - no upfront cache checking for dynamic caching!
         prompts = prompts_to_conversations(prompts)
-        ids = np.arange(len(prompts))
-
-        # if using cache, check for cached completions
-        if self.cache:
-            cached_results = [self.cache.get(prompt) for prompt in prompts]
-            cache_hit_ids = [
-                id for id, res in zip(ids, cached_results) if res is not None
-            ]
-            cache_hit_results = [res for res in cached_results if res is not None]
-            assert len(cache_hit_ids) == len(
-                cache_hit_results
-            ), "Cache hit ids and results must be the same length."
-            remaining_ids = np.array([i for i in ids if i not in cache_hit_ids])
-            remaining_prompts = [prompts[i] for i in remaining_ids]
-            print(
-                f"{len(cache_hit_ids)} cache hits; {len(remaining_ids)} prompts remaining."
-            )
-
-        else:
-            cache_hit_ids = []
-            cache_hit_results = []
-            remaining_prompts = prompts
-            remaining_ids = ids
-
+        ids = list(range(len(prompts)))
         results: list[APIResponse | None] = [None for _ in range(len(prompts))]
-        if len(remaining_prompts) > 0:
-            # Create StatusTracker with integrated progress bar
-            tracker = StatusTracker(
-                max_requests_per_minute=self.max_requests_per_minute,
-                max_tokens_per_minute=self.max_tokens_per_minute,
-                max_concurrent_requests=self.max_concurrent_requests,
-                use_progress_bar=show_progress,
-                progress_bar_total=len(prompts),
-                progress_bar_disable=not show_progress,
-                use_rich=show_progress,  # Disable Rich if progress is disabled
-            )
 
-            # Initialize progress bar and update with cache hits
-            tracker.init_progress_bar()
-            if len(cache_hit_ids) > 0:
-                tracker.update_pbar(len(cache_hit_ids))
-
-            if isinstance(ids, np.ndarray):
-                ids = ids.tolist()  # pyright: ignore
-
-            # calculate dynamically so we don't throttle RPM
-            seconds_to_sleep_each_loop = (60.0 * 0.9) / tracker.max_requests_per_minute
-            next_request = None  # variable to hold the next request to call
-            prompts_not_finished = True
-            prompts_iter = iter(zip(ids, prompts))
-            requests: list[APIRequestBase] = []
-            assert tracker.retry_queue, "retry queue not initialized"
-            while True:
-                # get next request (if one is not already waiting for capacity)
-                retry_request = False
-                if next_request is None:
-                    if not tracker.retry_queue.empty():
-                        next_request = tracker.retry_queue.get_nowait()
-                        retry_request = True
-                        print(f"Retrying request {next_request.task_id}.")
-                    elif prompts_not_finished:
+        # Create StatusTracker
+        tracker = StatusTracker(
+            max_requests_per_minute=self.max_requests_per_minute,
+            max_tokens_per_minute=self.max_tokens_per_minute,
+            max_concurrent_requests=self.max_concurrent_requests,
+            use_progress_bar=show_progress,
+            progress_bar_total=len(prompts),
+            progress_bar_disable=not show_progress,
+            use_rich=show_progress,
+        )
+
+        tracker.init_progress_bar()
+
+        # Create retry queue for failed requests
+        retry_queue: asyncio.Queue[RequestContext] = asyncio.Queue()
+
+        # Calculate sleep time for rate limiting
+        seconds_to_sleep_each_loop = (60.0 * 0.9) / tracker.max_requests_per_minute
+
+        # Main dispatch loop - using original pattern but with all prompts
+        next_context = None  # Persist across iterations like original
+        prompts_not_finished = True
+        prompts_iter = iter(zip(ids, prompts))
+
+        while True:
+            # Get next context (retry or new) - only if we don't already have one waiting
+            retry_request = False
+            if next_context is None:
+                if not retry_queue.empty():
+                    next_context = retry_queue.get_nowait()
+                    retry_request = True
+                    print(f"Retrying request {next_context.task_id}.")
+                elif prompts_not_finished:
+                    try:
+                        task_id, prompt = next(prompts_iter)
+                        model, sampling_params = self._select_model()
+                        assert isinstance(prompt, Conversation)
+                        next_context = RequestContext(
+                            task_id=task_id,
+                            model_name=model,
+                            prompt=prompt,
+                            sampling_params=sampling_params,
+                            attempts_left=self.max_attempts,
+                            request_timeout=self.request_timeout,
+                            status_tracker=tracker,
+                            tools=tools,
+                            cache=cache,
+                            use_responses_api=use_responses_api,
+                        )
+                    except StopIteration:
+                        prompts_not_finished = False
+
+            # Update capacity - original logic
+            tracker.update_capacity()
+
+            # Dispatch if capacity available - original logic
+            if next_context:
+                if tracker.check_capacity(next_context.num_tokens, retry=retry_request):
+                    tracker.set_limiting_factor(None)
+
+                    # Launch simplified request processing
+                    async def process_and_store(ctx: RequestContext):
                         try:
-                            # get new request
-                            id, prompt = next(prompts_iter)
-                            # select model
-                            model, sampling_params = self._select_model()
-
-                            next_request = create_api_request(
-                                task_id=id,
-                                model_name=model,
-                                prompt=prompt,  # type: ignore
-                                request_timeout=self.request_timeout,
-                                attempts_left=self.max_attempts,
-                                status_tracker=tracker,
-                                results_arr=requests,
-                                sampling_params=sampling_params,
-                                all_model_names=self.models,
-                                all_sampling_params=self.sampling_params,
-                                tools=tools,
-                                cache=cache,
-                                computer_use=computer_use,
-                                display_width=display_width,
-                                display_height=display_height,
-                                use_responses_api=use_responses_api,
+                            response = await self.process_single_request(
+                                ctx, retry_queue
                             )
-                            requests.append(next_request)
-
-                        except StopIteration:
-                            prompts_not_finished = False
-                            # print("API requests finished, only retries remain.")
-
-                # update available capacity
-                tracker.update_capacity()
-
-                # if enough capacity available, call API
-                if next_request:
-                    next_request_tokens = next_request.num_tokens
-                    if tracker.check_capacity(next_request_tokens, retry=retry_request):
-                        tracker.set_limiting_factor(None)
-                        # call API (attempts_left will be decremented in handle_error if it fails)
-                        asyncio.create_task(next_request.call_api())
-                        next_request = None  # reset next_request to empty
-                # update pbar status
-                tracker.update_pbar()
-
-                # if all tasks are finished, break
-                if tracker.num_tasks_in_progress == 0:
-                    break
-
-                # main loop sleeps briefly so concurrent tasks can run
-                await asyncio.sleep(seconds_to_sleep_each_loop)
-
-                # if a rate limit error was hit recently, pause to cool down
-                if tracker.seconds_to_pause > 0:
-                    await asyncio.sleep(tracker.seconds_to_pause)
-                    print(f"Pausing {tracker.seconds_to_pause}s to cool down.")
-
-            # after finishing, log final status
-            tracker.log_final_status()
-
-            # deduplicate results by id
-            api_results = deduplicate_responses(requests)
-            for res in api_results:
-                results[res.id] = res
-                # set to cache if result has a completion
-                if self.cache and res.completion:
-                    self.cache.put(prompts[res.id], res)
+                            results[ctx.task_id] = response
+                        except Exception as e:
+                            # Create an error response for validation errors and other exceptions
+                            from .api_requests.response import APIResponse
+
+                            error_response = APIResponse(
+                                id=ctx.task_id,
+                                model_internal=ctx.model_name,
+                                prompt=ctx.prompt,
+                                sampling_params=ctx.sampling_params,
+                                status_code=None,
+                                is_error=True,
+                                error_message=str(e),
+                            )
+                            results[ctx.task_id] = error_response
+                            # Mark task as completed so the main loop can finish
+                            if ctx.status_tracker:
+                                ctx.status_tracker.task_failed(ctx.task_id)
+
+                    asyncio.create_task(process_and_store(next_context))
+                    next_context = None  # Reset after successful dispatch
+
+            # Update progress - original logic
+            tracker.update_pbar()
+
+            # Check completion - original logic
+            if (
+                tracker.num_tasks_in_progress == 0
+                and not prompts_not_finished
+                and retry_queue.empty()
+            ):
+                break
 
-        # add cache hits back in
-        for id, res in zip(cache_hit_ids, cache_hit_results):
-            res.cache_hit = True
-            results[id] = res
+            # Sleep - original logic
+            await asyncio.sleep(seconds_to_sleep_each_loop + tracker.seconds_to_pause)
+            tracker.log_final_status()
 
         if return_completions_only:
             return [r.completion if r is not None else None for r in results]
@@ -363,7 +430,7 @@ class LLMClient(BaseModel):
         *,
         return_completions_only: bool = False,
         show_progress=True,
-        tools: list[Tool] | None = None,
+        tools: list[Tool | dict | MCPServer] | None = None,
         cache: CachePattern | None = None,
     ):
         return asyncio.run(
@@ -376,7 +443,11 @@ class LLMClient(BaseModel):
             )
         )
 
-    async def stream(self, prompt: str | Conversation, tools: list[Tool] | None = None):
+    async def stream(
+        self,
+        prompt: str | Conversation,
+        tools: list[Tool | dict | MCPServer] | None = None,
+    ):
         model, sampling_params = self._select_model()
         if isinstance(prompt, str):
             prompt = Conversation.user(prompt)
@@ -409,7 +480,7 @@ class LLMClient(BaseModel):
         if len(self.models) != 1:
             raise ValueError("Batch jobs can only be submitted with a single model.")
         model = self.models[0]
-        api_spec = registry[model].get("api_spec", None)
+        api_spec = registry[model].api_spec
 
         if api_spec == "openai":
             return await submit_batches_oa(model, self.sampling_params[0], prompts)

lm_deluge/file.py

@@ -141,8 +141,13 @@ class File:
         return filename, content, media_type
 
     def gemini(self) -> dict:
-        """For Gemini API - not yet supported."""
-        raise NotImplementedError("File support for Gemini is not yet implemented")
+        """For Gemini API - files are provided as inline data."""
+        return {
+            "inlineData": {
+                "mimeType": self._mime(),
+                "data": self._base64(include_header=False),
+            }
+        }
 
     def mistral(self) -> dict:
         """For Mistral API - not yet supported."""

lm_deluge/image.py

@@ -1,6 +1,5 @@
 import os
 from contextlib import contextmanager
-from functools import cached_property
 import io
 import requests
 from PIL import Image as PILImage  # type: ignore
@@ -18,6 +17,8 @@ class Image:
     media_type: str | None = None  # inferred if None
     detail: Literal["low", "high", "auto"] = "auto"
     type: str = field(init=False, default="image")
+    _fingerprint_cache: str | None = field(init=False, default=None)
+    _size_cache: tuple[int, int] | None = field(init=False, default=None)
 
     @classmethod
     def from_pdf(
@@ -95,12 +96,14 @@ class Image:
             if img:
                 img.close()
 
-    @cached_property
+    @property
     def size(self) -> tuple[int, int]:
-        with self._image() as img:
-            return img.size
+        if self._size_cache is None:
+            with self._image() as img:
+                self._size_cache = img.size
+        return self._size_cache
 
-    @cached_property
+    @property
     def num_pixels(self) -> int:
         return self.size[0] * self.size[1]
 
@@ -143,11 +146,13 @@ class Image:
             new_width = int(new_height / height * width)
         return self._resize((new_width, new_height))
 
-    @cached_property
+    @property
     def fingerprint(self) -> str:
         # return base64 of a very small version of the image
-        small_image = self._resize_longer(max_size=48)  # longer side = 48px
-        return base64.b64encode(small_image).decode("utf-8")
+        if self._fingerprint_cache is None:
+            small_image = self._resize_longer(max_size=48)  # longer side = 48px
+            self._fingerprint_cache = base64.b64encode(small_image).decode("utf-8")
+        return self._fingerprint_cache
 
     def resize(self, max_size: int) -> None:
         """