lm-deluge 0.0.13__tar.gz → 0.0.14__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.

Potentially problematic release.


This version of lm-deluge might be problematic. Click here for more details.

Files changed (79) hide show
  1. {lm_deluge-0.0.13/src/lm_deluge.egg-info → lm_deluge-0.0.14}/PKG-INFO +4 -1
  2. {lm_deluge-0.0.13 → lm_deluge-0.0.14}/README.md +3 -0
  3. {lm_deluge-0.0.13 → lm_deluge-0.0.14}/pyproject.toml +1 -1
  4. {lm_deluge-0.0.13 → lm_deluge-0.0.14}/src/lm_deluge/__init__.py +2 -0
  5. {lm_deluge-0.0.13 → lm_deluge-0.0.14}/src/lm_deluge/api_requests/base.py +2 -148
  6. {lm_deluge-0.0.13 → lm_deluge-0.0.14}/src/lm_deluge/api_requests/openai.py +72 -6
  7. lm_deluge-0.0.14/src/lm_deluge/api_requests/response.py +153 -0
  8. {lm_deluge-0.0.13 → lm_deluge-0.0.14}/src/lm_deluge/client.py +36 -48
  9. {lm_deluge-0.0.13 → lm_deluge-0.0.14}/src/lm_deluge/config.py +3 -2
  10. lm_deluge-0.0.14/src/lm_deluge/file.py +149 -0
  11. {lm_deluge-0.0.13 → lm_deluge-0.0.14}/src/lm_deluge/prompt.py +70 -9
  12. {lm_deluge-0.0.13 → lm_deluge-0.0.14}/src/lm_deluge/tracker.py +5 -3
  13. {lm_deluge-0.0.13 → lm_deluge-0.0.14/src/lm_deluge.egg-info}/PKG-INFO +4 -1
  14. {lm_deluge-0.0.13 → lm_deluge-0.0.14}/src/lm_deluge.egg-info/SOURCES.txt +5 -0
  15. {lm_deluge-0.0.13 → lm_deluge-0.0.14}/tests/test_all_models.py +24 -24
  16. lm_deluge-0.0.14/tests/test_file_integration.py +156 -0
  17. lm_deluge-0.0.14/tests/test_file_support.py +210 -0
  18. lm_deluge-0.0.14/tests/test_retry_fix.py +67 -0
  19. {lm_deluge-0.0.13 → lm_deluge-0.0.14}/LICENSE +0 -0
  20. {lm_deluge-0.0.13 → lm_deluge-0.0.14}/setup.cfg +0 -0
  21. {lm_deluge-0.0.13 → lm_deluge-0.0.14}/src/lm_deluge/agent.py +0 -0
  22. {lm_deluge-0.0.13 → lm_deluge-0.0.14}/src/lm_deluge/api_requests/__init__.py +0 -0
  23. {lm_deluge-0.0.13 → lm_deluge-0.0.14}/src/lm_deluge/api_requests/anthropic.py +0 -0
  24. {lm_deluge-0.0.13 → lm_deluge-0.0.14}/src/lm_deluge/api_requests/bedrock.py +0 -0
  25. {lm_deluge-0.0.13 → lm_deluge-0.0.14}/src/lm_deluge/api_requests/common.py +0 -0
  26. {lm_deluge-0.0.13 → lm_deluge-0.0.14}/src/lm_deluge/api_requests/deprecated/bedrock.py +0 -0
  27. {lm_deluge-0.0.13 → lm_deluge-0.0.14}/src/lm_deluge/api_requests/deprecated/cohere.py +0 -0
  28. {lm_deluge-0.0.13 → lm_deluge-0.0.14}/src/lm_deluge/api_requests/deprecated/deepseek.py +0 -0
  29. {lm_deluge-0.0.13 → lm_deluge-0.0.14}/src/lm_deluge/api_requests/deprecated/mistral.py +0 -0
  30. {lm_deluge-0.0.13 → lm_deluge-0.0.14}/src/lm_deluge/api_requests/deprecated/vertex.py +0 -0
  31. {lm_deluge-0.0.13 → lm_deluge-0.0.14}/src/lm_deluge/api_requests/mistral.py +0 -0
  32. {lm_deluge-0.0.13 → lm_deluge-0.0.14}/src/lm_deluge/batches.py +0 -0
  33. {lm_deluge-0.0.13 → lm_deluge-0.0.14}/src/lm_deluge/cache.py +0 -0
  34. {lm_deluge-0.0.13 → lm_deluge-0.0.14}/src/lm_deluge/computer_use/anthropic_tools.py +0 -0
  35. {lm_deluge-0.0.13 → lm_deluge-0.0.14}/src/lm_deluge/embed.py +0 -0
  36. {lm_deluge-0.0.13 → lm_deluge-0.0.14}/src/lm_deluge/errors.py +0 -0
  37. {lm_deluge-0.0.13 → lm_deluge-0.0.14}/src/lm_deluge/gemini_limits.py +0 -0
  38. {lm_deluge-0.0.13 → lm_deluge-0.0.14}/src/lm_deluge/image.py +0 -0
  39. {lm_deluge-0.0.13 → lm_deluge-0.0.14}/src/lm_deluge/llm_tools/__init__.py +0 -0
  40. {lm_deluge-0.0.13 → lm_deluge-0.0.14}/src/lm_deluge/llm_tools/extract.py +0 -0
  41. {lm_deluge-0.0.13 → lm_deluge-0.0.14}/src/lm_deluge/llm_tools/score.py +0 -0
  42. {lm_deluge-0.0.13 → lm_deluge-0.0.14}/src/lm_deluge/llm_tools/translate.py +0 -0
  43. {lm_deluge-0.0.13 → lm_deluge-0.0.14}/src/lm_deluge/models.py +0 -0
  44. {lm_deluge-0.0.13 → lm_deluge-0.0.14}/src/lm_deluge/rerank.py +0 -0
  45. {lm_deluge-0.0.13 → lm_deluge-0.0.14}/src/lm_deluge/tool.py +0 -0
  46. {lm_deluge-0.0.13 → lm_deluge-0.0.14}/src/lm_deluge/usage.py +0 -0
  47. {lm_deluge-0.0.13 → lm_deluge-0.0.14}/src/lm_deluge/util/json.py +0 -0
  48. {lm_deluge-0.0.13 → lm_deluge-0.0.14}/src/lm_deluge/util/logprobs.py +0 -0
  49. {lm_deluge-0.0.13 → lm_deluge-0.0.14}/src/lm_deluge/util/validation.py +0 -0
  50. {lm_deluge-0.0.13 → lm_deluge-0.0.14}/src/lm_deluge/util/xml.py +0 -0
  51. {lm_deluge-0.0.13 → lm_deluge-0.0.14}/src/lm_deluge.egg-info/dependency_links.txt +0 -0
  52. {lm_deluge-0.0.13 → lm_deluge-0.0.14}/src/lm_deluge.egg-info/requires.txt +0 -0
  53. {lm_deluge-0.0.13 → lm_deluge-0.0.14}/src/lm_deluge.egg-info/top_level.txt +0 -0
  54. {lm_deluge-0.0.13 → lm_deluge-0.0.14}/tests/test_batch_real.py +0 -0
  55. {lm_deluge-0.0.13 → lm_deluge-0.0.14}/tests/test_bedrock_computer_use.py +0 -0
  56. {lm_deluge-0.0.13 → lm_deluge-0.0.14}/tests/test_bedrock_models.py +0 -0
  57. {lm_deluge-0.0.13 → lm_deluge-0.0.14}/tests/test_cache.py +0 -0
  58. {lm_deluge-0.0.13 → lm_deluge-0.0.14}/tests/test_client_tracker_integration.py +0 -0
  59. {lm_deluge-0.0.13 → lm_deluge-0.0.14}/tests/test_computer_use.py +0 -0
  60. {lm_deluge-0.0.13 → lm_deluge-0.0.14}/tests/test_computer_use_integration.py +0 -0
  61. {lm_deluge-0.0.13 → lm_deluge-0.0.14}/tests/test_debug_format.py +0 -0
  62. {lm_deluge-0.0.13 → lm_deluge-0.0.14}/tests/test_image_models.py +0 -0
  63. {lm_deluge-0.0.13 → lm_deluge-0.0.14}/tests/test_image_utils.py +0 -0
  64. {lm_deluge-0.0.13 → lm_deluge-0.0.14}/tests/test_json_utils.py +0 -0
  65. {lm_deluge-0.0.13 → lm_deluge-0.0.14}/tests/test_logprobs_refactor.py +0 -0
  66. {lm_deluge-0.0.13 → lm_deluge-0.0.14}/tests/test_max_concurrent_requests.py +0 -0
  67. {lm_deluge-0.0.13 → lm_deluge-0.0.14}/tests/test_mcp_tools.py +0 -0
  68. {lm_deluge-0.0.13 → lm_deluge-0.0.14}/tests/test_openai_responses.py +0 -0
  69. {lm_deluge-0.0.13 → lm_deluge-0.0.14}/tests/test_prompt_caching.py +0 -0
  70. {lm_deluge-0.0.13 → lm_deluge-0.0.14}/tests/test_real_caching.py +0 -0
  71. {lm_deluge-0.0.13 → lm_deluge-0.0.14}/tests/test_real_caching_bedrock.py +0 -0
  72. {lm_deluge-0.0.13 → lm_deluge-0.0.14}/tests/test_rich_display.py +0 -0
  73. {lm_deluge-0.0.13 → lm_deluge-0.0.14}/tests/test_sampling_params.py +0 -0
  74. {lm_deluge-0.0.13 → lm_deluge-0.0.14}/tests/test_tool_calls.py +0 -0
  75. {lm_deluge-0.0.13 → lm_deluge-0.0.14}/tests/test_tool_from_function.py +0 -0
  76. {lm_deluge-0.0.13 → lm_deluge-0.0.14}/tests/test_tool_validation.py +0 -0
  77. {lm_deluge-0.0.13 → lm_deluge-0.0.14}/tests/test_tracker_refactor.py +0 -0
  78. {lm_deluge-0.0.13 → lm_deluge-0.0.14}/tests/test_translate.py +0 -0
  79. {lm_deluge-0.0.13 → lm_deluge-0.0.14}/tests/test_xml_utils.py +0 -0
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.4
2
2
  Name: lm_deluge
3
- Version: 0.0.13
3
+ Version: 0.0.14
4
4
  Summary: Python utility for using LLM API models.
5
5
  Author-email: Benjamin Anderson <ben@trytaylor.ai>
6
6
  Requires-Python: >=3.10
@@ -30,6 +30,7 @@ Dynamic: license-file
30
30
  `lm-deluge` is a lightweight helper library for maxing out your rate limits with LLM providers. It provides the following:
31
31
 
32
32
  - **Unified client** – Send prompts to all relevant models with a single client.
33
+ - **Files and Images** - Include images easily for multimodal models, and PDF files for models that support them (OpenAI and Anthropic).
33
34
  - **Massive concurrency with throttling** – Set `max_tokens_per_minute` and `max_requests_per_minute` and let it fly. The client will process as many requests as possible while respecting rate limits and retrying failures.
34
35
  - **Spray across models/providers** – Configure a client with multiple models from any provider(s), and sampling weights. The client samples a model for each request.
35
36
  - **Tool Use** – Unified API for defining tools for all providers, and creating tools automatically from python functions.
@@ -41,6 +42,8 @@ Dynamic: license-file
41
42
 
42
43
  **STREAMING IS NOT IN SCOPE.** There are plenty of packages that let you stream chat completions across providers. The sole purpose of this package is to do very fast batch inference using APIs. Sorry!
43
44
 
45
+ **Update 06/02/2025:** I lied, it supports (very basic) streaming now via client.stream(...). It will print tokens as they arrive, then return an APIResponse at the end. More sophisticated streaming may or may not be implemented later, don't count on it.
46
+
44
47
  ## Installation
45
48
 
46
49
  ```bash
@@ -3,6 +3,7 @@
3
3
  `lm-deluge` is a lightweight helper library for maxing out your rate limits with LLM providers. It provides the following:
4
4
 
5
5
  - **Unified client** – Send prompts to all relevant models with a single client.
6
+ - **Files and Images** - Include images easily for multimodal models, and PDF files for models that support them (OpenAI and Anthropic).
6
7
  - **Massive concurrency with throttling** – Set `max_tokens_per_minute` and `max_requests_per_minute` and let it fly. The client will process as many requests as possible while respecting rate limits and retrying failures.
7
8
  - **Spray across models/providers** – Configure a client with multiple models from any provider(s), and sampling weights. The client samples a model for each request.
8
9
  - **Tool Use** – Unified API for defining tools for all providers, and creating tools automatically from python functions.
@@ -14,6 +15,8 @@
14
15
 
15
16
  **STREAMING IS NOT IN SCOPE.** There are plenty of packages that let you stream chat completions across providers. The sole purpose of this package is to do very fast batch inference using APIs. Sorry!
16
17
 
18
+ **Update 06/02/2025:** I lied, it supports (very basic) streaming now via client.stream(...). It will print tokens as they arrive, then return an APIResponse at the end. More sophisticated streaming may or may not be implemented later, don't count on it.
19
+
17
20
  ## Installation
18
21
 
19
22
  ```bash
@@ -3,7 +3,7 @@ requires = ["setuptools", "wheel"]
3
3
 
4
4
  [project]
5
5
  name = "lm_deluge"
6
- version = "0.0.13"
6
+ version = "0.0.14"
7
7
  authors = [{ name = "Benjamin Anderson", email = "ben@trytaylor.ai" }]
8
8
  description = "Python utility for using LLM API models."
9
9
  readme = "README.md"
@@ -1,6 +1,7 @@
1
1
  from .client import LLMClient, SamplingParams, APIResponse
2
2
  from .prompt import Conversation, Message
3
3
  from .tool import Tool
4
+ from .file import File
4
5
  import dotenv
5
6
 
6
7
  dotenv.load_dotenv()
@@ -12,4 +13,5 @@ __all__ = [
12
13
  "Conversation",
13
14
  "Message",
14
15
  "Tool",
16
+ "File",
15
17
  ]
@@ -1,165 +1,19 @@
1
1
  import asyncio
2
- import json
3
2
  import random
4
3
  import traceback
5
4
  from abc import ABC, abstractmethod
6
- from dataclasses import dataclass
7
5
  from typing import Callable
8
6
 
9
7
  import aiohttp
10
8
  from aiohttp import ClientResponse
11
9
 
12
- from lm_deluge.prompt import CachePattern, Conversation, Message
13
- from lm_deluge.usage import Usage
10
+ from lm_deluge.prompt import CachePattern, Conversation
14
11
 
15
12
  from ..config import SamplingParams
16
13
  from ..errors import raise_if_modal_exception
17
14
  from ..models import APIModel
18
15
  from ..tracker import StatusTracker
19
-
20
-
21
- @dataclass
22
- class APIResponse:
23
- # request information
24
- id: int # should be unique to the request within a given prompt-processing call
25
- model_internal: str # our internal model tag
26
- prompt: Conversation
27
- sampling_params: SamplingParams
28
-
29
- # http response information
30
- status_code: int | None
31
- is_error: bool | None
32
- error_message: str | None
33
-
34
- # completion information - unified usage tracking
35
- usage: Usage | None = None
36
-
37
- # response content - structured format
38
- content: Message | None = None
39
-
40
- # optional or calculated automatically
41
- thinking: str | None = None # if model shows thinking tokens
42
- model_external: str | None = None # the model tag used by the API
43
- region: str | None = None
44
- logprobs: list | None = None
45
- finish_reason: str | None = None # make required later
46
- cost: float | None = None # calculated automatically
47
- cache_hit: bool = False # manually set if true
48
- # set to true if is_error and should be retried with a different model
49
- retry_with_different_model: bool | None = False
50
- # set to true if should NOT retry with the same model (unrecoverable error)
51
- give_up_if_no_other_models: bool | None = False
52
- # OpenAI Responses API specific - used for computer use continuation
53
- response_id: str | None = None
54
- # Raw API response for debugging
55
- raw_response: dict | None = None
56
-
57
- @property
58
- def completion(self) -> str | None:
59
- """Backward compatibility: extract text from content Message."""
60
- if self.content is not None:
61
- return self.content.completion
62
- return None
63
-
64
- @property
65
- def input_tokens(self) -> int | None:
66
- """Get input tokens from usage object."""
67
- return self.usage.input_tokens if self.usage else None
68
-
69
- @property
70
- def output_tokens(self) -> int | None:
71
- """Get output tokens from usage object."""
72
- return self.usage.output_tokens if self.usage else None
73
-
74
- @property
75
- def cache_read_tokens(self) -> int | None:
76
- """Get cache read tokens from usage object."""
77
- return self.usage.cache_read_tokens if self.usage else None
78
-
79
- @property
80
- def cache_write_tokens(self) -> int | None:
81
- """Get cache write tokens from usage object."""
82
- return self.usage.cache_write_tokens if self.usage else None
83
-
84
- def __post_init__(self):
85
- # calculate cost & get external model name
86
- self.id = int(self.id)
87
- api_model = APIModel.from_registry(self.model_internal)
88
- self.model_external = api_model.name
89
- self.cost = None
90
- if (
91
- self.usage is not None
92
- and api_model.input_cost is not None
93
- and api_model.output_cost is not None
94
- ):
95
- self.cost = (
96
- self.usage.input_tokens * api_model.input_cost / 1e6
97
- + self.usage.output_tokens * api_model.output_cost / 1e6
98
- )
99
- elif self.content is not None and self.completion is not None:
100
- print(
101
- f"Warning: Completion provided without token counts for model {self.model_internal}."
102
- )
103
-
104
- def to_dict(self):
105
- return {
106
- "id": self.id,
107
- "model_internal": self.model_internal,
108
- "model_external": self.model_external,
109
- "region": self.region,
110
- "prompt": self.prompt.to_log(), # destroys image if present
111
- "sampling_params": self.sampling_params.__dict__,
112
- "status_code": self.status_code,
113
- "is_error": self.is_error,
114
- "error_message": self.error_message,
115
- "completion": self.completion, # computed property
116
- "content": self.content.to_log() if self.content else None,
117
- "usage": self.usage.to_dict() if self.usage else None,
118
- "finish_reason": self.finish_reason,
119
- "cost": self.cost,
120
- }
121
-
122
- @classmethod
123
- def from_dict(cls, data: dict):
124
- # Handle backward compatibility for content/completion
125
- content = None
126
- if "content" in data and data["content"] is not None:
127
- # Reconstruct message from log format
128
- content = Message.from_log(data["content"])
129
- elif "completion" in data and data["completion"] is not None:
130
- # Backward compatibility: create a Message with just text
131
- content = Message.ai(data["completion"])
132
-
133
- usage = None
134
- if "usage" in data and data["usage"] is not None:
135
- usage = Usage.from_dict(data["usage"])
136
-
137
- return cls(
138
- id=data.get("id", random.randint(0, 1_000_000_000)),
139
- model_internal=data["model_internal"],
140
- prompt=Conversation.from_log(data["prompt"]),
141
- sampling_params=SamplingParams(**data["sampling_params"]),
142
- status_code=data["status_code"],
143
- is_error=data["is_error"],
144
- error_message=data["error_message"],
145
- usage=usage,
146
- content=content,
147
- thinking=data.get("thinking"),
148
- model_external=data.get("model_external"),
149
- region=data.get("region"),
150
- logprobs=data.get("logprobs"),
151
- finish_reason=data.get("finish_reason"),
152
- cost=data.get("cost"),
153
- cache_hit=data.get("cache_hit", False),
154
- )
155
-
156
- def write_to_file(self, filename):
157
- """
158
- Writes the APIResponse as a line to a file.
159
- If file exists, appends to it.
160
- """
161
- with open(filename, "a") as f:
162
- f.write(json.dumps(self.to_dict()) + "\n")
16
+ from .response import APIResponse
163
17
 
164
18
 
165
19
  class APIRequestBase(ABC):
@@ -1,17 +1,19 @@
1
- import warnings
2
- from aiohttp import ClientResponse
3
1
  import json
4
2
  import os
3
+ import warnings
5
4
  from typing import Callable
6
5
 
6
+ import aiohttp
7
+ from aiohttp import ClientResponse
8
+
7
9
  from lm_deluge.tool import Tool
8
10
 
9
- from .base import APIRequestBase, APIResponse
10
- from ..prompt import Conversation, Message, Text, ToolCall, Thinking, CachePattern
11
- from ..usage import Usage
12
- from ..tracker import StatusTracker
13
11
  from ..config import SamplingParams
14
12
  from ..models import APIModel
13
+ from ..prompt import CachePattern, Conversation, Message, Text, Thinking, ToolCall
14
+ from ..tracker import StatusTracker
15
+ from ..usage import Usage
16
+ from .base import APIRequestBase, APIResponse
15
17
 
16
18
 
17
19
  def _build_oa_chat_request(
@@ -111,6 +113,7 @@ class OpenAIRequest(APIRequestBase):
111
113
  status_code = http_response.status
112
114
  mimetype = http_response.headers.get("Content-Type", None)
113
115
  data = None
116
+ finish_reason = None
114
117
  if status_code >= 200 and status_code < 300:
115
118
  try:
116
119
  data = await http_response.json()
@@ -125,6 +128,7 @@ class OpenAIRequest(APIRequestBase):
125
128
  # Parse response into Message with parts
126
129
  parts = []
127
130
  message = data["choices"][0]["message"]
131
+ finish_reason = data["choices"][0]["finish_reason"]
128
132
 
129
133
  # Add text content if present
130
134
  if message.get("content"):
@@ -190,6 +194,7 @@ class OpenAIRequest(APIRequestBase):
190
194
  sampling_params=self.sampling_params,
191
195
  usage=usage,
192
196
  raw_response=data,
197
+ finish_reason=finish_reason,
193
198
  )
194
199
 
195
200
 
@@ -266,6 +271,13 @@ class OpenAIResponsesRequest(APIRequestBase):
266
271
  self.request_json["max_output_tokens"] = sampling_params.max_new_tokens
267
272
 
268
273
  if self.model.reasoning_model:
274
+ if sampling_params.reasoning_effort in [None, "none"]:
275
+ # gemini models can switch reasoning off
276
+ if "gemini" in self.model.id:
277
+ self.sampling_params.reasoning_effort = "none" # expects string
278
+ # openai models can only go down to "low"
279
+ else:
280
+ self.sampling_params.reasoning_effort = "low"
269
281
  self.request_json["temperature"] = 1.0
270
282
  self.request_json["top_p"] = 1.0
271
283
  self.request_json["reasoning"] = {
@@ -413,3 +425,57 @@ class OpenAIResponsesRequest(APIRequestBase):
413
425
  usage=usage,
414
426
  raw_response=data,
415
427
  )
428
+
429
+
430
+ async def stream_chat(
431
+ model_name: str, # must correspond to registry
432
+ prompt: Conversation,
433
+ sampling_params: SamplingParams = SamplingParams(),
434
+ tools: list | None = None,
435
+ cache: CachePattern | None = None,
436
+ ):
437
+ if cache is not None:
438
+ warnings.warn(
439
+ f"Cache parameter '{cache}' is only supported for Anthropic models, ignoring for {model_name}"
440
+ )
441
+
442
+ model = APIModel.from_registry(model_name)
443
+ if model.api_spec != "openai":
444
+ raise ValueError("streaming only supported on openai models for now")
445
+ url = f"{model.api_base}/chat/completions"
446
+ request_header = {"Authorization": f"Bearer {os.getenv(model.api_key_env_var)}"}
447
+ request_json = _build_oa_chat_request(model, prompt, tools, sampling_params)
448
+ request_json["stream"] = True
449
+
450
+ async with aiohttp.ClientSession() as s:
451
+ async with s.post(url, headers=request_header, json=request_json) as r:
452
+ r.raise_for_status() # bail on 4xx/5xx
453
+ content = ""
454
+ buf = ""
455
+ async for chunk in r.content.iter_any(): # raw bytes
456
+ buf += chunk.decode()
457
+ while "\n\n" in buf: # full SSE frame
458
+ event, buf = buf.split("\n\n", 1)
459
+ if not event.startswith("data:"):
460
+ continue # ignore comments
461
+ data = event[5:].strip() # after "data:"
462
+ if data == "[DONE]":
463
+ yield APIResponse(
464
+ id=0,
465
+ status_code=None,
466
+ is_error=False,
467
+ error_message=None,
468
+ prompt=prompt,
469
+ content=Message(
470
+ role="assistant", parts=[Text(text=content)]
471
+ ),
472
+ model_internal=model.id,
473
+ sampling_params=sampling_params,
474
+ usage=None,
475
+ raw_response=None,
476
+ )
477
+ msg = json.loads(data) # SSE payload
478
+ delta = msg["choices"][0]["delta"].get("content")
479
+ if delta:
480
+ content += delta
481
+ yield delta
@@ -0,0 +1,153 @@
1
+ import json
2
+ import random
3
+ from dataclasses import dataclass
4
+
5
+ from lm_deluge.prompt import Conversation, Message
6
+ from lm_deluge.usage import Usage
7
+
8
+ from ..config import SamplingParams
9
+ from ..models import APIModel
10
+
11
+
12
+ @dataclass
13
+ class APIResponse:
14
+ # request information
15
+ id: int # should be unique to the request within a given prompt-processing call
16
+ model_internal: str # our internal model tag
17
+ prompt: Conversation
18
+ sampling_params: SamplingParams
19
+
20
+ # http response information
21
+ status_code: int | None
22
+ is_error: bool | None
23
+ error_message: str | None
24
+
25
+ # completion information - unified usage tracking
26
+ usage: Usage | None = None
27
+
28
+ # response content - structured format
29
+ content: Message | None = None
30
+
31
+ # optional or calculated automatically
32
+ thinking: str | None = None # if model shows thinking tokens
33
+ model_external: str | None = None # the model tag used by the API
34
+ region: str | None = None
35
+ logprobs: list | None = None
36
+ finish_reason: str | None = None # make required later
37
+ cost: float | None = None # calculated automatically
38
+ cache_hit: bool = False # manually set if true
39
+ # set to true if is_error and should be retried with a different model
40
+ retry_with_different_model: bool | None = False
41
+ # set to true if should NOT retry with the same model (unrecoverable error)
42
+ give_up_if_no_other_models: bool | None = False
43
+ # OpenAI Responses API specific - used for computer use continuation
44
+ response_id: str | None = None
45
+ # Raw API response for debugging
46
+ raw_response: dict | None = None
47
+
48
+ @property
49
+ def completion(self) -> str | None:
50
+ """Backward compatibility: extract text from content Message."""
51
+ if self.content is not None:
52
+ return self.content.completion
53
+ return None
54
+
55
+ @property
56
+ def input_tokens(self) -> int | None:
57
+ """Get input tokens from usage object."""
58
+ return self.usage.input_tokens if self.usage else None
59
+
60
+ @property
61
+ def output_tokens(self) -> int | None:
62
+ """Get output tokens from usage object."""
63
+ return self.usage.output_tokens if self.usage else None
64
+
65
+ @property
66
+ def cache_read_tokens(self) -> int | None:
67
+ """Get cache read tokens from usage object."""
68
+ return self.usage.cache_read_tokens if self.usage else None
69
+
70
+ @property
71
+ def cache_write_tokens(self) -> int | None:
72
+ """Get cache write tokens from usage object."""
73
+ return self.usage.cache_write_tokens if self.usage else None
74
+
75
+ def __post_init__(self):
76
+ # calculate cost & get external model name
77
+ self.id = int(self.id)
78
+ api_model = APIModel.from_registry(self.model_internal)
79
+ self.model_external = api_model.name
80
+ self.cost = None
81
+ if (
82
+ self.usage is not None
83
+ and api_model.input_cost is not None
84
+ and api_model.output_cost is not None
85
+ ):
86
+ self.cost = (
87
+ self.usage.input_tokens * api_model.input_cost / 1e6
88
+ + self.usage.output_tokens * api_model.output_cost / 1e6
89
+ )
90
+ elif self.content is not None and self.completion is not None:
91
+ print(
92
+ f"Warning: Completion provided without token counts for model {self.model_internal}."
93
+ )
94
+
95
+ def to_dict(self):
96
+ return {
97
+ "id": self.id,
98
+ "model_internal": self.model_internal,
99
+ "model_external": self.model_external,
100
+ "region": self.region,
101
+ "prompt": self.prompt.to_log(), # destroys image if present
102
+ "sampling_params": self.sampling_params.__dict__,
103
+ "status_code": self.status_code,
104
+ "is_error": self.is_error,
105
+ "error_message": self.error_message,
106
+ "completion": self.completion, # computed property
107
+ "content": self.content.to_log() if self.content else None,
108
+ "usage": self.usage.to_dict() if self.usage else None,
109
+ "finish_reason": self.finish_reason,
110
+ "cost": self.cost,
111
+ }
112
+
113
+ @classmethod
114
+ def from_dict(cls, data: dict):
115
+ # Handle backward compatibility for content/completion
116
+ content = None
117
+ if "content" in data and data["content"] is not None:
118
+ # Reconstruct message from log format
119
+ content = Message.from_log(data["content"])
120
+ elif "completion" in data and data["completion"] is not None:
121
+ # Backward compatibility: create a Message with just text
122
+ content = Message.ai(data["completion"])
123
+
124
+ usage = None
125
+ if "usage" in data and data["usage"] is not None:
126
+ usage = Usage.from_dict(data["usage"])
127
+
128
+ return cls(
129
+ id=data.get("id", random.randint(0, 1_000_000_000)),
130
+ model_internal=data["model_internal"],
131
+ prompt=Conversation.from_log(data["prompt"]),
132
+ sampling_params=SamplingParams(**data["sampling_params"]),
133
+ status_code=data["status_code"],
134
+ is_error=data["is_error"],
135
+ error_message=data["error_message"],
136
+ usage=usage,
137
+ content=content,
138
+ thinking=data.get("thinking"),
139
+ model_external=data.get("model_external"),
140
+ region=data.get("region"),
141
+ logprobs=data.get("logprobs"),
142
+ finish_reason=data.get("finish_reason"),
143
+ cost=data.get("cost"),
144
+ cache_hit=data.get("cache_hit", False),
145
+ )
146
+
147
+ def write_to_file(self, filename):
148
+ """
149
+ Writes the APIResponse as a line to a file.
150
+ If file exists, appends to it.
151
+ """
152
+ with open(filename, "a") as f:
153
+ f.write(json.dumps(self.to_dict()) + "\n")
@@ -6,6 +6,7 @@ import yaml
6
6
  from pydantic import BaseModel
7
7
  from pydantic.functional_validators import model_validator
8
8
 
9
+ from lm_deluge.api_requests.openai import stream_chat
9
10
  from lm_deluge.batches import (
10
11
  submit_batches_anthropic,
11
12
  submit_batches_oa,
@@ -34,6 +35,12 @@ class LLMClient(BaseModel):
34
35
  """
35
36
 
36
37
  model_names: list[str] = ["gpt-4.1-mini"]
38
+
39
+ def __init__(self, model_name: str | list[str] | None = None, **kwargs):
40
+ if model_name is not None:
41
+ kwargs["model_names"] = model_name
42
+ super().__init__(**kwargs)
43
+
37
44
  max_requests_per_minute: int = 1_000
38
45
  max_tokens_per_minute: int = 100_000
39
46
  max_concurrent_requests: int = 225
@@ -81,7 +88,7 @@ class LLMClient(BaseModel):
81
88
  @model_validator(mode="before")
82
89
  @classmethod
83
90
  def fix_lists(cls, data) -> "LLMClient":
84
- if isinstance(data["model_names"], str):
91
+ if isinstance(data.get("model_names"), str):
85
92
  data["model_names"] = [data["model_names"]]
86
93
  if "sampling_params" not in data or len(data.get("sampling_params", [])) == 0:
87
94
  data["sampling_params"] = [
@@ -162,6 +169,11 @@ class LLMClient(BaseModel):
162
169
  kwargs["model_names"] = model
163
170
  return cls(**kwargs)
164
171
 
172
+ def _select_model(self):
173
+ assert isinstance(self.model_weights, list)
174
+ model_idx = np.random.choice(range(len(self.models)), p=self.model_weights)
175
+ return self.models[model_idx], self.sampling_params[model_idx]
176
+
165
177
  @overload
166
178
  async def process_prompts_async(
167
179
  self,
@@ -249,41 +261,6 @@ class LLMClient(BaseModel):
249
261
  if len(cache_hit_ids) > 0:
250
262
  tracker.update_pbar(len(cache_hit_ids))
251
263
 
252
- # api_task = asyncio.create_task(
253
- # process_api_prompts_async(
254
- # ids,
255
- # prompts, # type: ignore -- fix later for dry running conversations
256
- # self.models,
257
- # self.model_weights, # type: ignore
258
- # self.sampling_params, # type: ignore
259
- # max_attempts=self.max_attempts,
260
- # max_concurrent_requests=self.max_concurrent_requests,
261
- # request_timeout=self.request_timeout,
262
- # status_tracker=tracker,
263
- # tools=tools,
264
- # cache=cache,
265
- # computer_use=computer_use,
266
- # display_width=display_width,
267
- # display_height=display_height,
268
- # use_responses_api=use_responses_api,
269
- # )
270
- # )
271
- # async def process_api_prompts_async(
272
-
273
- # models: str | list[str],
274
- # model_weights: list[float],
275
- # sampling_params: list[SamplingParams],
276
- # max_attempts: int = 5,
277
- # max_concurrent_requests: int = 1_000,
278
- # request_timeout: int = 30,
279
- # status_tracker: StatusTracker | None = None,
280
- # tools: list[Tool] | None = None,
281
- # cache: CachePattern | None = None,
282
- # computer_use: bool = False,
283
- # display_width: int = 1024,
284
- # display_height: int = 768,
285
- # use_responses_api: bool = False,
286
- # ):
287
264
  if isinstance(ids, np.ndarray):
288
265
  ids = ids.tolist() # pyright: ignore
289
266
 
@@ -296,28 +273,28 @@ class LLMClient(BaseModel):
296
273
  assert tracker.retry_queue, "retry queue not initialized"
297
274
  while True:
298
275
  # get next request (if one is not already waiting for capacity)
276
+ retry_request = False
299
277
  if next_request is None:
300
278
  if not tracker.retry_queue.empty():
301
279
  next_request = tracker.retry_queue.get_nowait()
280
+ retry_request = True
302
281
  print(f"Retrying request {next_request.task_id}.")
303
282
  elif prompts_not_finished:
304
283
  try:
305
284
  # get new request
306
285
  id, prompt = next(prompts_iter)
307
286
  # select model
308
- assert isinstance(self.model_weights, list)
309
- model_idx = np.random.choice(
310
- range(len(self.models)), p=self.model_weights
311
- )
287
+ model, sampling_params = self._select_model()
288
+
312
289
  next_request = create_api_request(
313
290
  task_id=id,
314
- model_name=self.models[model_idx],
291
+ model_name=model,
315
292
  prompt=prompt, # type: ignore
316
293
  request_timeout=self.request_timeout,
317
294
  attempts_left=self.max_attempts,
318
295
  status_tracker=tracker,
319
296
  results_arr=requests,
320
- sampling_params=self.sampling_params[model_idx],
297
+ sampling_params=sampling_params,
321
298
  all_model_names=self.models,
322
299
  all_sampling_params=self.sampling_params,
323
300
  tools=tools,
@@ -339,10 +316,9 @@ class LLMClient(BaseModel):
339
316
  # if enough capacity available, call API
340
317
  if next_request:
341
318
  next_request_tokens = next_request.num_tokens
342
- if tracker.check_capacity(next_request_tokens):
319
+ if tracker.check_capacity(next_request_tokens, retry=retry_request):
343
320
  tracker.set_limiting_factor(None)
344
- next_request.attempts_left -= 1
345
- # call API
321
+ # call API (attempts_left will be decremented in handle_error if it fails)
346
322
  asyncio.create_task(next_request.call_api())
347
323
  next_request = None # reset next_request to empty
348
324
  # update pbar status
@@ -360,9 +336,10 @@ class LLMClient(BaseModel):
360
336
  await asyncio.sleep(tracker.seconds_to_pause)
361
337
  print(f"Pausing {tracker.seconds_to_pause}s to cool down.")
362
338
 
363
- # after finishing, log final status
364
- tracker.log_final_status()
365
- # deduplicate results by id
339
+ # after finishing, log final status
340
+ tracker.log_final_status()
341
+
342
+ # deduplicate results by id
366
343
  api_results = deduplicate_responses(requests)
367
344
  for res in api_results:
368
345
  results[res.id] = res
@@ -399,6 +376,17 @@ class LLMClient(BaseModel):
399
376
  )
400
377
  )
401
378
 
379
+ async def stream(self, prompt: str | Conversation, tools: list[Tool] | None = None):
380
+ model, sampling_params = self._select_model()
381
+ if isinstance(prompt, str):
382
+ prompt = Conversation.user(prompt)
383
+ async for item in stream_chat(model, prompt, sampling_params, tools, None):
384
+ if isinstance(item, str):
385
+ print(item, end="", flush=True)
386
+ else:
387
+ # final item
388
+ return item
389
+
402
390
  async def submit_batch_job(
403
391
  self,
404
392
  prompts: Sequence[str | list[dict] | Conversation],