84 lines
2.7 KiB
Python
84 lines
2.7 KiB
Python
from __future__ import annotations
|
|
|
|
import argparse
|
|
from collections.abc import Sequence
|
|
|
|
from google import genai
|
|
|
|
from copienator import configuration as config
|
|
from copienator import (
|
|
CliError,
|
|
EvaluationWorkspace,
|
|
ExitCode,
|
|
atomic_write_bytes,
|
|
evaluation_parser,
|
|
execute,
|
|
read_json,
|
|
workspace_from_args,
|
|
)
|
|
|
|
|
|
def run(workspace: EvaluationWorkspace, *, client=None) -> ExitCode:
|
|
if client is None:
|
|
if not config.API_KEY:
|
|
raise CliError("GEMINI_API_KEY is not configured")
|
|
client = genai.Client(api_key=config.API_KEY)
|
|
matching = []
|
|
if workspace.batch_jobs_file.is_file():
|
|
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}")
|
|
matching = [
|
|
client.batches.get(name=entry["name"])
|
|
for entry in jobs.values()
|
|
if isinstance(entry, dict) and isinstance(entry.get("name"), str)
|
|
]
|
|
else:
|
|
matching = [
|
|
job
|
|
for job in client.batches.list()
|
|
if workspace.name in str(getattr(job, "display_name", ""))
|
|
]
|
|
if not matching:
|
|
raise CliError(
|
|
f"No batch jobs found for evaluation {workspace.name!r}"
|
|
)
|
|
for job in matching:
|
|
state = job.state.name if hasattr(job.state, "name") else job.state
|
|
print(f"{job.display_name}: {state}")
|
|
if state != "JOB_STATE_SUCCEEDED":
|
|
print("Not all matching jobs have succeeded yet.")
|
|
return ExitCode.PARTIAL
|
|
|
|
chunks = []
|
|
incomplete = False
|
|
for job in matching:
|
|
destination = getattr(job, "dest", None)
|
|
file_name = getattr(destination, "file_name", None)
|
|
if not file_name:
|
|
print(f"Warning: {job.display_name} has no output file.")
|
|
incomplete = True
|
|
continue
|
|
payload = client.files.download(file=file_name)
|
|
chunks.append(payload.rstrip(b"\n"))
|
|
if not chunks:
|
|
return ExitCode.PARTIAL
|
|
output_path = workspace.batched_correction_result_file
|
|
atomic_write_bytes(output_path, b"\n".join(chunks) + b"\n")
|
|
print(f"Saved combined batch results to {output_path}")
|
|
return ExitCode.PARTIAL if incomplete else ExitCode.SUCCESS
|
|
|
|
|
|
def build_parser() -> argparse.ArgumentParser:
|
|
return evaluation_parser("Download and combine correction batch results")
|
|
|
|
|
|
def main(argv: Sequence[str] | None = None) -> int:
|
|
parser = build_parser()
|
|
return execute(parser, argv, lambda args: run(workspace_from_args(args)))
|
|
|
|
|
|
if __name__ == "__main__":
|
|
raise SystemExit(main())
|