bedrockhelper 1.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.
- bedrockhelper-1.1.0/LICENSE +21 -0
- bedrockhelper-1.1.0/PKG-INFO +262 -0
- bedrockhelper-1.1.0/README.md +241 -0
- bedrockhelper-1.1.0/pyproject.toml +45 -0
- bedrockhelper-1.1.0/src/bedrockhelper/__init__.py +17 -0
- bedrockhelper-1.1.0/src/bedrockhelper/batch/__init__.py +12 -0
- bedrockhelper-1.1.0/src/bedrockhelper/batch/models.py +41 -0
- bedrockhelper-1.1.0/src/bedrockhelper/batch/reconcile.py +118 -0
- bedrockhelper-1.1.0/src/bedrockhelper/batch/store.py +255 -0
- bedrockhelper-1.1.0/src/bedrockhelper/main.py +1057 -0
- bedrockhelper-1.1.0/src/bedrockhelper/models.py +176 -0
- bedrockhelper-1.1.0/src/bedrockhelper/types.py +15 -0
- bedrockhelper-1.1.0/src/bedrockhelper/utils.py +380 -0
|
@@ -0,0 +1,21 @@
|
|
|
1
|
+
MIT License
|
|
2
|
+
|
|
3
|
+
Copyright (c) 2025 Matthew Marion
|
|
4
|
+
|
|
5
|
+
Permission is hereby granted, free of charge, to any person obtaining a copy
|
|
6
|
+
of this software and associated documentation files (the "Software"), to deal
|
|
7
|
+
in the Software without restriction, including without limitation the rights
|
|
8
|
+
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
|
|
9
|
+
copies of the Software, and to permit persons to whom the Software is
|
|
10
|
+
furnished to do so, subject to the following conditions:
|
|
11
|
+
|
|
12
|
+
The above copyright notice and this permission notice shall be included in all
|
|
13
|
+
copies or substantial portions of the Software.
|
|
14
|
+
|
|
15
|
+
THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
|
|
16
|
+
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
|
|
17
|
+
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
|
|
18
|
+
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
|
|
19
|
+
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
|
|
20
|
+
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
|
|
21
|
+
SOFTWARE.
|
|
@@ -0,0 +1,262 @@
|
|
|
1
|
+
Metadata-Version: 2.1
|
|
2
|
+
Name: bedrockhelper
|
|
3
|
+
Version: 1.1.0
|
|
4
|
+
Summary: Helper utilities for Amazon Bedrock (RAG, embeddings, batch jobs)
|
|
5
|
+
License: MIT
|
|
6
|
+
Author: Matthew Marion
|
|
7
|
+
Author-email: matthew@endofline.ca
|
|
8
|
+
Requires-Python: >=3.10,<4.0
|
|
9
|
+
Classifier: License :: OSI Approved :: MIT License
|
|
10
|
+
Classifier: Programming Language :: Python :: 3
|
|
11
|
+
Classifier: Programming Language :: Python :: 3.10
|
|
12
|
+
Classifier: Programming Language :: Python :: 3.11
|
|
13
|
+
Classifier: Programming Language :: Python :: 3.12
|
|
14
|
+
Classifier: Programming Language :: Python :: 3.13
|
|
15
|
+
Requires-Dist: boto3 (>=1.42.25,<2.0.0)
|
|
16
|
+
Requires-Dist: boto3-stubs (>=1.42.42,<1.43.0)
|
|
17
|
+
Requires-Dist: botocore[crt] (>=1.42.42,<2.0.0)
|
|
18
|
+
Requires-Dist: ipykernel (>=7.1.0,<8.0.0)
|
|
19
|
+
Description-Content-Type: text/markdown
|
|
20
|
+
|
|
21
|
+
# BedrockHelper
|
|
22
|
+
|
|
23
|
+
A small, batteries-included helper for Amazon Bedrock:
|
|
24
|
+
|
|
25
|
+
- **RAG Q&A** using Bedrock Runtime `converse` / `converse_stream` when available, with automatic fallback.
|
|
26
|
+
- **Embeddings** using **Titan embed text v2**:
|
|
27
|
+
- Small batches inline
|
|
28
|
+
- Large batches via **Bedrock batch inference** (S3 JSONL)
|
|
29
|
+
- Works in **sync** and **async** codebases.
|
|
30
|
+
- Includes a **pluggable JobStore** abstraction for tracking batch jobs (in-memory, Redis, Postgres, etc.).
|
|
31
|
+
|
|
32
|
+
---
|
|
33
|
+
|
|
34
|
+
## Structure
|
|
35
|
+
``` plaintext
|
|
36
|
+
├── .github
|
|
37
|
+
│ └── workflows
|
|
38
|
+
│ └── mypy.yml
|
|
39
|
+
├── .gitignore
|
|
40
|
+
├── poetry.lock
|
|
41
|
+
├── pyproject.toml
|
|
42
|
+
├── README.md
|
|
43
|
+
├── src
|
|
44
|
+
│ └── bedrockhelper
|
|
45
|
+
│ ├── __init__.py
|
|
46
|
+
│ ├── batch
|
|
47
|
+
│ ├── main.py
|
|
48
|
+
│ ├── models.py
|
|
49
|
+
│ ├── types.py
|
|
50
|
+
│ └── utils.py
|
|
51
|
+
└── tests
|
|
52
|
+
├── test_batch_models.py
|
|
53
|
+
├── test_batch_store.py
|
|
54
|
+
├── test_main.py
|
|
55
|
+
├── test_models.py
|
|
56
|
+
├── test_reconcile.py
|
|
57
|
+
└── test_utils.py
|
|
58
|
+
```
|
|
59
|
+
## Installation
|
|
60
|
+
|
|
61
|
+
This module is intended to be vendored or installed directly from source.
|
|
62
|
+
|
|
63
|
+
Requirements:
|
|
64
|
+
|
|
65
|
+
- Python 3.10+
|
|
66
|
+
- boto3 / botocore
|
|
67
|
+
- AWS credentials configured for Bedrock + S3
|
|
68
|
+
|
|
69
|
+
---
|
|
70
|
+
|
|
71
|
+
## Quick Start
|
|
72
|
+
|
|
73
|
+
```python
|
|
74
|
+
from bedrockhelper import BedrockHelper
|
|
75
|
+
|
|
76
|
+
helper = BedrockHelper(
|
|
77
|
+
region_name="ca-central-1",
|
|
78
|
+
rag_model_id="global.anthropic.claude-sonnet-4-5-20250929-v1:0",
|
|
79
|
+
embedding_model_id="amazon.titan-embed-text-v2:0",
|
|
80
|
+
)
|
|
81
|
+
```
|
|
82
|
+
|
|
83
|
+
---
|
|
84
|
+
|
|
85
|
+
## RAG Usage
|
|
86
|
+
|
|
87
|
+
### Non-streaming
|
|
88
|
+
|
|
89
|
+
```python
|
|
90
|
+
resp = helper.generate_with_rag(
|
|
91
|
+
system_prompt="You are a helpful assistant.",
|
|
92
|
+
context=["Doc 1", "Doc 2"],
|
|
93
|
+
question="Summarize this",
|
|
94
|
+
)
|
|
95
|
+
|
|
96
|
+
print(resp.text)
|
|
97
|
+
```
|
|
98
|
+
|
|
99
|
+
### Streaming
|
|
100
|
+
|
|
101
|
+
```python
|
|
102
|
+
resp = helper.generate_with_rag(
|
|
103
|
+
system_prompt="You are concise.",
|
|
104
|
+
context="Some context",
|
|
105
|
+
question="Stream this",
|
|
106
|
+
stream=True,
|
|
107
|
+
)
|
|
108
|
+
|
|
109
|
+
for chunk in resp.stream:
|
|
110
|
+
print(chunk, end="")
|
|
111
|
+
```
|
|
112
|
+
|
|
113
|
+
---
|
|
114
|
+
|
|
115
|
+
## Async Usage
|
|
116
|
+
|
|
117
|
+
```python
|
|
118
|
+
resp = await helper.async_generate_with_rag(
|
|
119
|
+
system_prompt="You are helpful",
|
|
120
|
+
context="Context",
|
|
121
|
+
question="Async answer?",
|
|
122
|
+
)
|
|
123
|
+
|
|
124
|
+
print(resp.text)
|
|
125
|
+
```
|
|
126
|
+
|
|
127
|
+
Streaming:
|
|
128
|
+
|
|
129
|
+
```python
|
|
130
|
+
async for chunk in await helper.async_stream_generate_with_rag(
|
|
131
|
+
system_prompt="You are helpful",
|
|
132
|
+
context="Context",
|
|
133
|
+
question="Async stream",
|
|
134
|
+
):
|
|
135
|
+
print(chunk, end="")
|
|
136
|
+
```
|
|
137
|
+
|
|
138
|
+
---
|
|
139
|
+
|
|
140
|
+
## Embeddings
|
|
141
|
+
|
|
142
|
+
### Inline embeddings
|
|
143
|
+
|
|
144
|
+
```python
|
|
145
|
+
records = {
|
|
146
|
+
"a": "Hello world",
|
|
147
|
+
"b": "Another document",
|
|
148
|
+
}
|
|
149
|
+
|
|
150
|
+
resp = helper.embed_texts(records)
|
|
151
|
+
print(resp.embeddings["a"])
|
|
152
|
+
```
|
|
153
|
+
|
|
154
|
+
### Batch embeddings
|
|
155
|
+
|
|
156
|
+
```python
|
|
157
|
+
job = helper.embed_texts(
|
|
158
|
+
big_records,
|
|
159
|
+
s3_bucket="my-bucket",
|
|
160
|
+
s3_prefix="bedrock_batch",
|
|
161
|
+
role_arn="arn:aws:iam::123456789012:role/BedrockBatchRole",
|
|
162
|
+
)
|
|
163
|
+
|
|
164
|
+
print(job.job_id)
|
|
165
|
+
```
|
|
166
|
+
|
|
167
|
+
Wait for completion:
|
|
168
|
+
|
|
169
|
+
```python
|
|
170
|
+
info = helper.wait_for_batch_job(job.job_id)
|
|
171
|
+
print(info["status"])
|
|
172
|
+
```
|
|
173
|
+
|
|
174
|
+
Parse results:
|
|
175
|
+
|
|
176
|
+
```python
|
|
177
|
+
embeddings = helper.parse_batch_embeddings(job.output_s3_uri)
|
|
178
|
+
```
|
|
179
|
+
|
|
180
|
+
---
|
|
181
|
+
|
|
182
|
+
## Batch Job Reconciliation
|
|
183
|
+
|
|
184
|
+
The helper includes a small framework for tracking and reconciling batch jobs.
|
|
185
|
+
|
|
186
|
+
```python
|
|
187
|
+
from bedrockhelper import InMemoryJobStore
|
|
188
|
+
from bedrockhelper.batch.reconcile import reconcile_batch_embedding_jobs
|
|
189
|
+
|
|
190
|
+
store = InMemoryJobStore()
|
|
191
|
+
store.add(job)
|
|
192
|
+
|
|
193
|
+
|
|
194
|
+
def process_embeddings(job, embeddings):
|
|
195
|
+
print("Processed", len(embeddings))
|
|
196
|
+
|
|
197
|
+
|
|
198
|
+
reconcile_batch_embedding_jobs(
|
|
199
|
+
helper=helper,
|
|
200
|
+
store=store,
|
|
201
|
+
process_embeddings=process_embeddings,
|
|
202
|
+
)
|
|
203
|
+
```
|
|
204
|
+
|
|
205
|
+
---
|
|
206
|
+
|
|
207
|
+
## Implementing Your Own JobStore
|
|
208
|
+
|
|
209
|
+
To persist batch jobs across restarts or multiple workers, implement the `JobStore` interface.
|
|
210
|
+
|
|
211
|
+
### Redis Example
|
|
212
|
+
|
|
213
|
+
Use Redis keys per job plus a sorted set for ordering.
|
|
214
|
+
|
|
215
|
+
Key ideas:
|
|
216
|
+
|
|
217
|
+
- Store jobs as JSON
|
|
218
|
+
- Use TTL for terminal jobs
|
|
219
|
+
- Use a sorted set by `created_at`
|
|
220
|
+
|
|
221
|
+
This is ideal for horizontally scaled workers.
|
|
222
|
+
|
|
223
|
+
### Postgres Example
|
|
224
|
+
|
|
225
|
+
Suggested schema:
|
|
226
|
+
|
|
227
|
+
```sql
|
|
228
|
+
CREATE TABLE bedrock_batch_jobs (
|
|
229
|
+
job_id TEXT PRIMARY KEY,
|
|
230
|
+
job_name TEXT NOT NULL,
|
|
231
|
+
model_id TEXT NOT NULL,
|
|
232
|
+
input_s3_uri TEXT NOT NULL,
|
|
233
|
+
output_s3_uri TEXT NOT NULL,
|
|
234
|
+
state TEXT NOT NULL,
|
|
235
|
+
created_at TIMESTAMPTZ NOT NULL,
|
|
236
|
+
last_checked_at TIMESTAMPTZ,
|
|
237
|
+
attempts INTEGER NOT NULL DEFAULT 0,
|
|
238
|
+
last_error TEXT,
|
|
239
|
+
bedrock_status TEXT,
|
|
240
|
+
processed_at TIMESTAMPTZ,
|
|
241
|
+
embeddings_count INTEGER,
|
|
242
|
+
debug JSONB NOT NULL DEFAULT '{}'::jsonb
|
|
243
|
+
);
|
|
244
|
+
```
|
|
245
|
+
|
|
246
|
+
This approach is boring, reliable, and production-friendly.
|
|
247
|
+
|
|
248
|
+
---
|
|
249
|
+
|
|
250
|
+
## Recommended Patterns
|
|
251
|
+
|
|
252
|
+
- Submit batch jobs quickly and return job IDs
|
|
253
|
+
- Reconcile in a background worker or cron
|
|
254
|
+
- Make `process_embeddings()` idempotent
|
|
255
|
+
- Prefer Postgres if you need observability
|
|
256
|
+
|
|
257
|
+
---
|
|
258
|
+
|
|
259
|
+
## License
|
|
260
|
+
|
|
261
|
+
Internal / project-specific. Adapt as needed.
|
|
262
|
+
|
|
@@ -0,0 +1,241 @@
|
|
|
1
|
+
# BedrockHelper
|
|
2
|
+
|
|
3
|
+
A small, batteries-included helper for Amazon Bedrock:
|
|
4
|
+
|
|
5
|
+
- **RAG Q&A** using Bedrock Runtime `converse` / `converse_stream` when available, with automatic fallback.
|
|
6
|
+
- **Embeddings** using **Titan embed text v2**:
|
|
7
|
+
- Small batches inline
|
|
8
|
+
- Large batches via **Bedrock batch inference** (S3 JSONL)
|
|
9
|
+
- Works in **sync** and **async** codebases.
|
|
10
|
+
- Includes a **pluggable JobStore** abstraction for tracking batch jobs (in-memory, Redis, Postgres, etc.).
|
|
11
|
+
|
|
12
|
+
---
|
|
13
|
+
|
|
14
|
+
## Structure
|
|
15
|
+
``` plaintext
|
|
16
|
+
├── .github
|
|
17
|
+
│ └── workflows
|
|
18
|
+
│ └── mypy.yml
|
|
19
|
+
├── .gitignore
|
|
20
|
+
├── poetry.lock
|
|
21
|
+
├── pyproject.toml
|
|
22
|
+
├── README.md
|
|
23
|
+
├── src
|
|
24
|
+
│ └── bedrockhelper
|
|
25
|
+
│ ├── __init__.py
|
|
26
|
+
│ ├── batch
|
|
27
|
+
│ ├── main.py
|
|
28
|
+
│ ├── models.py
|
|
29
|
+
│ ├── types.py
|
|
30
|
+
│ └── utils.py
|
|
31
|
+
└── tests
|
|
32
|
+
├── test_batch_models.py
|
|
33
|
+
├── test_batch_store.py
|
|
34
|
+
├── test_main.py
|
|
35
|
+
├── test_models.py
|
|
36
|
+
├── test_reconcile.py
|
|
37
|
+
└── test_utils.py
|
|
38
|
+
```
|
|
39
|
+
## Installation
|
|
40
|
+
|
|
41
|
+
This module is intended to be vendored or installed directly from source.
|
|
42
|
+
|
|
43
|
+
Requirements:
|
|
44
|
+
|
|
45
|
+
- Python 3.10+
|
|
46
|
+
- boto3 / botocore
|
|
47
|
+
- AWS credentials configured for Bedrock + S3
|
|
48
|
+
|
|
49
|
+
---
|
|
50
|
+
|
|
51
|
+
## Quick Start
|
|
52
|
+
|
|
53
|
+
```python
|
|
54
|
+
from bedrockhelper import BedrockHelper
|
|
55
|
+
|
|
56
|
+
helper = BedrockHelper(
|
|
57
|
+
region_name="ca-central-1",
|
|
58
|
+
rag_model_id="global.anthropic.claude-sonnet-4-5-20250929-v1:0",
|
|
59
|
+
embedding_model_id="amazon.titan-embed-text-v2:0",
|
|
60
|
+
)
|
|
61
|
+
```
|
|
62
|
+
|
|
63
|
+
---
|
|
64
|
+
|
|
65
|
+
## RAG Usage
|
|
66
|
+
|
|
67
|
+
### Non-streaming
|
|
68
|
+
|
|
69
|
+
```python
|
|
70
|
+
resp = helper.generate_with_rag(
|
|
71
|
+
system_prompt="You are a helpful assistant.",
|
|
72
|
+
context=["Doc 1", "Doc 2"],
|
|
73
|
+
question="Summarize this",
|
|
74
|
+
)
|
|
75
|
+
|
|
76
|
+
print(resp.text)
|
|
77
|
+
```
|
|
78
|
+
|
|
79
|
+
### Streaming
|
|
80
|
+
|
|
81
|
+
```python
|
|
82
|
+
resp = helper.generate_with_rag(
|
|
83
|
+
system_prompt="You are concise.",
|
|
84
|
+
context="Some context",
|
|
85
|
+
question="Stream this",
|
|
86
|
+
stream=True,
|
|
87
|
+
)
|
|
88
|
+
|
|
89
|
+
for chunk in resp.stream:
|
|
90
|
+
print(chunk, end="")
|
|
91
|
+
```
|
|
92
|
+
|
|
93
|
+
---
|
|
94
|
+
|
|
95
|
+
## Async Usage
|
|
96
|
+
|
|
97
|
+
```python
|
|
98
|
+
resp = await helper.async_generate_with_rag(
|
|
99
|
+
system_prompt="You are helpful",
|
|
100
|
+
context="Context",
|
|
101
|
+
question="Async answer?",
|
|
102
|
+
)
|
|
103
|
+
|
|
104
|
+
print(resp.text)
|
|
105
|
+
```
|
|
106
|
+
|
|
107
|
+
Streaming:
|
|
108
|
+
|
|
109
|
+
```python
|
|
110
|
+
async for chunk in await helper.async_stream_generate_with_rag(
|
|
111
|
+
system_prompt="You are helpful",
|
|
112
|
+
context="Context",
|
|
113
|
+
question="Async stream",
|
|
114
|
+
):
|
|
115
|
+
print(chunk, end="")
|
|
116
|
+
```
|
|
117
|
+
|
|
118
|
+
---
|
|
119
|
+
|
|
120
|
+
## Embeddings
|
|
121
|
+
|
|
122
|
+
### Inline embeddings
|
|
123
|
+
|
|
124
|
+
```python
|
|
125
|
+
records = {
|
|
126
|
+
"a": "Hello world",
|
|
127
|
+
"b": "Another document",
|
|
128
|
+
}
|
|
129
|
+
|
|
130
|
+
resp = helper.embed_texts(records)
|
|
131
|
+
print(resp.embeddings["a"])
|
|
132
|
+
```
|
|
133
|
+
|
|
134
|
+
### Batch embeddings
|
|
135
|
+
|
|
136
|
+
```python
|
|
137
|
+
job = helper.embed_texts(
|
|
138
|
+
big_records,
|
|
139
|
+
s3_bucket="my-bucket",
|
|
140
|
+
s3_prefix="bedrock_batch",
|
|
141
|
+
role_arn="arn:aws:iam::123456789012:role/BedrockBatchRole",
|
|
142
|
+
)
|
|
143
|
+
|
|
144
|
+
print(job.job_id)
|
|
145
|
+
```
|
|
146
|
+
|
|
147
|
+
Wait for completion:
|
|
148
|
+
|
|
149
|
+
```python
|
|
150
|
+
info = helper.wait_for_batch_job(job.job_id)
|
|
151
|
+
print(info["status"])
|
|
152
|
+
```
|
|
153
|
+
|
|
154
|
+
Parse results:
|
|
155
|
+
|
|
156
|
+
```python
|
|
157
|
+
embeddings = helper.parse_batch_embeddings(job.output_s3_uri)
|
|
158
|
+
```
|
|
159
|
+
|
|
160
|
+
---
|
|
161
|
+
|
|
162
|
+
## Batch Job Reconciliation
|
|
163
|
+
|
|
164
|
+
The helper includes a small framework for tracking and reconciling batch jobs.
|
|
165
|
+
|
|
166
|
+
```python
|
|
167
|
+
from bedrockhelper import InMemoryJobStore
|
|
168
|
+
from bedrockhelper.batch.reconcile import reconcile_batch_embedding_jobs
|
|
169
|
+
|
|
170
|
+
store = InMemoryJobStore()
|
|
171
|
+
store.add(job)
|
|
172
|
+
|
|
173
|
+
|
|
174
|
+
def process_embeddings(job, embeddings):
|
|
175
|
+
print("Processed", len(embeddings))
|
|
176
|
+
|
|
177
|
+
|
|
178
|
+
reconcile_batch_embedding_jobs(
|
|
179
|
+
helper=helper,
|
|
180
|
+
store=store,
|
|
181
|
+
process_embeddings=process_embeddings,
|
|
182
|
+
)
|
|
183
|
+
```
|
|
184
|
+
|
|
185
|
+
---
|
|
186
|
+
|
|
187
|
+
## Implementing Your Own JobStore
|
|
188
|
+
|
|
189
|
+
To persist batch jobs across restarts or multiple workers, implement the `JobStore` interface.
|
|
190
|
+
|
|
191
|
+
### Redis Example
|
|
192
|
+
|
|
193
|
+
Use Redis keys per job plus a sorted set for ordering.
|
|
194
|
+
|
|
195
|
+
Key ideas:
|
|
196
|
+
|
|
197
|
+
- Store jobs as JSON
|
|
198
|
+
- Use TTL for terminal jobs
|
|
199
|
+
- Use a sorted set by `created_at`
|
|
200
|
+
|
|
201
|
+
This is ideal for horizontally scaled workers.
|
|
202
|
+
|
|
203
|
+
### Postgres Example
|
|
204
|
+
|
|
205
|
+
Suggested schema:
|
|
206
|
+
|
|
207
|
+
```sql
|
|
208
|
+
CREATE TABLE bedrock_batch_jobs (
|
|
209
|
+
job_id TEXT PRIMARY KEY,
|
|
210
|
+
job_name TEXT NOT NULL,
|
|
211
|
+
model_id TEXT NOT NULL,
|
|
212
|
+
input_s3_uri TEXT NOT NULL,
|
|
213
|
+
output_s3_uri TEXT NOT NULL,
|
|
214
|
+
state TEXT NOT NULL,
|
|
215
|
+
created_at TIMESTAMPTZ NOT NULL,
|
|
216
|
+
last_checked_at TIMESTAMPTZ,
|
|
217
|
+
attempts INTEGER NOT NULL DEFAULT 0,
|
|
218
|
+
last_error TEXT,
|
|
219
|
+
bedrock_status TEXT,
|
|
220
|
+
processed_at TIMESTAMPTZ,
|
|
221
|
+
embeddings_count INTEGER,
|
|
222
|
+
debug JSONB NOT NULL DEFAULT '{}'::jsonb
|
|
223
|
+
);
|
|
224
|
+
```
|
|
225
|
+
|
|
226
|
+
This approach is boring, reliable, and production-friendly.
|
|
227
|
+
|
|
228
|
+
---
|
|
229
|
+
|
|
230
|
+
## Recommended Patterns
|
|
231
|
+
|
|
232
|
+
- Submit batch jobs quickly and return job IDs
|
|
233
|
+
- Reconcile in a background worker or cron
|
|
234
|
+
- Make `process_embeddings()` idempotent
|
|
235
|
+
- Prefer Postgres if you need observability
|
|
236
|
+
|
|
237
|
+
---
|
|
238
|
+
|
|
239
|
+
## License
|
|
240
|
+
|
|
241
|
+
Internal / project-specific. Adapt as needed.
|
|
@@ -0,0 +1,45 @@
|
|
|
1
|
+
[tool.poetry]
|
|
2
|
+
name = "bedrockhelper"
|
|
3
|
+
version = "1.1.0"
|
|
4
|
+
description = "Helper utilities for Amazon Bedrock (RAG, embeddings, batch jobs)"
|
|
5
|
+
authors = ["Matthew Marion <matthew@endofline.ca>"]
|
|
6
|
+
readme = "README.md"
|
|
7
|
+
license = "MIT"
|
|
8
|
+
packages = [{ include = "bedrockhelper", from = "src" }]
|
|
9
|
+
include = ["README.md", "LICENSE*"]
|
|
10
|
+
|
|
11
|
+
[tool.poetry.dependencies]
|
|
12
|
+
python = ">=3.10,<4.0"
|
|
13
|
+
boto3 = "^1.42.25"
|
|
14
|
+
botocore = { extras = ["crt"], version = "^1.42.42" }
|
|
15
|
+
boto3-stubs = "~=1.42.42"
|
|
16
|
+
ipykernel = "^7.1.0"
|
|
17
|
+
|
|
18
|
+
[tool.poetry.group.dev.dependencies]
|
|
19
|
+
pytest = "^9.0.2"
|
|
20
|
+
pytest-cov = "^7.0.0"
|
|
21
|
+
mypy = "^1.14.0"
|
|
22
|
+
ruff = "^0.15.0"
|
|
23
|
+
pre-commit = "^4.0.0"
|
|
24
|
+
|
|
25
|
+
[tool.pytest.ini_options]
|
|
26
|
+
testpaths = ["tests"]
|
|
27
|
+
|
|
28
|
+
[tool.ruff]
|
|
29
|
+
line-length = 100
|
|
30
|
+
target-version = "py310"
|
|
31
|
+
exclude = [".venv", "dist", "build"]
|
|
32
|
+
|
|
33
|
+
[tool.mypy]
|
|
34
|
+
python_version = "3.10"
|
|
35
|
+
mypy_path = ["src"]
|
|
36
|
+
packages = ["bedrockhelper"]
|
|
37
|
+
strict = true
|
|
38
|
+
warn_unused_ignores = true
|
|
39
|
+
warn_return_any = true
|
|
40
|
+
no_implicit_optional = true
|
|
41
|
+
exclude = "(^dist/|^build/|^\\.venv/)"
|
|
42
|
+
|
|
43
|
+
[build-system]
|
|
44
|
+
requires = ["poetry-core>=1.9.0"]
|
|
45
|
+
build-backend = "poetry.core.masonry.api"
|
|
@@ -0,0 +1,17 @@
|
|
|
1
|
+
from __future__ import annotations
|
|
2
|
+
|
|
3
|
+
from .main import BedrockHelper
|
|
4
|
+
from .models import InvocationMetrics, RAGResponse, EmbeddingResponse, BatchJobResponse
|
|
5
|
+
from .batch import InMemoryJobStore, reconcile_batch_embedding_jobs
|
|
6
|
+
|
|
7
|
+
__all__ = [
|
|
8
|
+
'BedrockHelper',
|
|
9
|
+
'InvocationMetrics',
|
|
10
|
+
'RAGResponse',
|
|
11
|
+
'EmbeddingResponse',
|
|
12
|
+
'BatchJobResponse',
|
|
13
|
+
'InMemoryJobStore',
|
|
14
|
+
'reconcile_batch_embedding_jobs',
|
|
15
|
+
]
|
|
16
|
+
|
|
17
|
+
__version__ = '1.1.0' # x-release-please-version
|
|
@@ -0,0 +1,12 @@
|
|
|
1
|
+
from __future__ import annotations
|
|
2
|
+
|
|
3
|
+
from .store import JobStore, InMemoryJobStore
|
|
4
|
+
from .models import StoredBatchJob
|
|
5
|
+
from .reconcile import reconcile_batch_embedding_jobs
|
|
6
|
+
|
|
7
|
+
__all__ = [
|
|
8
|
+
'JobStore',
|
|
9
|
+
'InMemoryJobStore',
|
|
10
|
+
'StoredBatchJob',
|
|
11
|
+
'reconcile_batch_embedding_jobs',
|
|
12
|
+
]
|
|
@@ -0,0 +1,41 @@
|
|
|
1
|
+
from __future__ import annotations
|
|
2
|
+
|
|
3
|
+
from dataclasses import dataclass, field
|
|
4
|
+
from datetime import datetime
|
|
5
|
+
from typing import Optional, Dict, Any
|
|
6
|
+
|
|
7
|
+
|
|
8
|
+
@dataclass(frozen=True, slots=True)
|
|
9
|
+
class BatchJobRef:
|
|
10
|
+
"""
|
|
11
|
+
Minimal, durable information needed to reconcile a Bedrock batch job later.
|
|
12
|
+
|
|
13
|
+
Intentionally excludes large/unstable provider payloads.
|
|
14
|
+
"""
|
|
15
|
+
|
|
16
|
+
job_id: str
|
|
17
|
+
job_name: str
|
|
18
|
+
model_id: str
|
|
19
|
+
input_s3_uri: str
|
|
20
|
+
output_s3_uri: str
|
|
21
|
+
|
|
22
|
+
|
|
23
|
+
@dataclass(frozen=True, slots=True)
|
|
24
|
+
class StoredBatchJob:
|
|
25
|
+
"""
|
|
26
|
+
A persisted representation of a submitted batch job.
|
|
27
|
+
|
|
28
|
+
`job` is the minimal durable reference; store/reconciler owns `state`.
|
|
29
|
+
"""
|
|
30
|
+
|
|
31
|
+
job: BatchJobRef
|
|
32
|
+
state: str # "SUBMITTED" | "RUNNING" | "COMPLETED" | "FAILED" | "PROCESSED"
|
|
33
|
+
created_at: datetime
|
|
34
|
+
last_checked_at: Optional[datetime] = None
|
|
35
|
+
attempts: int = 0
|
|
36
|
+
last_error: Optional[str] = None
|
|
37
|
+
bedrock_status: Optional[str] = None
|
|
38
|
+
processed_at: Optional[datetime] = None
|
|
39
|
+
embeddings_count: Optional[int] = None
|
|
40
|
+
|
|
41
|
+
debug: Dict[str, Any] = field(default_factory=dict)
|