batchwizard 0.1.0__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
@@ -0,0 +1,21 @@
1
+ MIT License
2
+
3
+ Copyright (c) 2024 Carl Kugblenu
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,172 @@
1
+ Metadata-Version: 2.1
2
+ Name: batchwizard
3
+ Version: 0.1.0
4
+ Summary: BatchWizard: Manage OpenAI batch processing jobs with ease
5
+ Home-page: https://github.com/cmakafui/batchwizard
6
+ License: MIT
7
+ Keywords: openai,batch,cli,async
8
+ Author: Carl Kugblenu
9
+ Requires-Python: >=3.9,<4.0
10
+ Classifier: License :: OSI Approved :: MIT License
11
+ Classifier: Programming Language :: Python :: 3
12
+ Classifier: Programming Language :: Python :: 3.9
13
+ Classifier: Programming Language :: Python :: 3.10
14
+ Classifier: Programming Language :: Python :: 3.11
15
+ Classifier: Programming Language :: Python :: 3.12
16
+ Requires-Dist: aiofiles (>=24.1.0,<25.0.0)
17
+ Requires-Dist: loguru (>=0.7.2,<0.8.0)
18
+ Requires-Dist: openai (>=1.37.0,<2.0.0)
19
+ Requires-Dist: pydantic (>=2.8.2,<3.0.0)
20
+ Requires-Dist: pydantic-settings (>=2.3.4,<3.0.0)
21
+ Requires-Dist: python-dotenv (>=1.0.1,<2.0.0)
22
+ Requires-Dist: rich (>=13.7.1,<14.0.0)
23
+ Requires-Dist: typer (>=0.12.3,<0.13.0)
24
+ Project-URL: Repository, https://github.com/cmakafui/batchwizard
25
+ Description-Content-Type: text/markdown
26
+
27
+ # BatchWizard
28
+
29
+ BatchWizard is a powerful CLI tool for managing OpenAI batch processing jobs with ease. It provides functionalities to upload files, create batch jobs, check their status, and download the results. The tool uses asynchronous processing to efficiently handle multiple jobs concurrently.
30
+
31
+ ## Table of Contents
32
+
33
+ - [Installation](#installation)
34
+ - [Usage](#usage)
35
+ - [Configuration](#configuration)
36
+ - [Commands](#commands)
37
+ - [Features](#features)
38
+ - [Contributing](#contributing)
39
+ - [License](#license)
40
+
41
+ ## Installation
42
+
43
+ You can install BatchWizard using `pipx` for an isolated environment or directly via `pip`.
44
+
45
+ ### Using pipx (recommended)
46
+
47
+ ```bash
48
+ pipx install batchwizard
49
+ ```
50
+
51
+ ### Using pip
52
+
53
+ ```bash
54
+ pip install batchwizard
55
+ ```
56
+
57
+ Ensure you have `pipx` or `pip` installed on your system. For `pipx`, you can follow the installation instructions [here](https://pipx.pypa.io/stable/installation/).
58
+
59
+ ## Usage
60
+
61
+ BatchWizard provides a command-line interface (CLI) for managing batch jobs. Here are some example commands:
62
+
63
+ ### Process Batch Jobs
64
+
65
+ To process a directory containing input files:
66
+
67
+ ```bash
68
+ batchwizard process <input_directory> [--output-directory OUTPUT_DIR] [--max-concurrent-jobs NUM] [--check-interval SECONDS]
69
+ ```
70
+
71
+ ### List Recent Jobs
72
+
73
+ To list recent batch jobs:
74
+
75
+ ```bash
76
+ batchwizard list-jobs [--limit NUM] [--all]
77
+ ```
78
+
79
+ ### Cancel a Job
80
+
81
+ To cancel a specific batch job:
82
+
83
+ ```bash
84
+ batchwizard cancel <job_id>
85
+ ```
86
+
87
+ ### Download Job Results
88
+
89
+ To download results for a completed batch job:
90
+
91
+ ```bash
92
+ batchwizard download <job_id> [--output-file FILE_PATH]
93
+ ```
94
+
95
+ ## Configuration
96
+
97
+ ### Setting up the OpenAI API Key
98
+
99
+ To set the OpenAI API key:
100
+
101
+ ```bash
102
+ batchwizard configure --set-key YOUR_API_KEY
103
+ ```
104
+
105
+ ### Show Current Configuration
106
+
107
+ To show the current configuration:
108
+
109
+ ```bash
110
+ batchwizard configure --show
111
+ ```
112
+
113
+ ### Reset Configuration
114
+
115
+ To reset the configuration to default values:
116
+
117
+ ```bash
118
+ batchwizard configure --reset
119
+ ```
120
+
121
+ ## Commands
122
+
123
+ BatchWizard supports the following commands:
124
+
125
+ - `process`: Process batch jobs from input files in the specified directory.
126
+ - `configure`: Manage BatchWizard configuration.
127
+ - `list-jobs`: List recent batch jobs.
128
+ - `cancel`: Cancel a specific batch job.
129
+ - `download`: Download results for a completed batch job.
130
+
131
+ For detailed information on each command, use the `--help` option:
132
+
133
+ ```bash
134
+ batchwizard <command> --help
135
+ ```
136
+
137
+ ## Features
138
+
139
+ - **Asynchronous Processing**: Efficiently handle multiple batch jobs concurrently.
140
+ - **Rich UI**: Display progress and job status using a rich, interactive interface.
141
+ - **Flexible Configuration**: Easily manage API keys and other settings.
142
+ - **Job Management**: List, cancel, and download results for batch jobs.
143
+ - **Error Handling**: Robust error handling and informative error messages.
144
+
145
+ ## Contributing
146
+
147
+ We welcome contributions to BatchWizard! To contribute, follow these steps:
148
+
149
+ 1. Fork the repository.
150
+ 2. Create a new branch: `git checkout -b feature/your-feature-name`.
151
+ 3. Make your changes and commit them: `git commit -m 'Add some feature'`.
152
+ 4. Push to the branch: `git push origin feature/your-feature-name`.
153
+ 5. Open a pull request.
154
+
155
+ ### Running Tests
156
+
157
+ To run tests, use `pytest`:
158
+
159
+ ```bash
160
+ pytest --cov=batchwizard tests/
161
+ ```
162
+
163
+ Ensure your code passes all tests and meets the coding standards before opening a pull request.
164
+
165
+ ## License
166
+
167
+ This project is licensed under the MIT License. See the [LICENSE](LICENSE) file for details.
168
+
169
+ ## Contact
170
+
171
+ For any questions or feedback, feel free to open an issue on the [GitHub repository](https://github.com/cmakafui/batchwizard).
172
+
@@ -0,0 +1,145 @@
1
+ # BatchWizard
2
+
3
+ BatchWizard is a powerful CLI tool for managing OpenAI batch processing jobs with ease. It provides functionalities to upload files, create batch jobs, check their status, and download the results. The tool uses asynchronous processing to efficiently handle multiple jobs concurrently.
4
+
5
+ ## Table of Contents
6
+
7
+ - [Installation](#installation)
8
+ - [Usage](#usage)
9
+ - [Configuration](#configuration)
10
+ - [Commands](#commands)
11
+ - [Features](#features)
12
+ - [Contributing](#contributing)
13
+ - [License](#license)
14
+
15
+ ## Installation
16
+
17
+ You can install BatchWizard using `pipx` for an isolated environment or directly via `pip`.
18
+
19
+ ### Using pipx (recommended)
20
+
21
+ ```bash
22
+ pipx install batchwizard
23
+ ```
24
+
25
+ ### Using pip
26
+
27
+ ```bash
28
+ pip install batchwizard
29
+ ```
30
+
31
+ Ensure you have `pipx` or `pip` installed on your system. For `pipx`, you can follow the installation instructions [here](https://pipx.pypa.io/stable/installation/).
32
+
33
+ ## Usage
34
+
35
+ BatchWizard provides a command-line interface (CLI) for managing batch jobs. Here are some example commands:
36
+
37
+ ### Process Batch Jobs
38
+
39
+ To process a directory containing input files:
40
+
41
+ ```bash
42
+ batchwizard process <input_directory> [--output-directory OUTPUT_DIR] [--max-concurrent-jobs NUM] [--check-interval SECONDS]
43
+ ```
44
+
45
+ ### List Recent Jobs
46
+
47
+ To list recent batch jobs:
48
+
49
+ ```bash
50
+ batchwizard list-jobs [--limit NUM] [--all]
51
+ ```
52
+
53
+ ### Cancel a Job
54
+
55
+ To cancel a specific batch job:
56
+
57
+ ```bash
58
+ batchwizard cancel <job_id>
59
+ ```
60
+
61
+ ### Download Job Results
62
+
63
+ To download results for a completed batch job:
64
+
65
+ ```bash
66
+ batchwizard download <job_id> [--output-file FILE_PATH]
67
+ ```
68
+
69
+ ## Configuration
70
+
71
+ ### Setting up the OpenAI API Key
72
+
73
+ To set the OpenAI API key:
74
+
75
+ ```bash
76
+ batchwizard configure --set-key YOUR_API_KEY
77
+ ```
78
+
79
+ ### Show Current Configuration
80
+
81
+ To show the current configuration:
82
+
83
+ ```bash
84
+ batchwizard configure --show
85
+ ```
86
+
87
+ ### Reset Configuration
88
+
89
+ To reset the configuration to default values:
90
+
91
+ ```bash
92
+ batchwizard configure --reset
93
+ ```
94
+
95
+ ## Commands
96
+
97
+ BatchWizard supports the following commands:
98
+
99
+ - `process`: Process batch jobs from input files in the specified directory.
100
+ - `configure`: Manage BatchWizard configuration.
101
+ - `list-jobs`: List recent batch jobs.
102
+ - `cancel`: Cancel a specific batch job.
103
+ - `download`: Download results for a completed batch job.
104
+
105
+ For detailed information on each command, use the `--help` option:
106
+
107
+ ```bash
108
+ batchwizard <command> --help
109
+ ```
110
+
111
+ ## Features
112
+
113
+ - **Asynchronous Processing**: Efficiently handle multiple batch jobs concurrently.
114
+ - **Rich UI**: Display progress and job status using a rich, interactive interface.
115
+ - **Flexible Configuration**: Easily manage API keys and other settings.
116
+ - **Job Management**: List, cancel, and download results for batch jobs.
117
+ - **Error Handling**: Robust error handling and informative error messages.
118
+
119
+ ## Contributing
120
+
121
+ We welcome contributions to BatchWizard! To contribute, follow these steps:
122
+
123
+ 1. Fork the repository.
124
+ 2. Create a new branch: `git checkout -b feature/your-feature-name`.
125
+ 3. Make your changes and commit them: `git commit -m 'Add some feature'`.
126
+ 4. Push to the branch: `git push origin feature/your-feature-name`.
127
+ 5. Open a pull request.
128
+
129
+ ### Running Tests
130
+
131
+ To run tests, use `pytest`:
132
+
133
+ ```bash
134
+ pytest --cov=batchwizard tests/
135
+ ```
136
+
137
+ Ensure your code passes all tests and meets the coding standards before opening a pull request.
138
+
139
+ ## License
140
+
141
+ This project is licensed under the MIT License. See the [LICENSE](LICENSE) file for details.
142
+
143
+ ## Contact
144
+
145
+ For any questions or feedback, feel free to open an issue on the [GitHub repository](https://github.com/cmakafui/batchwizard).
@@ -0,0 +1,37 @@
1
+ [tool.poetry]
2
+ name = "batchwizard"
3
+ version = "0.1.0"
4
+ description = "BatchWizard: Manage OpenAI batch processing jobs with ease"
5
+ authors = ["Carl Kugblenu"]
6
+ license = "MIT"
7
+ readme = "README.md"
8
+ homepage = "https://github.com/cmakafui/batchwizard"
9
+ repository = "https://github.com/cmakafui/batchwizard"
10
+ keywords = ["openai", "batch", "cli", "async"]
11
+
12
+ [tool.poetry.dependencies]
13
+ python = "^3.9"
14
+ typer = "^0.12.3"
15
+ rich = "^13.7.1"
16
+ openai = "^1.37.0"
17
+ pydantic = "^2.8.2"
18
+ aiofiles = "^24.1.0"
19
+ loguru = "^0.7.2"
20
+ pydantic-settings = "^2.3.4"
21
+ python-dotenv = "^1.0.1"
22
+
23
+
24
+ [tool.poetry.group.dev.dependencies]
25
+ pytest = "^8.3.2"
26
+ pytest-asyncio = "^0.23.8"
27
+ isort = "^5.13.2"
28
+ black = "^24.4.2"
29
+ ruff = "^0.5.5"
30
+ pytest-mock = "^3.14.0"
31
+
32
+ [tool.poetry.scripts]
33
+ batchwizard = "batchwizard.cli:app"
34
+
35
+ [build-system]
36
+ requires = ["poetry-core"]
37
+ build-backend = "poetry.core.masonry.api"
@@ -0,0 +1,6 @@
1
+ # __init__.py
2
+ from .cli import app
3
+ from .models import BatchJob, BatchJobResult
4
+ from .processor import BatchProcessor
5
+
6
+ __all__ = ["BatchProcessor", "BatchJob", "BatchJobResult", "app"]
@@ -0,0 +1,191 @@
1
+ # cli.py
2
+ import asyncio
3
+ from datetime import datetime
4
+ from pathlib import Path
5
+ from typing import Optional
6
+
7
+ import typer
8
+ from rich.console import Console
9
+ from rich.table import Table
10
+
11
+ from .config import BatchWizardSettings, config
12
+ from .processor import BatchProcessor
13
+ from .ui import BatchWizardUI
14
+ from .utils import get_api_key, set_api_key, setup_logger
15
+
16
+ app = typer.Typer(help="BatchWizard: Manage OpenAI batch processing jobs with ease")
17
+ console = Console()
18
+ logger = setup_logger(console)
19
+
20
+
21
+ @app.command()
22
+ def process(
23
+ input_directory: Path = typer.Argument(
24
+ ..., help="Directory containing input files for processing"
25
+ ),
26
+ output_directory: Optional[Path] = typer.Option(
27
+ None, help="Directory to store output files"
28
+ ),
29
+ max_concurrent_jobs: int = typer.Option(
30
+ 5, help="Maximum number of concurrent jobs"
31
+ ),
32
+ check_interval: int = typer.Option(
33
+ 5, help="Initial interval (in seconds) between job status checks"
34
+ ),
35
+ ):
36
+ """Process batch jobs from input files in the specified directory."""
37
+ if not output_directory:
38
+ output_directory = input_directory / "results"
39
+ output_directory.mkdir(parents=True, exist_ok=True)
40
+
41
+ api_key = get_api_key()
42
+ if not api_key:
43
+ logger.error(
44
+ "API key not set. Please set it using 'batchwizard configure --set-key YOUR_API_KEY'"
45
+ )
46
+ raise typer.Exit(code=1)
47
+
48
+ config.settings.max_concurrent_jobs = max_concurrent_jobs
49
+ config.settings.check_interval = check_interval
50
+ config.save()
51
+
52
+ processor = BatchProcessor()
53
+ ui = BatchWizardUI(Console())
54
+
55
+ async def run_and_close():
56
+ try:
57
+ await ui.run_processing(processor, input_directory, output_directory)
58
+ finally:
59
+ await processor.close()
60
+
61
+ asyncio.run(run_and_close())
62
+
63
+
64
+ @app.command()
65
+ def configure(
66
+ set_key: Optional[str] = typer.Option(
67
+ None, "--set-key", help="Set the OpenAI API key"
68
+ ),
69
+ show: bool = typer.Option(False, "--show", help="Show the current configuration"),
70
+ reset: bool = typer.Option(
71
+ False, "--reset", help="Reset the configuration to default values"
72
+ ),
73
+ ):
74
+ """Manage BatchWizard configuration."""
75
+ if set_key:
76
+ set_api_key(set_key)
77
+ console.print("[green]API key set successfully.[/green]")
78
+ elif show:
79
+ api_key = get_api_key()
80
+ masked_key = f"{api_key[:4]}...{api_key[-4:]}" if api_key else "Not set"
81
+ console.print(f"API Key: {masked_key}")
82
+ console.print(f"Max Concurrent Jobs: {config.settings.max_concurrent_jobs}")
83
+ console.print(f"Check Interval: {config.settings.check_interval} seconds")
84
+ elif reset:
85
+ config.settings = BatchWizardSettings()
86
+ config.save()
87
+ console.print("[yellow]Configuration reset to default values.[/yellow]")
88
+ else:
89
+ console.print(
90
+ "Use --set-key, --show, or --reset options to manage configuration."
91
+ )
92
+
93
+
94
+ @app.command()
95
+ def list_jobs(
96
+ limit: int = typer.Option(100, help="Number of jobs to display"),
97
+ all: bool = typer.Option(False, "--all", help="Display all jobs"),
98
+ ):
99
+ """List recent batch jobs."""
100
+
101
+ async def fetch_jobs():
102
+ processor = BatchProcessor()
103
+ console = Console() # Create a Console object directly
104
+ try:
105
+ jobs = await processor.client.batches.list(limit=None if all else limit)
106
+ table = Table(title="Batch Jobs")
107
+ table.add_column("Job ID", style="cyan")
108
+ table.add_column("Status", style="magenta")
109
+ table.add_column("Created At", style="green")
110
+ table.add_column("Completed", style="blue")
111
+ table.add_column("Failed", style="red")
112
+
113
+ for job in jobs.data:
114
+ created_at = datetime.fromtimestamp(job.created_at).strftime(
115
+ "%Y-%m-%d %H:%M:%S"
116
+ )
117
+ table.add_row(
118
+ job.id,
119
+ job.status,
120
+ created_at,
121
+ str(job.request_counts.completed),
122
+ str(job.request_counts.failed),
123
+ )
124
+ console.print(table) # Use the console object to print the table
125
+ finally:
126
+ await processor.close()
127
+
128
+ asyncio.run(fetch_jobs())
129
+
130
+
131
+ @app.command()
132
+ def cancel(
133
+ job_id: str = typer.Argument(..., help="ID of the batch job to cancel"),
134
+ ):
135
+ """Cancel a specific batch job."""
136
+
137
+ async def cancel_job():
138
+ processor = BatchProcessor()
139
+ try:
140
+ await processor.client.batches.cancel(job_id)
141
+ console.print(f"[green]Job {job_id} cancelled successfully.[/green]")
142
+ except Exception as e:
143
+ console.print(f"[red]Error cancelling job {job_id}: {str(e)}[/red]")
144
+ finally:
145
+ await processor.close()
146
+
147
+ asyncio.run(cancel_job())
148
+
149
+
150
+ @app.command()
151
+ def download(
152
+ job_id: str = typer.Argument(
153
+ ..., help="ID of the batch job to download results for"
154
+ ),
155
+ output_file: Path = typer.Option(
156
+ None, help="Path to save the output file (default: <job_id>_results.jsonl)"
157
+ ),
158
+ ):
159
+ """Download results for a completed batch job."""
160
+ if not output_file:
161
+ output_file = Path(f"{job_id}_results.jsonl")
162
+
163
+ async def download_results():
164
+ processor = BatchProcessor()
165
+ try:
166
+ batch_job = await processor.client.batches.retrieve(job_id)
167
+ if batch_job.status != "completed":
168
+ console.print(
169
+ f"[yellow]Job {job_id} is not completed (status: {batch_job.status}). Cannot download results.[/yellow]"
170
+ )
171
+ return
172
+
173
+ success = await processor.download_batch_results(batch_job, output_file)
174
+ if success:
175
+ console.print(
176
+ f"[green]Results for job {job_id} downloaded successfully to {output_file}[/green]"
177
+ )
178
+ else:
179
+ console.print(f"[red]Failed to download results for job {job_id}[/red]")
180
+ except Exception as e:
181
+ console.print(
182
+ f"[red]Error downloading results for job {job_id}: {str(e)}[/red]"
183
+ )
184
+ finally:
185
+ await processor.close()
186
+
187
+ asyncio.run(download_results())
188
+
189
+
190
+ if __name__ == "__main__":
191
+ app()
@@ -0,0 +1,48 @@
1
+ # config.py
2
+ import os
3
+ from pathlib import Path
4
+ from typing import Optional
5
+
6
+ import typer
7
+ from pydantic import BaseModel, Field
8
+ from pydantic_settings import BaseSettings, SettingsConfigDict
9
+
10
+
11
+ class BatchWizardSettings(BaseSettings):
12
+ api_key: Optional[str] = Field(None, env="OPENAI_API_KEY")
13
+ max_concurrent_jobs: int = 5
14
+ check_interval: int = 5
15
+ model_config = SettingsConfigDict(env_file=".env", env_file_encoding="utf-8")
16
+
17
+
18
+ class Config(BaseModel):
19
+ settings: BatchWizardSettings = BatchWizardSettings()
20
+
21
+ @property
22
+ def config_dir(self) -> Path:
23
+ return Path(typer.get_app_dir("BatchWizard"))
24
+
25
+ @property
26
+ def config_file(self) -> Path:
27
+ return self.config_dir / "config.json"
28
+
29
+ def load(self) -> None:
30
+ if self.config_file.exists():
31
+ self.settings = BatchWizardSettings.model_validate_json(
32
+ self.config_file.read_text()
33
+ )
34
+
35
+ def save(self) -> None:
36
+ self.config_dir.mkdir(parents=True, exist_ok=True)
37
+ self.config_file.write_text(self.settings.model_dump_json())
38
+
39
+ def get_api_key(self) -> Optional[str]:
40
+ return self.settings.api_key or os.getenv("OPENAI_API_KEY")
41
+
42
+ def set_api_key(self, api_key: str) -> None:
43
+ self.settings.api_key = api_key
44
+ self.save()
45
+
46
+
47
+ config = Config()
48
+ config.load()
@@ -0,0 +1,18 @@
1
+ # models.py
2
+ from pathlib import Path
3
+ from typing import Optional
4
+
5
+ from pydantic import BaseModel
6
+
7
+
8
+ class BatchJob(BaseModel):
9
+ id: str
10
+ status: str
11
+ input_file_id: str
12
+ output_file_id: Optional[str] = None
13
+
14
+
15
+ class BatchJobResult(BaseModel):
16
+ job_id: str
17
+ success: bool
18
+ output_file_path: Optional[Path] = None
@@ -0,0 +1,145 @@
1
+ # processor.py
2
+ import asyncio
3
+ from pathlib import Path
4
+ from typing import List, Optional
5
+
6
+ import aiofiles
7
+ from loguru import logger
8
+ from openai import AsyncOpenAI
9
+
10
+ from .config import config
11
+ from .models import BatchJob, BatchJobResult
12
+
13
+
14
+ class BatchProcessor:
15
+ def __init__(self):
16
+ self.client = AsyncOpenAI(api_key=config.get_api_key())
17
+ self.settings = config.settings
18
+
19
+ async def upload_file(self, file_path: Path) -> Optional[str]:
20
+ try:
21
+ async with aiofiles.open(file_path, "rb") as file:
22
+ file_content = await file.read()
23
+
24
+ response = await self.client.files.create(
25
+ file=(file_path.name, file_content), purpose="batch"
26
+ )
27
+ logger.info(
28
+ f"File uploaded successfully: {response.id}, Filename: {file_path.name}"
29
+ )
30
+ return response.id
31
+ except Exception as e:
32
+ logger.error(f"Error uploading file {file_path.name}: {str(e)}")
33
+ return None
34
+
35
+ async def create_batch_job(self, input_file_id: str) -> Optional[BatchJob]:
36
+ try:
37
+ batch_job = await self.client.batches.create(
38
+ input_file_id=input_file_id,
39
+ endpoint="/v1/chat/completions",
40
+ completion_window="24h",
41
+ )
42
+ logger.info(f"Created batch job with ID: {batch_job.id}")
43
+ return BatchJob(
44
+ id=batch_job.id,
45
+ status=self.normalize_status(batch_job.status),
46
+ input_file_id=input_file_id,
47
+ )
48
+ except Exception as e:
49
+ logger.error(f"Error creating batch job: {str(e)}")
50
+ return None
51
+
52
+ async def check_batch_status(self, batch_id: str) -> Optional[str]:
53
+ try:
54
+ batch_job = await self.client.batches.retrieve(batch_id)
55
+ return self.normalize_status(batch_job.status)
56
+ except Exception as e:
57
+ logger.error(f"Error checking batch status: {str(e)}")
58
+ return None
59
+
60
+ def normalize_status(self, status: str) -> str:
61
+ """Normalize the status string to lowercase with underscores."""
62
+ return status.lower().replace(" ", "_")
63
+
64
+ async def download_batch_results(self, batch_job, output_file_path: Path) -> bool:
65
+ try:
66
+ if batch_job.status == "completed" and batch_job.output_file_id:
67
+ result = await self.client.files.content(batch_job.output_file_id)
68
+ async with aiofiles.open(output_file_path, "wb") as file:
69
+ await file.write(result.content)
70
+ logger.info(f"Downloaded results to {output_file_path}")
71
+ return True
72
+ else:
73
+ logger.warning(
74
+ f"Batch job not completed or missing output file. Status: {batch_job.status}"
75
+ )
76
+ return False
77
+ except Exception as e:
78
+ logger.error(f"Error downloading batch results: {str(e)}")
79
+ return False
80
+
81
+ async def process_batch_job(
82
+ self, batch_job: BatchJob, output_dir: Path
83
+ ) -> BatchJobResult:
84
+ check_interval = self.settings.check_interval
85
+ while True:
86
+ status = await self.check_batch_status(batch_job.id)
87
+ if status == "completed":
88
+ try:
89
+ batch_job = await self.client.batches.retrieve(batch_job.id)
90
+ if batch_job.output_file_id:
91
+ output_file = output_dir / f"{batch_job.id}_results.jsonl"
92
+ if await self.download_batch_results(batch_job, output_file):
93
+ logger.info(
94
+ f"Successfully processed batch job {batch_job.id}"
95
+ )
96
+ return BatchJobResult(
97
+ job_id=batch_job.id,
98
+ success=True,
99
+ output_file_path=output_file,
100
+ )
101
+ else:
102
+ logger.error(
103
+ f"No output file ID found for completed batch job {batch_job.id}"
104
+ )
105
+ except Exception as e:
106
+ logger.error(
107
+ f"Error processing completed batch job {batch_job.id}: {str(e)}"
108
+ )
109
+ elif status in ["failed", "expired", "cancelled"]:
110
+ logger.error(f"Batch job {batch_job.id} {status}")
111
+ return BatchJobResult(job_id=batch_job.id, success=False)
112
+ elif status is None:
113
+ logger.error(f"Failed to retrieve status for batch job {batch_job.id}")
114
+ return BatchJobResult(job_id=batch_job.id, success=False)
115
+
116
+ await asyncio.sleep(check_interval)
117
+ check_interval = min(
118
+ check_interval * 1.5, 60
119
+ ) # Implement exponential backoff
120
+
121
+ async def process_directory(
122
+ self, input_dir: Path, output_dir: Path
123
+ ) -> List[BatchJobResult]:
124
+ input_files = list(input_dir.glob("*_batch_*.jsonl"))
125
+ if not input_files:
126
+ logger.warning(f"No input files found in {input_dir}")
127
+ return []
128
+
129
+ semaphore = asyncio.Semaphore(self.settings.max_concurrent_jobs)
130
+
131
+ async def process_file(input_file: Path) -> Optional[BatchJobResult]:
132
+ async with semaphore:
133
+ file_id = await self.upload_file(input_file)
134
+ if file_id:
135
+ batch_job = await self.create_batch_job(file_id)
136
+ if batch_job:
137
+ return await self.process_batch_job(batch_job, output_dir)
138
+ return None
139
+
140
+ tasks = [process_file(file) for file in input_files]
141
+ results = await asyncio.gather(*tasks)
142
+ return [result for result in results if result is not None]
143
+
144
+ async def close(self):
145
+ await self.client.close()
@@ -0,0 +1,250 @@
1
+ # ui.py
2
+ import asyncio
3
+ from datetime import datetime
4
+ from pathlib import Path
5
+ from typing import List
6
+
7
+ from loguru import logger
8
+ from rich.console import Console
9
+ from rich.layout import Layout
10
+ from rich.live import Live
11
+ from rich.panel import Panel
12
+ from rich.progress import (BarColumn, Progress, SpinnerColumn,
13
+ TaskProgressColumn, TextColumn, TimeElapsedColumn)
14
+ from rich.table import Table
15
+ from rich.text import Text
16
+
17
+ from .processor import BatchProcessor
18
+
19
+
20
+ class BatchWizardUI:
21
+ def __init__(self, console: Console):
22
+ self.console = console
23
+ self.job_table = self.create_job_table()
24
+ self.stats_table = self.create_stats_table()
25
+ self.log_messages = []
26
+ self.jobs = {} # Dictionary to store job data
27
+ self.total_jobs = 0
28
+ self.completed_jobs = 0
29
+ self.failed_jobs = 0
30
+
31
+ def create_layout(self) -> Layout:
32
+ layout = Layout()
33
+ layout.split_column(Layout(name="header", size=3), Layout(name="body"))
34
+ layout["body"].split_row(
35
+ Layout(name="main", ratio=2), Layout(name="sidebar", ratio=1)
36
+ )
37
+ layout["body"]["main"].split_column(
38
+ Layout(name="progress", ratio=1), Layout(name="job_status", ratio=2)
39
+ )
40
+ layout["body"]["sidebar"].split_column(
41
+ Layout(name="stats", ratio=1), Layout(name="logs", ratio=2)
42
+ )
43
+ return layout
44
+
45
+ def create_progress_bars(self):
46
+ overall_progress = Progress(
47
+ SpinnerColumn(),
48
+ TextColumn("[progress.description]{task.description}"),
49
+ BarColumn(),
50
+ TaskProgressColumn(),
51
+ TimeElapsedColumn(),
52
+ )
53
+ job_progress = Progress(
54
+ TextColumn("[progress.description]{task.description}"),
55
+ BarColumn(),
56
+ TaskProgressColumn(),
57
+ )
58
+ return overall_progress, job_progress
59
+
60
+ def create_job_table(self):
61
+ job_table = Table(show_header=True, header_style="bold magenta", expand=True)
62
+ job_table.add_column("Job ID", style="dim", no_wrap=True)
63
+ job_table.add_column("Status", style="bold")
64
+ job_table.add_column("Progress", justify="right")
65
+ return job_table
66
+
67
+ def update_job_status(self, job_id: str, status: str, progress: str):
68
+ color = (
69
+ "green"
70
+ if status == "completed"
71
+ else "red" if status in ["failed", "expired", "cancelled"] else "yellow"
72
+ )
73
+ self.jobs[job_id] = (f"[{color}]{status}", progress)
74
+ self.job_table = self.create_job_table()
75
+ for job_id, (status, progress) in self.jobs.items():
76
+ self.job_table.add_row(job_id, status, progress)
77
+
78
+ def create_stats_table(self):
79
+ stats_table = Table(show_header=False, expand=True)
80
+ stats_table.add_column("Metric", style="cyan")
81
+ stats_table.add_column("Value", justify="right")
82
+ stats_table.add_row("Total Jobs", "0")
83
+ stats_table.add_row("Completed", "0")
84
+ stats_table.add_row("In Progress", "0")
85
+ stats_table.add_row("Failed", "0")
86
+ return stats_table
87
+
88
+ def update_stats(self):
89
+ self.stats_table.rows.clear()
90
+ self.stats_table.add_row("Total Jobs", str(self.total_jobs))
91
+ self.stats_table.add_row("Completed", f"[green]{self.completed_jobs}[/green]")
92
+ self.stats_table.add_row(
93
+ "In Progress",
94
+ f"[yellow]{self.total_jobs - self.completed_jobs - self.failed_jobs}[/yellow]",
95
+ )
96
+ self.stats_table.add_row("Failed", f"[red]{self.failed_jobs}[/red]")
97
+
98
+ def add_log(self, message: str):
99
+ self.log_messages.append(message)
100
+ if len(self.log_messages) > 10: # Keep only the last 10 log messages
101
+ self.log_messages.pop(0)
102
+ logger.info(message)
103
+
104
+ def get_log_panel(self):
105
+ return Panel("\n".join(self.log_messages), title="Logs", border_style="green")
106
+
107
+ def update_layout(
108
+ self,
109
+ layout: Layout,
110
+ header: str,
111
+ overall_progress: Progress,
112
+ job_table: Table,
113
+ stats_table: Table,
114
+ log_panel: Panel,
115
+ ):
116
+ layout["header"].update(
117
+ Panel(Text(header, style="bold blue"), border_style="blue")
118
+ )
119
+ layout["body"]["main"]["progress"].update(
120
+ Panel(overall_progress, title="Overall Progress", border_style="green")
121
+ )
122
+ layout["body"]["main"]["job_status"].update(
123
+ Panel(job_table, title="Job Status", border_style="magenta")
124
+ )
125
+ layout["body"]["sidebar"]["stats"].update(
126
+ Panel(stats_table, title="Statistics", border_style="cyan")
127
+ )
128
+ layout["body"]["sidebar"]["logs"].update(log_panel)
129
+
130
+ async def run_processing(
131
+ self, processor: BatchProcessor, input_dir: Path, output_dir: Path
132
+ ):
133
+ layout = self.create_layout()
134
+ overall_progress, job_progress = self.create_progress_bars()
135
+
136
+ overall_task = overall_progress.add_task("[bold blue]Processing", total=100)
137
+ upload_task = overall_progress.add_task("[green]Uploading files", visible=False)
138
+ process_task = overall_progress.add_task("[cyan]Processing jobs", visible=False)
139
+
140
+ async def update_ui():
141
+ self.update_layout(
142
+ layout,
143
+ "BatchWizard Processing",
144
+ overall_progress,
145
+ self.job_table,
146
+ self.stats_table,
147
+ self.get_log_panel(),
148
+ )
149
+
150
+ with Live(layout, console=self.console, screen=True, refresh_per_second=4):
151
+ await update_ui()
152
+
153
+ input_files = list(input_dir.glob("*_batch_*.jsonl"))
154
+ if not input_files:
155
+ self.add_log(f"No input files found in {input_dir}")
156
+ await update_ui()
157
+ return
158
+
159
+ overall_progress.update(upload_task, total=len(input_files), visible=True)
160
+ overall_progress.update(process_task, total=len(input_files), visible=True)
161
+
162
+ self.total_jobs = len(input_files)
163
+
164
+ async def process_file(input_file: Path):
165
+ file_id = await processor.upload_file(input_file)
166
+ overall_progress.update(upload_task, advance=1)
167
+ if file_id:
168
+ batch_job = await processor.create_batch_job(file_id)
169
+ if batch_job:
170
+ self.update_job_status(batch_job.id, batch_job.status, "0%")
171
+ self.add_log(f"Job created: {batch_job.id}")
172
+ await update_ui()
173
+
174
+ result = await processor.process_batch_job(
175
+ batch_job, output_dir
176
+ )
177
+ if result.success:
178
+ self.completed_jobs += 1
179
+ self.update_job_status(batch_job.id, "completed", "100%")
180
+ self.add_log(f"Job completed: {batch_job.id}")
181
+ if result.output_file_path:
182
+ self.add_log(
183
+ f"Results saved to: {result.output_file_path}"
184
+ )
185
+ else:
186
+ self.failed_jobs += 1
187
+ self.update_job_status(batch_job.id, "failed", "100%")
188
+ self.add_log(f"Job failed: {batch_job.id}")
189
+
190
+ overall_progress.update(process_task, advance=1)
191
+ self.update_stats()
192
+ await update_ui()
193
+ else:
194
+ self.add_log(
195
+ f"Failed to create batch job for {input_file.name}"
196
+ )
197
+ else:
198
+ self.add_log(f"Failed to upload file {input_file.name}")
199
+ await update_ui()
200
+
201
+ tasks = [process_file(file) for file in input_files]
202
+ await asyncio.gather(*tasks)
203
+
204
+ overall_progress.update(overall_task, completed=100)
205
+ await update_ui()
206
+
207
+ self.console.print("[bold green]Processing completed![/bold green]")
208
+ self.console.print(f"Total jobs: {self.total_jobs}")
209
+ self.console.print(f"Completed jobs: {self.completed_jobs}")
210
+ self.console.print(f"Failed jobs: {self.failed_jobs}")
211
+ if self.completed_jobs > 0:
212
+ self.console.print(f"Results saved in: {output_dir}")
213
+
214
+ def display_job_list(self, jobs: List[dict]):
215
+ table = Table(title="Batch Jobs")
216
+ table.add_column("Job ID", style="cyan")
217
+ table.add_column("Status", style="magenta")
218
+ table.add_column("Created At", style="green")
219
+ table.add_column("Completed", style="blue")
220
+ table.add_column("Failed", style="red")
221
+
222
+ for job in jobs:
223
+ created_at = datetime.fromtimestamp(job["created_at"]).strftime(
224
+ "%Y-%m-%d %H:%M:%S"
225
+ )
226
+ table.add_row(
227
+ job["id"],
228
+ job["status"],
229
+ created_at,
230
+ str(job["request_counts"]["completed"]),
231
+ str(job["request_counts"]["failed"]),
232
+ )
233
+
234
+ self.console.print(table)
235
+
236
+ def display_cancel_result(self, job_id: str, success: bool):
237
+ if success:
238
+ self.console.print(f"[green]Job {job_id} cancelled successfully.[/green]")
239
+ else:
240
+ self.console.print(f"[red]Failed to cancel job {job_id}.[/red]")
241
+
242
+ def display_download_result(self, job_id: str, output_file: Path, success: bool):
243
+ if success:
244
+ self.console.print(
245
+ f"[green]Results for job {job_id} downloaded successfully to {output_file}[/green]"
246
+ )
247
+ else:
248
+ self.console.print(
249
+ f"[red]Failed to download results for job {job_id}[/red]"
250
+ )
@@ -0,0 +1,36 @@
1
+ # utils.py
2
+ from typing import Optional
3
+
4
+ from loguru import logger
5
+ from rich.console import Console
6
+
7
+ from .config import config
8
+
9
+
10
+ def setup_logger(console: Console = None):
11
+ logger.remove()
12
+ if console is None:
13
+ console = Console()
14
+ logger.add(
15
+ console.file,
16
+ format="<red>{time:YYYY-MM-DD HH:mm:ss}</red> | <level>{level: <8}</level> | <cyan>{name}</cyan>:<cyan>{function}</cyan>:<cyan>{line}</cyan> - <level>{message}</level>",
17
+ level="ERROR",
18
+ backtrace=True,
19
+ diagnose=True,
20
+ )
21
+ return logger
22
+
23
+
24
+ def get_api_key() -> Optional[str]:
25
+ """Get the API key from config or environment variable."""
26
+ return config.get_api_key()
27
+
28
+
29
+ def set_api_key(api_key: str) -> None:
30
+ """Set the API key in the config."""
31
+ config.set_api_key(api_key)
32
+
33
+
34
+ def get_settings():
35
+ """Get the current settings."""
36
+ return config.settings