Standardisation 8
This commit is contained in:
+79
-74
@@ -1,88 +1,93 @@
|
||||
import os
|
||||
import sys
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
from collections.abc import Sequence
|
||||
from pathlib import Path
|
||||
|
||||
from google import genai
|
||||
|
||||
if "GEMINI_API_KEY" not in os.environ:
|
||||
sys.exit("Error: GEMINI_API_KEY environment variable not set.")
|
||||
|
||||
client = genai.Client()
|
||||
|
||||
def list_jobs():
|
||||
print("Fetching recent batch jobs...\n")
|
||||
try:
|
||||
batch_jobs = client.batches.list()
|
||||
jobs_found = False
|
||||
|
||||
for job in batch_jobs:
|
||||
jobs_found = True
|
||||
state = job.state.name if hasattr(job.state, 'name') else job.state
|
||||
|
||||
print("-" * 60)
|
||||
print(f"Job Name: {job.name}")
|
||||
|
||||
if hasattr(job, 'display_name') and job.display_name:
|
||||
print(f"Display Name: {job.display_name}")
|
||||
|
||||
print(f"State: {state}")
|
||||
|
||||
if state == 'JOB_STATE_FAILED' and hasattr(job, 'error'):
|
||||
print(f"Error: {job.error}")
|
||||
|
||||
if state == 'JOB_STATE_SUCCEEDED' and hasattr(job, 'dest') and job.dest:
|
||||
if hasattr(job.dest, 'file_name') and job.dest.file_name:
|
||||
print(f"Output File: {job.dest.file_name}")
|
||||
|
||||
if not jobs_found:
|
||||
print("No batch jobs found.")
|
||||
else:
|
||||
print("-" * 60)
|
||||
print("\nTo download a completed job, run:")
|
||||
print("python batch_status.py --download batches/<YOUR_BATCH_ID>")
|
||||
|
||||
except Exception as e:
|
||||
sys.exit(f"An error occurred while listing jobs: {e}")
|
||||
import config
|
||||
from copienator import (
|
||||
CliError,
|
||||
ExitCode,
|
||||
atomic_write_bytes,
|
||||
execute,
|
||||
standard_parser,
|
||||
)
|
||||
|
||||
|
||||
def download_job(job_name):
|
||||
print(f"Checking status for {job_name}...\n")
|
||||
try:
|
||||
job = client.batches.get(name=job_name)
|
||||
state = job.state.name if hasattr(job.state, 'name') else job.state
|
||||
def _client():
|
||||
if not config.API_KEY:
|
||||
raise CliError("GEMINI_API_KEY is not configured")
|
||||
return genai.Client(api_key=config.API_KEY)
|
||||
|
||||
print(f"State: {state}")
|
||||
|
||||
if state != 'JOB_STATE_SUCCEEDED':
|
||||
print("Job is not ready yet or has failed.")
|
||||
if state == 'JOB_STATE_FAILED' and hasattr(job, 'error'):
|
||||
print(f"Error: {job.error}")
|
||||
return
|
||||
def list_jobs(*, client=None) -> ExitCode:
|
||||
client = client or _client()
|
||||
print("Fetching recent batch jobs...")
|
||||
jobs = list(client.batches.list())
|
||||
for job in jobs:
|
||||
state = job.state.name if hasattr(job.state, "name") else job.state
|
||||
print(f"{job.name}: {state}")
|
||||
if getattr(job, "display_name", None):
|
||||
print(f" Display name: {job.display_name}")
|
||||
if state == "JOB_STATE_FAILED" and getattr(job, "error", None):
|
||||
print(f" Error: {job.error}")
|
||||
destination = getattr(job, "dest", None)
|
||||
if state == "JOB_STATE_SUCCEEDED" and getattr(
|
||||
destination, "file_name", None
|
||||
):
|
||||
print(f" Output file: {destination.file_name}")
|
||||
if not jobs:
|
||||
print("No batch jobs found.")
|
||||
return ExitCode.SUCCESS
|
||||
|
||||
if hasattr(job, 'dest') and job.dest and hasattr(job.dest, 'file_name') and job.dest.file_name:
|
||||
result_file_name = job.dest.file_name
|
||||
print(f"Downloading results from {result_file_name}...")
|
||||
|
||||
file_content_bytes = client.files.download(file=result_file_name)
|
||||
output_path = f"results_{job_name.replace('/', '_')}.jsonl"
|
||||
def download_job(
|
||||
job_name: str,
|
||||
*,
|
||||
output: Path | None = None,
|
||||
client=None,
|
||||
) -> ExitCode:
|
||||
client = client or _client()
|
||||
job = client.batches.get(name=job_name)
|
||||
state = job.state.name if hasattr(job.state, "name") else job.state
|
||||
print(f"State: {state}")
|
||||
if state != "JOB_STATE_SUCCEEDED":
|
||||
if state == "JOB_STATE_FAILED" and getattr(job, "error", None):
|
||||
print(f"Error: {job.error}")
|
||||
return ExitCode.PARTIAL
|
||||
destination = getattr(job, "dest", None)
|
||||
file_name = getattr(destination, "file_name", None)
|
||||
if not file_name:
|
||||
print("Job succeeded but no output file was found.")
|
||||
return ExitCode.PARTIAL
|
||||
payload = client.files.download(file=file_name)
|
||||
output_path = output or Path(f"results_{job_name.replace('/', '_')}.jsonl")
|
||||
atomic_write_bytes(output_path, payload)
|
||||
print(f"Saved batch results to {output_path}")
|
||||
return ExitCode.SUCCESS
|
||||
|
||||
with open(output_path, "wb") as f:
|
||||
f.write(file_content_bytes)
|
||||
|
||||
print(f"Success! Saved to {output_path}")
|
||||
print(f"You can now feed this to your correction script using: --deal-with-batched {output_path}")
|
||||
else:
|
||||
print("Job succeeded but no output file was found.")
|
||||
def build_parser() -> argparse.ArgumentParser:
|
||||
parser = standard_parser("List or download Gemini correction batch jobs")
|
||||
parser.add_argument("--download", metavar="JOB_NAME")
|
||||
parser.add_argument("--output", type=Path, help="Downloaded JSONL destination")
|
||||
return parser
|
||||
|
||||
|
||||
def main(argv: Sequence[str] | None = None) -> int:
|
||||
parser = build_parser()
|
||||
|
||||
def handle(args: argparse.Namespace) -> ExitCode:
|
||||
if args.output is not None and not args.download:
|
||||
raise CliError("--output requires --download", ExitCode.INVALID_ARGUMENTS)
|
||||
if args.download:
|
||||
return download_job(args.download, output=args.output)
|
||||
return list_jobs()
|
||||
|
||||
return execute(parser, argv, handle)
|
||||
|
||||
except Exception as e:
|
||||
sys.exit(f"An error occurred while fetching the job: {e}")
|
||||
|
||||
if __name__ == "__main__":
|
||||
parser = argparse.ArgumentParser(description="Manage Gemini Batch Jobs")
|
||||
parser.add_argument("--download", type=str, metavar="JOB_NAME",
|
||||
help="Download the results for a specific batch job (e.g. batches/123456)")
|
||||
args = parser.parse_args()
|
||||
|
||||
if args.download:
|
||||
download_job(args.download)
|
||||
else:
|
||||
list_jobs()
|
||||
raise SystemExit(main())
|
||||
|
||||
Reference in New Issue
Block a user