llm-batch-runner 0.1.0__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.
- llm_batch_runner-0.1.0/PKG-INFO +121 -0
- llm_batch_runner-0.1.0/README.md +105 -0
- llm_batch_runner-0.1.0/llm_batch_runner/__init__.py +5 -0
- llm_batch_runner-0.1.0/llm_batch_runner/processor.py +378 -0
- llm_batch_runner-0.1.0/llm_batch_runner.egg-info/PKG-INFO +121 -0
- llm_batch_runner-0.1.0/llm_batch_runner.egg-info/SOURCES.txt +9 -0
- llm_batch_runner-0.1.0/llm_batch_runner.egg-info/dependency_links.txt +1 -0
- llm_batch_runner-0.1.0/llm_batch_runner.egg-info/requires.txt +2 -0
- llm_batch_runner-0.1.0/llm_batch_runner.egg-info/top_level.txt +1 -0
- llm_batch_runner-0.1.0/pyproject.toml +30 -0
- llm_batch_runner-0.1.0/setup.cfg +4 -0
|
@@ -0,0 +1,121 @@
|
|
|
1
|
+
Metadata-Version: 2.4
|
|
2
|
+
Name: llm-batch-runner
|
|
3
|
+
Version: 0.1.0
|
|
4
|
+
Summary: Concurrent, caching, retrying OpenAI-compatible batch chat-completion runner with structured output validation.
|
|
5
|
+
Author-email: Your Name <you@example.com>
|
|
6
|
+
License: MIT
|
|
7
|
+
Project-URL: Homepage, https://github.com/yourusername/llm-batch-runner
|
|
8
|
+
Keywords: openai,llm,batch,structured-output,pydantic,concurrency,cache
|
|
9
|
+
Classifier: Programming Language :: Python :: 3
|
|
10
|
+
Classifier: License :: OSI Approved :: MIT License
|
|
11
|
+
Classifier: Operating System :: OS Independent
|
|
12
|
+
Requires-Python: >=3.9
|
|
13
|
+
Description-Content-Type: text/markdown
|
|
14
|
+
Requires-Dist: openai>=1.40.0
|
|
15
|
+
Requires-Dist: pydantic>=2.0.0
|
|
16
|
+
|
|
17
|
+
# llm-batch-runner
|
|
18
|
+
|
|
19
|
+
Concurrent, caching, retrying batch runner for OpenAI-compatible chat
|
|
20
|
+
completion endpoints, with optional structured-output validation against
|
|
21
|
+
a Pydantic model.
|
|
22
|
+
|
|
23
|
+
## Features
|
|
24
|
+
|
|
25
|
+
- Fires requests concurrently (configurable `max_workers`) through the
|
|
26
|
+
official `openai` SDK, so it works with OpenAI itself or any
|
|
27
|
+
OpenAI-compatible `base_url` (vLLM, Together, Groq, local servers, ...).
|
|
28
|
+
- If you pass a Pydantic `base_model`, requests are made with a
|
|
29
|
+
JSON-schema `response_format` and every response is validated against
|
|
30
|
+
the model. Invalid output is retried, up to `n_retries` times.
|
|
31
|
+
- If no `base_model` is passed, it just does a normal completion and
|
|
32
|
+
returns the raw text.
|
|
33
|
+
- Every sample is cached to disk as its own JSON file, keyed by a SHA-256
|
|
34
|
+
hash of its input `messages`. On a re-run, already-cached samples are
|
|
35
|
+
skipped unless `overwrite_existing=True`.
|
|
36
|
+
- Failures are logged (via the standard `logging` module) and returned
|
|
37
|
+
in the result list rather than raised, so one bad sample never kills
|
|
38
|
+
the whole batch.
|
|
39
|
+
- `LLMBatchRunner.load_cache(cache_dir)` is a static method that loads
|
|
40
|
+
every cached record back into a dict, keyed by hash.
|
|
41
|
+
|
|
42
|
+
|
|
43
|
+
## Usage
|
|
44
|
+
|
|
45
|
+
```python
|
|
46
|
+
from pydantic import BaseModel
|
|
47
|
+
from llm_batch_runner import LLMBatchRunner
|
|
48
|
+
|
|
49
|
+
class Answer(BaseModel):
|
|
50
|
+
reasoning: str
|
|
51
|
+
value: int
|
|
52
|
+
|
|
53
|
+
conversations = [
|
|
54
|
+
[{"role": "user", "content": "What is 12 * 7?"}],
|
|
55
|
+
[{"role": "user", "content": "What is 9 * 9?"}],
|
|
56
|
+
# ... as many as you like
|
|
57
|
+
]
|
|
58
|
+
|
|
59
|
+
runner = LLMBatchRunner(
|
|
60
|
+
base_model=Answer, # or None for plain text completions
|
|
61
|
+
messages=conversations, # List[List[dict]] — one conversation per sample
|
|
62
|
+
max_workers=16,
|
|
63
|
+
client_config={"api_key": "wow_very secret", "base_url": None},
|
|
64
|
+
model_name="gpt-4o-mini",
|
|
65
|
+
cache_dir="./llm_cache",
|
|
66
|
+
n_retries=3,
|
|
67
|
+
overwrite_existing=False,
|
|
68
|
+
conversations_meta=None # or list of dicts: additional values to store
|
|
69
|
+
)
|
|
70
|
+
|
|
71
|
+
results = runner.run()
|
|
72
|
+
print(results[0])
|
|
73
|
+
# OUTPUT:
|
|
74
|
+
```
|
|
75
|
+
|
|
76
|
+
Output:
|
|
77
|
+
|
|
78
|
+
```json
|
|
79
|
+
{
|
|
80
|
+
"attempts": 1,
|
|
81
|
+
"conversation": [
|
|
82
|
+
{
|
|
83
|
+
"role": "user",
|
|
84
|
+
"content": "What is 12 * 7?"
|
|
85
|
+
}
|
|
86
|
+
],
|
|
87
|
+
"error": "None",
|
|
88
|
+
"from_cache": "False",
|
|
89
|
+
"hash": "f699170cc22a5e871fccdf3a414f35606d100c943d18387e58d75c0a1006b6e6",
|
|
90
|
+
"index": 0,
|
|
91
|
+
"output": {
|
|
92
|
+
"reasoning": "12 * 7 can be calculated as (10 * 7) + (2 * 7) = 70 + 14 = 84.",
|
|
93
|
+
"value": 84
|
|
94
|
+
},
|
|
95
|
+
"raw_output": "{\"reasoning\": \"12 * 7 can be calculated as (10 * 7) + (2 * 7) = 70 + 14 = 84.\", \"value\": 84}",
|
|
96
|
+
"success": "True",
|
|
97
|
+
"meta": null
|
|
98
|
+
}
|
|
99
|
+
```
|
|
100
|
+
|
|
101
|
+
### Loading cached results later
|
|
102
|
+
|
|
103
|
+
```python
|
|
104
|
+
from llm_batch_runner import LLMBatchRunner
|
|
105
|
+
|
|
106
|
+
cache = LLMBatchRunner.load_cache("./llm_cache")
|
|
107
|
+
# List of dicts - loaded data from the cache
|
|
108
|
+
# [
|
|
109
|
+
# {"success": True, "output": {...}, "error": None, ...},
|
|
110
|
+
# ...
|
|
111
|
+
# ]
|
|
112
|
+
```
|
|
113
|
+
|
|
114
|
+
## Notes
|
|
115
|
+
|
|
116
|
+
- `messages` is a list of conversations — each conversation is itself the
|
|
117
|
+
standard OpenAI `messages` list (`[{"role": ..., "content": ...}, ...]`).
|
|
118
|
+
- The cache key is a hash of the conversation only (not the model name),
|
|
119
|
+
so if you change `model_name` and want fresh results, pass a different
|
|
120
|
+
`cache_dir` or use `overwrite_existing=True`.
|
|
121
|
+
- Failed samples (that exhausted all retries) are also cached.
|
|
@@ -0,0 +1,105 @@
|
|
|
1
|
+
# llm-batch-runner
|
|
2
|
+
|
|
3
|
+
Concurrent, caching, retrying batch runner for OpenAI-compatible chat
|
|
4
|
+
completion endpoints, with optional structured-output validation against
|
|
5
|
+
a Pydantic model.
|
|
6
|
+
|
|
7
|
+
## Features
|
|
8
|
+
|
|
9
|
+
- Fires requests concurrently (configurable `max_workers`) through the
|
|
10
|
+
official `openai` SDK, so it works with OpenAI itself or any
|
|
11
|
+
OpenAI-compatible `base_url` (vLLM, Together, Groq, local servers, ...).
|
|
12
|
+
- If you pass a Pydantic `base_model`, requests are made with a
|
|
13
|
+
JSON-schema `response_format` and every response is validated against
|
|
14
|
+
the model. Invalid output is retried, up to `n_retries` times.
|
|
15
|
+
- If no `base_model` is passed, it just does a normal completion and
|
|
16
|
+
returns the raw text.
|
|
17
|
+
- Every sample is cached to disk as its own JSON file, keyed by a SHA-256
|
|
18
|
+
hash of its input `messages`. On a re-run, already-cached samples are
|
|
19
|
+
skipped unless `overwrite_existing=True`.
|
|
20
|
+
- Failures are logged (via the standard `logging` module) and returned
|
|
21
|
+
in the result list rather than raised, so one bad sample never kills
|
|
22
|
+
the whole batch.
|
|
23
|
+
- `LLMBatchRunner.load_cache(cache_dir)` is a static method that loads
|
|
24
|
+
every cached record back into a dict, keyed by hash.
|
|
25
|
+
|
|
26
|
+
|
|
27
|
+
## Usage
|
|
28
|
+
|
|
29
|
+
```python
|
|
30
|
+
from pydantic import BaseModel
|
|
31
|
+
from llm_batch_runner import LLMBatchRunner
|
|
32
|
+
|
|
33
|
+
class Answer(BaseModel):
|
|
34
|
+
reasoning: str
|
|
35
|
+
value: int
|
|
36
|
+
|
|
37
|
+
conversations = [
|
|
38
|
+
[{"role": "user", "content": "What is 12 * 7?"}],
|
|
39
|
+
[{"role": "user", "content": "What is 9 * 9?"}],
|
|
40
|
+
# ... as many as you like
|
|
41
|
+
]
|
|
42
|
+
|
|
43
|
+
runner = LLMBatchRunner(
|
|
44
|
+
base_model=Answer, # or None for plain text completions
|
|
45
|
+
messages=conversations, # List[List[dict]] — one conversation per sample
|
|
46
|
+
max_workers=16,
|
|
47
|
+
client_config={"api_key": "wow_very secret", "base_url": None},
|
|
48
|
+
model_name="gpt-4o-mini",
|
|
49
|
+
cache_dir="./llm_cache",
|
|
50
|
+
n_retries=3,
|
|
51
|
+
overwrite_existing=False,
|
|
52
|
+
conversations_meta=None # or list of dicts: additional values to store
|
|
53
|
+
)
|
|
54
|
+
|
|
55
|
+
results = runner.run()
|
|
56
|
+
print(results[0])
|
|
57
|
+
# OUTPUT:
|
|
58
|
+
```
|
|
59
|
+
|
|
60
|
+
Output:
|
|
61
|
+
|
|
62
|
+
```json
|
|
63
|
+
{
|
|
64
|
+
"attempts": 1,
|
|
65
|
+
"conversation": [
|
|
66
|
+
{
|
|
67
|
+
"role": "user",
|
|
68
|
+
"content": "What is 12 * 7?"
|
|
69
|
+
}
|
|
70
|
+
],
|
|
71
|
+
"error": "None",
|
|
72
|
+
"from_cache": "False",
|
|
73
|
+
"hash": "f699170cc22a5e871fccdf3a414f35606d100c943d18387e58d75c0a1006b6e6",
|
|
74
|
+
"index": 0,
|
|
75
|
+
"output": {
|
|
76
|
+
"reasoning": "12 * 7 can be calculated as (10 * 7) + (2 * 7) = 70 + 14 = 84.",
|
|
77
|
+
"value": 84
|
|
78
|
+
},
|
|
79
|
+
"raw_output": "{\"reasoning\": \"12 * 7 can be calculated as (10 * 7) + (2 * 7) = 70 + 14 = 84.\", \"value\": 84}",
|
|
80
|
+
"success": "True",
|
|
81
|
+
"meta": null
|
|
82
|
+
}
|
|
83
|
+
```
|
|
84
|
+
|
|
85
|
+
### Loading cached results later
|
|
86
|
+
|
|
87
|
+
```python
|
|
88
|
+
from llm_batch_runner import LLMBatchRunner
|
|
89
|
+
|
|
90
|
+
cache = LLMBatchRunner.load_cache("./llm_cache")
|
|
91
|
+
# List of dicts - loaded data from the cache
|
|
92
|
+
# [
|
|
93
|
+
# {"success": True, "output": {...}, "error": None, ...},
|
|
94
|
+
# ...
|
|
95
|
+
# ]
|
|
96
|
+
```
|
|
97
|
+
|
|
98
|
+
## Notes
|
|
99
|
+
|
|
100
|
+
- `messages` is a list of conversations — each conversation is itself the
|
|
101
|
+
standard OpenAI `messages` list (`[{"role": ..., "content": ...}, ...]`).
|
|
102
|
+
- The cache key is a hash of the conversation only (not the model name),
|
|
103
|
+
so if you change `model_name` and want fresh results, pass a different
|
|
104
|
+
`cache_dir` or use `overwrite_existing=True`.
|
|
105
|
+
- Failed samples (that exhausted all retries) are also cached.
|
|
@@ -0,0 +1,378 @@
|
|
|
1
|
+
"""
|
|
2
|
+
Core implementation of LLMBatchRunner: a concurrent, caching, retrying
|
|
3
|
+
wrapper around the OpenAI-compatible chat completions API, with optional
|
|
4
|
+
structured-output validation against a Pydantic model.
|
|
5
|
+
"""
|
|
6
|
+
|
|
7
|
+
from __future__ import annotations
|
|
8
|
+
|
|
9
|
+
import hashlib
|
|
10
|
+
import json
|
|
11
|
+
import logging
|
|
12
|
+
import threading
|
|
13
|
+
from concurrent.futures import ThreadPoolExecutor, as_completed
|
|
14
|
+
from dataclasses import dataclass, asdict
|
|
15
|
+
from pathlib import Path
|
|
16
|
+
from typing import Any, Dict, List, Optional, Type, Union, Tuple
|
|
17
|
+
|
|
18
|
+
from openai import OpenAI
|
|
19
|
+
from pydantic import BaseModel, ValidationError
|
|
20
|
+
|
|
21
|
+
logger = logging.getLogger("llm_batch_runner")
|
|
22
|
+
if not logger.handlers:
|
|
23
|
+
# Give the library a sane default handler so errors are visible
|
|
24
|
+
# out of the box, without clobbering a user's own logging config.
|
|
25
|
+
_handler = logging.StreamHandler()
|
|
26
|
+
_handler.setFormatter(
|
|
27
|
+
logging.Formatter("[%(levelname)s] %(name)s: %(message)s")
|
|
28
|
+
)
|
|
29
|
+
logger.addHandler(_handler)
|
|
30
|
+
logger.setLevel(logging.INFO)
|
|
31
|
+
|
|
32
|
+
Message = Dict[str, Any]
|
|
33
|
+
Conversation = List[Message]
|
|
34
|
+
|
|
35
|
+
|
|
36
|
+
@dataclass
|
|
37
|
+
class SampleResult:
|
|
38
|
+
"""Result of processing a single conversation."""
|
|
39
|
+
|
|
40
|
+
index: int
|
|
41
|
+
hash: str
|
|
42
|
+
success: bool
|
|
43
|
+
output: Optional[Union[dict, str]] = None
|
|
44
|
+
raw_output: str = None
|
|
45
|
+
error: Optional[str] = None
|
|
46
|
+
from_cache: bool = False
|
|
47
|
+
attempts: int = 0
|
|
48
|
+
conversation: Conversation = None
|
|
49
|
+
meta: List[Any] = None
|
|
50
|
+
|
|
51
|
+
def to_dict(self) -> dict:
|
|
52
|
+
return asdict(self)
|
|
53
|
+
|
|
54
|
+
|
|
55
|
+
def hash_messages(messages: Conversation) -> str:
|
|
56
|
+
"""Deterministically hash a conversation (list of message dicts)."""
|
|
57
|
+
payload = json.dumps(messages, sort_keys=True, ensure_ascii=False)
|
|
58
|
+
return hashlib.sha256(payload.encode("utf-8")).hexdigest()
|
|
59
|
+
|
|
60
|
+
|
|
61
|
+
class LLMBatchRunner:
|
|
62
|
+
"""
|
|
63
|
+
Send a batch of chat-completion requests concurrently to an
|
|
64
|
+
OpenAI-compatible endpoint, optionally enforcing structured output
|
|
65
|
+
against a Pydantic model, with retries, per-sample disk caching
|
|
66
|
+
(keyed by a hash of the input messages), and error reporting.
|
|
67
|
+
|
|
68
|
+
Example
|
|
69
|
+
-------
|
|
70
|
+
>>> from pydantic import BaseModel
|
|
71
|
+
>>> class Answer(BaseModel):
|
|
72
|
+
... value: int
|
|
73
|
+
>>> runner = LLMBatchRunner(
|
|
74
|
+
... base_model=Answer,
|
|
75
|
+
... messages=[[{"role": "user", "content": "2+2=?"}]],
|
|
76
|
+
... max_workers=8,
|
|
77
|
+
... client_config={"api_key": "wow_very secret", "base_url": None},
|
|
78
|
+
... model_name="gpt-4o-mini",
|
|
79
|
+
... cache_dir="./cache",
|
|
80
|
+
... n_retries=3,
|
|
81
|
+
... conversations_meta=None
|
|
82
|
+
... )
|
|
83
|
+
>>> results = runner.run()
|
|
84
|
+
"""
|
|
85
|
+
|
|
86
|
+
def __init__(
|
|
87
|
+
self,
|
|
88
|
+
base_model: Optional[Type[BaseModel]],
|
|
89
|
+
messages: List[Conversation],
|
|
90
|
+
max_workers: int,
|
|
91
|
+
client_config: dict,
|
|
92
|
+
model_name: str,
|
|
93
|
+
cache_dir: Union[str, Path],
|
|
94
|
+
n_retries: int = 3,
|
|
95
|
+
overwrite_existing: bool = False,
|
|
96
|
+
temperature: Optional[float] = None,
|
|
97
|
+
extra_create_kwargs: Optional[Dict[str, Any]] = None,
|
|
98
|
+
show_progress: bool = True,
|
|
99
|
+
conversations_meta: List[Any]|None = None,
|
|
100
|
+
additional_validation_func = None,
|
|
101
|
+
postprocessing_func=None
|
|
102
|
+
) -> None:
|
|
103
|
+
if n_retries < 1:
|
|
104
|
+
raise ValueError("n_retries must be >= 1")
|
|
105
|
+
if max_workers < 1:
|
|
106
|
+
raise ValueError("max_workers must be >= 1")
|
|
107
|
+
|
|
108
|
+
self.base_model = base_model
|
|
109
|
+
self.messages = messages
|
|
110
|
+
self.max_workers = max_workers
|
|
111
|
+
self.model_name = model_name
|
|
112
|
+
self.n_retries = n_retries
|
|
113
|
+
self.overwrite_existing = overwrite_existing
|
|
114
|
+
self.temperature = temperature
|
|
115
|
+
self.extra_create_kwargs = extra_create_kwargs or {}
|
|
116
|
+
self.show_progress = show_progress
|
|
117
|
+
|
|
118
|
+
self.cache_dir = Path(cache_dir)
|
|
119
|
+
self.cache_dir.mkdir(parents=True, exist_ok=True)
|
|
120
|
+
|
|
121
|
+
self.client = OpenAI(**client_config)
|
|
122
|
+
|
|
123
|
+
self._progress_lock = threading.Lock()
|
|
124
|
+
self._completed = 0
|
|
125
|
+
|
|
126
|
+
self.additional_validation_func = additional_validation_func
|
|
127
|
+
self.postprocessing_func = postprocessing_func
|
|
128
|
+
|
|
129
|
+
self.conversations_meta = [None for _ in range(len(self.messages))]
|
|
130
|
+
if conversations_meta is not None:
|
|
131
|
+
assert isinstance(conversations_meta, list)
|
|
132
|
+
assert len(conversations_meta) == len(messages)
|
|
133
|
+
self.conversations_meta = conversations_meta
|
|
134
|
+
|
|
135
|
+
|
|
136
|
+
# ------------------------------------------------------------------ #
|
|
137
|
+
# Public API
|
|
138
|
+
# ------------------------------------------------------------------ #
|
|
139
|
+
|
|
140
|
+
def run(self) -> List[dict]:
|
|
141
|
+
"""
|
|
142
|
+
Process all conversations concurrently and return a list of
|
|
143
|
+
SampleResult objects, ordered the same way as the input.
|
|
144
|
+
"""
|
|
145
|
+
total = len(self.messages)
|
|
146
|
+
results: List[Optional[SampleResult]] = [None] * total
|
|
147
|
+
|
|
148
|
+
with ThreadPoolExecutor(max_workers=self.max_workers) as pool:
|
|
149
|
+
futures = {
|
|
150
|
+
pool.submit(self._process_one, idx, conv): idx
|
|
151
|
+
for idx, conv in enumerate(self.messages)
|
|
152
|
+
}
|
|
153
|
+
for future in as_completed(futures):
|
|
154
|
+
idx = futures[future]
|
|
155
|
+
try:
|
|
156
|
+
result = future.result()
|
|
157
|
+
except Exception as exc: # pragma: no cover - safety net
|
|
158
|
+
logger.error(
|
|
159
|
+
"Unhandled exception while processing sample %d: %s",
|
|
160
|
+
idx,
|
|
161
|
+
exc,
|
|
162
|
+
)
|
|
163
|
+
result = SampleResult(
|
|
164
|
+
index=idx,
|
|
165
|
+
hash=hash_messages(self.messages[idx]),
|
|
166
|
+
success=False,
|
|
167
|
+
error=str(exc),
|
|
168
|
+
conversation=self.messages[idx],
|
|
169
|
+
meta=self.conversations_meta[idx]
|
|
170
|
+
)
|
|
171
|
+
results[idx] = result
|
|
172
|
+
self._report_progress(total)
|
|
173
|
+
|
|
174
|
+
results = [
|
|
175
|
+
x.to_dict() for x in results if x is not None
|
|
176
|
+
]
|
|
177
|
+
return results
|
|
178
|
+
|
|
179
|
+
@staticmethod
|
|
180
|
+
def load_cache(cache_dir: Union[str, Path]) -> List[Dict]:
|
|
181
|
+
"""
|
|
182
|
+
Load every cached JSON result from `cache_dir`.
|
|
183
|
+
|
|
184
|
+
Returns a dict mapping {messages_hash: cached_record}, where
|
|
185
|
+
cached_record has keys like "output", "success", "model_name", etc.
|
|
186
|
+
"""
|
|
187
|
+
cache_dir = Path(cache_dir)
|
|
188
|
+
if not cache_dir.exists():
|
|
189
|
+
return []
|
|
190
|
+
|
|
191
|
+
|
|
192
|
+
cache = []
|
|
193
|
+
for path in cache_dir.glob("*.json"):
|
|
194
|
+
try:
|
|
195
|
+
with open(path, "r", encoding="utf-8") as f:
|
|
196
|
+
cache.append(json.load(f))
|
|
197
|
+
except Exception as exc:
|
|
198
|
+
logger.warning("Skipping unreadable cache file %s: %s", path, exc)
|
|
199
|
+
return cache
|
|
200
|
+
|
|
201
|
+
# ------------------------------------------------------------------ #
|
|
202
|
+
# Internal helpers
|
|
203
|
+
# ------------------------------------------------------------------ #
|
|
204
|
+
|
|
205
|
+
def _cache_path(self, key: str) -> Path:
|
|
206
|
+
return self.cache_dir / f"{key}.json"
|
|
207
|
+
|
|
208
|
+
def _load_cached(self, key: str) -> Optional[dict]:
|
|
209
|
+
path = self._cache_path(key)
|
|
210
|
+
if not path.exists():
|
|
211
|
+
return None
|
|
212
|
+
try:
|
|
213
|
+
with open(path, "r", encoding="utf-8") as f:
|
|
214
|
+
return json.load(f)
|
|
215
|
+
except Exception as exc:
|
|
216
|
+
logger.warning("Cache file %s is corrupt, will reprocess: %s", path, exc)
|
|
217
|
+
return None
|
|
218
|
+
|
|
219
|
+
def _save_cache(self, key: str, record: dict) -> None:
|
|
220
|
+
path = self._cache_path(key)
|
|
221
|
+
tmp_path = path.with_suffix(".json.tmp")
|
|
222
|
+
with open(tmp_path, "w", encoding="utf-8") as f:
|
|
223
|
+
json.dump(record, f, ensure_ascii=False, indent=2)
|
|
224
|
+
tmp_path.replace(path)
|
|
225
|
+
|
|
226
|
+
def _process_one(self, index: int, conversation: Conversation) -> SampleResult:
|
|
227
|
+
key = hash_messages(conversation)
|
|
228
|
+
|
|
229
|
+
if not self.overwrite_existing:
|
|
230
|
+
cached = self._load_cached(key)
|
|
231
|
+
if cached is not None:
|
|
232
|
+
if cached.get("success", False):
|
|
233
|
+
return SampleResult(
|
|
234
|
+
index=index,
|
|
235
|
+
hash=key,
|
|
236
|
+
success=cached.get("success", True),
|
|
237
|
+
output=cached.get("output"),
|
|
238
|
+
error=cached.get("error"),
|
|
239
|
+
from_cache=True,
|
|
240
|
+
attempts=0,
|
|
241
|
+
conversation=conversation,
|
|
242
|
+
raw_output=cached.get("raw_output"),
|
|
243
|
+
meta=cached.get("meta")
|
|
244
|
+
)
|
|
245
|
+
|
|
246
|
+
last_error: Optional[str] = None
|
|
247
|
+
for attempt in range(1, self.n_retries + 1):
|
|
248
|
+
try:
|
|
249
|
+
output, raw_output = self._call_and_validate(conversation)
|
|
250
|
+
if self.postprocessing_func is not None:
|
|
251
|
+
output = self.postprocessing_func(output, conversation)
|
|
252
|
+
record = {
|
|
253
|
+
"success": True,
|
|
254
|
+
"output": output,
|
|
255
|
+
"error": None,
|
|
256
|
+
"model_name": self.model_name,
|
|
257
|
+
"attempts": attempt,
|
|
258
|
+
"conversation": conversation,
|
|
259
|
+
"raw_output": raw_output,
|
|
260
|
+
"meta": self.conversations_meta[index]
|
|
261
|
+
}
|
|
262
|
+
self._save_cache(key, record)
|
|
263
|
+
return SampleResult(
|
|
264
|
+
index=index,
|
|
265
|
+
hash=key,
|
|
266
|
+
success=True,
|
|
267
|
+
output=output,
|
|
268
|
+
attempts=attempt,
|
|
269
|
+
conversation=conversation,
|
|
270
|
+
raw_output=raw_output
|
|
271
|
+
)
|
|
272
|
+
except Exception as exc:
|
|
273
|
+
last_error = f"{type(exc).__name__}: {exc}"
|
|
274
|
+
logger.warning(
|
|
275
|
+
"Sample %d (hash=%s) attempt %d/%d failed: %s",
|
|
276
|
+
index,
|
|
277
|
+
key[:8],
|
|
278
|
+
attempt,
|
|
279
|
+
self.n_retries,
|
|
280
|
+
last_error,
|
|
281
|
+
)
|
|
282
|
+
|
|
283
|
+
logger.error(
|
|
284
|
+
"Sample %d (hash=%s) failed after %d attempts: %s",
|
|
285
|
+
index,
|
|
286
|
+
key[:8],
|
|
287
|
+
self.n_retries,
|
|
288
|
+
last_error,
|
|
289
|
+
)
|
|
290
|
+
# # Cache the failure too, so a re-run without overwrite_existing
|
|
291
|
+
# doesn't silently skip a sample that never succeeded.
|
|
292
|
+
# record = {
|
|
293
|
+
# "success": False,
|
|
294
|
+
# "output": None,
|
|
295
|
+
# "error": last_error,
|
|
296
|
+
# "model_name": self.model_name,
|
|
297
|
+
# "attempts": self.n_retries,
|
|
298
|
+
# "conversation": conversation,
|
|
299
|
+
# "raw_output": None,
|
|
300
|
+
# "meta": self.conversations_meta[index]
|
|
301
|
+
# }
|
|
302
|
+
# self._save_cache(key, record)
|
|
303
|
+
return SampleResult(
|
|
304
|
+
index=index,
|
|
305
|
+
hash=key,
|
|
306
|
+
success=False,
|
|
307
|
+
error=last_error,
|
|
308
|
+
attempts=self.n_retries,
|
|
309
|
+
raw_output=None,
|
|
310
|
+
conversation=conversation,
|
|
311
|
+
output=None
|
|
312
|
+
)
|
|
313
|
+
|
|
314
|
+
def _call_and_validate(self, conversation: Conversation) -> Tuple[Union[dict, str], str]:
|
|
315
|
+
|
|
316
|
+
# Structured output
|
|
317
|
+
if self.base_model is not None:
|
|
318
|
+
create_kwargs: Dict[str, Any] = dict(
|
|
319
|
+
model=self.model_name,
|
|
320
|
+
messages=conversation,
|
|
321
|
+
response_format=self.base_model,
|
|
322
|
+
**self.extra_create_kwargs,
|
|
323
|
+
)
|
|
324
|
+
if self.temperature is not None:
|
|
325
|
+
create_kwargs["temperature"] = self.temperature
|
|
326
|
+
|
|
327
|
+
response = self.client.beta.chat.completions.parse(**create_kwargs)
|
|
328
|
+
message = response.choices[0].message
|
|
329
|
+
|
|
330
|
+
if hasattr(message, "refusal"):
|
|
331
|
+
if message.refusal:
|
|
332
|
+
raise ValueError(f"Model refused to answer: {message.refusal}")
|
|
333
|
+
|
|
334
|
+
if message.parsed is None:
|
|
335
|
+
raise ValueError("Model did not return a parsed structured output")
|
|
336
|
+
|
|
337
|
+
json_model_output = message.parsed.model_dump()
|
|
338
|
+
validated_json_model_output = self._validate(json_model_output, conversation)
|
|
339
|
+
return validated_json_model_output, message.content
|
|
340
|
+
|
|
341
|
+
# Regular (non-structured) completion
|
|
342
|
+
create_kwargs = dict(
|
|
343
|
+
model=self.model_name,
|
|
344
|
+
messages=conversation,
|
|
345
|
+
**self.extra_create_kwargs,
|
|
346
|
+
)
|
|
347
|
+
if self.temperature is not None:
|
|
348
|
+
create_kwargs["temperature"] = self.temperature
|
|
349
|
+
|
|
350
|
+
response = self.client.chat.completions.create(**create_kwargs)
|
|
351
|
+
return response.choices[0].message.content, response.choices[0].message.content
|
|
352
|
+
|
|
353
|
+
def _validate(self, extracted_json, conversation) -> dict:
|
|
354
|
+
|
|
355
|
+
try:
|
|
356
|
+
instance = self.base_model.model_validate_json(json.dumps(extracted_json))
|
|
357
|
+
except Exception as exc:
|
|
358
|
+
raise ValueError(f"Structured output validation failed: {exc}") from exc
|
|
359
|
+
|
|
360
|
+
if self.additional_validation_func is not None:
|
|
361
|
+
try:
|
|
362
|
+
is_valid = self.additional_validation_func(extracted_json, conversation)
|
|
363
|
+
if not is_valid:
|
|
364
|
+
raise ValueError(f"Structured output validation failed via CUSTOM FUNCTION: {extracted_json}")
|
|
365
|
+
except Exception as exc:
|
|
366
|
+
raise ValueError(f"Structured output validation failed: {exc}") from exc
|
|
367
|
+
|
|
368
|
+
|
|
369
|
+
return instance.model_dump()
|
|
370
|
+
|
|
371
|
+
def _report_progress(self, total: int) -> None:
|
|
372
|
+
if not self.show_progress:
|
|
373
|
+
return
|
|
374
|
+
with self._progress_lock:
|
|
375
|
+
self._completed += 1
|
|
376
|
+
print(f"\rProcessed {self._completed}/{total}", end="", flush=True)
|
|
377
|
+
if self._completed == total:
|
|
378
|
+
print()
|
|
@@ -0,0 +1,121 @@
|
|
|
1
|
+
Metadata-Version: 2.4
|
|
2
|
+
Name: llm-batch-runner
|
|
3
|
+
Version: 0.1.0
|
|
4
|
+
Summary: Concurrent, caching, retrying OpenAI-compatible batch chat-completion runner with structured output validation.
|
|
5
|
+
Author-email: Your Name <you@example.com>
|
|
6
|
+
License: MIT
|
|
7
|
+
Project-URL: Homepage, https://github.com/yourusername/llm-batch-runner
|
|
8
|
+
Keywords: openai,llm,batch,structured-output,pydantic,concurrency,cache
|
|
9
|
+
Classifier: Programming Language :: Python :: 3
|
|
10
|
+
Classifier: License :: OSI Approved :: MIT License
|
|
11
|
+
Classifier: Operating System :: OS Independent
|
|
12
|
+
Requires-Python: >=3.9
|
|
13
|
+
Description-Content-Type: text/markdown
|
|
14
|
+
Requires-Dist: openai>=1.40.0
|
|
15
|
+
Requires-Dist: pydantic>=2.0.0
|
|
16
|
+
|
|
17
|
+
# llm-batch-runner
|
|
18
|
+
|
|
19
|
+
Concurrent, caching, retrying batch runner for OpenAI-compatible chat
|
|
20
|
+
completion endpoints, with optional structured-output validation against
|
|
21
|
+
a Pydantic model.
|
|
22
|
+
|
|
23
|
+
## Features
|
|
24
|
+
|
|
25
|
+
- Fires requests concurrently (configurable `max_workers`) through the
|
|
26
|
+
official `openai` SDK, so it works with OpenAI itself or any
|
|
27
|
+
OpenAI-compatible `base_url` (vLLM, Together, Groq, local servers, ...).
|
|
28
|
+
- If you pass a Pydantic `base_model`, requests are made with a
|
|
29
|
+
JSON-schema `response_format` and every response is validated against
|
|
30
|
+
the model. Invalid output is retried, up to `n_retries` times.
|
|
31
|
+
- If no `base_model` is passed, it just does a normal completion and
|
|
32
|
+
returns the raw text.
|
|
33
|
+
- Every sample is cached to disk as its own JSON file, keyed by a SHA-256
|
|
34
|
+
hash of its input `messages`. On a re-run, already-cached samples are
|
|
35
|
+
skipped unless `overwrite_existing=True`.
|
|
36
|
+
- Failures are logged (via the standard `logging` module) and returned
|
|
37
|
+
in the result list rather than raised, so one bad sample never kills
|
|
38
|
+
the whole batch.
|
|
39
|
+
- `LLMBatchRunner.load_cache(cache_dir)` is a static method that loads
|
|
40
|
+
every cached record back into a dict, keyed by hash.
|
|
41
|
+
|
|
42
|
+
|
|
43
|
+
## Usage
|
|
44
|
+
|
|
45
|
+
```python
|
|
46
|
+
from pydantic import BaseModel
|
|
47
|
+
from llm_batch_runner import LLMBatchRunner
|
|
48
|
+
|
|
49
|
+
class Answer(BaseModel):
|
|
50
|
+
reasoning: str
|
|
51
|
+
value: int
|
|
52
|
+
|
|
53
|
+
conversations = [
|
|
54
|
+
[{"role": "user", "content": "What is 12 * 7?"}],
|
|
55
|
+
[{"role": "user", "content": "What is 9 * 9?"}],
|
|
56
|
+
# ... as many as you like
|
|
57
|
+
]
|
|
58
|
+
|
|
59
|
+
runner = LLMBatchRunner(
|
|
60
|
+
base_model=Answer, # or None for plain text completions
|
|
61
|
+
messages=conversations, # List[List[dict]] — one conversation per sample
|
|
62
|
+
max_workers=16,
|
|
63
|
+
client_config={"api_key": "wow_very secret", "base_url": None},
|
|
64
|
+
model_name="gpt-4o-mini",
|
|
65
|
+
cache_dir="./llm_cache",
|
|
66
|
+
n_retries=3,
|
|
67
|
+
overwrite_existing=False,
|
|
68
|
+
conversations_meta=None # or list of dicts: additional values to store
|
|
69
|
+
)
|
|
70
|
+
|
|
71
|
+
results = runner.run()
|
|
72
|
+
print(results[0])
|
|
73
|
+
# OUTPUT:
|
|
74
|
+
```
|
|
75
|
+
|
|
76
|
+
Output:
|
|
77
|
+
|
|
78
|
+
```json
|
|
79
|
+
{
|
|
80
|
+
"attempts": 1,
|
|
81
|
+
"conversation": [
|
|
82
|
+
{
|
|
83
|
+
"role": "user",
|
|
84
|
+
"content": "What is 12 * 7?"
|
|
85
|
+
}
|
|
86
|
+
],
|
|
87
|
+
"error": "None",
|
|
88
|
+
"from_cache": "False",
|
|
89
|
+
"hash": "f699170cc22a5e871fccdf3a414f35606d100c943d18387e58d75c0a1006b6e6",
|
|
90
|
+
"index": 0,
|
|
91
|
+
"output": {
|
|
92
|
+
"reasoning": "12 * 7 can be calculated as (10 * 7) + (2 * 7) = 70 + 14 = 84.",
|
|
93
|
+
"value": 84
|
|
94
|
+
},
|
|
95
|
+
"raw_output": "{\"reasoning\": \"12 * 7 can be calculated as (10 * 7) + (2 * 7) = 70 + 14 = 84.\", \"value\": 84}",
|
|
96
|
+
"success": "True",
|
|
97
|
+
"meta": null
|
|
98
|
+
}
|
|
99
|
+
```
|
|
100
|
+
|
|
101
|
+
### Loading cached results later
|
|
102
|
+
|
|
103
|
+
```python
|
|
104
|
+
from llm_batch_runner import LLMBatchRunner
|
|
105
|
+
|
|
106
|
+
cache = LLMBatchRunner.load_cache("./llm_cache")
|
|
107
|
+
# List of dicts - loaded data from the cache
|
|
108
|
+
# [
|
|
109
|
+
# {"success": True, "output": {...}, "error": None, ...},
|
|
110
|
+
# ...
|
|
111
|
+
# ]
|
|
112
|
+
```
|
|
113
|
+
|
|
114
|
+
## Notes
|
|
115
|
+
|
|
116
|
+
- `messages` is a list of conversations — each conversation is itself the
|
|
117
|
+
standard OpenAI `messages` list (`[{"role": ..., "content": ...}, ...]`).
|
|
118
|
+
- The cache key is a hash of the conversation only (not the model name),
|
|
119
|
+
so if you change `model_name` and want fresh results, pass a different
|
|
120
|
+
`cache_dir` or use `overwrite_existing=True`.
|
|
121
|
+
- Failed samples (that exhausted all retries) are also cached.
|
|
@@ -0,0 +1,9 @@
|
|
|
1
|
+
README.md
|
|
2
|
+
pyproject.toml
|
|
3
|
+
llm_batch_runner/__init__.py
|
|
4
|
+
llm_batch_runner/processor.py
|
|
5
|
+
llm_batch_runner.egg-info/PKG-INFO
|
|
6
|
+
llm_batch_runner.egg-info/SOURCES.txt
|
|
7
|
+
llm_batch_runner.egg-info/dependency_links.txt
|
|
8
|
+
llm_batch_runner.egg-info/requires.txt
|
|
9
|
+
llm_batch_runner.egg-info/top_level.txt
|
|
@@ -0,0 +1 @@
|
|
|
1
|
+
|
|
@@ -0,0 +1 @@
|
|
|
1
|
+
llm_batch_runner
|
|
@@ -0,0 +1,30 @@
|
|
|
1
|
+
[build-system]
|
|
2
|
+
requires = ["setuptools>=68", "wheel"]
|
|
3
|
+
build-backend = "setuptools.build_meta"
|
|
4
|
+
|
|
5
|
+
[project]
|
|
6
|
+
name = "llm-batch-runner"
|
|
7
|
+
version = "0.1.0"
|
|
8
|
+
description = "Concurrent, caching, retrying OpenAI-compatible batch chat-completion runner with structured output validation."
|
|
9
|
+
readme = "README.md"
|
|
10
|
+
requires-python = ">=3.9"
|
|
11
|
+
license = { text = "MIT" }
|
|
12
|
+
authors = [
|
|
13
|
+
{ name = "Your Name", email = "you@example.com" }
|
|
14
|
+
]
|
|
15
|
+
keywords = ["openai", "llm", "batch", "structured-output", "pydantic", "concurrency", "cache"]
|
|
16
|
+
classifiers = [
|
|
17
|
+
"Programming Language :: Python :: 3",
|
|
18
|
+
"License :: OSI Approved :: MIT License",
|
|
19
|
+
"Operating System :: OS Independent",
|
|
20
|
+
]
|
|
21
|
+
dependencies = [
|
|
22
|
+
"openai>=1.40.0",
|
|
23
|
+
"pydantic>=2.0.0",
|
|
24
|
+
]
|
|
25
|
+
|
|
26
|
+
[project.urls]
|
|
27
|
+
Homepage = "https://github.com/yourusername/llm-batch-runner"
|
|
28
|
+
|
|
29
|
+
[tool.setuptools.packages.find]
|
|
30
|
+
include = ["llm_batch_runner*"]
|