Restructuration de l'application

copienator/commands/fetch_batched_results.py

@@ -0,0 +1,83 @@
+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())