lm-deluge 0.0.31__tar.gz → 0.0.32__tar.gz

{lm_deluge-0.0.31/src/lm_deluge.egg-info → lm_deluge-0.0.32}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: lm_deluge
-Version: 0.0.31
+Version: 0.0.32
 Summary: Python utility for using LLM API models.
 Author-email: Benjamin Anderson <ben@trytaylor.ai>
 Requires-Python: >=3.10

{lm_deluge-0.0.31 → lm_deluge-0.0.32}/pyproject.toml

@@ -3,7 +3,7 @@ requires = ["setuptools", "wheel"]
 
 [project]
 name = "lm_deluge"
-version = "0.0.31"
+version = "0.0.32"
 authors = [{ name = "Benjamin Anderson", email = "ben@trytaylor.ai" }]
 description = "Python utility for using LLM API models."
 readme = "README.md"

{lm_deluge-0.0.31 → lm_deluge-0.0.32}/src/lm_deluge/api_requests/anthropic.py

@@ -38,8 +38,8 @@ def _build_anthropic_request(
     tools = context.tools
     sampling_params = context.sampling_params
     system_message, messages = prompt.to_anthropic(cache_pattern=cache_pattern)
-    if not system_message:
-        print("WARNING: system_message is None")
+    # if not system_message:
+    #     print("WARNING: system_message is None")
     base_headers = {
         "x-api-key": os.getenv(model.api_key_env_var),
         "anthropic-version": "2023-06-01",

{lm_deluge-0.0.31 → lm_deluge-0.0.32}/src/lm_deluge/client.py

@@ -4,7 +4,7 @@ from typing import Any, Literal, Self, Sequence, overload
 
 import numpy as np
 import yaml
-from pydantic import BaseModel
+from pydantic import BaseModel, PrivateAttr
 from pydantic.functional_validators import model_validator
 
 from lm_deluge.api_requests.openai import stream_chat
@@ -32,13 +32,7 @@ class LLMClient(BaseModel):
     Handles models, sampling params for each model, model weights, rate limits, etc.
     """
 
-    model_names: list[str] = ["gpt-4.1-mini"]
-
-    def __init__(self, model_name: str | list[str] | None = None, **kwargs):
-        if model_name is not None:
-            kwargs["model_names"] = model_name
-        super().__init__(**kwargs)
-
+    model_names: str | list[str] = ["gpt-4.1-mini"]
     max_requests_per_minute: int = 1_000
     max_tokens_per_minute: int = 100_000
     max_concurrent_requests: int = 225
@@ -59,6 +53,13 @@ class LLMClient(BaseModel):
     top_logprobs: int | None = None
     force_local_mcp: bool = False
 
+    # Internal state for async task handling
+    _next_task_id: int = PrivateAttr(default=0)
+    _tasks: dict[int, asyncio.Task] = PrivateAttr(default_factory=dict)
+    _results: dict[int, APIResponse] = PrivateAttr(default_factory=dict)
+    _tracker: StatusTracker | None = PrivateAttr(default=None)
+    _capacity_lock: asyncio.Lock = PrivateAttr(default_factory=asyncio.Lock)
+
     # NEW! Builder methods
     def with_model(self, model: str):
         self.model_names = [model]
@@ -81,6 +82,18 @@ class LLMClient(BaseModel):
         if max_concurrent_requests:
             self.max_concurrent_requests = max_concurrent_requests
 
+    def _get_tracker(self) -> StatusTracker:
+        if self._tracker is None:
+            self._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=False,
+                progress_bar_disable=True,
+                use_rich=False,
+            )
+        return self._tracker
+
     @property
     def models(self):
         return self.model_names  # why? idk
@@ -90,6 +103,8 @@ class LLMClient(BaseModel):
     def fix_lists(cls, data) -> "LLMClient":
         if isinstance(data.get("model_names"), str):
             data["model_names"] = [data["model_names"]]
+        if not isinstance(data.get("sampling_params", []), list):
+            data["sampling_params"] = [data["sampling_params"]]
         if "sampling_params" not in data or len(data.get("sampling_params", [])) == 0:
             data["sampling_params"] = [
                 SamplingParams(
@@ -192,14 +207,19 @@ class LLMClient(BaseModel):
         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):
+    async def _wait_for_capacity(
+        self, num_tokens: int, tracker: StatusTracker, *, retry: bool = False
+    ):
         while True:
-            if tracker.check_capacity(num_tokens):
-                tracker.set_limiting_factor(None)
-                return
+            async with self._capacity_lock:
+                tracker.update_capacity()
+                if tracker.check_capacity(num_tokens, retry=retry):
+                    tracker.set_limiting_factor(None)
+                    return
+                seconds_to_pause = tracker.seconds_to_pause
 
-            if tracker.seconds_to_pause > 0:
-                await asyncio.sleep(tracker.seconds_to_pause)
+            if seconds_to_pause > 0:
+                await asyncio.sleep(seconds_to_pause)
             else:
                 await asyncio.sleep(random.random())
 
@@ -446,6 +466,78 @@ class LLMClient(BaseModel):
             )
         )
 
+    async def _run_context(self, context: RequestContext) -> APIResponse:
+        tracker = self._get_tracker()
+        retry = False
+        retry_queue: asyncio.Queue[RequestContext] = asyncio.Queue()
+        current = context
+        while True:
+            await self._wait_for_capacity(current.num_tokens, tracker, retry=retry)
+            response = await self.process_single_request(current, retry_queue)
+            if not response.is_error or retry_queue.empty():
+                self._results[context.task_id] = response
+                return response
+            current = await retry_queue.get()
+            retry = True
+
+    def start_nowait(
+        self,
+        prompt: str | Conversation,
+        *,
+        tools: list[Tool | dict | MCPServer] | None = None,
+        cache: CachePattern | None = None,
+        use_responses_api: bool = False,
+    ) -> int:
+        tracker = self._get_tracker()
+        task_id = self._next_task_id
+        self._next_task_id += 1
+        model, sampling_params = self._select_model()
+        if isinstance(prompt, str):
+            prompt = Conversation.user(prompt)
+        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,
+            extra_headers=self.extra_headers,
+            force_local_mcp=self.force_local_mcp,
+        )
+        task = asyncio.create_task(self._run_context(context))
+        self._tasks[task_id] = task
+        return task_id
+
+    async def start(
+        self,
+        prompt: str | Conversation,
+        *,
+        tools: list[Tool | dict | MCPServer] | None = None,
+        cache: CachePattern | None = None,
+        use_responses_api: bool = False,
+    ) -> APIResponse | None:
+        task_id = self.start_nowait(
+            prompt, tools=tools, cache=cache, use_responses_api=use_responses_api
+        )
+        return await self.wait_for(task_id)
+
+    async def wait_for(self, task_id: int) -> APIResponse | None:
+        task = self._tasks.get(task_id)
+        if task:
+            return await task
+        return self._results.get(task_id)
+
+    async def wait_for_all(
+        self, task_ids: Sequence[int] | None = None
+    ) -> list[APIResponse | None]:
+        if task_ids is None:
+            task_ids = list(self._tasks.keys())
+        return [await self.wait_for(tid) for tid in task_ids]
+
     async def stream(
         self,
         prompt: str | Conversation,

{lm_deluge-0.0.31 → lm_deluge-0.0.32/src/lm_deluge.egg-info}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: lm_deluge
-Version: 0.0.31
+Version: 0.0.32
 Summary: Python utility for using LLM API models.
 Author-email: Benjamin Anderson <ben@trytaylor.ai>
 Requires-Python: >=3.10