forgeoptimizer 0.1.0__py3-none-any.whl

forgecli/optimizer/ponytail/__init__.py

@@ -0,0 +1,181 @@
+"""Ponytail prompt-optimizer integration.
+
+ForgeCLI integrates the [Ponytail](https://ponytail.dev/) ruleset
+behind a small :class:`PromptOptimizer` interface and applies it
+transparently to every chat call. Two implementations are shipped:
+
+* :class:`PonytailRulesetOptimizer` — a self-contained Python
+  implementation of the Ponytail "ladder" (lite / full / ultra).
+  Always available, no external dependencies.
+* :class:`PonytailCLIOptimizer` — an optional adapter that shells out
+  to an external ``ponytail`` binary if one is on ``PATH``.
+
+The :class:`CompositeOptimizer` picks between them at runtime.
+"""
+
+from __future__ import annotations
+
+from abc import ABC, abstractmethod
+from collections.abc import Sequence
+from dataclasses import dataclass
+from enum import Enum
+from typing import Any
+
+from forgecli.providers.base import ChatMessage, ChatRequest
+
+
+class Intensity(str, Enum):
+    """How aggressively to optimize prompts.
+
+    The semantics match the Ponytail commands exactly:
+
+    * ``off``   - no rewriting; pass prompts through unchanged.
+    * ``lite``  - default; appends a one-line hint naming the lazier
+                  alternative and lets the model decide.
+    * ``full``  - applies the full "ladder" ruleset and instructs the
+                  model to ship the shortest diff.
+    * ``ultra`` - aggressive YAGNI; instructs the model to challenge
+                  the rest of the requirement in the same breath.
+    """
+
+    OFF = "off"
+    LITE = "lite"
+    FULL = "full"
+    ULTRA = "ultra"
+
+    @classmethod
+    def parse(cls, value: str | Intensity | None) -> Intensity:
+        """Parse a string into an :class:`Intensity`, falling back to LITE."""
+        if value is None or value == "":
+            return cls.LITE
+        if isinstance(value, cls):
+            return value
+        try:
+            return cls(value.lower())
+        except ValueError as exc:
+            raise ValueError(
+                f"Unknown intensity {value!r}; expected one of {', '.join(i.value for i in cls)}"
+            ) from exc
+
+
+@dataclass(frozen=True)
+class OptimizedRequest:
+    """The output of :meth:`PromptOptimizer.optimize_chat`."""
+
+    request: ChatRequest
+    notes: tuple[str, ...] = ()
+    intensity: Any = Intensity.LITE
+    source: str = "ruleset"  # "ruleset" | "external" | "passthrough"
+
+
+class PromptOptimizer(ABC):
+    """Strategy interface for prompt optimization.
+
+    Implementations must be deterministic and side-effect free
+    (other than logging), so they can be invoked transparently before
+    every model call.
+    """
+
+    name: str = "abstract"
+
+    @abstractmethod
+    async def optimize_chat(self, request: ChatRequest) -> OptimizedRequest:
+        """Return an optimized copy of ``request``."""
+
+    async def is_available(self) -> bool:
+        """Return whether this optimizer is available to use."""
+        return True
+
+
+class CompositeOptimizer(PromptOptimizer):
+    """Pick the right optimizer for the configured :class:`Intensity`."""
+
+    name = "composite"
+
+    def __init__(
+        self,
+        *,
+        intensity: Intensity = Intensity.LITE,
+        ruleset: PromptOptimizer | None = None,
+        external: PromptOptimizer | None = None,
+    ) -> None:
+        self._intensity = intensity
+        self._ruleset = ruleset
+        self._external = external
+        self._sync_ruleset_intensity()
+
+    def _sync_ruleset_intensity(self) -> None:
+        """Propagate the composite's intensity to the ruleset, if compatible."""
+        ruleset = self._ruleset
+        if isinstance(ruleset, PonytailRulesetOptimizer):
+            ruleset.set_intensity(self._intensity)
+
+    @property
+    def intensity(self) -> Intensity:
+        return self._intensity
+
+    def set_intensity(self, intensity: Intensity | str) -> None:
+        self._intensity = Intensity.parse(intensity)
+        # Keep the ruleset in sync so its output reflects the new level
+        # when we fall back to it.
+        if isinstance(self._ruleset, PonytailRulesetOptimizer):
+            self._ruleset.set_intensity(self._intensity)
+
+    async def optimize_chat(self, request: ChatRequest) -> OptimizedRequest:
+        if self._intensity is Intensity.OFF:
+            return OptimizedRequest(
+                request=request,
+                notes=("optimizer off",),
+                intensity=Intensity.OFF,
+                source="passthrough",
+            )
+
+        if self._external is not None and await self._external.is_available():
+            result = await self._external.optimize_chat(request)
+            return OptimizedRequest(
+                request=result.request,
+                notes=result.notes,
+                intensity=self._intensity,
+                source="external",
+            )
+
+        if self._ruleset is None:
+            return OptimizedRequest(
+                request=request,
+                notes=("no ruleset registered",),
+                intensity=self._intensity,
+                source="passthrough",
+            )
+
+        result = await self._ruleset.optimize_chat(request)
+        return OptimizedRequest(
+            request=result.request,
+            notes=result.notes,
+            intensity=self._intensity,
+            source="ruleset",
+        )
+
+
+def _ensure_user_message(messages: Sequence[ChatMessage]) -> bool:
+    return any(m.role.value == "user" for m in messages)
+
+
+def _clone_request(request: ChatRequest, messages: Sequence[ChatMessage]) -> ChatRequest:
+    """Return a copy of ``request`` with ``messages`` replaced."""
+    return request.model_copy(update={"messages": list(messages)})
+
+
+__all__ = [
+    "CompositeOptimizer",
+    "Intensity",
+    "OptimizedProvider",
+    "OptimizedRequest",
+    "PonytailCLIOptimizer",
+    "PonytailRulesetOptimizer",
+    "PromptOptimizer",
+]
+
+# Re-export the two concrete implementations and the decorator.
+from forgecli.optimizer.ponytail.cli import PonytailCLIOptimizer  # noqa: E402
+from forgecli.optimizer.ponytail.decorator import OptimizedProvider  # noqa: E402
+from forgecli.optimizer.ponytail.ruleset import PonytailRulesetOptimizer  # noqa: E402

forgecli/optimizer/ponytail/cli.py

@@ -0,0 +1,159 @@
+"""Optional adapter for an external ``ponytail`` CLI.
+
+Some Ponytail installations (or future refactors) ship a binary on
+``PATH`` that can rewrite a prompt itself. This adapter detects the
+binary and forwards the conversation to it. When the binary is not
+installed, :meth:`is_available` returns ``False`` and the
+:class:`CompositeOptimizer` falls back to the in-process ruleset.
+"""
+
+from __future__ import annotations
+
+import asyncio
+import contextlib
+import os
+import shutil
+import tempfile
+from pathlib import Path
+
+from forgecli.core.errors import ForgeCLIError
+from forgecli.optimizer.ponytail import (
+    Intensity,
+    OptimizedRequest,
+    PromptOptimizer,
+)
+from forgecli.providers.base import ChatMessage, ChatRequest, Role
+
+
+class PonytailCLIError(ForgeCLIError):
+    """Raised when the external ``ponytail`` binary fails."""
+
+
+class PonytailCLIOptimizer(PromptOptimizer):
+    """Forward prompts to an external ``ponytail`` binary via stdin/stdout.
+
+    The CLI is expected to accept a JSON payload on stdin and emit a
+    JSON payload on stdout. When such a binary is not present the
+    adapter is harmless — :meth:`is_available` returns ``False`` and
+    the composite falls back to the in-process ruleset.
+    """
+
+    name = "ponytail.cli"
+
+    def __init__(
+        self,
+        *,
+        executable: str | None = None,
+        timeout: float = 30.0,
+    ) -> None:
+        self._executable = executable or os.environ.get("FORGECLI_PONYTAIL_BIN", "ponytail")
+        self._timeout = timeout
+
+    @property
+    def executable(self) -> str:
+        return self._executable
+
+    async def is_available(self) -> bool:
+        return shutil.which(self._executable) is not None
+
+    async def optimize_chat(self, request: ChatRequest) -> OptimizedRequest:
+        if not await self.is_available():
+            raise PonytailCLIError(f"Ponytail executable {self._executable!r} not found on PATH")
+
+        binary = shutil.which(self._executable) or self._executable
+        payload = _request_to_json(request)
+        with tempfile.NamedTemporaryFile(
+            "w", suffix=".json", delete=False, encoding="utf-8"
+        ) as src:
+            src.write(payload)
+            src_path = Path(src.name)
+        try:
+            proc = await asyncio.create_subprocess_exec(
+                binary,
+                "optimize",
+                "--stdin",
+                str(src_path),
+                stdout=asyncio.subprocess.PIPE,
+                stderr=asyncio.subprocess.PIPE,
+            )
+            try:
+                stdout, stderr = await asyncio.wait_for(proc.communicate(), timeout=self._timeout)
+            except TimeoutError as exc:
+                proc.kill()
+                raise PonytailCLIError(
+                    f"ponytail optimize timed out after {self._timeout}s"
+                ) from exc
+        finally:
+            with contextlib.suppress(OSError):
+                src_path.unlink()
+
+        if proc.returncode != 0:
+            raise PonytailCLIError(
+                f"ponytail exited with {proc.returncode}: {stderr.decode(errors='replace').strip()}"
+            )
+
+        return _json_to_optimized(stdout.decode(errors="replace"))
+
+
+def _request_to_json(request: ChatRequest) -> str:
+    """Serialize a :class:`ChatRequest` for the CLI."""
+    import json
+
+    return json.dumps(
+        {
+            "model": request.model,
+            "temperature": request.temperature,
+            "max_tokens": request.max_tokens,
+            "messages": [
+                {
+                    "role": m.role.value,
+                    "content": m.content,
+                    "name": m.name,
+                    "tool_call_id": m.tool_call_id,
+                }
+                for m in request.messages
+            ],
+        },
+        ensure_ascii=False,
+    )
+
+
+def _json_to_optimized(raw: str) -> OptimizedRequest:
+    """Parse the JSON the CLI emitted and rebuild an :class:`OptimizedRequest`."""
+    import json
+
+    try:
+        payload = json.loads(raw)
+    except json.JSONDecodeError as exc:
+        raise PonytailCLIError(f"ponytail returned invalid JSON: {exc}") from exc
+
+    try:
+        intensity = Intensity.parse(payload.get("intensity", "lite"))
+        messages = [
+            ChatMessage(
+                role=Role(m["role"]),
+                content=m["content"],
+                name=m.get("name"),
+                tool_call_id=m.get("tool_call_id"),
+            )
+            for m in payload.get("messages", [])
+        ]
+    except (KeyError, ValueError) as exc:
+        raise PonytailCLIError(f"ponytail returned an unexpected payload: {exc}") from exc
+
+    new_request = ChatRequest(
+        model=payload.get("model"),
+        messages=messages,
+        temperature=payload.get("temperature"),
+        max_tokens=payload.get("max_tokens"),
+    )
+    notes = tuple(payload.get("notes", ()))
+    return OptimizedRequest(
+        request=new_request,
+        notes=notes,
+        intensity=intensity,
+        source="external",
+    )
+
+
+__all__ = ["PonytailCLIError", "PonytailCLIOptimizer"]

forgecli/optimizer/ponytail/decorator.py

@@ -0,0 +1,70 @@
+"""A :class:`Provider` decorator that runs a :class:`PromptOptimizer`
+on every chat request before delegating to the wrapped provider.
+
+The decorator is transparent for every other operation (``stream``,
+``embed``, ``list_models``) — those either pass through unchanged or
+inherit the base behaviour.
+"""
+
+from __future__ import annotations
+
+from collections.abc import AsyncIterator
+from typing import Any
+
+from forgecli.optimizer.ponytail import OptimizedRequest, PromptOptimizer
+from forgecli.providers.base import (
+    ChatRequest,
+    ChatResponse,
+    EmbeddingRequest,
+    EmbeddingResponse,
+    ModelInfo,
+    Provider,
+    StreamChunk,
+)
+
+
+class OptimizedProvider(Provider[Any]):
+    """Wrap a :class:`Provider` so every chat call is pre-optimized."""
+
+    name = "optimized"
+
+    def __init__(
+        self,
+        base: Provider[Any],
+        optimizer: PromptOptimizer,
+    ) -> None:
+        # The wrapped provider already has its own concrete config; we
+        # expose it via ``.config`` for compatibility.
+        super().__init__(base.config)
+        self._base = base
+        self._optimizer = optimizer
+
+    @property
+    def base(self) -> Provider[Any]:
+        return self._base
+
+    @property
+    def optimizer(self) -> PromptOptimizer:
+        return self._optimizer
+
+    async def chat(self, request: ChatRequest) -> ChatResponse:
+        optimized: OptimizedRequest = await self._optimizer.optimize_chat(request)
+        return await self._base.chat(optimized.request)
+
+    async def stream(self, request: ChatRequest) -> AsyncIterator[StreamChunk]:
+        optimized: OptimizedRequest = await self._optimizer.optimize_chat(request)
+        async for chunk in self._base.stream(optimized.request):
+            yield chunk
+
+    async def embed(self, request: EmbeddingRequest) -> EmbeddingResponse:
+        # Embedding calls don't carry chat messages; pass through.
+        return await self._base.embed(request)
+
+    async def list_models(self) -> list[ModelInfo]:
+        return await self._base.list_models()
+
+    def validate(self) -> None:
+        return self._base.validate()
+
+
+__all__ = ["OptimizedProvider"]

forgecli/optimizer/ponytail/factory.py

@@ -0,0 +1,51 @@
+"""Factory for building the live :class:`PromptOptimizer`."""
+
+from __future__ import annotations
+
+from forgecli.config.settings import ForgeSettings
+from forgecli.optimizer.ponytail import (
+    CompositeOptimizer,
+    Intensity,
+    PonytailCLIOptimizer,
+    PonytailRulesetOptimizer,
+    PromptOptimizer,
+)
+from forgecli.optimizer.ponytail.state import OptimizerState
+
+
+def build_optimizer(
+    state: OptimizerState,
+    settings: ForgeSettings | None = None,
+) -> PromptOptimizer:
+    """Build the composite optimizer from the runtime state + config.
+
+    The composite always falls back to the in-process ruleset; the
+    external CLI adapter is added on top whenever the configured
+    backend asks for it and the binary is available.
+    """
+    ruleset = PonytailRulesetOptimizer(intensity=state.intensity)
+
+    external: PromptOptimizer | None = None
+    if state.backend in {"cli", "auto"}:
+        binary = state.binary
+        if settings is not None and settings.prompt_optimizer.binary:
+            binary = settings.prompt_optimizer.binary
+        external = PonytailCLIOptimizer(executable=binary)
+    elif state.backend == "ruleset":
+        external = None
+
+    if settings is not None and not settings.prompt_optimizer.enabled:
+        return CompositeOptimizer(
+            intensity=Intensity.OFF,
+            ruleset=ruleset,
+            external=external,
+        )
+
+    return CompositeOptimizer(
+        intensity=state.intensity,
+        ruleset=ruleset,
+        external=external,
+    )
+
+
+__all__ = ["build_optimizer"]

forgecli/optimizer/ponytail/ruleset.py

@@ -0,0 +1,168 @@
+"""Self-contained Python implementation of the Ponytail ruleset.
+
+The ruleset mirrors the behavior described on https://ponytail.dev/ and
+the project's GitHub README (DietrichGebert/ponytail). It does not
+shell out to anything; it rewrites the system message of a
+:class:`ChatRequest` to bias the model toward the lazier correct
+solution.
+
+The four intensity levels match the official ``/ponytail`` command:
+
+* ``off``   - pass-through (handled by the composite, not this class).
+* ``lite``  - default; appends a single sentence reminding the model to
+              name the lazier alternative and lets the user pick.
+* ``full``  - prepends the full "ladder" instruction and tells the
+              model to ship the shortest diff.
+* ``ultra`` - aggressive YAGNI; tells the model to ship the one-liner
+              and challenge the rest of the requirement in the same
+              breath.
+
+The output is deterministic for a given (intensity, prompt) pair, so
+tests can pin it.
+"""
+
+from __future__ import annotations
+
+from forgecli.optimizer.ponytail import (
+    CompositeOptimizer,
+    Intensity,
+    OptimizedRequest,
+    PromptOptimizer,
+    _clone_request,
+    _ensure_user_message,
+)
+from forgecli.providers.base import ChatMessage, ChatRequest, Role
+
+LADDER_INSTRUCTION = (
+    "Stop at the first rung that holds. Apply the Ponytail ladder before "
+    "writing any code:\n"
+    "  1. Does this need to exist? Speculative need = skip it (YAGNI).\n"
+    "  2. Already in this codebase? Reuse the helper, util, or pattern "
+    "that already lives here.\n"
+    "  3. Does the standard library do it? Use it.\n"
+    "  4. Native platform feature covers it? (e.g. <input type='date'> "
+    "over a picker lib.)\n"
+    "  5. Already-installed dependency solves it? Use it. Don't add a "
+    "new one.\n"
+    "  6. Can it be one line? One line.\n"
+    "  7. Only then: the minimum code that works."
+)
+
+
+_INTENSITY_GUIDANCE: dict[Intensity, str] = {
+    Intensity.OFF: "",
+    Intensity.LITE: (
+        "Ponytail (lite): when you reach for code, also name the lazier "
+        "correct alternative in a single sentence. The user picks."
+    ),
+    Intensity.FULL: (
+        "Ponytail (full): apply the Ponytail ladder before writing any code. "
+        "Ship the shortest diff and the shortest explanation. Never "
+        "simplify away validation, error handling, security, or accessibility.\n\n"
+        f"{LADDER_INSTRUCTION}"
+    ),
+    Intensity.ULTRA: (
+        "Ponytail (ultra): YAGNI extremist. Ship the one-liner. In the "
+        "same breath, name what is being *cut* from the original "
+        "requirement and why the cut is safe."
+    ),
+}
+
+
+class PonytailRulesetOptimizer(PromptOptimizer):
+    """Rewrite system messages to apply the Ponytail ruleset."""
+
+    name = "ponytail.ruleset"
+
+    def __init__(self, *, intensity: Intensity = Intensity.LITE) -> None:
+        self._intensity = Intensity.parse(intensity)
+
+    @property
+    def intensity(self) -> Intensity:
+        return self._intensity
+
+    def set_intensity(self, intensity: Intensity | str) -> None:
+        self._intensity = Intensity.parse(intensity)
+
+    async def optimize_chat(self, request: ChatRequest) -> OptimizedRequest:
+        if self._intensity is Intensity.OFF:
+            return OptimizedRequest(
+                request=request,
+                notes=("ponytail off",),
+                intensity=Intensity.OFF,
+                source="ruleset",
+            )
+
+        guidance = _INTENSITY_GUIDANCE[self._intensity]
+        messages = list(request.messages)
+
+        if not _ensure_user_message(messages):
+            # The model must see a user turn; without one the request is
+            # malformed. Pass through unchanged and let the caller decide.
+            return OptimizedRequest(
+                request=request,
+                notes=("no user message; passthrough",),
+                intensity=self._intensity,
+                source="ruleset",
+            )
+
+        rewritten = _rewrite_messages(messages, guidance)
+        notes = _build_notes(self._intensity, len(messages), len(rewritten))
+        return OptimizedRequest(
+            request=_clone_request(request, rewritten),
+            notes=notes,
+            intensity=self._intensity,
+            source="ruleset",
+        )
+
+
+def _rewrite_messages(
+    messages: list[ChatMessage],
+    guidance: str,
+) -> list[ChatMessage]:
+    """Insert Ponytail guidance at the head of the conversation.
+
+    Rules:
+    * If the first message is a system message, prepend the guidance
+      to it (separated by two newlines).
+    * If there is no system message, insert one.
+    * If the first system message is *empty*, replace it.
+    * Otherwise the user/assistant turns are left untouched.
+    """
+    if not messages:
+        return [ChatMessage(role=Role.SYSTEM, content=guidance)]
+
+    first = messages[0]
+    if first.role is Role.SYSTEM:
+        existing = first.content.strip()
+        if not existing:
+            return [ChatMessage(role=Role.SYSTEM, content=guidance), *messages[1:]]
+        new_content = f"{guidance}\n\n{first.content}" if guidance else first.content
+        return [ChatMessage(role=Role.SYSTEM, content=new_content), *messages[1:]]
+
+    return [ChatMessage(role=Role.SYSTEM, content=guidance), *messages]
+
+
+def _build_notes(
+    intensity: Intensity, original_count: int, rewritten_count: int
+) -> tuple[str, ...]:
+    notes: list[str] = [f"ponytail intensity={intensity.value}"]
+    if intensity is Intensity.FULL:
+        notes.append("ladder enforced")
+    elif intensity is Intensity.ULTRA:
+        notes.append("yagni extremist mode")
+    elif intensity is Intensity.LITE:
+        notes.append("named the lazier alternative")
+    if rewritten_count != original_count:
+        notes.append(f"inserted system message ({rewritten_count} > {original_count})")
+    return tuple(notes)
+
+
+__all__ = [
+    "LADDER_INSTRUCTION",
+    "PonytailRulesetOptimizer",
+]
+
+
+# Re-export the composite here so the public namespace stays compact.
+_ = CompositeOptimizer

forgecli/optimizer/ponytail/state.py

@@ -0,0 +1,51 @@
+"""Live runtime state for the prompt optimizer.
+
+The optimizer intensity can be flipped from the CLI via
+``forge optimizer lite|full|ultra|off``. Because the setting is held
+in :class:`AppContext.extras`, every subcommand that resolves the
+optimizer sees the latest value. Tests can override it directly.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+from forgecli.optimizer.ponytail import Intensity
+
+
+@dataclass
+class OptimizerState:
+    """Mutable intensity for the running CLI process."""
+
+    intensity: Intensity = Intensity.LITE
+    backend: str = "ruleset"  # "ruleset" | "cli" | "auto"
+    binary: str = "ponytail"
+
+    @classmethod
+    def from_extras(cls, extras: dict[str, object]) -> OptimizerState:
+        state = cls()
+        intensity = extras.get("optimizer.intensity") or extras.get("optimizer_intensity")
+        backend = extras.get("optimizer.backend") or extras.get("optimizer_backend")
+        binary = extras.get("optimizer.binary") or extras.get("optimizer_binary")
+        if isinstance(intensity, str):
+            try:
+                state.intensity = Intensity.parse(intensity)
+            except ValueError:
+                state.intensity = Intensity.LITE
+        elif isinstance(intensity, Intensity):
+            state.intensity = intensity
+        if isinstance(backend, str):
+            state.backend = backend
+        if isinstance(binary, str):
+            state.binary = binary
+        return state
+
+    def to_extras(self) -> dict[str, str]:
+        return {
+            "optimizer.intensity": self.intensity.value,
+            "optimizer.backend": self.backend,
+            "optimizer.binary": self.binary,
+        }
+
+
+__all__ = ["OptimizerState"]