lm-deluge 0.0.8__py3-none-any.whl → 0.0.10__py3-none-any.whl

lm_deluge/client.py

@@ -9,6 +9,7 @@ from typing import Sequence, overload, Literal, Any
 from tqdm.auto import tqdm
 
 from lm_deluge.prompt import Conversation
+from lm_deluge.tool import Tool
 
 from .tracker import StatusTracker
 from .sampling_params import SamplingParams
@@ -225,39 +226,47 @@ class LLMClient:
     async def process_prompts_async(
         self,
         prompts: Sequence[str | list[dict] | Conversation],
+        *,
         return_completions_only: bool,
         show_progress: bool = ...,
-        dry_run: Literal[True] = ...,
+        dry_run: Literal[True],
         verbose: bool = ...,
+        tools: list[Tool] | None = ...,
     ) -> dict[str, int]: ...
 
     @overload
     async def process_prompts_async(
         self,
         prompts: Sequence[str | list[dict] | Conversation],
+        *,
         return_completions_only: Literal[True],
         show_progress: bool = ...,
-        dry_run: Literal[False] = ...,
+        dry_run: bool = ...,
         verbose: bool = ...,
+        tools: list[Tool] | None = ...,
     ) -> list[str | None]: ...
 
     @overload
     async def process_prompts_async(
         self,
         prompts: Sequence[str | list[dict] | Conversation],
+        *,
         return_completions_only: Literal[False] = ...,
         show_progress: bool = ...,
-        dry_run: Literal[False] = ...,
+        dry_run: bool = ...,
         verbose: bool = ...,
+        tools: list[Tool] | None = ...,
     ) -> list[APIResponse | None]: ...
 
     async def process_prompts_async(
         self,
         prompts: Sequence[str | list[dict] | Conversation],
+        *,
         return_completions_only: bool = False,
         show_progress: bool = True,
         dry_run: bool = False,
         verbose: bool = False,
+        tools: list[Tool] | None = None,
     ) -> list[APIResponse | None] | list[str | None] | dict[str, int]:
         # if prompts are not Conversations, convert them.
         # can only handle strings for now
@@ -335,6 +344,7 @@ class LLMClient:
                     progress_bar=pbar,
                     use_qps=self.use_qps,
                     verbose=verbose,
+                    tools=tools,
                 )
             )
             api_results: list[APIResponse] = await api_task
@@ -357,10 +367,12 @@ class LLMClient:
     def process_prompts_sync(
         self,
         prompts: Sequence[str | list[dict] | Conversation],
+        *,
         return_completions_only: bool = False,
         show_progress=True,
         dry_run: bool = False,
         verbose: bool = False,
+        tools: list[Tool] | None = None,
     ):
         return asyncio.run(
             self.process_prompts_async(
@@ -369,6 +381,7 @@ class LLMClient:
                 show_progress=show_progress,
                 dry_run=dry_run,
                 verbose=verbose,
+                tools=tools,
             )
         )
 
@@ -556,6 +569,7 @@ async def process_api_prompts_async(
     progress_bar: tqdm | None = None,
     use_qps: bool = False,
     verbose: bool = False,
+    tools: list[Tool] | None = None,
 ):
     """Processes API requests in parallel, throttling to stay under rate limits."""
     # change ids to integer list
@@ -639,6 +653,7 @@ async def process_api_prompts_async(
                         pbar=progress_bar,
                         all_model_names=models,
                         all_sampling_params=sampling_params,
+                        tools=tools,
                     )
                     status_tracker.num_tasks_started += 1
                     status_tracker.num_tasks_in_progress += 1

lm_deluge/models.py

@@ -629,30 +629,95 @@ registry = {
     #  ░███    ░███░███░░░  ░███ ░███  ░███     ░███ ░███░███  ███ ░███░░███
     #  ███████████ ░░██████ ░░████████ █████    ░░██████ ░░██████  ████ █████
     # ░░░░░░░░░░░   ░░░░░░   ░░░░░░░░ ░░░░░      ░░░░░░   ░░░░░░  ░░░░ ░░░░░
-    # "claude-haiku-bedrock": {
-    #     "id": "claude-haiku-bedrock",
-    #     "name": "anthropic.claude-3-haiku-20240307-v1:0",
-    #     "regions": ["us-east-1", "us-west-2", "ap-southeast-2", "eu-west-3"],
-    #     "api_base": "",
-    #     "api_key_env_var": "",
-    #     "api_spec": "bedrock_anthropic",
-    #     "input_cost": 0.25,
-    #     "output_cost": 1.25,
-    #     "requests_per_minute": 4_000,
-    #     "tokens_per_minute": 8_000_000,
-    # },
-    # "claude-sonnet-bedrock": {
-    #     "id": "claude-sonnet-bedrock",
-    #     "name": "anthropic.claude-3-sonnet-20240229-v1:0",
-    #     "regions": ["us-east-1", "us-west-2", "ap-southeast-2", "eu-west-3"],
-    #     "api_base": "",
-    #     "api_key_env_var": "",
-    #     "api_spec": "bedrock_anthropic",
-    #     "input_cost": 3.0,
-    #     "output_cost": 15.0,
-    #     "requests_per_minute": 2_000,
-    #     "tokens_per_minute": 4_000_000,
-    # },
+    "claude-3-haiku-bedrock": {
+        "id": "claude-3-haiku-bedrock",
+        "name": "us.anthropic.claude-3-haiku-20240307-v1:0",
+        "regions": ["us-east-1", "us-west-2", "ap-southeast-2", "eu-west-3"],
+        "api_base": "",
+        "api_key_env_var": "",
+        "api_spec": "bedrock",
+        "input_cost": 0.25,
+        "output_cost": 1.25,
+        "requests_per_minute": 4_000,
+        "tokens_per_minute": 8_000_000,
+    },
+    "claude-3.5-haiku-bedrock": {
+        "id": "claude-3.5-haiku-bedrock",
+        "name": "us.anthropic.claude-3-5-haiku-20241022-v1:0",
+        "regions": ["us-east-1", "us-west-2", "ap-southeast-2", "eu-west-3"],
+        "api_base": "",
+        "api_key_env_var": "",
+        "api_spec": "bedrock",
+        "input_cost": 0.25,
+        "output_cost": 1.25,
+        "requests_per_minute": 4_000,
+        "tokens_per_minute": 8_000_000,
+    },
+    "claude-3.5-sonnet-bedrock": {
+        "id": "claude-3.5-sonnet-bedrock",
+        "name": "us.anthropic.claude-3-5-sonnet-20240620-v1:0",
+        "regions": ["us-east-1", "us-west-2"],
+        "api_base": "",
+        "api_key_env_var": "",
+        "api_spec": "bedrock",
+        "input_cost": 3.0,
+        "output_cost": 15.0,
+        "requests_per_minute": 4_000,
+        "tokens_per_minute": 400_000,
+        "reasoning_model": False,
+    },
+    "claude-3.6-sonnet-bedrock": {
+        "id": "claude-3.6-sonnet-bedrock",
+        "name": "us.anthropic.claude-3-5-sonnet-20241022-v2:0",
+        "regions": ["us-east-1", "us-west-2", "us-east-2"],
+        "api_base": "",
+        "api_key_env_var": "",
+        "api_spec": "bedrock",
+        "input_cost": 3.0,
+        "output_cost": 15.0,
+        "requests_per_minute": 4_000,
+        "tokens_per_minute": 400_000,
+        "reasoning_model": False,
+    },
+    "claude-3.7-sonnet-bedrock": {
+        "id": "claude-3.7-sonnet-bedrock",
+        "name": "us.anthropic.claude-3-7-sonnet-20250219-v1:0",
+        "regions": ["us-east-1", "us-west-2", "eu-west-1"],
+        "api_base": "",
+        "api_key_env_var": "",
+        "api_spec": "bedrock",
+        "input_cost": 3.0,
+        "output_cost": 15.0,
+        "requests_per_minute": 4_000,
+        "tokens_per_minute": 400_000,
+        "reasoning_model": True,
+    },
+    "claude-4-sonnet-bedrock": {
+        "id": "claude-4-sonnet-bedrock",
+        "name": "us.anthropic.claude-sonnet-4-20250514-v1:0",
+        "regions": ["us-east-1", "us-west-2", "us-east-2"],
+        "api_base": "",
+        "api_key_env_var": "",
+        "api_spec": "bedrock",
+        "input_cost": 3.0,
+        "output_cost": 15.0,
+        "requests_per_minute": 4_000,
+        "tokens_per_minute": 400_000,
+        "reasoning_model": True,
+    },
+    "claude-4-opus-bedrock": {
+        "id": "claude-4-opus-bedrock",
+        "name": "us.anthropic.claude-opus-4-20250514-v1:0",
+        "regions": ["us-east-1", "us-west-2", "us-east-2"],
+        "api_base": "",
+        "api_key_env_var": "",
+        "api_spec": "bedrock",
+        "input_cost": 3.0,
+        "output_cost": 15.0,
+        "requests_per_minute": 4_000,
+        "tokens_per_minute": 400_000,
+        "reasoning_model": True,
+    },
     # "mistral-7b-bedrock": {
     #     "id": "mistral-7b-bedrock",
     #     "name": "mistral.mistral-7b-instruct-v0:2",