orizu 0.3.2 → 0.3.4

package/vendor/orizu-gepa-python/src/orizu_gepa/cli.py

@@ -6,6 +6,7 @@ from pathlib import Path
 from typing import Any
 
 from .client import OrizuClient, OrizuEventSink
+from .local_log import LocalOptimizationLogger
 from .optimizer import TextGepaConfig, optimize_loaded_text_candidate
 from .reflection import reflect_with_provider
 from .runner import make_candidate_runner, make_scorer_runner
@@ -74,6 +75,8 @@ def main() -> None:
     parser.add_argument("--auto-promote", action="store_true")
     parser.add_argument("--promotion-label")
     parser.add_argument("--log-row-snapshots", action="store_true")
+    parser.add_argument("--log-dir", default="logs")
+    parser.add_argument("--no-local-log", action="store_true")
     parser.add_argument("--metadata", default="{}")
     args = parser.parse_args()
 
@@ -123,6 +126,19 @@ def main() -> None:
         split=args.val_split,
     )
 
+    run_metadata = {
+        **metadata,
+        "optimizer_package": "orizu-gepa-python",
+        "optimizer_family": "gepa",
+        "mode": "text-candidate",
+        "inference_lm": prompt_context.provider_settings.get("model"),
+        "reflection_lm": config.reflection_model,
+        "scorer_lm": scorer_context.provider_settings.get("model"),
+        "dataset_size": len(trainset) + len(valset),
+        "train_count": len(trainset),
+        "validation_count": len(valset),
+    }
+
     run_id = client.start_run(
         project=args.project,
         optimizer_version_id=args.optimizer_version_id,
@@ -132,21 +148,29 @@ def main() -> None:
         split_set_id=args.split_set_id,
         train_split=args.train_split,
         validation_split=args.val_split,
-        metadata={
-            **metadata,
-            "optimizer_package": "orizu-gepa-python",
-            "optimizer_family": "gepa",
-            "mode": "text-candidate",
-            "inference_lm": prompt_context.provider_settings.get("model"),
-            "reflection_lm": config.reflection_model,
-            "scorer_lm": scorer_context.provider_settings.get("model"),
-            "dataset_size": len(trainset) + len(valset),
-            "train_count": len(trainset),
-            "validation_count": len(valset),
-        },
+        metadata=run_metadata,
     )
     print(f"[orizu-gepa] started optimization run {run_id}", flush=True)
 
+    local_logger = None
+    if not args.no_local_log:
+        local_logger = LocalOptimizationLogger.create(args.log_dir, run_id)
+        local_logger.write_context(
+            project=args.project,
+            run_id=run_id,
+            args={
+                key: value
+                for key, value in vars(args).items()
+                if key not in {"metadata"}
+            },
+            prompt_context=prompt_context,
+            scorer_context=scorer_context,
+            trainset=trainset,
+            valset=valset,
+            metadata=run_metadata,
+        )
+        print(f"[orizu-gepa] local log: {local_logger.path}", flush=True)
+
     result = optimize_loaded_text_candidate(
         run_id=run_id,
         prompt_context=prompt_context,
@@ -156,9 +180,17 @@ def main() -> None:
         candidate_runner=make_candidate_runner(Path(args.candidate_runner_dir), run_id),
         scorer_runner=make_scorer_runner(Path(args.scorer_runner_dir), run_id),
         reflector=reflect_with_provider,
-        event_sink=OrizuEventSink(client, run_id, fail_on_log_error=config.fail_on_log_error),
+        event_sink=OrizuEventSink(
+            client,
+            run_id,
+            fail_on_log_error=config.fail_on_log_error,
+            local_logger=local_logger,
+        ),
         config=config,
+        local_logger=local_logger,
     )
+    if local_logger is not None:
+        local_logger.write_result(result)
     print(json.dumps({
         "optimization_run_id": result.run_id,
         "best_candidate_id": result.best_candidate_id,
@@ -166,6 +198,7 @@ def main() -> None:
         "seed_score": result.seed_score,
         "promoted_prompt_version_id": result.promoted_prompt_version_id,
         "budget": result.budget.to_payload(),
+        "local_log_dir": local_logger.path if local_logger is not None else None,
     }))
 
 

package/vendor/orizu-gepa-python/src/orizu_gepa/client.py

@@ -255,12 +255,14 @@ class OrizuEventSink:
         *,
         fail_on_log_error: bool = True,
         max_log_retries: int = 2,
+        local_logger: Any | None = None,
     ):
         self.client = client
         self.run_id = run_id
         self.sequence = 0
         self.fail_on_log_error = fail_on_log_error
         self.max_log_retries = max(0, max_log_retries)
+        self.local_logger = local_logger
 
     def log_event(
         self,
@@ -275,6 +277,20 @@ class OrizuEventSink:
         child_candidate_id: str | None = None,
     ) -> None:
         self.sequence += 1
+        local_event = {
+            "run_id": self.run_id,
+            "sequence": self.sequence,
+            "event_type": event_type,
+            "event_layer": event_layer,
+            "optimizer_family": optimizer_family,
+            "iteration": iteration,
+            "candidate_id": candidate_id,
+            "parent_candidate_id": parent_candidate_id,
+            "child_candidate_id": child_candidate_id,
+            "payload": payload or {},
+        }
+        if self.local_logger is not None:
+            self.local_logger.append_event(local_event)
         last_error: Exception | None = None
         for attempt in range(self.max_log_retries + 1):
             try:
@@ -308,7 +324,37 @@ class OrizuEventSink:
         prompt_version_id = self.client.promote_candidate(self.run_id, **kwargs)
         # Promotion writes a system event at max(sequence)+1 in Postgres.
         self.sequence += 1
+        if self.local_logger is not None:
+            self.local_logger.append_event({
+                "run_id": self.run_id,
+                "sequence": self.sequence,
+                "event_type": "candidate_promoted",
+                "event_layer": "system",
+                "optimizer_family": "gepa",
+                "iteration": None,
+                "candidate_id": kwargs.get("candidate_id"),
+                "parent_candidate_id": None,
+                "child_candidate_id": None,
+                "payload": {
+                    **kwargs,
+                    "prompt_version_id": prompt_version_id,
+                },
+            })
         return prompt_version_id
 
     def finish_run(self, **kwargs: Any) -> None:
+        if self.local_logger is not None:
+            self.sequence += 1
+            self.local_logger.append_event({
+                "run_id": self.run_id,
+                "sequence": self.sequence,
+                "event_type": "run_status_updated",
+                "event_layer": "system",
+                "optimizer_family": "gepa",
+                "iteration": None,
+                "candidate_id": kwargs.get("best_candidate_id"),
+                "parent_candidate_id": None,
+                "child_candidate_id": None,
+                "payload": kwargs,
+            })
         self.client.update_run(self.run_id, **kwargs)

package/vendor/orizu-gepa-python/src/orizu_gepa/local_log.py

@@ -0,0 +1,123 @@
+from __future__ import annotations
+
+import dataclasses
+import json
+from pathlib import Path
+from typing import Any
+
+from .optimizer import DatasetRow, PromptContext, RowEvaluation, TextGepaResult
+
+
+def _json_default(value: Any) -> Any:
+    if dataclasses.is_dataclass(value):
+        return dataclasses.asdict(value)
+    return str(value)
+
+
+def _write_json(path: Path, payload: dict[str, Any]) -> None:
+    path.write_text(
+        json.dumps(payload, indent=2, ensure_ascii=False, default=_json_default) + "\n",
+        encoding="utf-8",
+    )
+
+
+def _append_jsonl(path: Path, payload: dict[str, Any]) -> None:
+    with path.open("a", encoding="utf-8") as handle:
+        handle.write(json.dumps(payload, ensure_ascii=False, default=_json_default) + "\n")
+
+
+class LocalOptimizationLogger:
+    def __init__(self, directory: Path):
+        self.directory = directory
+        self.directory.mkdir(parents=True, exist_ok=True)
+
+    @classmethod
+    def create(cls, root: str | Path, run_id: str) -> "LocalOptimizationLogger":
+        return cls(Path(root) / run_id)
+
+    @property
+    def path(self) -> str:
+        return str(self.directory)
+
+    def write_context(
+        self,
+        *,
+        project: str,
+        run_id: str,
+        args: dict[str, Any],
+        prompt_context: PromptContext,
+        scorer_context: PromptContext,
+        trainset: list[DatasetRow],
+        valset: list[DatasetRow],
+        metadata: dict[str, Any],
+    ) -> None:
+        _write_json(self.directory / "run.json", {
+            "schema_version": "orizu.optimization-local-log.v1",
+            "optimization_run_id": run_id,
+            "project": project,
+            "args": args,
+            "metadata": metadata,
+        })
+        _write_json(self.directory / "prompt_context.json", dataclasses.asdict(prompt_context))
+        _write_json(self.directory / "scorer_context.json", dataclasses.asdict(scorer_context))
+        _write_json(self.directory / "trainset.json", {
+            "rows": [dataclasses.asdict(row) for row in trainset],
+        })
+        _write_json(self.directory / "valset.json", {
+            "rows": [dataclasses.asdict(row) for row in valset],
+        })
+
+    def append_event(self, event: dict[str, Any]) -> None:
+        _append_jsonl(self.directory / "events.jsonl", event)
+
+    def append_evaluations(
+        self,
+        *,
+        stage: str,
+        split: str,
+        iteration: int | None,
+        candidate_id: str,
+        parent_candidate_id: str | None = None,
+        results: list[RowEvaluation],
+    ) -> None:
+        for result in results:
+            _append_jsonl(self.directory / "evaluations.jsonl", {
+                "stage": stage,
+                "split": split,
+                "iteration": iteration,
+                "candidate_id": candidate_id,
+                "parent_candidate_id": parent_candidate_id,
+                **result.to_payload(),
+            })
+
+    def append_reflection(
+        self,
+        *,
+        iteration: int,
+        parent_candidate_id: str,
+        child_candidate_id: str,
+        row_ids: list[str],
+        prompt: str,
+        response: str,
+        candidate_text: str,
+    ) -> None:
+        _append_jsonl(self.directory / "reflections.jsonl", {
+            "iteration": iteration,
+            "parent_candidate_id": parent_candidate_id,
+            "child_candidate_id": child_candidate_id,
+            "row_ids": row_ids,
+            "prompt": prompt,
+            "response": response,
+            "candidate_text": candidate_text,
+        })
+
+    def write_result(self, result: TextGepaResult) -> None:
+        _write_json(self.directory / "result.json", {
+            "optimization_run_id": result.run_id,
+            "best_candidate_id": result.best_candidate_id,
+            "best_candidate_text": result.best_candidate_text,
+            "best_score": result.best_score,
+            "seed_score": result.seed_score,
+            "promoted_prompt_version_id": result.promoted_prompt_version_id,
+            "budget": result.budget.to_payload(),
+        })

package/vendor/orizu-gepa-python/src/orizu_gepa/optimizer.py

@@ -658,6 +658,7 @@ def optimize_loaded_text_candidate(
     reflector: Reflector,
     event_sink: EventSink,
     config: TextGepaConfig,
+    local_logger: Any | None = None,
 ) -> TextGepaResult:
     seed_text = prompt_context.body or ""
     budget = Budget.from_config(config)
@@ -709,6 +710,14 @@ def optimize_loaded_text_candidate(
             budget=budget,
             evaluation_cache=evaluation_cache,
         )
+        if local_logger is not None:
+            local_logger.append_evaluations(
+                stage="seed_val_set",
+                split="validation",
+                iteration=None,
+                candidate_id="seed",
+                results=seed_results,
+            )
         if not _completed_all_rows(seed_results, valset):
             raise RuntimeError("Budget exhausted before seed validation completed")
         budget.used_full_evals += 1
@@ -804,6 +813,14 @@ def optimize_loaded_text_candidate(
                 budget=budget,
                 evaluation_cache=evaluation_cache,
             )
+            if local_logger is not None:
+                local_logger.append_evaluations(
+                    stage="parent_minibatch",
+                    split="train",
+                    iteration=iteration,
+                    candidate_id=parent_candidate_id,
+                    results=parent_results,
+                )
             if not _completed_all_rows(parent_results, minibatch):
                 event_sink.log_event(
                     "budget_exhausted",
@@ -852,11 +869,21 @@ def optimize_loaded_text_candidate(
             child_text = reflection.candidate_text
             candidate_text_by_id[child_id] = child_text
             budget.used_candidate_proposals += 1
+            if local_logger is not None:
+                local_logger.append_reflection(
+                    iteration=iteration,
+                    parent_candidate_id=parent_candidate_id,
+                    child_candidate_id=child_id,
+                    row_ids=[row.id for row in minibatch],
+                    prompt=reflection.prompt,
+                    response=reflection.response,
+                    candidate_text=reflection.candidate_text,
+                )
             event_sink.log_event(
                 "reflection_completed",
                 {
                     **_text_log_fields("prompt", reflection.prompt, include_text=config.log_row_snapshots),
-                    **_text_log_fields("response", reflection.response, include_text=config.log_row_snapshots),
+                    "response": reflection.response,
                     "child_candidate_id": child_id,
                     "candidate_text": child_text,
                 },
@@ -908,6 +935,15 @@ def optimize_loaded_text_candidate(
                 budget=budget,
                 evaluation_cache=evaluation_cache,
             )
+            if local_logger is not None:
+                local_logger.append_evaluations(
+                    stage="child_minibatch",
+                    split="train",
+                    iteration=iteration,
+                    candidate_id=child_id,
+                    parent_candidate_id=parent_candidate_id,
+                    results=child_results,
+                )
             if not _completed_all_rows(child_results, minibatch):
                 event_sink.log_event(
                     "budget_exhausted",
@@ -981,6 +1017,15 @@ def optimize_loaded_text_candidate(
                     budget=budget,
                     evaluation_cache=evaluation_cache,
                 )
+                if local_logger is not None:
+                    local_logger.append_evaluations(
+                        stage="child_val_set",
+                        split="validation",
+                        iteration=iteration,
+                        candidate_id=child_id,
+                        parent_candidate_id=parent_candidate_id,
+                        results=child_val_results,
+                    )
                 if not _completed_all_rows(child_val_results, valset):
                     event_sink.log_event(
                         "budget_exhausted",