139 lines
5.2 KiB
Python
139 lines
5.2 KiB
Python
from __future__ import annotations
|
||
|
||
import argparse
|
||
from collections.abc import Sequence
|
||
from pathlib import Path
|
||
|
||
from google import genai
|
||
|
||
from copienator import configuration as config
|
||
from copienator import (
|
||
CliError,
|
||
EvaluationWorkspace,
|
||
ExitCode,
|
||
atomic_write_bytes,
|
||
execute,
|
||
read_json,
|
||
standard_parser,
|
||
workspace_from_args,
|
||
)
|
||
|
||
|
||
def _client():
|
||
if not config.API_KEY:
|
||
raise CliError("GEMINI_API_KEY is not configured")
|
||
return genai.Client(api_key=config.API_KEY)
|
||
|
||
|
||
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
|
||
|
||
|
||
def check_evaluation_jobs(workspace: EvaluationWorkspace, *, client=None) -> ExitCode:
|
||
"""Only report readiness after checking every job recorded for this evaluation."""
|
||
if not workspace.batch_jobs_file.is_file():
|
||
print("Impossible de vérifier les batchs : batch_jobs.json est absent.")
|
||
return ExitCode.PARTIAL
|
||
manifest = read_json(workspace.batch_jobs_file)
|
||
jobs = manifest.get("jobs") if isinstance(manifest, dict) else None
|
||
if not isinstance(jobs, dict):
|
||
raise CliError(f"Invalid batch manifest: {workspace.batch_jobs_file}")
|
||
if not jobs:
|
||
print("Aucun job enregistré pour cette évaluation.")
|
||
return ExitCode.PARTIAL
|
||
if any(not isinstance(entry, dict) or not isinstance(entry.get("name"), str)
|
||
or not entry["name"].strip() for entry in jobs.values()):
|
||
raise CliError(f"Invalid batch job in {workspace.batch_jobs_file}")
|
||
client = client or _client()
|
||
ready = True
|
||
for tier, entry in jobs.items():
|
||
job = client.batches.get(name=entry["name"])
|
||
state = job.state.name if hasattr(job.state, "name") else job.state
|
||
print(f"{tier} — {entry['name']}: {state}")
|
||
if state != "JOB_STATE_SUCCEEDED":
|
||
ready = False
|
||
if getattr(job, "error", None):
|
||
print(f" Erreur : {job.error}")
|
||
elif not getattr(getattr(job, "dest", None), "file_name", None):
|
||
ready = False
|
||
print(" Le fichier de résultats n’est pas encore disponible.")
|
||
if ready:
|
||
print("Tous les batchs de l’évaluation ont réussi. Les résultats sont prêts à récupérer.")
|
||
return ExitCode.SUCCESS
|
||
print("Les résultats ne sont pas tous prêts. Consultez à nouveau cette étape plus tard.")
|
||
return ExitCode.PARTIAL
|
||
|
||
|
||
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
|
||
|
||
|
||
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")
|
||
parser.add_argument("--evaluation", type=Path,
|
||
help="Check readiness of jobs recorded in this evaluation's batch_jobs.json")
|
||
return parser
|
||
|
||
|
||
def main(argv: Sequence[str] | None = None) -> int:
|
||
parser = build_parser()
|
||
|
||
def handle(args: argparse.Namespace) -> ExitCode:
|
||
if args.evaluation is not None:
|
||
if args.download or args.output is not None:
|
||
raise CliError("--evaluation cannot be combined with --download or --output",
|
||
ExitCode.INVALID_ARGUMENTS)
|
||
return check_evaluation_jobs(workspace_from_args(args))
|
||
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)
|
||
|
||
|
||
if __name__ == "__main__":
|
||
raise SystemExit(main())
|