lm-deluge 0.0.15__py3-none-any.whl → 0.0.17__py3-none-any.whl

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)
@@ -387,6 +458,89 @@ class LLMClient(BaseModel):
                 # final item
                 return item
 
+    async def run_agent_loop(
+        self,
+        conversation: str | Conversation,
+        *,
+        tools: list[Tool | dict] | None = None,
+        max_rounds: int = 5,
+        show_progress: bool = False,
+    ) -> tuple[Conversation, APIResponse]:
+        """Run a simple agent loop until no more tool calls are returned.
+
+        The provided ``conversation`` will be mutated and returned alongside the
+        final ``APIResponse`` from the model. ``tools`` may include ``Tool``
+        instances or built‑in tool dictionaries.
+        """
+
+        if isinstance(conversation, str):
+            conversation = Conversation.user(conversation)
+
+        last_response: APIResponse | None = None
+
+        for _ in range(max_rounds):
+            responses = await self.process_prompts_async(
+                [conversation],
+                tools=tools,  # type: ignore
+                return_completions_only=False,
+                show_progress=show_progress,
+            )
+
+            last_response = responses[0]
+            if last_response is None or last_response.content is None:
+                break
+
+            conversation.add(last_response.content)
+
+            tool_calls = last_response.content.tool_calls
+            if not tool_calls:
+                break
+
+            for call in tool_calls:
+                tool_obj = None
+                if tools:
+                    for t in tools:
+                        if isinstance(t, Tool) and t.name == call.name:
+                            tool_obj = t
+                            break
+
+                if isinstance(tool_obj, Tool) and tool_obj.run is not None:
+                    try:
+                        result = await tool_obj.acall(**call.arguments)
+                    except Exception as e:  # pragma: no cover - best effort
+                        result = f"Error: {e}"
+                else:
+                    result = f"Tool {call.name} not found"
+
+                if not isinstance(result, (str, dict, list)):
+                    result = str(result)
+
+                conversation.add_tool_result(call.id, result)  # type: ignore
+
+        if last_response is None:
+            raise RuntimeError("model did not return a response")
+
+        return conversation, last_response
+
+    def run_agent_loop_sync(
+        self,
+        conversation: str | Conversation,
+        *,
+        tools: list[Tool | dict | MCPServer] | None = None,
+        max_rounds: int = 5,
+        show_progress: bool = False,
+    ) -> tuple[Conversation, APIResponse]:
+        """Synchronous wrapper for :meth:`run_agent_loop`."""
+
+        return asyncio.run(
+            self.run_agent_loop(
+                conversation,
+                tools=tools,  # type: ignore
+                max_rounds=max_rounds,
+                show_progress=show_progress,
+            )
+        )
+
     async def submit_batch_job(
         self,
         prompts: Sequence[str | list[dict] | Conversation],
@@ -409,7 +563,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/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:
         """

lm_deluge/llm_tools/extract.py

@@ -1,9 +1,12 @@
+import asyncio
 import io
 import json
-from ..prompt import Conversation
-import asyncio
-from ..client import LLMClient
 from typing import Any
+
+from lm_deluge.file import File
+
+from ..client import LLMClient
+from ..prompt import Conversation
 from ..util.json import load_json
 
 try:
@@ -62,6 +65,15 @@ async def extract_async(
         + 'like `{"error": "The document is not relevant to the schema."}`.'
     )
 
+    file_prompt = (
+        f"Given the attached {document_name} file, extract the {object_name}information "
+        + "from it according to the following JSON schema:\n\n```json\n"
+        + json.dumps(schema_dict, indent=2)
+        + "Return the extracted information as JSON, no explanation required. "
+        + f"If the {document_name} seems to be totally irrelevant to the schema (not just incomplete), you may return a JSON object "
+        + 'like `{"error": "The document is not relevant to the schema."}`.'
+    )
+
     prompts = []
     for input in inputs:
         if isinstance(input, str):
@@ -74,8 +86,15 @@ async def extract_async(
             prompts.append(
                 Conversation.user(text=image_only_prompt, image=buffer.getvalue())
             )
+        elif isinstance(input, File):
+            data = input.data
+            if isinstance(data, io.BytesIO):
+                data = data.getvalue()
+            prompts.append(Conversation.user(text=file_prompt, file=data))
         else:
-            raise ValueError("inputs must be a list of strings or PIL images.")
+            raise ValueError(
+                "inputs must be a list of strings or PIL images or a File object."
+            )
 
     if return_prompts:
         return prompts

lm_deluge/llm_tools/ocr.py

@@ -0,0 +1 @@
+# NOT IMPLEMENTED YET!