Restructuration de l'application

copienator/commands/submit_batches.py

@@ -0,0 +1,96 @@
+from __future__ import annotations
+
+import argparse
+from collections.abc import Sequence
+
+from google import genai
+from google.genai import types
+
+from copienator import configuration as config
+from copienator import (
+    CliError,
+    EvaluationWorkspace,
+    ExitCode,
+    atomic_write_json,
+    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)
+    batches = (
+        (
+            "flash",
+            workspace.root / "batch_requests_flash.jsonl",
+            config.MODEL_FLASH_ID,
+            f"flash-correction-{workspace.name}",
+        ),
+        (
+            "pro",
+            workspace.root / "batch_requests_pro.jsonl",
+            config.MODEL_PRO_ID,
+            f"pro-correction-{workspace.name}",
+        ),
+    )
+    manifest = {
+        "version": 1,
+        "evaluation": workspace.name,
+        "jobs": {},
+    }
+    if workspace.batch_jobs_file.is_file():
+        previous = read_json(workspace.batch_jobs_file)
+        if isinstance(previous, dict) and isinstance(previous.get("jobs"), dict):
+            manifest["jobs"] = previous["jobs"]
+    started = 0
+    for tier, file_path, model_id, display_name in batches:
+        if not file_path.is_file():
+            print(f"Skipping {model_id}: {file_path.name} does not exist.")
+            continue
+        if file_path.stat().st_size == 0:
+            print(f"Skipping {model_id}: {file_path.name} is empty.")
+            continue
+        print(f"Uploading {file_path.name} for model {model_id}...")
+        uploaded = client.files.upload(
+            file=str(file_path),
+            config=types.UploadFileConfig(
+                display_name=f"{display_name}-input",
+                mime_type="jsonl",
+            ),
+        )
+        job = client.batches.create(
+            model=model_id,
+            src=uploaded.name,
+            config={"display_name": display_name},
+        )
+        started += 1
+        manifest["jobs"][tier] = {
+            "name": job.name,
+            "display_name": display_name,
+            "model": model_id,
+            "request_file": file_path.name,
+        }
+        atomic_write_json(workspace.batch_jobs_file, manifest)
+        print(f"Started batch job: {job.name}")
+    if not started:
+        print("No non-empty batch request files were found.")
+        return ExitCode.PARTIAL
+    return ExitCode.SUCCESS
+
+
+def build_parser() -> argparse.ArgumentParser:
+    return evaluation_parser("Upload correction JSONL files and start Gemini batches")
+
+
+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())