batchgrid 0.1.0b1__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.
- batchgrid-0.1.0b1/.gitignore +12 -0
- batchgrid-0.1.0b1/PKG-INFO +159 -0
- batchgrid-0.1.0b1/README.md +133 -0
- batchgrid-0.1.0b1/pyproject.toml +38 -0
- batchgrid-0.1.0b1/src/batchgrid/__init__.py +30 -0
- batchgrid-0.1.0b1/src/batchgrid/_api.py +457 -0
- batchgrid-0.1.0b1/src/batchgrid/_cli.py +146 -0
- batchgrid-0.1.0b1/src/batchgrid/errors.py +55 -0
- batchgrid-0.1.0b1/src/batchgrid/py.typed +0 -0
- batchgrid-0.1.0b1/tests/conftest.py +154 -0
- batchgrid-0.1.0b1/tests/test_api.py +157 -0
|
@@ -0,0 +1,159 @@
|
|
|
1
|
+
Metadata-Version: 2.5
|
|
2
|
+
Name: batchgrid
|
|
3
|
+
Version: 0.1.0b1
|
|
4
|
+
Summary: Run a prompt, or a whole pipeline, over every row of a DataFrame or spreadsheet - with cost estimates, retries and resumable runs.
|
|
5
|
+
Project-URL: Homepage, https://github.com/mertguvencli/batchgrid
|
|
6
|
+
Project-URL: Source, https://github.com/mertguvencli/batchgrid/tree/main/python
|
|
7
|
+
Project-URL: Issues, https://github.com/mertguvencli/batchgrid/issues
|
|
8
|
+
Author: Mert Guvencli
|
|
9
|
+
License-Expression: MIT
|
|
10
|
+
Keywords: anthropic,batch,csv,dataframe,gemini,llm,openai,pandas
|
|
11
|
+
Classifier: Development Status :: 4 - Beta
|
|
12
|
+
Classifier: Intended Audience :: Developers
|
|
13
|
+
Classifier: Intended Audience :: Science/Research
|
|
14
|
+
Classifier: License :: OSI Approved :: MIT License
|
|
15
|
+
Classifier: Programming Language :: Python :: 3
|
|
16
|
+
Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
|
|
17
|
+
Requires-Python: >=3.9
|
|
18
|
+
Provides-Extra: dev
|
|
19
|
+
Requires-Dist: pandas>=1.5; extra == 'dev'
|
|
20
|
+
Requires-Dist: pytest>=7; extra == 'dev'
|
|
21
|
+
Provides-Extra: pandas
|
|
22
|
+
Requires-Dist: pandas>=1.5; extra == 'pandas'
|
|
23
|
+
Provides-Extra: progress
|
|
24
|
+
Requires-Dist: tqdm>=4.60; extra == 'progress'
|
|
25
|
+
Description-Content-Type: text/markdown
|
|
26
|
+
|
|
27
|
+
# batchgrid for Python
|
|
28
|
+
|
|
29
|
+
Run a prompt, or a whole pipeline, over every row of a DataFrame or spreadsheet. You get a cost
|
|
30
|
+
estimate before anything runs, parallel calls with retries and rate-limit backoff, and runs that
|
|
31
|
+
resume after an interruption.
|
|
32
|
+
|
|
33
|
+
```python
|
|
34
|
+
import pandas as pd
|
|
35
|
+
import batchgrid
|
|
36
|
+
|
|
37
|
+
df = pd.read_csv("reviews.csv")
|
|
38
|
+
result = batchgrid.run(df, "classify the sentiment and extract keywords", max_cost=5)
|
|
39
|
+
result.data # df with the new columns
|
|
40
|
+
```
|
|
41
|
+
|
|
42
|
+
> **Beta.** This package drives the [batchgrid CLI](https://www.npmjs.com/package/batchgrid), so it
|
|
43
|
+
> needs **Node.js 22 or newer**. It uses a `batchgrid` on your PATH, or fetches the CLI with `npx`
|
|
44
|
+
> on first use.
|
|
45
|
+
|
|
46
|
+
## Install
|
|
47
|
+
|
|
48
|
+
```bash
|
|
49
|
+
pip install "batchgrid[pandas,progress]"
|
|
50
|
+
```
|
|
51
|
+
|
|
52
|
+
`pandas` is needed for DataFrame input and `progress` adds a tqdm progress bar. File paths work
|
|
53
|
+
without either.
|
|
54
|
+
|
|
55
|
+
Set the key of the provider you use, either as an environment variable (`OPENAI_API_KEY`,
|
|
56
|
+
`ANTHROPIC_API_KEY`, `GOOGLE_API_KEY`, `OPENROUTER_API_KEY`, `XAI_API_KEY`) or once with
|
|
57
|
+
`npx batchgrid config --set-key openai:sk-…`.
|
|
58
|
+
|
|
59
|
+
## Plan, check, run
|
|
60
|
+
|
|
61
|
+
Planning asks a model, so the same request can produce a slightly different plan each time. Look at
|
|
62
|
+
the plan and its cost first, then run it:
|
|
63
|
+
|
|
64
|
+
```python
|
|
65
|
+
plan = batchgrid.plan(df, "classify the sentiment of each review")
|
|
66
|
+
print(plan) # the steps, the columns they write, the estimated cost
|
|
67
|
+
plan.cost_usd # 0.08
|
|
68
|
+
|
|
69
|
+
result = batchgrid.run(df, plan=plan)
|
|
70
|
+
```
|
|
71
|
+
|
|
72
|
+
A pipeline that must behave the same on every run should save the plan once and run that file.
|
|
73
|
+
Running a saved plan does not ask the planner again:
|
|
74
|
+
|
|
75
|
+
```python
|
|
76
|
+
plan.save("sentiment.json")
|
|
77
|
+
|
|
78
|
+
# later, in the pipeline
|
|
79
|
+
result = batchgrid.run("next_week.csv", plan="sentiment.json", max_cost=10)
|
|
80
|
+
```
|
|
81
|
+
|
|
82
|
+
The file is the same one `batchgrid --save-plan` writes and `batchgrid --plan` reads.
|
|
83
|
+
|
|
84
|
+
## Reference
|
|
85
|
+
|
|
86
|
+
### `batchgrid.run(data, prompt=None, *, plan=None, ...) -> Result`
|
|
87
|
+
|
|
88
|
+
| Argument | Meaning |
|
|
89
|
+
|---|---|
|
|
90
|
+
| `data` | A CSV, TSV or Excel path, a pandas DataFrame, or `None` when the request brings its own rows |
|
|
91
|
+
| `prompt` | What to do, in plain language. Pass either this or `plan` |
|
|
92
|
+
| `plan` | A `Plan`, its dict, or a saved JSON file |
|
|
93
|
+
| `model` | `"gpt-5.6"` or `"provider:model"`, e.g. `"anthropic:claude-sonnet-5"`. The default is the saved choice |
|
|
94
|
+
| `rows` | Run only the first N rows, a cheap way to try a plan |
|
|
95
|
+
| `concurrency` | The most requests running at once |
|
|
96
|
+
| `output` | Where to write the result. The extension picks the format |
|
|
97
|
+
| `max_cost` | In US dollars. The plan is priced first and nothing runs above it, or when the model has no price list |
|
|
98
|
+
| `progress` | Show a progress bar when tqdm is installed. Defaults to on |
|
|
99
|
+
| `on_event` | Called with each event the CLI reports: plan, progress, result… |
|
|
100
|
+
|
|
101
|
+
`Result` has `status` (`done` or `stopped`), `total`, `success`, `failed`, `input_tokens`,
|
|
102
|
+
`output_tokens`, `duration_ms`, `output_path`, `error_log_path`, and `data`. For DataFrame input,
|
|
103
|
+
`data` is a copy of your DataFrame with the new columns added, with its index and dtypes kept. If a
|
|
104
|
+
plan filters rows or drops columns, `data` holds the output as it was written.
|
|
105
|
+
|
|
106
|
+
Rows that still fail after their retries do not raise an error. Check `result.failed`, and look in
|
|
107
|
+
`result.error_log_path` for the reasons.
|
|
108
|
+
|
|
109
|
+
### `batchgrid.plan(data, prompt, *, model=None, rows=None) -> Plan`
|
|
110
|
+
|
|
111
|
+
Plans the request and prices it without running anything. `Plan` has `title`, `summary`,
|
|
112
|
+
`cost_usd`, `cost_formatted`, `rows`, `missing_secrets`, the raw `steps`, and `save(path)` /
|
|
113
|
+
`Plan.load(path)`.
|
|
114
|
+
|
|
115
|
+
### `batchgrid.resume(data) -> Result`
|
|
116
|
+
|
|
117
|
+
Ctrl+C (or a Jupyter interrupt) stops a run and keeps the rows that finished. Call `resume` with the
|
|
118
|
+
same file or DataFrame to finish it. Only the unfinished rows are sent to the model.
|
|
119
|
+
|
|
120
|
+
```python
|
|
121
|
+
try:
|
|
122
|
+
result = batchgrid.run(df, plan="sentiment.json")
|
|
123
|
+
except KeyboardInterrupt:
|
|
124
|
+
result = batchgrid.resume(df)
|
|
125
|
+
```
|
|
126
|
+
|
|
127
|
+
### Errors
|
|
128
|
+
|
|
129
|
+
Everything raises `batchgrid.BatchgridError`, and its `code` says what went wrong:
|
|
130
|
+
|
|
131
|
+
- `no_api_key`: no provider key is set.
|
|
132
|
+
- `invalid_plan`: the plan does not fit the data.
|
|
133
|
+
- `no_plan`: the planner asked a question instead of returning a plan. This raises `NoPlanError`,
|
|
134
|
+
and its `reply` holds the question.
|
|
135
|
+
- `missing_secrets`: the plan needs keys that are not saved. This raises `MissingSecretsError`,
|
|
136
|
+
and its `names` lists them.
|
|
137
|
+
- `cost_limit`: the run would go over `max_cost`. This raises `CostLimitError`, and its `plan`
|
|
138
|
+
holds the priced plan.
|
|
139
|
+
- `run_failed`: the run itself failed.
|
|
140
|
+
|
|
141
|
+
### Choosing the CLI
|
|
142
|
+
|
|
143
|
+
The package looks for the CLI in this order:
|
|
144
|
+
|
|
145
|
+
1. the `cli=[...]` argument
|
|
146
|
+
2. the `BATCHGRID_CLI` environment variable, e.g. `node /path/to/cli/dist/index.js`
|
|
147
|
+
3. a `batchgrid` on the PATH
|
|
148
|
+
4. `npx "batchgrid@>=0.2.0 <1"`
|
|
149
|
+
|
|
150
|
+
## Development
|
|
151
|
+
|
|
152
|
+
```bash
|
|
153
|
+
pnpm --filter batchgrid build # the tests drive the real CLI
|
|
154
|
+
cd python
|
|
155
|
+
uv venv && uv pip install -e ".[dev,progress]"
|
|
156
|
+
.venv/bin/python -m pytest
|
|
157
|
+
```
|
|
158
|
+
|
|
159
|
+
The tests replace only the model: a local server stands in for the OpenAI API.
|
|
@@ -0,0 +1,133 @@
|
|
|
1
|
+
# batchgrid for Python
|
|
2
|
+
|
|
3
|
+
Run a prompt, or a whole pipeline, over every row of a DataFrame or spreadsheet. You get a cost
|
|
4
|
+
estimate before anything runs, parallel calls with retries and rate-limit backoff, and runs that
|
|
5
|
+
resume after an interruption.
|
|
6
|
+
|
|
7
|
+
```python
|
|
8
|
+
import pandas as pd
|
|
9
|
+
import batchgrid
|
|
10
|
+
|
|
11
|
+
df = pd.read_csv("reviews.csv")
|
|
12
|
+
result = batchgrid.run(df, "classify the sentiment and extract keywords", max_cost=5)
|
|
13
|
+
result.data # df with the new columns
|
|
14
|
+
```
|
|
15
|
+
|
|
16
|
+
> **Beta.** This package drives the [batchgrid CLI](https://www.npmjs.com/package/batchgrid), so it
|
|
17
|
+
> needs **Node.js 22 or newer**. It uses a `batchgrid` on your PATH, or fetches the CLI with `npx`
|
|
18
|
+
> on first use.
|
|
19
|
+
|
|
20
|
+
## Install
|
|
21
|
+
|
|
22
|
+
```bash
|
|
23
|
+
pip install "batchgrid[pandas,progress]"
|
|
24
|
+
```
|
|
25
|
+
|
|
26
|
+
`pandas` is needed for DataFrame input and `progress` adds a tqdm progress bar. File paths work
|
|
27
|
+
without either.
|
|
28
|
+
|
|
29
|
+
Set the key of the provider you use, either as an environment variable (`OPENAI_API_KEY`,
|
|
30
|
+
`ANTHROPIC_API_KEY`, `GOOGLE_API_KEY`, `OPENROUTER_API_KEY`, `XAI_API_KEY`) or once with
|
|
31
|
+
`npx batchgrid config --set-key openai:sk-…`.
|
|
32
|
+
|
|
33
|
+
## Plan, check, run
|
|
34
|
+
|
|
35
|
+
Planning asks a model, so the same request can produce a slightly different plan each time. Look at
|
|
36
|
+
the plan and its cost first, then run it:
|
|
37
|
+
|
|
38
|
+
```python
|
|
39
|
+
plan = batchgrid.plan(df, "classify the sentiment of each review")
|
|
40
|
+
print(plan) # the steps, the columns they write, the estimated cost
|
|
41
|
+
plan.cost_usd # 0.08
|
|
42
|
+
|
|
43
|
+
result = batchgrid.run(df, plan=plan)
|
|
44
|
+
```
|
|
45
|
+
|
|
46
|
+
A pipeline that must behave the same on every run should save the plan once and run that file.
|
|
47
|
+
Running a saved plan does not ask the planner again:
|
|
48
|
+
|
|
49
|
+
```python
|
|
50
|
+
plan.save("sentiment.json")
|
|
51
|
+
|
|
52
|
+
# later, in the pipeline
|
|
53
|
+
result = batchgrid.run("next_week.csv", plan="sentiment.json", max_cost=10)
|
|
54
|
+
```
|
|
55
|
+
|
|
56
|
+
The file is the same one `batchgrid --save-plan` writes and `batchgrid --plan` reads.
|
|
57
|
+
|
|
58
|
+
## Reference
|
|
59
|
+
|
|
60
|
+
### `batchgrid.run(data, prompt=None, *, plan=None, ...) -> Result`
|
|
61
|
+
|
|
62
|
+
| Argument | Meaning |
|
|
63
|
+
|---|---|
|
|
64
|
+
| `data` | A CSV, TSV or Excel path, a pandas DataFrame, or `None` when the request brings its own rows |
|
|
65
|
+
| `prompt` | What to do, in plain language. Pass either this or `plan` |
|
|
66
|
+
| `plan` | A `Plan`, its dict, or a saved JSON file |
|
|
67
|
+
| `model` | `"gpt-5.6"` or `"provider:model"`, e.g. `"anthropic:claude-sonnet-5"`. The default is the saved choice |
|
|
68
|
+
| `rows` | Run only the first N rows, a cheap way to try a plan |
|
|
69
|
+
| `concurrency` | The most requests running at once |
|
|
70
|
+
| `output` | Where to write the result. The extension picks the format |
|
|
71
|
+
| `max_cost` | In US dollars. The plan is priced first and nothing runs above it, or when the model has no price list |
|
|
72
|
+
| `progress` | Show a progress bar when tqdm is installed. Defaults to on |
|
|
73
|
+
| `on_event` | Called with each event the CLI reports: plan, progress, result… |
|
|
74
|
+
|
|
75
|
+
`Result` has `status` (`done` or `stopped`), `total`, `success`, `failed`, `input_tokens`,
|
|
76
|
+
`output_tokens`, `duration_ms`, `output_path`, `error_log_path`, and `data`. For DataFrame input,
|
|
77
|
+
`data` is a copy of your DataFrame with the new columns added, with its index and dtypes kept. If a
|
|
78
|
+
plan filters rows or drops columns, `data` holds the output as it was written.
|
|
79
|
+
|
|
80
|
+
Rows that still fail after their retries do not raise an error. Check `result.failed`, and look in
|
|
81
|
+
`result.error_log_path` for the reasons.
|
|
82
|
+
|
|
83
|
+
### `batchgrid.plan(data, prompt, *, model=None, rows=None) -> Plan`
|
|
84
|
+
|
|
85
|
+
Plans the request and prices it without running anything. `Plan` has `title`, `summary`,
|
|
86
|
+
`cost_usd`, `cost_formatted`, `rows`, `missing_secrets`, the raw `steps`, and `save(path)` /
|
|
87
|
+
`Plan.load(path)`.
|
|
88
|
+
|
|
89
|
+
### `batchgrid.resume(data) -> Result`
|
|
90
|
+
|
|
91
|
+
Ctrl+C (or a Jupyter interrupt) stops a run and keeps the rows that finished. Call `resume` with the
|
|
92
|
+
same file or DataFrame to finish it. Only the unfinished rows are sent to the model.
|
|
93
|
+
|
|
94
|
+
```python
|
|
95
|
+
try:
|
|
96
|
+
result = batchgrid.run(df, plan="sentiment.json")
|
|
97
|
+
except KeyboardInterrupt:
|
|
98
|
+
result = batchgrid.resume(df)
|
|
99
|
+
```
|
|
100
|
+
|
|
101
|
+
### Errors
|
|
102
|
+
|
|
103
|
+
Everything raises `batchgrid.BatchgridError`, and its `code` says what went wrong:
|
|
104
|
+
|
|
105
|
+
- `no_api_key`: no provider key is set.
|
|
106
|
+
- `invalid_plan`: the plan does not fit the data.
|
|
107
|
+
- `no_plan`: the planner asked a question instead of returning a plan. This raises `NoPlanError`,
|
|
108
|
+
and its `reply` holds the question.
|
|
109
|
+
- `missing_secrets`: the plan needs keys that are not saved. This raises `MissingSecretsError`,
|
|
110
|
+
and its `names` lists them.
|
|
111
|
+
- `cost_limit`: the run would go over `max_cost`. This raises `CostLimitError`, and its `plan`
|
|
112
|
+
holds the priced plan.
|
|
113
|
+
- `run_failed`: the run itself failed.
|
|
114
|
+
|
|
115
|
+
### Choosing the CLI
|
|
116
|
+
|
|
117
|
+
The package looks for the CLI in this order:
|
|
118
|
+
|
|
119
|
+
1. the `cli=[...]` argument
|
|
120
|
+
2. the `BATCHGRID_CLI` environment variable, e.g. `node /path/to/cli/dist/index.js`
|
|
121
|
+
3. a `batchgrid` on the PATH
|
|
122
|
+
4. `npx "batchgrid@>=0.2.0 <1"`
|
|
123
|
+
|
|
124
|
+
## Development
|
|
125
|
+
|
|
126
|
+
```bash
|
|
127
|
+
pnpm --filter batchgrid build # the tests drive the real CLI
|
|
128
|
+
cd python
|
|
129
|
+
uv venv && uv pip install -e ".[dev,progress]"
|
|
130
|
+
.venv/bin/python -m pytest
|
|
131
|
+
```
|
|
132
|
+
|
|
133
|
+
The tests replace only the model: a local server stands in for the OpenAI API.
|
|
@@ -0,0 +1,38 @@
|
|
|
1
|
+
[build-system]
|
|
2
|
+
requires = ["hatchling>=1.24"]
|
|
3
|
+
build-backend = "hatchling.build"
|
|
4
|
+
|
|
5
|
+
[project]
|
|
6
|
+
name = "batchgrid"
|
|
7
|
+
version = "0.1.0b1"
|
|
8
|
+
description = "Run a prompt, or a whole pipeline, over every row of a DataFrame or spreadsheet - with cost estimates, retries and resumable runs."
|
|
9
|
+
readme = "README.md"
|
|
10
|
+
license = "MIT"
|
|
11
|
+
requires-python = ">=3.9"
|
|
12
|
+
authors = [{ name = "Mert Guvencli" }]
|
|
13
|
+
keywords = ["llm", "batch", "pandas", "dataframe", "csv", "openai", "anthropic", "gemini"]
|
|
14
|
+
classifiers = [
|
|
15
|
+
"Development Status :: 4 - Beta",
|
|
16
|
+
"Intended Audience :: Developers",
|
|
17
|
+
"Intended Audience :: Science/Research",
|
|
18
|
+
"License :: OSI Approved :: MIT License",
|
|
19
|
+
"Programming Language :: Python :: 3",
|
|
20
|
+
"Topic :: Scientific/Engineering :: Artificial Intelligence",
|
|
21
|
+
]
|
|
22
|
+
dependencies = []
|
|
23
|
+
|
|
24
|
+
[project.optional-dependencies]
|
|
25
|
+
pandas = ["pandas>=1.5"]
|
|
26
|
+
progress = ["tqdm>=4.60"]
|
|
27
|
+
dev = ["pytest>=7", "pandas>=1.5"]
|
|
28
|
+
|
|
29
|
+
[project.urls]
|
|
30
|
+
Homepage = "https://github.com/mertguvencli/batchgrid"
|
|
31
|
+
Source = "https://github.com/mertguvencli/batchgrid/tree/main/python"
|
|
32
|
+
Issues = "https://github.com/mertguvencli/batchgrid/issues"
|
|
33
|
+
|
|
34
|
+
[tool.hatch.build.targets.wheel]
|
|
35
|
+
packages = ["src/batchgrid"]
|
|
36
|
+
|
|
37
|
+
[tool.pytest.ini_options]
|
|
38
|
+
testpaths = ["tests"]
|
|
@@ -0,0 +1,30 @@
|
|
|
1
|
+
"""batchgrid - run a prompt, or a whole pipeline, over every row of a DataFrame or spreadsheet.
|
|
2
|
+
|
|
3
|
+
import batchgrid
|
|
4
|
+
|
|
5
|
+
result = batchgrid.run(df, "classify the sentiment of each review", max_cost=5)
|
|
6
|
+
result.data # df with the new columns
|
|
7
|
+
|
|
8
|
+
The work is done by the batchgrid CLI (Node.js 22+), which this package starts
|
|
9
|
+
and follows; see find_cli() for how it is located.
|
|
10
|
+
"""
|
|
11
|
+
|
|
12
|
+
from ._api import Plan, Result, plan, resume, run
|
|
13
|
+
from ._cli import CLI_SPEC, find_cli
|
|
14
|
+
from .errors import BatchgridError, CostLimitError, MissingSecretsError, NoPlanError
|
|
15
|
+
|
|
16
|
+
__version__ = "0.1.0b1"
|
|
17
|
+
|
|
18
|
+
__all__ = [
|
|
19
|
+
"BatchgridError",
|
|
20
|
+
"CLI_SPEC",
|
|
21
|
+
"CostLimitError",
|
|
22
|
+
"MissingSecretsError",
|
|
23
|
+
"NoPlanError",
|
|
24
|
+
"Plan",
|
|
25
|
+
"Result",
|
|
26
|
+
"find_cli",
|
|
27
|
+
"plan",
|
|
28
|
+
"resume",
|
|
29
|
+
"run",
|
|
30
|
+
]
|
|
@@ -0,0 +1,457 @@
|
|
|
1
|
+
"""plan(), run() and resume() - the Python face of the batchgrid CLI."""
|
|
2
|
+
|
|
3
|
+
from __future__ import annotations
|
|
4
|
+
|
|
5
|
+
import json
|
|
6
|
+
import os
|
|
7
|
+
import tempfile
|
|
8
|
+
from dataclasses import dataclass, field
|
|
9
|
+
from pathlib import Path
|
|
10
|
+
from typing import TYPE_CHECKING, Any, Callable, Dict, List, Mapping, Optional, Sequence, Tuple, Union
|
|
11
|
+
|
|
12
|
+
from ._cli import Event, find_cli, stream_events
|
|
13
|
+
from .errors import BatchgridError, CostLimitError, MissingSecretsError, NoPlanError
|
|
14
|
+
|
|
15
|
+
if TYPE_CHECKING:
|
|
16
|
+
import pandas as pd
|
|
17
|
+
|
|
18
|
+
PathLike = Union[str, "os.PathLike[str]"]
|
|
19
|
+
#: A file path, or a pandas DataFrame
|
|
20
|
+
Data = Union[PathLike, "pd.DataFrame"]
|
|
21
|
+
EventHandler = Callable[[Event], None]
|
|
22
|
+
|
|
23
|
+
|
|
24
|
+
@dataclass
|
|
25
|
+
class Plan:
|
|
26
|
+
"""A workflow the planner wrote, with what it will cost.
|
|
27
|
+
|
|
28
|
+
Save it and pass it to ``run(plan=...)`` to repeat the exact same steps
|
|
29
|
+
later - running a saved plan does not ask the planner again.
|
|
30
|
+
"""
|
|
31
|
+
|
|
32
|
+
steps: Dict[str, Any]
|
|
33
|
+
"""The plan itself, as the engine reads it."""
|
|
34
|
+
summary: str = ""
|
|
35
|
+
"""One readable paragraph: the steps and the columns they write."""
|
|
36
|
+
cost_usd: Optional[float] = None
|
|
37
|
+
"""Estimated cost of the run in US dollars, or None when the model has no price list."""
|
|
38
|
+
cost_formatted: str = ""
|
|
39
|
+
rows: int = 0
|
|
40
|
+
"""How many rows the run will process."""
|
|
41
|
+
missing_secrets: List[str] = field(default_factory=list)
|
|
42
|
+
"""Service keys the plan needs that are not saved yet."""
|
|
43
|
+
cost: Dict[str, Any] = field(default_factory=dict)
|
|
44
|
+
"""The full estimate: input and output tokens, model, pricing availability."""
|
|
45
|
+
|
|
46
|
+
@property
|
|
47
|
+
def title(self) -> str:
|
|
48
|
+
return str(self.steps.get("title") or self.steps.get("intent") or "")
|
|
49
|
+
|
|
50
|
+
def save(self, path: PathLike) -> Path:
|
|
51
|
+
"""Write the plan as JSON - the same file ``batchgrid --plan`` reads."""
|
|
52
|
+
target = Path(path)
|
|
53
|
+
target.write_text(json.dumps(self.steps, indent=2) + "\n", encoding="utf-8")
|
|
54
|
+
return target
|
|
55
|
+
|
|
56
|
+
@classmethod
|
|
57
|
+
def load(cls, path: PathLike) -> "Plan":
|
|
58
|
+
"""Read a plan saved with ``save()`` or ``batchgrid --save-plan``."""
|
|
59
|
+
data = json.loads(Path(path).read_text(encoding="utf-8"))
|
|
60
|
+
# A saved workflow exported from the app wraps the plan: { name, plan }
|
|
61
|
+
if isinstance(data, dict) and "steps" not in data and isinstance(data.get("plan"), dict):
|
|
62
|
+
data = data["plan"]
|
|
63
|
+
return cls(steps=data)
|
|
64
|
+
|
|
65
|
+
@classmethod
|
|
66
|
+
def _from_event(cls, event: Event) -> "Plan":
|
|
67
|
+
cost = event.get("cost") or {}
|
|
68
|
+
assert isinstance(cost, dict)
|
|
69
|
+
priced = bool(cost.get("pricingAvailable"))
|
|
70
|
+
return cls(
|
|
71
|
+
steps=event["plan"], # type: ignore[arg-type]
|
|
72
|
+
summary=str(event.get("summary") or ""),
|
|
73
|
+
cost_usd=float(cost["estimatedCostUsd"]) if priced and "estimatedCostUsd" in cost else None,
|
|
74
|
+
cost_formatted=str(event.get("costFormatted") or ""),
|
|
75
|
+
rows=int(event.get("rows") or 0), # type: ignore[arg-type]
|
|
76
|
+
missing_secrets=list(event.get("missingSecrets") or []), # type: ignore[arg-type]
|
|
77
|
+
cost=cost,
|
|
78
|
+
)
|
|
79
|
+
|
|
80
|
+
def __str__(self) -> str:
|
|
81
|
+
parts = [self.summary or self.title]
|
|
82
|
+
if self.cost_formatted:
|
|
83
|
+
parts.append(self.cost_formatted)
|
|
84
|
+
return "\n\n".join(parts)
|
|
85
|
+
|
|
86
|
+
def _repr_markdown_(self) -> str:
|
|
87
|
+
cost = self.cost_formatted or "no estimate"
|
|
88
|
+
return f"**{self.title}** · {self.rows:,} rows · {cost}\n\n```\n{self.summary}\n```"
|
|
89
|
+
|
|
90
|
+
|
|
91
|
+
@dataclass
|
|
92
|
+
class Result:
|
|
93
|
+
"""How a run ended."""
|
|
94
|
+
|
|
95
|
+
status: str
|
|
96
|
+
"""``done``, or ``stopped`` when it was interrupted - resume() picks it up."""
|
|
97
|
+
total: int
|
|
98
|
+
success: int
|
|
99
|
+
failed: int
|
|
100
|
+
"""Rows that failed after their retries; ``error_log_path`` has the reasons."""
|
|
101
|
+
input_tokens: int
|
|
102
|
+
output_tokens: int
|
|
103
|
+
duration_ms: int
|
|
104
|
+
output_path: Optional[Path]
|
|
105
|
+
"""The file the run wrote: the input with the new columns, or a zip when it produced files."""
|
|
106
|
+
error_log_path: Optional[Path]
|
|
107
|
+
data: Optional["pd.DataFrame"] = None
|
|
108
|
+
"""The result as a DataFrame, when the input was one."""
|
|
109
|
+
|
|
110
|
+
@classmethod
|
|
111
|
+
def _from_event(cls, event: Event) -> "Result":
|
|
112
|
+
tokens = event.get("tokens") or {}
|
|
113
|
+
assert isinstance(tokens, dict)
|
|
114
|
+
output = event.get("outputPath")
|
|
115
|
+
errors = event.get("errorLogPath")
|
|
116
|
+
return cls(
|
|
117
|
+
status=str(event.get("status")),
|
|
118
|
+
total=int(event.get("total") or 0), # type: ignore[arg-type]
|
|
119
|
+
success=int(event.get("success") or 0), # type: ignore[arg-type]
|
|
120
|
+
failed=int(event.get("failed") or 0), # type: ignore[arg-type]
|
|
121
|
+
input_tokens=int(tokens.get("input") or 0),
|
|
122
|
+
output_tokens=int(tokens.get("output") or 0),
|
|
123
|
+
duration_ms=int(event.get("durationMs") or 0), # type: ignore[arg-type]
|
|
124
|
+
output_path=Path(str(output)) if output else None,
|
|
125
|
+
error_log_path=Path(str(errors)) if errors else None,
|
|
126
|
+
)
|
|
127
|
+
|
|
128
|
+
|
|
129
|
+
def plan(
|
|
130
|
+
data: Optional[Data],
|
|
131
|
+
prompt: str,
|
|
132
|
+
*,
|
|
133
|
+
model: Optional[str] = None,
|
|
134
|
+
rows: Optional[int] = None,
|
|
135
|
+
on_event: Optional[EventHandler] = None,
|
|
136
|
+
cli: Optional[Sequence[str]] = None,
|
|
137
|
+
env: Optional[Mapping[str, str]] = None,
|
|
138
|
+
) -> Plan:
|
|
139
|
+
"""Ask the planner for a workflow and price it, without running anything.
|
|
140
|
+
|
|
141
|
+
``data`` is a CSV/TSV/Excel path or a pandas DataFrame, or None when the
|
|
142
|
+
request brings its own rows ("find 50 AI startups in Berlin").
|
|
143
|
+
``model`` is a model name or ``provider:model``, e.g.
|
|
144
|
+
``"anthropic:claude-sonnet-5"``; the default is the one saved in
|
|
145
|
+
``batchgrid config``.
|
|
146
|
+
"""
|
|
147
|
+
with _input_file(data) as (path, _):
|
|
148
|
+
events = _run_cli(
|
|
149
|
+
[*_input_args(path), "--prompt", prompt, *_common_args(model, rows, None)],
|
|
150
|
+
on_event,
|
|
151
|
+
cli,
|
|
152
|
+
env,
|
|
153
|
+
)
|
|
154
|
+
return _expect_plan(events)
|
|
155
|
+
|
|
156
|
+
|
|
157
|
+
def run(
|
|
158
|
+
data: Optional[Data],
|
|
159
|
+
prompt: Optional[str] = None,
|
|
160
|
+
*,
|
|
161
|
+
plan: Union[Plan, Mapping[str, Any], PathLike, None] = None,
|
|
162
|
+
model: Optional[str] = None,
|
|
163
|
+
rows: Optional[int] = None,
|
|
164
|
+
concurrency: Optional[int] = None,
|
|
165
|
+
output: Optional[PathLike] = None,
|
|
166
|
+
max_cost: Optional[float] = None,
|
|
167
|
+
progress: bool = True,
|
|
168
|
+
on_event: Optional[EventHandler] = None,
|
|
169
|
+
cli: Optional[Sequence[str]] = None,
|
|
170
|
+
env: Optional[Mapping[str, str]] = None,
|
|
171
|
+
) -> Result:
|
|
172
|
+
"""Plan a request and run it over every row - or run a plan you already have.
|
|
173
|
+
|
|
174
|
+
Pass ``prompt`` to have the planner write the workflow, or ``plan`` (a
|
|
175
|
+
``Plan``, its dict, or a saved JSON file) to repeat one exactly. With
|
|
176
|
+
``max_cost``, the plan is priced first and nothing runs when the estimate
|
|
177
|
+
is higher - or when the model has no price list to estimate with.
|
|
178
|
+
|
|
179
|
+
For a DataFrame, the result's ``data`` is that DataFrame with the new
|
|
180
|
+
columns added. For a file, the output is written next to it (or to
|
|
181
|
+
``output``) and ``output_path`` says where.
|
|
182
|
+
|
|
183
|
+
Ctrl+C stops the run and keeps the rows that finished; ``resume()`` with
|
|
184
|
+
the same data carries on from there.
|
|
185
|
+
"""
|
|
186
|
+
if (prompt is None) == (plan is None):
|
|
187
|
+
raise ValueError("Pass either a prompt or a plan.")
|
|
188
|
+
|
|
189
|
+
with _input_file(data) as (path, frame):
|
|
190
|
+
if max_cost is not None:
|
|
191
|
+
# Price it on this data first, without running; the run then repeats this exact plan
|
|
192
|
+
with _plan_file(plan) as plan_path:
|
|
193
|
+
request = ["--plan", plan_path] if plan_path else ["--prompt", prompt or ""]
|
|
194
|
+
priced = _expect_plan(
|
|
195
|
+
_run_cli([*_input_args(path), *request, *_common_args(model, rows, None)], on_event, cli, env)
|
|
196
|
+
)
|
|
197
|
+
_check_cost(priced, max_cost)
|
|
198
|
+
plan = priced
|
|
199
|
+
|
|
200
|
+
with _plan_file(plan) as plan_path:
|
|
201
|
+
request = ["--plan", plan_path] if plan_path else ["--prompt", prompt or ""]
|
|
202
|
+
args = [*_input_args(path), *request, "--yes", *_common_args(model, rows, concurrency)]
|
|
203
|
+
out = output if output is not None else (_frame_output(path) if frame is not None else None)
|
|
204
|
+
if out is not None:
|
|
205
|
+
args += ["--output", os.fspath(out)]
|
|
206
|
+
result = _expect_result(_run_cli(args, on_event, cli, env, progress=progress))
|
|
207
|
+
if frame is not None:
|
|
208
|
+
result.data = _merge_output(frame, result.output_path)
|
|
209
|
+
return result
|
|
210
|
+
|
|
211
|
+
|
|
212
|
+
def resume(
|
|
213
|
+
data: Data,
|
|
214
|
+
*,
|
|
215
|
+
output: Optional[PathLike] = None,
|
|
216
|
+
progress: bool = True,
|
|
217
|
+
on_event: Optional[EventHandler] = None,
|
|
218
|
+
cli: Optional[Sequence[str]] = None,
|
|
219
|
+
env: Optional[Mapping[str, str]] = None,
|
|
220
|
+
) -> Result:
|
|
221
|
+
"""Finish a run that was stopped, on the same file or DataFrame.
|
|
222
|
+
|
|
223
|
+
Only the rows that had not finished are sent to the model.
|
|
224
|
+
"""
|
|
225
|
+
with _input_file(data) as (path, frame):
|
|
226
|
+
out = output if output is not None else (_frame_output(path) if frame is not None else None)
|
|
227
|
+
args = [*_input_args(path), "--resume"]
|
|
228
|
+
if out is not None:
|
|
229
|
+
args += ["--output", os.fspath(out)]
|
|
230
|
+
result = _expect_result(_run_cli(args, on_event, cli, env, progress=progress))
|
|
231
|
+
if frame is not None:
|
|
232
|
+
result.data = _merge_output(frame, result.output_path)
|
|
233
|
+
return result
|
|
234
|
+
|
|
235
|
+
|
|
236
|
+
# ── Running the CLI ────────────────────────────────────────────────────
|
|
237
|
+
|
|
238
|
+
|
|
239
|
+
def _run_cli(
|
|
240
|
+
args: List[str],
|
|
241
|
+
on_event: Optional[EventHandler],
|
|
242
|
+
cli: Optional[Sequence[str]],
|
|
243
|
+
env: Optional[Mapping[str, str]],
|
|
244
|
+
progress: bool = False,
|
|
245
|
+
) -> List[Event]:
|
|
246
|
+
bar = _ProgressBar() if progress else None
|
|
247
|
+
events: List[Event] = []
|
|
248
|
+
try:
|
|
249
|
+
for event in stream_events(find_cli(cli), [*args, "--json"], env):
|
|
250
|
+
events.append(event)
|
|
251
|
+
if bar is not None and event.get("type") == "progress":
|
|
252
|
+
bar.update(event)
|
|
253
|
+
if on_event is not None and event.get("type") != "exit":
|
|
254
|
+
on_event(event)
|
|
255
|
+
finally:
|
|
256
|
+
if bar is not None:
|
|
257
|
+
bar.close()
|
|
258
|
+
return events
|
|
259
|
+
|
|
260
|
+
|
|
261
|
+
def _check_cost(p: Plan, max_cost: float) -> None:
|
|
262
|
+
if p.cost_usd is None:
|
|
263
|
+
# A plan without model calls costs nothing; one without a price list cannot be held to a limit
|
|
264
|
+
if p.cost and not p.cost.get("pricingAvailable", True):
|
|
265
|
+
raise CostLimitError(
|
|
266
|
+
f"There is no price list for {p.cost.get('model')}, so the run cannot be held to max_cost={max_cost}.",
|
|
267
|
+
p,
|
|
268
|
+
)
|
|
269
|
+
return
|
|
270
|
+
if p.cost_usd > max_cost:
|
|
271
|
+
raise CostLimitError(
|
|
272
|
+
f"The run is estimated at ${p.cost_usd:.4f}, above max_cost=${max_cost:.4f}. Nothing was run.",
|
|
273
|
+
p,
|
|
274
|
+
)
|
|
275
|
+
|
|
276
|
+
|
|
277
|
+
def _failure(events: List[Event]) -> BatchgridError:
|
|
278
|
+
"""The error the CLI reported - or, without one, what it printed on stderr"""
|
|
279
|
+
exit_event = events[-1] if events and events[-1].get("type") == "exit" else {}
|
|
280
|
+
stderr = str(exit_event.get("stderr") or "")
|
|
281
|
+
error = next((e for e in reversed(events) if e.get("type") == "error"), None)
|
|
282
|
+
if error is None:
|
|
283
|
+
hint = " The installed batchgrid CLI is too old - update it." if "unknown option" in stderr else ""
|
|
284
|
+
return BatchgridError("cli_failed", f"The batchgrid CLI stopped unexpectedly.{hint}\n{stderr}".rstrip(), stderr)
|
|
285
|
+
|
|
286
|
+
code, message = str(error.get("code")), str(error.get("message"))
|
|
287
|
+
if code == "no_plan":
|
|
288
|
+
reply = next((str(e["text"]) for e in reversed(events) if e.get("type") == "message"), None)
|
|
289
|
+
return NoPlanError(message, reply)
|
|
290
|
+
if code == "missing_secrets":
|
|
291
|
+
planned = next((e for e in events if e.get("type") == "plan"), {})
|
|
292
|
+
return MissingSecretsError(message, list(planned.get("missingSecrets") or [])) # type: ignore[arg-type]
|
|
293
|
+
return BatchgridError(code, message, stderr)
|
|
294
|
+
|
|
295
|
+
|
|
296
|
+
def _expect_plan(events: List[Event]) -> Plan:
|
|
297
|
+
planned = next((e for e in events if e.get("type") == "plan"), None)
|
|
298
|
+
if planned is None or any(e.get("type") == "error" for e in events):
|
|
299
|
+
raise _failure(events)
|
|
300
|
+
return Plan._from_event(planned)
|
|
301
|
+
|
|
302
|
+
|
|
303
|
+
def _expect_result(events: List[Event]) -> Result:
|
|
304
|
+
result = next((e for e in reversed(events) if e.get("type") == "result"), None)
|
|
305
|
+
if result is None or result.get("status") == "error":
|
|
306
|
+
raise _failure(events)
|
|
307
|
+
parsed = Result._from_event(result)
|
|
308
|
+
if any(e.get("type") == "error" for e in events) and parsed.status == "done":
|
|
309
|
+
# The rows ran, but something after them failed - writing the output
|
|
310
|
+
raise _failure(events)
|
|
311
|
+
return parsed
|
|
312
|
+
|
|
313
|
+
|
|
314
|
+
def _input_args(path: Optional[str]) -> List[str]:
|
|
315
|
+
return [path] if path else []
|
|
316
|
+
|
|
317
|
+
|
|
318
|
+
def _common_args(model: Optional[str], rows: Optional[int], concurrency: Optional[int]) -> List[str]:
|
|
319
|
+
args: List[str] = []
|
|
320
|
+
if model:
|
|
321
|
+
args += ["--model", model]
|
|
322
|
+
if rows is not None:
|
|
323
|
+
args += ["--rows", str(rows)]
|
|
324
|
+
if concurrency is not None:
|
|
325
|
+
args += ["--concurrency", str(concurrency)]
|
|
326
|
+
return args
|
|
327
|
+
|
|
328
|
+
|
|
329
|
+
# ── Files in and out ───────────────────────────────────────────────────
|
|
330
|
+
|
|
331
|
+
|
|
332
|
+
class _TempDir:
|
|
333
|
+
"""One scratch directory per call, removed on exit"""
|
|
334
|
+
|
|
335
|
+
def __init__(self) -> None:
|
|
336
|
+
self.path: Optional[str] = None
|
|
337
|
+
|
|
338
|
+
def get(self) -> str:
|
|
339
|
+
if self.path is None:
|
|
340
|
+
self.path = tempfile.mkdtemp(prefix="batchgrid-")
|
|
341
|
+
return self.path
|
|
342
|
+
|
|
343
|
+
def cleanup(self) -> None:
|
|
344
|
+
if self.path is not None:
|
|
345
|
+
import shutil
|
|
346
|
+
|
|
347
|
+
shutil.rmtree(self.path, ignore_errors=True)
|
|
348
|
+
|
|
349
|
+
|
|
350
|
+
class _input_file:
|
|
351
|
+
"""The data as a path the CLI can open: a DataFrame is written to a CSV first"""
|
|
352
|
+
|
|
353
|
+
def __init__(self, data: Optional[Data]) -> None:
|
|
354
|
+
self.data = data
|
|
355
|
+
self.tmp = _TempDir()
|
|
356
|
+
|
|
357
|
+
def __enter__(self) -> Tuple[Optional[str], Optional["pd.DataFrame"]]:
|
|
358
|
+
if self.data is None:
|
|
359
|
+
return None, None
|
|
360
|
+
if _is_dataframe(self.data):
|
|
361
|
+
frame = self.data
|
|
362
|
+
path = os.path.join(self.tmp.get(), "data.csv")
|
|
363
|
+
# The same DataFrame writes the same bytes, so its checkpoint is found again on resume()
|
|
364
|
+
frame.rename(columns=str).to_csv(path, index=False) # type: ignore[union-attr]
|
|
365
|
+
return path, frame # type: ignore[return-value]
|
|
366
|
+
path = os.fspath(self.data) # type: ignore[arg-type]
|
|
367
|
+
if not os.path.exists(path):
|
|
368
|
+
raise FileNotFoundError(path)
|
|
369
|
+
return os.path.abspath(path), None
|
|
370
|
+
|
|
371
|
+
def __exit__(self, *exc: object) -> None:
|
|
372
|
+
self.tmp.cleanup()
|
|
373
|
+
|
|
374
|
+
|
|
375
|
+
class _plan_file:
|
|
376
|
+
"""A plan as a JSON file for --plan, or None when the planner writes one"""
|
|
377
|
+
|
|
378
|
+
def __init__(self, value: Union[Plan, Mapping[str, Any], PathLike, None]) -> None:
|
|
379
|
+
self.value = value
|
|
380
|
+
self.tmp = _TempDir()
|
|
381
|
+
|
|
382
|
+
def __enter__(self) -> Optional[str]:
|
|
383
|
+
if self.value is None:
|
|
384
|
+
return None
|
|
385
|
+
if isinstance(self.value, Plan):
|
|
386
|
+
steps: Mapping[str, Any] = self.value.steps
|
|
387
|
+
elif isinstance(self.value, Mapping):
|
|
388
|
+
steps = self.value
|
|
389
|
+
else:
|
|
390
|
+
return os.path.abspath(os.fspath(self.value))
|
|
391
|
+
path = os.path.join(self.tmp.get(), "plan.json")
|
|
392
|
+
with open(path, "w", encoding="utf-8") as fh:
|
|
393
|
+
json.dump(steps, fh)
|
|
394
|
+
return path
|
|
395
|
+
|
|
396
|
+
def __exit__(self, *exc: object) -> None:
|
|
397
|
+
self.tmp.cleanup()
|
|
398
|
+
|
|
399
|
+
|
|
400
|
+
def _frame_output(input_path: Optional[str]) -> str:
|
|
401
|
+
# Next to the temp input, so it is removed with it once read back
|
|
402
|
+
return os.path.join(os.path.dirname(input_path or ""), "output.csv")
|
|
403
|
+
|
|
404
|
+
|
|
405
|
+
def _is_dataframe(value: object) -> bool:
|
|
406
|
+
return type(value).__module__.startswith("pandas") and type(value).__name__ == "DataFrame"
|
|
407
|
+
|
|
408
|
+
|
|
409
|
+
def _merge_output(frame: "pd.DataFrame", output_path: Optional[Path]) -> Optional["pd.DataFrame"]:
|
|
410
|
+
"""The input DataFrame with the run's new columns added.
|
|
411
|
+
|
|
412
|
+
The original columns keep their dtypes and the index is kept. When the
|
|
413
|
+
run reshaped the sheet - filtered rows, dropped columns - the output is
|
|
414
|
+
returned as read instead.
|
|
415
|
+
"""
|
|
416
|
+
if output_path is None or output_path.suffix.lower() != ".csv" or not output_path.exists():
|
|
417
|
+
return None
|
|
418
|
+
import pandas as pd
|
|
419
|
+
|
|
420
|
+
out = pd.read_csv(output_path, dtype=str, keep_default_na=False)
|
|
421
|
+
names = [str(c) for c in frame.columns]
|
|
422
|
+
if len(out) != len(frame) or not set(names).issubset(out.columns):
|
|
423
|
+
return out
|
|
424
|
+
merged = frame.copy()
|
|
425
|
+
for column in out.columns:
|
|
426
|
+
if column not in names:
|
|
427
|
+
merged[column] = out[column].to_numpy()
|
|
428
|
+
return merged
|
|
429
|
+
|
|
430
|
+
|
|
431
|
+
class _ProgressBar:
|
|
432
|
+
"""A tqdm bar when tqdm is installed, otherwise nothing"""
|
|
433
|
+
|
|
434
|
+
def __init__(self) -> None:
|
|
435
|
+
self.bar: Any = None
|
|
436
|
+
try:
|
|
437
|
+
from tqdm.auto import tqdm
|
|
438
|
+
|
|
439
|
+
self._tqdm: Any = tqdm
|
|
440
|
+
except ImportError:
|
|
441
|
+
self._tqdm = None
|
|
442
|
+
|
|
443
|
+
def update(self, event: Event) -> None:
|
|
444
|
+
if self._tqdm is None:
|
|
445
|
+
return
|
|
446
|
+
total = int(event.get("total") or 0) # type: ignore[arg-type]
|
|
447
|
+
if self.bar is None:
|
|
448
|
+
self.bar = self._tqdm(total=total, unit="row", desc="batchgrid")
|
|
449
|
+
if total and self.bar.total != total:
|
|
450
|
+
self.bar.total = total
|
|
451
|
+
self.bar.n = int(event.get("completed") or 0) # type: ignore[arg-type]
|
|
452
|
+
self.bar.set_postfix(failed=int(event.get("failed") or 0), refresh=False) # type: ignore[arg-type]
|
|
453
|
+
self.bar.refresh()
|
|
454
|
+
|
|
455
|
+
def close(self) -> None:
|
|
456
|
+
if self.bar is not None:
|
|
457
|
+
self.bar.close()
|
|
@@ -0,0 +1,146 @@
|
|
|
1
|
+
"""Running the batchgrid CLI and reading the JSON events it prints.
|
|
2
|
+
|
|
3
|
+
The engine is the Node CLI; this module only starts it with ``--json`` and
|
|
4
|
+
turns its stdout into Python dicts. The event format is documented in the
|
|
5
|
+
CLI README under "JSON output".
|
|
6
|
+
"""
|
|
7
|
+
|
|
8
|
+
from __future__ import annotations
|
|
9
|
+
|
|
10
|
+
import json
|
|
11
|
+
import os
|
|
12
|
+
import shlex
|
|
13
|
+
import shutil
|
|
14
|
+
import signal
|
|
15
|
+
import subprocess
|
|
16
|
+
import sys
|
|
17
|
+
import threading
|
|
18
|
+
from collections import deque
|
|
19
|
+
from typing import Callable, Deque, Dict, Iterator, List, Mapping, Optional, Sequence
|
|
20
|
+
|
|
21
|
+
#: The CLI releases this package speaks to: 0.2.0 was the first with ``--json``, and its events only gain fields
|
|
22
|
+
CLI_SPEC = "batchgrid@>=0.2.0 <1"
|
|
23
|
+
|
|
24
|
+
Event = Dict[str, object]
|
|
25
|
+
|
|
26
|
+
|
|
27
|
+
def find_cli(cli: Optional[Sequence[str]] = None) -> List[str]:
|
|
28
|
+
"""The command that starts the CLI.
|
|
29
|
+
|
|
30
|
+
In order: the ``cli`` argument, the ``BATCHGRID_CLI`` environment variable
|
|
31
|
+
(a command line, e.g. ``node /path/to/dist/index.js``), a ``batchgrid`` on
|
|
32
|
+
the PATH, then ``npx`` fetching the matching release.
|
|
33
|
+
"""
|
|
34
|
+
if cli:
|
|
35
|
+
return list(cli)
|
|
36
|
+
from_env = os.environ.get("BATCHGRID_CLI", "").strip()
|
|
37
|
+
if from_env:
|
|
38
|
+
return shlex.split(from_env, posix=os.name != "nt")
|
|
39
|
+
installed = shutil.which("batchgrid")
|
|
40
|
+
if installed:
|
|
41
|
+
return [installed]
|
|
42
|
+
npx = shutil.which("npx")
|
|
43
|
+
if npx:
|
|
44
|
+
return [npx, "--yes", CLI_SPEC]
|
|
45
|
+
from .errors import BatchgridError
|
|
46
|
+
|
|
47
|
+
raise BatchgridError(
|
|
48
|
+
"cli_not_found",
|
|
49
|
+
"The batchgrid CLI was not found. Install Node.js 22 or newer "
|
|
50
|
+
"(https://nodejs.org), or point BATCHGRID_CLI at a batchgrid command.",
|
|
51
|
+
)
|
|
52
|
+
|
|
53
|
+
|
|
54
|
+
def stream_events(
|
|
55
|
+
command: Sequence[str],
|
|
56
|
+
args: Sequence[str],
|
|
57
|
+
env: Optional[Mapping[str, str]] = None,
|
|
58
|
+
) -> Iterator[Event]:
|
|
59
|
+
"""Run the CLI and yield each JSON event as it is printed.
|
|
60
|
+
|
|
61
|
+
The last event is always a synthetic ``{"type": "exit", "code": ...,
|
|
62
|
+
"stderr": ...}``, so the caller sees how the process ended.
|
|
63
|
+
|
|
64
|
+
Ctrl+C (KeyboardInterrupt, also Jupyter's interrupt) is passed on to the
|
|
65
|
+
CLI as one SIGINT: the run stops, keeps its finished rows in the
|
|
66
|
+
checkpoint, and reports a ``stopped`` result before the interrupt is
|
|
67
|
+
raised again here.
|
|
68
|
+
"""
|
|
69
|
+
full_env = {**os.environ, **(env or {}), "FORCE_COLOR": "0", "NO_COLOR": "1"}
|
|
70
|
+
popen_kwargs: Dict[str, object] = {}
|
|
71
|
+
if os.name == "nt":
|
|
72
|
+
popen_kwargs["creationflags"] = subprocess.CREATE_NEW_PROCESS_GROUP # type: ignore[attr-defined]
|
|
73
|
+
else:
|
|
74
|
+
# Its own session, so a terminal's Ctrl+C reaches the CLI only through
|
|
75
|
+
# us - two SIGINTs would make it quit without saving
|
|
76
|
+
popen_kwargs["start_new_session"] = True
|
|
77
|
+
|
|
78
|
+
proc = subprocess.Popen(
|
|
79
|
+
[*command, *args],
|
|
80
|
+
stdin=subprocess.DEVNULL,
|
|
81
|
+
stdout=subprocess.PIPE,
|
|
82
|
+
stderr=subprocess.PIPE,
|
|
83
|
+
env=full_env,
|
|
84
|
+
text=True,
|
|
85
|
+
encoding="utf-8",
|
|
86
|
+
bufsize=1,
|
|
87
|
+
**popen_kwargs, # type: ignore[arg-type]
|
|
88
|
+
)
|
|
89
|
+
|
|
90
|
+
# stderr is drained on its own thread so a chatty run cannot fill the pipe
|
|
91
|
+
stderr_tail: Deque[str] = deque(maxlen=40)
|
|
92
|
+
drain = threading.Thread(target=_drain, args=(proc.stderr, stderr_tail.append), daemon=True)
|
|
93
|
+
drain.start()
|
|
94
|
+
|
|
95
|
+
interrupted = False
|
|
96
|
+
try:
|
|
97
|
+
assert proc.stdout is not None
|
|
98
|
+
while True:
|
|
99
|
+
try:
|
|
100
|
+
line = proc.stdout.readline()
|
|
101
|
+
except KeyboardInterrupt:
|
|
102
|
+
if interrupted:
|
|
103
|
+
proc.kill()
|
|
104
|
+
raise
|
|
105
|
+
interrupted = True
|
|
106
|
+
_interrupt(proc)
|
|
107
|
+
continue
|
|
108
|
+
if not line:
|
|
109
|
+
break
|
|
110
|
+
line = line.strip()
|
|
111
|
+
if not line:
|
|
112
|
+
continue
|
|
113
|
+
try:
|
|
114
|
+
event = json.loads(line)
|
|
115
|
+
except ValueError:
|
|
116
|
+
# Not an event - an old CLI without --json, or stray output
|
|
117
|
+
stderr_tail.append(line)
|
|
118
|
+
continue
|
|
119
|
+
if isinstance(event, dict):
|
|
120
|
+
yield event
|
|
121
|
+
code = proc.wait()
|
|
122
|
+
drain.join(timeout=2)
|
|
123
|
+
yield {"type": "exit", "code": code, "stderr": "\n".join(stderr_tail)}
|
|
124
|
+
finally:
|
|
125
|
+
if proc.poll() is None:
|
|
126
|
+
proc.kill()
|
|
127
|
+
proc.wait()
|
|
128
|
+
|
|
129
|
+
if interrupted:
|
|
130
|
+
raise KeyboardInterrupt
|
|
131
|
+
|
|
132
|
+
|
|
133
|
+
def _drain(pipe: Optional[object], sink: Callable[[str], None]) -> None:
|
|
134
|
+
if pipe is None:
|
|
135
|
+
return
|
|
136
|
+
for line in pipe: # type: ignore[attr-defined]
|
|
137
|
+
sink(line.rstrip("\n"))
|
|
138
|
+
|
|
139
|
+
|
|
140
|
+
def _interrupt(proc: "subprocess.Popen[str]") -> None:
|
|
141
|
+
if os.name == "nt":
|
|
142
|
+
# Windows has no SIGINT for another process group; CTRL_BREAK stops it at once
|
|
143
|
+
proc.send_signal(signal.CTRL_BREAK_EVENT) # type: ignore[attr-defined]
|
|
144
|
+
else:
|
|
145
|
+
proc.send_signal(signal.SIGINT)
|
|
146
|
+
print("batchgrid: stopping the run and keeping finished rows…", file=sys.stderr)
|
|
@@ -0,0 +1,55 @@
|
|
|
1
|
+
"""Exceptions raised by batchgrid. Every one carries the CLI's error code."""
|
|
2
|
+
|
|
3
|
+
from __future__ import annotations
|
|
4
|
+
|
|
5
|
+
from typing import List, Optional
|
|
6
|
+
|
|
7
|
+
|
|
8
|
+
class BatchgridError(Exception):
|
|
9
|
+
"""A plan or a run could not go ahead.
|
|
10
|
+
|
|
11
|
+
``code`` is the error code the CLI reported (``no_api_key``,
|
|
12
|
+
``invalid_plan``, ``run_failed``...), or ``cli_failed`` when the CLI
|
|
13
|
+
stopped without saying why - ``stderr`` then holds its last output.
|
|
14
|
+
"""
|
|
15
|
+
|
|
16
|
+
def __init__(self, code: str, message: str, stderr: str = "") -> None:
|
|
17
|
+
super().__init__(message)
|
|
18
|
+
self.code = code
|
|
19
|
+
self.message = message
|
|
20
|
+
self.stderr = stderr
|
|
21
|
+
|
|
22
|
+
|
|
23
|
+
class NoPlanError(BatchgridError):
|
|
24
|
+
"""The planner answered with a question instead of a plan.
|
|
25
|
+
|
|
26
|
+
``reply`` is what it said - usually what it needs to know to make one.
|
|
27
|
+
"""
|
|
28
|
+
|
|
29
|
+
def __init__(self, message: str, reply: Optional[str]) -> None:
|
|
30
|
+
super().__init__("no_plan", f"{message} The planner said: {reply}" if reply else message)
|
|
31
|
+
self.reply = reply
|
|
32
|
+
|
|
33
|
+
|
|
34
|
+
class MissingSecretsError(BatchgridError):
|
|
35
|
+
"""The plan calls a service whose key is not saved.
|
|
36
|
+
|
|
37
|
+
Save each one with ``batchgrid config --set-secret <name>:<value>``, or
|
|
38
|
+
set it in the environment, and run again.
|
|
39
|
+
"""
|
|
40
|
+
|
|
41
|
+
def __init__(self, message: str, names: List[str]) -> None:
|
|
42
|
+
super().__init__("missing_secrets", message)
|
|
43
|
+
self.names = names
|
|
44
|
+
|
|
45
|
+
|
|
46
|
+
class CostLimitError(BatchgridError):
|
|
47
|
+
"""The estimated cost is above ``max_cost``, so nothing was run.
|
|
48
|
+
|
|
49
|
+
``plan`` is the plan that was priced - inspect it, then run it with a
|
|
50
|
+
higher limit or on fewer rows.
|
|
51
|
+
"""
|
|
52
|
+
|
|
53
|
+
def __init__(self, message: str, plan: object) -> None:
|
|
54
|
+
super().__init__("cost_limit", message)
|
|
55
|
+
self.plan = plan
|
|
File without changes
|
|
@@ -0,0 +1,154 @@
|
|
|
1
|
+
"""Test fixtures: the real batchgrid CLI against a local stand-in for the OpenAI API.
|
|
2
|
+
|
|
3
|
+
Nothing in the engine is stubbed. The CLI is the one built in packages/cli;
|
|
4
|
+
only the model is replaced, by a server that answers the three kinds of call a
|
|
5
|
+
run makes - column analysis, planning and one call per row.
|
|
6
|
+
"""
|
|
7
|
+
|
|
8
|
+
from __future__ import annotations
|
|
9
|
+
|
|
10
|
+
import json
|
|
11
|
+
import os
|
|
12
|
+
import shutil
|
|
13
|
+
import threading
|
|
14
|
+
import time
|
|
15
|
+
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
|
|
16
|
+
from pathlib import Path
|
|
17
|
+
from typing import Any, Dict, Iterator, List
|
|
18
|
+
|
|
19
|
+
import pytest
|
|
20
|
+
|
|
21
|
+
REPO = Path(__file__).resolve().parents[2]
|
|
22
|
+
CLI_ENTRY = REPO / "packages" / "cli" / "dist" / "index.js"
|
|
23
|
+
|
|
24
|
+
PLAN: Dict[str, Any] = {
|
|
25
|
+
"title": "Classify sentiment",
|
|
26
|
+
"intent": "Label each review",
|
|
27
|
+
"sourceColumns": ["review"],
|
|
28
|
+
"steps": [
|
|
29
|
+
{
|
|
30
|
+
"id": "s1",
|
|
31
|
+
"kind": "llm",
|
|
32
|
+
"name": "Classify",
|
|
33
|
+
"promptTemplate": "Review: {{review}}",
|
|
34
|
+
"systemPrompt": "Classify the sentiment.",
|
|
35
|
+
"outputColumns": [{"name": "sentiment", "description": "positive or negative"}],
|
|
36
|
+
}
|
|
37
|
+
],
|
|
38
|
+
"output": {"format": "csv"},
|
|
39
|
+
}
|
|
40
|
+
|
|
41
|
+
SUMMARY = {"description": "Product reviews", "columnDescriptions": [], "suggestions": []}
|
|
42
|
+
|
|
43
|
+
|
|
44
|
+
class FakeModel:
|
|
45
|
+
"""What the stand-in server answers, and what it was asked"""
|
|
46
|
+
|
|
47
|
+
def __init__(self) -> None:
|
|
48
|
+
self.planner_reply: Any = PLAN
|
|
49
|
+
self.row_delay = 0.0
|
|
50
|
+
self.calls: List[str] = []
|
|
51
|
+
self.lock = threading.Lock()
|
|
52
|
+
|
|
53
|
+
def answer(self, system: str, user: str) -> Any:
|
|
54
|
+
with self.lock:
|
|
55
|
+
self.calls.append(system[:40])
|
|
56
|
+
if system.startswith("You are a data analyst"):
|
|
57
|
+
return SUMMARY
|
|
58
|
+
if system.startswith("You are a batch workflow planner"):
|
|
59
|
+
return self.planner_reply
|
|
60
|
+
if self.row_delay:
|
|
61
|
+
time.sleep(self.row_delay)
|
|
62
|
+
return {"sentiment": "negative" if "Terrible" in user else "positive"}
|
|
63
|
+
|
|
64
|
+
@property
|
|
65
|
+
def row_calls(self) -> int:
|
|
66
|
+
return sum(1 for c in self.calls if c.startswith("Classify the sentiment."))
|
|
67
|
+
|
|
68
|
+
@property
|
|
69
|
+
def planner_calls(self) -> int:
|
|
70
|
+
return sum(1 for c in self.calls if c.startswith("You are a batch workflow planner"))
|
|
71
|
+
|
|
72
|
+
|
|
73
|
+
def _handler(model: FakeModel) -> type:
|
|
74
|
+
class Handler(BaseHTTPRequestHandler):
|
|
75
|
+
def log_message(self, *args: object) -> None:
|
|
76
|
+
pass
|
|
77
|
+
|
|
78
|
+
def _send(self, body: Any) -> None:
|
|
79
|
+
data = json.dumps(body).encode()
|
|
80
|
+
self.send_response(200)
|
|
81
|
+
self.send_header("Content-Type", "application/json")
|
|
82
|
+
self.send_header("Content-Length", str(len(data)))
|
|
83
|
+
self.end_headers()
|
|
84
|
+
self.wfile.write(data)
|
|
85
|
+
|
|
86
|
+
def do_GET(self) -> None:
|
|
87
|
+
self._send({"object": "list", "data": [{"id": "gpt-5.6", "object": "model"}]})
|
|
88
|
+
|
|
89
|
+
def do_POST(self) -> None:
|
|
90
|
+
request = json.loads(self.rfile.read(int(self.headers["Content-Length"])))
|
|
91
|
+
messages = request["messages"]
|
|
92
|
+
system = next((m["content"] for m in messages if m["role"] == "system"), "")
|
|
93
|
+
user = "\n".join(m["content"] for m in messages if m["role"] == "user")
|
|
94
|
+
content = json.dumps(model.answer(system, user))
|
|
95
|
+
self._send(
|
|
96
|
+
{
|
|
97
|
+
"id": "x",
|
|
98
|
+
"object": "chat.completion",
|
|
99
|
+
"created": 0,
|
|
100
|
+
"model": request["model"],
|
|
101
|
+
"choices": [{"index": 0, "finish_reason": "stop", "message": {"role": "assistant", "content": content}}],
|
|
102
|
+
"usage": {"prompt_tokens": 10, "completion_tokens": 4, "total_tokens": 14},
|
|
103
|
+
}
|
|
104
|
+
)
|
|
105
|
+
|
|
106
|
+
return Handler
|
|
107
|
+
|
|
108
|
+
|
|
109
|
+
@pytest.fixture
|
|
110
|
+
def fake_model() -> Iterator[FakeModel]:
|
|
111
|
+
yield FakeModel()
|
|
112
|
+
|
|
113
|
+
|
|
114
|
+
@pytest.fixture
|
|
115
|
+
def cli_env(tmp_path: Path, fake_model: FakeModel) -> Iterator[Dict[str, str]]:
|
|
116
|
+
"""Environment for the CLI: an isolated home, a key, and the stand-in model"""
|
|
117
|
+
# BATCHGRID_TEST_CLI tests an installed CLI instead, e.g. one from a packed tarball
|
|
118
|
+
installed = os.environ.get("BATCHGRID_TEST_CLI")
|
|
119
|
+
node = shutil.which("node")
|
|
120
|
+
if not installed and (node is None or not CLI_ENTRY.exists()):
|
|
121
|
+
pytest.skip("needs Node.js and a built CLI (pnpm --filter batchgrid build)")
|
|
122
|
+
|
|
123
|
+
server = ThreadingHTTPServer(("127.0.0.1", 0), _handler(fake_model))
|
|
124
|
+
thread = threading.Thread(target=server.serve_forever, daemon=True)
|
|
125
|
+
thread.start()
|
|
126
|
+
|
|
127
|
+
env = {
|
|
128
|
+
"BATCHGRID_CLI": f'"{installed}"' if installed else f'"{node}" "{CLI_ENTRY}"',
|
|
129
|
+
"BATCHGRID_HOME": str(tmp_path / "home"),
|
|
130
|
+
"OPENAI_API_KEY": "sk-test",
|
|
131
|
+
"OPENAI_BASE_URL": f"http://127.0.0.1:{server.server_port}/v1",
|
|
132
|
+
}
|
|
133
|
+
# Keys from the developer's own environment must not pick another provider
|
|
134
|
+
cleared = ["ANTHROPIC_API_KEY", "GOOGLE_API_KEY", "GEMINI_API_KEY", "OPENROUTER_API_KEY", "XAI_API_KEY"]
|
|
135
|
+
saved = {k: os.environ.get(k) for k in [*env, *cleared]}
|
|
136
|
+
os.environ.update(env)
|
|
137
|
+
for key in cleared:
|
|
138
|
+
os.environ.pop(key, None)
|
|
139
|
+
try:
|
|
140
|
+
yield env
|
|
141
|
+
finally:
|
|
142
|
+
server.shutdown()
|
|
143
|
+
for key, value in saved.items():
|
|
144
|
+
if value is None:
|
|
145
|
+
os.environ.pop(key, None)
|
|
146
|
+
else:
|
|
147
|
+
os.environ[key] = value
|
|
148
|
+
|
|
149
|
+
|
|
150
|
+
@pytest.fixture
|
|
151
|
+
def reviews_csv(tmp_path: Path) -> Path:
|
|
152
|
+
path = tmp_path / "reviews.csv"
|
|
153
|
+
path.write_text("review\nGreat product\nTerrible service\nOkay I guess\n", encoding="utf-8")
|
|
154
|
+
return path
|
|
@@ -0,0 +1,157 @@
|
|
|
1
|
+
from __future__ import annotations
|
|
2
|
+
|
|
3
|
+
import _thread
|
|
4
|
+
import sys
|
|
5
|
+
import threading
|
|
6
|
+
import time
|
|
7
|
+
from pathlib import Path
|
|
8
|
+
|
|
9
|
+
import pandas as pd
|
|
10
|
+
import pytest
|
|
11
|
+
|
|
12
|
+
import batchgrid
|
|
13
|
+
from batchgrid import BatchgridError, CostLimitError, NoPlanError, Plan
|
|
14
|
+
|
|
15
|
+
from conftest import PLAN, FakeModel
|
|
16
|
+
|
|
17
|
+
pytestmark = pytest.mark.usefixtures("cli_env")
|
|
18
|
+
|
|
19
|
+
|
|
20
|
+
def test_plan_prices_without_running(reviews_csv: Path, fake_model: FakeModel) -> None:
|
|
21
|
+
plan = batchgrid.plan(reviews_csv, "classify the sentiment of each review")
|
|
22
|
+
|
|
23
|
+
assert plan.title == "Classify sentiment"
|
|
24
|
+
assert plan.rows == 3
|
|
25
|
+
assert plan.missing_secrets == []
|
|
26
|
+
assert "sentiment" in plan.summary
|
|
27
|
+
assert fake_model.row_calls == 0
|
|
28
|
+
|
|
29
|
+
|
|
30
|
+
def test_run_on_a_file_writes_the_output_next_to_it(reviews_csv: Path) -> None:
|
|
31
|
+
result = batchgrid.run(reviews_csv, "classify the sentiment of each review", progress=False)
|
|
32
|
+
|
|
33
|
+
assert (result.status, result.total, result.success, result.failed) == ("done", 3, 3, 0)
|
|
34
|
+
assert result.output_path == reviews_csv.with_name("reviews_output.csv")
|
|
35
|
+
assert "Terrible service,negative" in result.output_path.read_text(encoding="utf-8")
|
|
36
|
+
assert result.data is None
|
|
37
|
+
|
|
38
|
+
|
|
39
|
+
def test_run_on_a_dataframe_adds_the_new_columns() -> None:
|
|
40
|
+
df = pd.DataFrame(
|
|
41
|
+
{"review": ["Great product", "Terrible service", "Okay I guess"], "stars": [5, 1, 3]},
|
|
42
|
+
index=["a", "b", "c"],
|
|
43
|
+
)
|
|
44
|
+
|
|
45
|
+
result = batchgrid.run(df, "classify the sentiment of each review", progress=False)
|
|
46
|
+
|
|
47
|
+
assert result.data is not None
|
|
48
|
+
assert list(result.data.columns) == ["review", "stars", "sentiment"]
|
|
49
|
+
assert list(result.data.index) == ["a", "b", "c"]
|
|
50
|
+
# The columns it had keep their types
|
|
51
|
+
assert result.data["stars"].dtype == df["stars"].dtype
|
|
52
|
+
assert result.data.loc["b", "sentiment"] == "negative"
|
|
53
|
+
# The caller's DataFrame is left as it was
|
|
54
|
+
assert list(df.columns) == ["review", "stars"]
|
|
55
|
+
|
|
56
|
+
|
|
57
|
+
def test_a_saved_plan_runs_without_the_planner(reviews_csv: Path, tmp_path: Path, fake_model: FakeModel) -> None:
|
|
58
|
+
saved = batchgrid.plan(reviews_csv, "classify the sentiment of each review").save(tmp_path / "plan.json")
|
|
59
|
+
planner_calls = fake_model.planner_calls
|
|
60
|
+
|
|
61
|
+
result = batchgrid.run(reviews_csv, plan=Plan.load(saved), progress=False)
|
|
62
|
+
|
|
63
|
+
assert result.success == 3
|
|
64
|
+
assert fake_model.planner_calls == planner_calls
|
|
65
|
+
|
|
66
|
+
|
|
67
|
+
def test_a_plan_dict_runs_as_is(reviews_csv: Path, fake_model: FakeModel) -> None:
|
|
68
|
+
result = batchgrid.run(reviews_csv, plan=PLAN, progress=False)
|
|
69
|
+
assert result.success == 3
|
|
70
|
+
assert fake_model.planner_calls == 0
|
|
71
|
+
|
|
72
|
+
|
|
73
|
+
def test_max_cost_stops_before_any_row_runs(reviews_csv: Path, fake_model: FakeModel) -> None:
|
|
74
|
+
with pytest.raises(CostLimitError) as caught:
|
|
75
|
+
batchgrid.run(reviews_csv, "classify the sentiment of each review", max_cost=0.0000001, progress=False)
|
|
76
|
+
|
|
77
|
+
assert isinstance(caught.value.plan, Plan)
|
|
78
|
+
assert fake_model.row_calls == 0
|
|
79
|
+
|
|
80
|
+
|
|
81
|
+
def test_max_cost_within_budget_runs_the_priced_plan(reviews_csv: Path, fake_model: FakeModel) -> None:
|
|
82
|
+
result = batchgrid.run(reviews_csv, "classify the sentiment of each review", max_cost=100, progress=False)
|
|
83
|
+
|
|
84
|
+
assert result.success == 3
|
|
85
|
+
# Priced once, then the same plan ran - the planner was asked once
|
|
86
|
+
assert fake_model.planner_calls == 1
|
|
87
|
+
|
|
88
|
+
|
|
89
|
+
def test_a_plan_that_does_not_fit_the_sheet_is_refused(reviews_csv: Path, fake_model: FakeModel) -> None:
|
|
90
|
+
wrong = {**PLAN, "sourceColumns": ["body"], "steps": [{**PLAN["steps"][0], "promptTemplate": "{{body}}"}]}
|
|
91
|
+
|
|
92
|
+
with pytest.raises(BatchgridError) as caught:
|
|
93
|
+
batchgrid.run(reviews_csv, plan=wrong, progress=False)
|
|
94
|
+
|
|
95
|
+
assert caught.value.code == "invalid_plan"
|
|
96
|
+
assert fake_model.row_calls == 0
|
|
97
|
+
|
|
98
|
+
|
|
99
|
+
def test_a_question_from_the_planner_is_raised_with_its_text(reviews_csv: Path, fake_model: FakeModel) -> None:
|
|
100
|
+
fake_model.planner_reply = {"kind": "message", "message": "Which column holds the review?"}
|
|
101
|
+
|
|
102
|
+
with pytest.raises(NoPlanError) as caught:
|
|
103
|
+
batchgrid.plan(reviews_csv, "do the thing")
|
|
104
|
+
|
|
105
|
+
assert caught.value.reply == "Which column holds the review?"
|
|
106
|
+
|
|
107
|
+
|
|
108
|
+
def test_no_key_is_reported_by_code(reviews_csv: Path, monkeypatch: pytest.MonkeyPatch) -> None:
|
|
109
|
+
monkeypatch.delenv("OPENAI_API_KEY")
|
|
110
|
+
with pytest.raises(BatchgridError) as caught:
|
|
111
|
+
batchgrid.plan(reviews_csv, "classify")
|
|
112
|
+
assert caught.value.code == "no_api_key"
|
|
113
|
+
|
|
114
|
+
|
|
115
|
+
def test_events_reach_the_callback(reviews_csv: Path) -> None:
|
|
116
|
+
seen = []
|
|
117
|
+
batchgrid.run(reviews_csv, "classify the sentiment of each review", progress=False, on_event=seen.append)
|
|
118
|
+
|
|
119
|
+
types = [e["type"] for e in seen]
|
|
120
|
+
assert "plan" in types
|
|
121
|
+
assert types[-1] == "result"
|
|
122
|
+
|
|
123
|
+
|
|
124
|
+
def test_either_a_prompt_or_a_plan(reviews_csv: Path) -> None:
|
|
125
|
+
with pytest.raises(ValueError):
|
|
126
|
+
batchgrid.run(reviews_csv)
|
|
127
|
+
with pytest.raises(ValueError):
|
|
128
|
+
batchgrid.run(reviews_csv, "x", plan=PLAN)
|
|
129
|
+
|
|
130
|
+
|
|
131
|
+
@pytest.mark.skipif(sys.platform == "win32", reason="SIGINT forwarding is POSIX-only")
|
|
132
|
+
def test_ctrl_c_keeps_finished_rows_and_resume_finishes(tmp_path: Path, fake_model: FakeModel) -> None:
|
|
133
|
+
rows = "\n".join(f"Review number {i}" for i in range(40))
|
|
134
|
+
path = tmp_path / "many.csv"
|
|
135
|
+
path.write_text(f"review\n{rows}\n", encoding="utf-8")
|
|
136
|
+
fake_model.row_delay = 0.2
|
|
137
|
+
|
|
138
|
+
def interrupt_once_rows_run() -> None:
|
|
139
|
+
# Past the first few rows, so some have finished and are checkpointed
|
|
140
|
+
while fake_model.row_calls < 10:
|
|
141
|
+
time.sleep(0.02)
|
|
142
|
+
_thread.interrupt_main()
|
|
143
|
+
|
|
144
|
+
threading.Thread(target=interrupt_once_rows_run, daemon=True).start()
|
|
145
|
+
with pytest.raises(KeyboardInterrupt):
|
|
146
|
+
batchgrid.run(path, plan=PLAN, concurrency=2, progress=False)
|
|
147
|
+
stopped_after = fake_model.row_calls
|
|
148
|
+
assert stopped_after < 40
|
|
149
|
+
|
|
150
|
+
fake_model.row_delay = 0
|
|
151
|
+
result = batchgrid.resume(path, progress=False)
|
|
152
|
+
|
|
153
|
+
assert result.status == "done"
|
|
154
|
+
assert result.total == 40
|
|
155
|
+
# Only the unfinished rows went to the model again
|
|
156
|
+
assert fake_model.row_calls - stopped_after < 40
|
|
157
|
+
assert len(pd.read_csv(result.output_path)) == 40
|