omegaprompt 0.1.0__py3-none-any.whl

omegaprompt/__init__.py

@@ -0,0 +1,51 @@
+"""omegaprompt - Calibration discipline for Claude API prompts.
+
+Sensitivity-driven coordinate descent and walk-forward validation, ported
+from the omega-lock parameter-calibration framework to the prompt-engineering
+setting. Treats prompt parameters (system variant, few-shot count, effort
+level, max_tokens bucket, thinking on/off) as a calibratable parameter
+space; measures stress on each axis; runs grid search in the top-K unlock
+subspace; validates on a held-out walk-forward slice. LLM-as-judge scoring
+with user-supplied rubrics; hard gates on refusal / format / safety
+collapse the fitness to zero, so overfitting a single metric does not pass.
+
+Public API:
+    from omegaprompt import (
+        PromptTarget, PromptSpace, ParamVariants,
+        Dataset, DatasetItem,
+        JudgeRubric, Dimension, HardGate, JudgeResult,
+        CompositeFitness,
+    )
+
+Depends on:
+    omega-lock (>=0.1.4)  - provides CalibrableTarget, run_p1, stress,
+                            walk-forward, kill criteria, and benchmark.
+    anthropic (>=0.40.0)  - Claude API SDK.
+"""
+
+from omegaprompt.dataset import Dataset, DatasetItem
+from omegaprompt.fitness import CompositeFitness
+from omegaprompt.judge import JudgeRubric, Dimension, HardGate, JudgeResult
+from omegaprompt.schema import (
+    ParamVariants,
+    PromptSpace,
+    CalibrationOutcome,
+)
+from omegaprompt.target import PromptTarget
+
+__version__ = "0.1.0"
+
+__all__ = [
+    "PromptTarget",
+    "PromptSpace",
+    "ParamVariants",
+    "Dataset",
+    "DatasetItem",
+    "JudgeRubric",
+    "Dimension",
+    "HardGate",
+    "JudgeResult",
+    "CompositeFitness",
+    "CalibrationOutcome",
+    "__version__",
+]

omegaprompt/__main__.py

@@ -0,0 +1,7 @@
+"""Module entry point: ``python -m omegaprompt``."""
+
+from omegaprompt.cli import app
+
+
+if __name__ == "__main__":
+    app()

omegaprompt/api.py

@@ -0,0 +1,192 @@
+"""Anthropic API wrappers for both target and judge calls.
+
+Two separate client boundaries:
+
+- ``call_target(...)`` - issues the user's prompt to the target model and
+  returns the raw response text. No schema enforcement; the target is
+  allowed to respond however it wants (that is what the judge grades).
+- ``call_judge(...)`` - issues the rubric + response to the judge model
+  via ``messages.parse(output_format=JudgeResult)``. Schema enforcement
+  at the SDK boundary guarantees the judge cannot return malformed data.
+
+Both wrappers accept the client via a duck-typed Protocol so tests mock
+without importing ``anthropic``. Prompt caching is applied to the judge's
+system prompt - the same prompt is used thousands of times per run so
+cache hits dominate the judge cost.
+"""
+
+from __future__ import annotations
+
+from typing import Any, Protocol
+
+from omegaprompt.judge import JudgeResult, JudgeRubric
+from omegaprompt.prompts import JUDGE_SYSTEM_PROMPT
+
+
+class _MessagesNamespace(Protocol):
+    def create(self, **kwargs: Any) -> Any: ...  # noqa: D401
+
+    def parse(self, **kwargs: Any) -> Any: ...  # noqa: D401
+
+
+class _AnthropicLike(Protocol):
+    messages: _MessagesNamespace
+
+
+# Effort levels the target is willing to accept. The calibration searcher
+# emits an integer in {0, 1, 2}; these labels are the corresponding string
+# values the Anthropic SDK expects.
+_EFFORT_LABELS = ("low", "medium", "high")
+
+# Max-tokens buckets the searcher picks between. Picked to exercise short
+# vs. long output behavior of typical frontier models; the calibration
+# will surface which bucket is appropriate for the task.
+_MAX_TOKENS_BUCKETS = (1024, 4096, 16000)
+
+
+def effort_from_int(idx: int) -> str:
+    """Map a calibration integer index to the vendor's effort label."""
+    idx = max(0, min(len(_EFFORT_LABELS) - 1, idx))
+    return _EFFORT_LABELS[idx]
+
+
+def max_tokens_from_int(idx: int) -> int:
+    """Map a calibration integer index to a concrete max_tokens value."""
+    idx = max(0, min(len(_MAX_TOKENS_BUCKETS) - 1, idx))
+    return _MAX_TOKENS_BUCKETS[idx]
+
+
+def call_target(
+    client: _AnthropicLike,
+    *,
+    model: str,
+    system_prompt: str,
+    user_message: str,
+    few_shots: list[dict[str, str]],
+    effort: str,
+    max_tokens: int,
+    thinking_enabled: bool,
+) -> tuple[str, dict[str, int]]:
+    """Call the target model.
+
+    Returns ``(response_text, usage_dict)``. ``response_text`` is the
+    concatenated text of every text content block in the response.
+    """
+    messages: list[dict[str, Any]] = []
+    for shot in few_shots:
+        messages.append({"role": "user", "content": shot["input"]})
+        messages.append({"role": "assistant", "content": shot["output"]})
+    messages.append({"role": "user", "content": user_message})
+
+    kwargs: dict[str, Any] = {
+        "model": model,
+        "max_tokens": max_tokens,
+        "system": [
+            {
+                "type": "text",
+                "text": system_prompt,
+                "cache_control": {"type": "ephemeral"},
+            }
+        ],
+        "messages": messages,
+    }
+    if thinking_enabled:
+        kwargs["thinking"] = {"type": "adaptive"}
+        kwargs["output_config"] = {"effort": effort}
+
+    response = client.messages.create(**kwargs)
+
+    text_parts: list[str] = []
+    for block in getattr(response, "content", []) or []:
+        if getattr(block, "type", None) == "text":
+            chunk = getattr(block, "text", "")
+            if chunk:
+                text_parts.append(chunk)
+    return "\n".join(text_parts).strip(), _usage_to_dict(getattr(response, "usage", None))
+
+
+def call_judge(
+    client: _AnthropicLike,
+    *,
+    judge_model: str,
+    rubric: JudgeRubric,
+    task_input: str,
+    task_reference: str | None,
+    target_response: str,
+) -> tuple[JudgeResult, dict[str, int]]:
+    """Call the judge model with schema enforcement.
+
+    Returns ``(JudgeResult, usage_dict)``. Raises ``RuntimeError`` if the
+    judge refuses or returns no parsed output.
+    """
+    user_payload = _build_judge_payload(rubric, task_input, task_reference, target_response)
+
+    response = client.messages.parse(
+        model=judge_model,
+        max_tokens=2048,
+        system=[
+            {
+                "type": "text",
+                "text": JUDGE_SYSTEM_PROMPT,
+                "cache_control": {"type": "ephemeral"},
+            }
+        ],
+        messages=[{"role": "user", "content": user_payload}],
+        output_format=JudgeResult,
+    )
+
+    if getattr(response, "stop_reason", None) == "refusal":
+        raise RuntimeError(
+            "Judge refused to score the response. This usually means the target "
+            "response contained content the judge's safety layer flagged. "
+            "Inspect the task input and response for anything that could trigger "
+            "the judge's refusal path."
+        )
+
+    parsed: JudgeResult | None = getattr(response, "parsed_output", None)
+    if parsed is None:
+        raise RuntimeError(
+            "Judge returned no parsed_output. The judge response did not conform "
+            "to the JudgeResult schema. This indicates a prompt-caching drift, a "
+            "judge model downgrade, or a transient SDK issue. "
+            f"stop_reason={getattr(response, 'stop_reason', None)!r}."
+        )
+    if not isinstance(parsed, JudgeResult):
+        parsed = JudgeResult.model_validate(parsed)
+
+    return parsed, _usage_to_dict(getattr(response, "usage", None))
+
+
+def _build_judge_payload(
+    rubric: JudgeRubric,
+    task_input: str,
+    task_reference: str | None,
+    target_response: str,
+) -> str:
+    """Render the judge-user-turn payload the system prompt expects."""
+    import json as _json
+
+    rubric_json = _json.dumps(
+        {
+            "dimensions": [d.model_dump() for d in rubric.dimensions],
+            "hard_gates": [g.model_dump() for g in rubric.hard_gates],
+        },
+        ensure_ascii=False,
+        indent=2,
+    )
+    ref_block = f"<reference>\n{task_reference}\n</reference>\n\n" if task_reference else ""
+    return (
+        f"<rubric>\n{rubric_json}\n</rubric>\n\n"
+        f"<input>\n{task_input}\n</input>\n\n"
+        f"{ref_block}"
+        f"<response>\n{target_response}\n</response>"
+    )
+
+
+def _usage_to_dict(usage: Any) -> dict[str, int]:
+    return {
+        "input_tokens": getattr(usage, "input_tokens", 0) or 0,
+        "output_tokens": getattr(usage, "output_tokens", 0) or 0,
+        "cache_creation_input_tokens": getattr(usage, "cache_creation_input_tokens", 0) or 0,
+        "cache_read_input_tokens": getattr(usage, "cache_read_input_tokens", 0) or 0,
+    }

omegaprompt/cli.py

@@ -0,0 +1,43 @@
+"""Top-level Typer application for omegaprompt."""
+
+import typer
+
+from omegaprompt import __version__
+from omegaprompt.commands import calibrate as calibrate_cmd
+
+app = typer.Typer(
+    name="omegaprompt",
+    help="Calibration discipline for Claude API prompts.",
+    no_args_is_help=True,
+    add_completion=False,
+)
+
+
+def _version_callback(value: bool) -> None:
+    if value:
+        typer.echo(f"omegaprompt {__version__}")
+        raise typer.Exit()
+
+
+@app.callback()
+def _root(
+    version: bool = typer.Option(  # noqa: B008
+        False,
+        "--version",
+        "-V",
+        callback=_version_callback,
+        is_eager=True,
+        help="Show version and exit.",
+    ),
+) -> None:
+    """omegaprompt - apply omega-lock's calibration discipline to prompts."""
+
+
+app.command(
+    name="calibrate",
+    help="Calibrate a prompt configuration against a dataset (stress + grid + walk-forward).",
+)(calibrate_cmd.calibrate)
+
+
+if __name__ == "__main__":
+    app()

omegaprompt/commands/__init__.py

@@ -0,0 +1 @@
+"""Subcommand modules for the omegaprompt CLI."""

omegaprompt/commands/calibrate.py

@@ -0,0 +1,232 @@
+"""`omegaprompt calibrate` - end-to-end calibration from a single command.
+
+Takes a training dataset, an optional test dataset, a judge rubric, and
+a ParamVariants config; runs omega-lock's P1 pipeline (stress ->
+top-K unlock -> grid -> walk-forward); writes a ``CalibrationOutcome``
+JSON artifact with the winning parameters and the generalization gap.
+
+This command requires omega-lock to be installed (``pip install omega-lock``)
+and ``ANTHROPIC_API_KEY`` to be set.
+"""
+
+from __future__ import annotations
+
+import json
+import os
+from pathlib import Path
+
+import typer
+
+from omegaprompt.dataset import Dataset
+from omegaprompt.judge import JudgeRubric
+from omegaprompt.schema import CalibrationOutcome, ParamVariants, PromptSpace
+
+
+def calibrate(
+    dataset_path: Path = typer.Argument(  # noqa: B008
+        ...,
+        help="Path to the training dataset (.jsonl).",
+        exists=True,
+        file_okay=True,
+        dir_okay=False,
+        readable=True,
+    ),
+    rubric_path: Path = typer.Option(  # noqa: B008
+        ...,
+        "--rubric",
+        "-r",
+        help="Path to the JudgeRubric JSON.",
+        exists=True,
+        file_okay=True,
+        dir_okay=False,
+        readable=True,
+    ),
+    variants_path: Path = typer.Option(  # noqa: B008
+        ...,
+        "--variants",
+        "-v",
+        help="Path to the ParamVariants JSON (system_prompts + few_shot_examples).",
+        exists=True,
+        file_okay=True,
+        dir_okay=False,
+        readable=True,
+    ),
+    test_path: Path | None = typer.Option(  # noqa: B008
+        None,
+        "--test",
+        "-t",
+        help="Path to the walk-forward test dataset (.jsonl). Recommended.",
+        exists=True,
+        file_okay=True,
+        dir_okay=False,
+        readable=True,
+    ),
+    output_path: Path = typer.Option(  # noqa: B008
+        Path("calibration_outcome.json"),
+        "--output",
+        "-o",
+        help="Where to write the CalibrationOutcome JSON artifact.",
+        file_okay=True,
+        dir_okay=False,
+    ),
+    target_model: str = typer.Option(  # noqa: B008
+        ...,
+        "--target-model",
+        help="Model string to calibrate (e.g. 'claude-haiku-4-5').",
+    ),
+    judge_model: str = typer.Option(  # noqa: B008
+        ...,
+        "--judge-model",
+        help="Model string to use as judge. Can equal --target-model.",
+    ),
+    method: str = typer.Option(  # noqa: B008
+        "p1",
+        "--method",
+        "-m",
+        help="Calibration method: 'p1' (grid + KC-4) or 'grid' (grid only).",
+        case_sensitive=False,
+    ),
+    unlock_k: int = typer.Option(  # noqa: B008
+        3,
+        "--unlock-k",
+        min=1,
+        help="How many top-stress parameters to unlock for grid search.",
+    ),
+    space_path: Path | None = typer.Option(  # noqa: B008
+        None,
+        "--space",
+        help="Optional PromptSpace JSON to override the default parameter bounds.",
+        exists=True,
+        file_okay=True,
+        dir_okay=False,
+        readable=True,
+    ),
+) -> None:
+    """Calibrate a Claude API prompt configuration against a dataset."""
+
+    if not os.getenv("ANTHROPIC_API_KEY"):
+        typer.secho(
+            "ANTHROPIC_API_KEY is not set. Export it (see https://console.anthropic.com) "
+            "before running `omegaprompt calibrate`.",
+            fg=typer.colors.RED,
+            err=True,
+        )
+        raise typer.Exit(code=2)
+
+    # Lazy imports so the CLI can at least print --help without omega-lock
+    # / anthropic installed.
+    try:
+        from anthropic import Anthropic  # noqa: F401
+    except ImportError as exc:
+        typer.secho(
+            f"The 'anthropic' package is required for `calibrate`. Install with "
+            f"`pip install omegaprompt[anthropic]` or `pip install anthropic`. ({exc})",
+            fg=typer.colors.RED,
+            err=True,
+        )
+        raise typer.Exit(code=2) from exc
+    try:
+        from omega_lock import run_p1, P1Config  # noqa: F401
+    except ImportError as exc:
+        typer.secho(
+            f"The 'omega-lock' package is required for `calibrate`. Install with "
+            f"`pip install omega-lock`. ({exc})",
+            fg=typer.colors.RED,
+            err=True,
+        )
+        raise typer.Exit(code=2) from exc
+
+    from anthropic import Anthropic
+    from omega_lock import run_p1, P1Config
+
+    from omegaprompt.target import PromptTarget
+
+    typer.secho(f"Loading dataset from {dataset_path} ...", fg=typer.colors.BRIGHT_BLACK)
+    train_ds = Dataset.from_jsonl(dataset_path)
+    typer.secho(f"  {len(train_ds)} items", fg=typer.colors.BRIGHT_BLACK)
+
+    test_ds: Dataset | None = None
+    if test_path is not None:
+        typer.secho(f"Loading test set from {test_path} ...", fg=typer.colors.BRIGHT_BLACK)
+        test_ds = Dataset.from_jsonl(test_path)
+        typer.secho(f"  {len(test_ds)} items", fg=typer.colors.BRIGHT_BLACK)
+
+    rubric = JudgeRubric.from_json(rubric_path)
+    variants_payload = json.loads(rubric_path.read_text(encoding="utf-8")) if False else None
+    variants = ParamVariants.model_validate_json(variants_path.read_text(encoding="utf-8"))
+    space: PromptSpace | None = None
+    if space_path is not None:
+        space = PromptSpace.model_validate_json(space_path.read_text(encoding="utf-8"))
+
+    client = Anthropic()
+    train_target = PromptTarget(
+        target_client=client,
+        judge_client=client,
+        dataset=train_ds,
+        rubric=rubric,
+        variants=variants,
+        space=space,
+        target_model=target_model,
+        judge_model=judge_model,
+    )
+    test_target: PromptTarget | None = None
+    if test_ds is not None:
+        test_target = PromptTarget(
+            target_client=client,
+            judge_client=client,
+            dataset=test_ds,
+            rubric=rubric,
+            variants=variants,
+            space=space,
+            target_model=target_model,
+            judge_model=judge_model,
+        )
+
+    typer.secho(
+        f"Starting omega-lock run_p1 calibration (unlock_k={unlock_k}, method={method}) ...",
+        fg=typer.colors.BRIGHT_BLACK,
+    )
+    typer.secho(
+        "This issues Claude API calls. Budget accordingly.",
+        fg=typer.colors.YELLOW,
+    )
+
+    config = P1Config(unlock_k=unlock_k)
+    result = run_p1(
+        train_target=train_target,
+        test_target=test_target,
+        config=config,
+    )
+
+    best_params = getattr(result, "grid_best", None) or {}
+    best_params_dict = best_params.get("unlocked", {}) if isinstance(best_params, dict) else {}
+    best_fitness = float(best_params.get("fitness", 0.0)) if isinstance(best_params, dict) else 0.0
+    test_fitness = None
+    if test_target is not None and isinstance(best_params, dict):
+        test_fitness = best_params.get("test_fitness")
+    gen_gap = None
+    if test_fitness is not None and best_fitness != 0:
+        gen_gap = abs(best_fitness - float(test_fitness)) / abs(best_fitness)
+
+    outcome = CalibrationOutcome(
+        best_params=best_params_dict,
+        best_fitness=best_fitness,
+        test_fitness=float(test_fitness) if test_fitness is not None else None,
+        generalization_gap=gen_gap,
+        hard_gate_pass_rate=train_target._fitness.pass_rate(),
+        n_candidates_evaluated=train_target.total_api_calls
+        // max(len(train_ds) * 2, 1),
+        total_api_calls=train_target.total_api_calls
+        + (test_target.total_api_calls if test_target else 0),
+        method=method,
+        usage_summary=dict(train_target.last_usage),
+    )
+
+    output_path.write_text(outcome.model_dump_json(indent=2) + "\n", encoding="utf-8")
+    typer.secho(
+        f"Calibration complete. best_fitness={best_fitness:.4f}"
+        + (f", test_fitness={test_fitness:.4f}" if test_fitness is not None else "")
+        + (f", gen_gap={gen_gap:.2%}" if gen_gap is not None else ""),
+        fg=typer.colors.GREEN,
+    )
+    typer.secho(f"Artifact: {output_path}", fg=typer.colors.GREEN)

omegaprompt/dataset.py

@@ -0,0 +1,87 @@
+"""Dataset loading for omegaprompt calibration.
+
+The expected format is JSON Lines (``.jsonl``) with one item per line:
+
+.. code-block:: json
+
+   {"id": "t1", "input": "...", "reference": "...", "metadata": {...}}
+
+``reference`` is optional — judges can score without a reference when the
+rubric is purely qualitative. ``metadata`` is a free-form dict that the
+judge may consult (e.g. difficulty tier, target domain).
+"""
+
+from __future__ import annotations
+
+import json
+from collections.abc import Iterable
+from pathlib import Path
+
+from pydantic import BaseModel, ConfigDict, Field
+
+
+class DatasetItem(BaseModel):
+    """A single calibration example."""
+
+    model_config = ConfigDict(extra="allow")
+
+    id: str = Field(..., min_length=1)
+    input: str = Field(..., min_length=1, description="Input given to the target model.")
+    reference: str | None = Field(
+        default=None,
+        description="Optional reference / expected output for the judge.",
+    )
+    metadata: dict = Field(default_factory=dict)
+
+
+class Dataset(BaseModel):
+    """A set of ``DatasetItem`` s loaded from a JSONL file."""
+
+    items: list[DatasetItem]
+
+    @classmethod
+    def from_jsonl(cls, path: str | Path) -> Dataset:
+        """Read one item per line from ``path``."""
+        p = Path(path)
+        if not p.exists():
+            raise FileNotFoundError(f"Dataset not found: {p}")
+        items: list[DatasetItem] = []
+        with p.open("r", encoding="utf-8") as fh:
+            for lineno, raw in enumerate(fh, start=1):
+                line = raw.strip()
+                if not line:
+                    continue
+                try:
+                    payload = json.loads(line)
+                except json.JSONDecodeError as exc:
+                    raise ValueError(
+                        f"{p}:{lineno} - not valid JSON: {exc.msg}"
+                    ) from exc
+                try:
+                    items.append(DatasetItem.model_validate(payload))
+                except Exception as exc:
+                    raise ValueError(f"{p}:{lineno} - schema invalid: {exc}") from exc
+        if not items:
+            raise ValueError(f"{p} contained zero items (empty or blank-only lines).")
+        ids = [it.id for it in items]
+        if len(set(ids)) != len(ids):
+            dupes = sorted({i for i in ids if ids.count(i) > 1})
+            raise ValueError(f"{p} has duplicate ids: {dupes}")
+        return cls(items=items)
+
+    @classmethod
+    def from_items(cls, items: Iterable[DatasetItem | dict]) -> Dataset:
+        """Construct a Dataset from an in-memory iterable. Useful in tests."""
+        parsed: list[DatasetItem] = []
+        for it in items:
+            if isinstance(it, DatasetItem):
+                parsed.append(it)
+            else:
+                parsed.append(DatasetItem.model_validate(it))
+        return cls(items=parsed)
+
+    def __len__(self) -> int:
+        return len(self.items)
+
+    def __iter__(self):
+        return iter(self.items)

omegaprompt/fitness.py

@@ -0,0 +1,75 @@
+"""Composite fitness: ``hard_gate × soft_score``.
+
+The omegaprompt fitness rule is deliberately non-negotiable: if any hard
+gate fails on a dataset item, that item's contribution is zero, period.
+This is the structural defense against prompts that score beautifully on
+the soft rubric but refuse on a subset of inputs (a common failure mode
+with aggressive system prompts) or emit malformed output.
+"""
+
+from __future__ import annotations
+
+from collections.abc import Iterable
+from dataclasses import dataclass
+
+from omegaprompt.judge import JudgeResult, JudgeRubric
+
+
+@dataclass(frozen=True)
+class PerItemScore:
+    """Fitness breakdown for a single dataset item."""
+
+    item_id: str
+    soft_score: float
+    gates_passed: bool
+    final_score: float  # soft_score if gates_passed else 0.0
+    notes: str = ""
+
+
+class CompositeFitness:
+    """Aggregate per-item judge results into a scalar fitness.
+
+    The aggregation policy (mean by default) is intentionally simple and
+    boring. Fancy aggregations (robust mean, quantile-weighted, etc.) live
+    downstream of the raw per-item record preserved in ``last_per_item``.
+    """
+
+    def __init__(self, rubric: JudgeRubric) -> None:
+        self.rubric = rubric
+        self.last_per_item: list[PerItemScore] = []
+
+    def evaluate(
+        self,
+        judge_results: Iterable[tuple[str, JudgeResult]],
+    ) -> float:
+        """Score a batch of (item_id, judge_result) pairs.
+
+        Returns the mean final score across all items in the batch. The
+        per-item breakdown is retained on ``self.last_per_item`` for the
+        caller to inspect (useful for populating ``metadata`` on the
+        omega-lock EvalResult).
+        """
+        per_item: list[PerItemScore] = []
+        for item_id, jr in judge_results:
+            soft = jr.weighted_score(self.rubric)
+            passed = not jr.any_gate_failed()
+            per_item.append(
+                PerItemScore(
+                    item_id=item_id,
+                    soft_score=soft,
+                    gates_passed=passed,
+                    final_score=soft if passed else 0.0,
+                    notes=jr.notes,
+                )
+            )
+
+        self.last_per_item = per_item
+        if not per_item:
+            return 0.0
+        return sum(p.final_score for p in per_item) / len(per_item)
+
+    def pass_rate(self) -> float:
+        """Fraction of items in the last batch that cleared every gate."""
+        if not self.last_per_item:
+            return 0.0
+        return sum(1 for p in self.last_per_item if p.gates_passed) / len(self.last_per_item)