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.
@@ -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)