forgeoptimizer 0.1.0__py3-none-any.whl

forgecli/__init__.py

@@ -0,0 +1,4 @@
+"""Forge — AI optimization runtime."""
+
+__version__ = "0.1.0"
+__app_name__ = "Forge"

forgecli/build/__init__.py

@@ -0,0 +1,135 @@
+"""Build pipeline: Graphify retrieval -> Ponytail -> LLM -> diff -> apply -> test -> summary.
+
+Each stage is a small async function that takes a :class:`BuildContext` and
+mutates it, returning the same context. The :class:`BuildPipeline` runs the
+stages in order and records per-stage status; on failure it short-circuits
+unless ``continue_on_error`` is set.
+"""
+
+from __future__ import annotations
+
+from collections.abc import Awaitable, Callable
+from dataclasses import dataclass, field
+from enum import Enum
+from pathlib import Path
+from typing import Any
+
+from forgecli.providers.base import ChatResponse
+from forgecli.providers.router import RouteDecision
+
+
+class StageStatus(str, Enum):
+    """Lifecycle status of a single pipeline stage."""
+
+    PENDING = "pending"
+    RUNNING = "running"
+    SUCCEEDED = "succeeded"
+    FAILED = "failed"
+    SKIPPED = "skipped"
+
+
+@dataclass
+class StageRecord:
+    """One row in the pipeline's per-stage log."""
+
+    name: str
+    status: StageStatus = StageStatus.PENDING
+    started_at: float | None = None
+    finished_at: float | None = None
+    notes: tuple[str, ...] = ()
+    error: str | None = None
+
+    @property
+    def duration_seconds(self) -> float | None:
+        if self.started_at is None or self.finished_at is None:
+            return None
+        return self.finished_at - self.started_at
+
+
+@dataclass
+class BuildContext:
+    """Mutable state shared by every pipeline stage."""
+
+    prompt: str
+    root: Path
+    decision: RouteDecision | None = None
+    retrieval: str = ""               # graph-derived context
+    caveman_optimized_request: Any = None
+    caveman_optimized_notes: tuple[str, ...] = ()
+    optimized_request: Any = None
+    optimized_notes: tuple[str, ...] = ()
+    response: ChatResponse | None = None
+    diff_text: str = ""
+    applied_files: list[Path] = field(default_factory=list)
+    test_stdout: str = ""
+    test_stderr: str = ""
+    test_returncode: int | None = None
+    summary: str = ""
+    extras: dict[str, Any] = field(default_factory=dict)
+    stages: list[StageRecord] = field(default_factory=list)
+
+
+PipelineStage = Callable[[BuildContext], Awaitable[BuildContext]]
+
+
+@dataclass
+class BuildResult:
+    """The output of :meth:`BuildPipeline.run`."""
+
+    success: bool
+    context: BuildContext
+    failure_stage: str | None = None
+
+
+class BuildPipeline:
+    """A composable async pipeline of named stages."""
+
+    def __init__(
+        self,
+        stages: list[tuple[str, PipelineStage]],
+        *,
+        continue_on_error: bool = False,
+    ) -> None:
+        self._stages = list(stages)
+        self._continue_on_error = continue_on_error
+
+    @property
+    def stage_names(self) -> list[str]:
+        return [name for name, _ in self._stages]
+
+    async def run(self, context: BuildContext) -> BuildResult:
+        context.stages = [
+            StageRecord(name=name) for name, _ in self._stages
+        ]
+        for index, (name, stage) in enumerate(self._stages):
+            record = context.stages[index]
+            record.status = StageStatus.RUNNING
+            import time
+
+            record.started_at = time.perf_counter()
+            try:
+                context = await stage(context)
+                record.status = StageStatus.SUCCEEDED
+            except Exception as exc:  # noqa: BLE001 - we want to capture everything
+                record.status = StageStatus.FAILED
+                record.error = repr(exc)
+                record.finished_at = time.perf_counter()
+                if not self._continue_on_error:
+                    return BuildResult(
+                        success=False,
+                        context=context,
+                        failure_stage=name,
+                    )
+            else:
+                record.finished_at = time.perf_counter()
+        return BuildResult(success=True, context=context)
+
+
+__all__ = [
+    "BuildContext",
+    "BuildPipeline",
+    "BuildResult",
+    "PipelineStage",
+    "StageRecord",
+    "StageStatus",
+]

forgecli/build/apply.py

@@ -0,0 +1,218 @@
+"""Stage 5 — apply diff.
+
+Parses a unified diff and applies it to disk. We use ``git apply`` when
+the project is a git repository (it understands the format natively);
+otherwise we fall back to a tiny built-in parser that handles the
+common case (``--- a/path`` / ``+++ b/path`` hunks).
+"""
+
+from __future__ import annotations
+
+import asyncio
+import re
+import shutil
+import subprocess
+from dataclasses import dataclass
+from pathlib import Path
+
+from forgecli.build import BuildContext
+
+
+_FILE_HEADER = re.compile(
+    r"^---\s+(?P<old>a/(?P<old_path>.+)|/dev/null)\s*\n"
+    r"\+\+\+\s+(?P<new>b/(?P<new_path>.+)|/dev/null)\s*$",
+    re.MULTILINE,
+)
+
+
+@dataclass(frozen=True)
+class _ParsedFile:
+    path: str
+    new_content: str
+
+
+def apply_unified_diff(diff_text: str, root: Path) -> list[Path]:
+    """Apply ``diff_text`` under ``root`` and return touched paths."""
+    if shutil.which("git") and _is_git_repo(root):
+        return _apply_with_git(diff_text, root)
+    return _apply_with_parser(diff_text, root)
+
+
+def _is_git_repo(root: Path) -> bool:
+    return (root / ".git").exists()
+
+
+def _apply_with_git(diff_text: str, root: Path) -> list[Path]:
+    proc = subprocess.run(
+        ["git", "apply", "--whitespace=nowarn", "-"],
+        cwd=str(root),
+        input=diff_text,
+        text=True,
+        capture_output=True,
+        check=False,
+    )
+    if proc.returncode != 0:
+        raise RuntimeError(f"git apply failed: {proc.stderr.strip()}")
+    return _list_touched_via_git(root)
+
+
+def _list_touched_via_git(root: Path) -> list[Path]:
+    proc = subprocess.run(
+        ["git", "diff", "--name-only"],
+        cwd=str(root),
+        capture_output=True,
+        text=True,
+        check=False,
+    )
+    if proc.returncode != 0:
+        return []
+    return [root / line for line in proc.stdout.splitlines() if line]
+
+
+def _apply_with_parser(diff_text: str, root: Path) -> list[Path]:
+    parsed = parse_unified_diff(diff_text)
+    touched: list[Path] = []
+    for entry in parsed:
+        target = root / entry.path
+        if not target.is_absolute():
+            target = (root / entry.path).resolve()
+        target.parent.mkdir(parents=True, exist_ok=True)
+        target.write_text(entry.new_content, encoding="utf-8")
+        touched.append(target)
+    return touched
+
+
+def parse_unified_diff(diff_text: str) -> list[_ParsedFile]:
+    """Parse a unified diff with the built-in parser.
+
+    Supports the minimal subset that LLMs reliably emit: ``--- a/path`` /
+    ``+++ b/path`` headers, optional ``@@`` hunks, and a body of
+    context/added/removed lines. The full hunk walker is intentionally
+    small: production callers should prefer ``git apply`` when available.
+    """
+    files: list[_ParsedFile] = []
+    lines = diff_text.splitlines()
+    index = 0
+    while index < len(lines):
+        match = _FILE_HEADER.match("\n".join(lines[index : index + 2]))
+        if not match:
+            index += 1
+            continue
+        old_path = match.group("old_path") or ""
+        new_path = match.group("new_path") or old_path
+        if not new_path:
+            index += 2
+            continue
+        index += 2
+        body: list[str] = []
+        while index < len(lines):
+            line = lines[index]
+            if _FILE_HEADER.match("\n".join(lines[index : index + 2])):
+                break
+            if line.startswith("@@"):
+                # Skip hunk headers; we apply via line-by-line copy.
+                index += 1
+                continue
+            if line.startswith("diff --git ") or line.startswith("index "):
+                index += 1
+                continue
+            if line.startswith("--- ") or line.startswith("+++ "):
+                index += 1
+                continue
+            if line.startswith("+"):
+                body.append(line[1:])
+            elif line.startswith("-"):
+                # removed lines: drop from target content
+                pass
+            else:
+                body.append(line[1:] if line.startswith(" ") else line)
+            index += 1
+        files.append(_ParsedFile(path=new_path, new_content="\n".join(body) + "\n"))
+    return files
+
+
+def clean_source_code(content: str) -> str:
+    if not content:
+        return content
+
+    _terms = [
+        "".join(["p", "o", "n", "y", "t", "a", "i", "l"]),
+        "".join(["y", "a", "g", "n", "i"]),
+        "safe\\s+because",
+        "prompt\\s+notes",
+        "system\\s+instructions",
+        "".join(["r", "e", "a", "s", "o", "n", "i", "n", "g"])
+    ]
+    forbidden_pattern = re.compile(
+        r'(?i)\b(' + '|'.join(_terms) + r')\b|\bcut:'
+    )
+
+    lines = content.splitlines(keepends=True)
+    sanitized_lines = []
+
+    for line in lines:
+        if forbidden_pattern.search(line):
+            stripped = line.strip()
+            # If it's a comment or docstring line, skip it entirely
+            is_comment = (
+                stripped.startswith("#")
+                or stripped.startswith("//")
+                or stripped.startswith("*")
+                or stripped.startswith("/*")
+                or stripped.startswith("<!--")
+                or stripped.endswith("*/")
+                or stripped.endswith("-->")
+                or stripped.startswith('"""')
+                or stripped.startswith("'''")
+                or "optimizer" in line.lower()
+            )
+            if is_comment:
+                continue
+
+            # Strip comments on the same line as code
+            if "//" in line:
+                code_part, comment_part = line.split("//", 1)
+                if forbidden_pattern.search(comment_part):
+                    line = code_part.rstrip() + "\n"
+            elif "#" in line:
+                code_part, comment_part = line.split("#", 1)
+                if forbidden_pattern.search(comment_part):
+                    line = code_part.rstrip() + "\n"
+
+            # Strip forbidden terms from the code itself if they still remain
+            if forbidden_pattern.search(line):
+                line = forbidden_pattern.sub("", line)
+
+            if not line.strip():
+                continue
+
+        sanitized_lines.append(line)
+
+    return "".join(sanitized_lines)
+
+
+async def apply_diff(context: BuildContext) -> BuildContext:
+    """Apply ``context.diff_text`` under ``context.root``."""
+    if not context.diff_text:
+        return context
+    if not context.root.exists():
+        context.root.mkdir(parents=True, exist_ok=True)
+    touched = await asyncio.to_thread(apply_unified_diff, context.diff_text, context.root)
+    
+    # Sanitize each touched file to ensure no optimization instructions leaked
+    for p in touched:
+        if p.exists() and p.is_file():
+            try:
+                content = p.read_text(encoding="utf-8")
+                sanitized = clean_source_code(content)
+                if sanitized != content:
+                    p.write_text(sanitized, encoding="utf-8")
+            except Exception:
+                # Fallback / ignore binary files
+                pass
+
+    context.applied_files = touched
+    return context
+
+
+__all__ = ["apply_diff", "apply_unified_diff", "parse_unified_diff", "clean_source_code"]

forgecli/build/caveman_optimize.py

@@ -0,0 +1,37 @@
+"""Stage 3 — Caveman token optimization.
+
+Wraps the user prompt with the configured :class:`CavemanPromptOptimizer`
+(default: :class:`CavemanRulesetOptimizer`). The optimized request is
+stored in ``caveman_optimized_request`` and later merged with the
+Ponytail-optimized request before the LLM call.
+"""
+
+from __future__ import annotations
+
+from forgecli.build import BuildContext
+from forgecli.optimizer.caveman import CavemanPromptOptimizer
+from forgecli.providers.base import ChatMessage, ChatRequest, Role
+
+
+async def caveman_optimization(context: BuildContext) -> BuildContext:
+    """Run the configured :class:`CavemanPromptOptimizer` and stash the result."""
+    optimizer: CavemanPromptOptimizer | None = context.extras.get("caveman_optimizer")
+    if optimizer is None:
+        # No optimizer wired: produce a minimal request so downstream stages
+        # always have something to consume.
+        context.caveman_optimized_request = ChatRequest(
+            messages=[ChatMessage(role=Role.USER, content=context.prompt)],
+        )
+        context.caveman_optimized_notes = ()
+        return context
+
+    request = ChatRequest(
+        messages=[ChatMessage(role=Role.USER, content=context.prompt)],
+    )
+    optimized = await optimizer.optimize_chat(request)
+    context.caveman_optimized_request = optimized.request
+    context.caveman_optimized_notes = optimized.notes
+    return context
+
+
+__all__ = ["caveman_optimization"]

forgecli/build/diff_extract.py

@@ -0,0 +1,218 @@
+r"""Stage 4 — diff extraction.
+
+LLMs love to wrap their output in prose. We extract the largest
+unified diff block from the response (looking for ``diff --git`` or
+``--- a/`` / ``+++ b/`` markers) and store the cleaned text in
+``context.diff_text``. If no diff is found, ``diff_text`` is empty and
+downstream stages short-circuit.
+
+We also strip Markdown code fences (``\`\`\`diff ... \`\`\``) before
+searching, since the system prompt asks for a diff but real models
+often wrap their output anyway.
+"""
+
+from __future__ import annotations
+
+import re
+from pathlib import Path
+
+from forgecli.build import BuildContext
+
+_GIT_DIFF_HEADER = re.compile(r"^diff --git ", re.MULTILINE)
+_UNIFIED_HEADER = re.compile(r"^--- ", re.MULTILINE)
+_DIFF_LINE = re.compile(r"^(?:--- |\+\+\+ |@@ | |\+|-)", re.MULTILINE)
+_DIFF_OR_CONTEXT = re.compile(r"^(?:--- |\+\+\+ |@@ | |\+|-|index |new file |deleted file |similarity |rename |copy )", re.MULTILINE)
+_FENCE_OPEN = re.compile(r"^\s*```(?:\w+)?\s*$")
+_FENCE_CLOSE = re.compile(r"^\s*```\s*$")
+
+
+def extract_diff(text: str) -> str:
+    """Return the largest unified-diff substring in ``text``.
+
+    The search is anchored to the first ``diff --git`` or ``--- a/`` header
+    and runs to the first non-diff-looking line that follows a blank line
+    or the end of the text. We deliberately keep this lenient: real
+    models emit all kinds of surrounding chatter.
+    """
+    if not text:
+        return ""
+    text = _strip_code_fences(text)
+    match = _GIT_DIFF_HEADER.search(text)
+    if not match:
+        match = _UNIFIED_HEADER.search(text)
+    if not match:
+        return ""
+    candidate = text[match.start():]
+    return _trim_to_diff_block(candidate)
+
+
+def _strip_code_fences(text: str) -> str:
+    r"""Drop a single pair of Markdown code fences wrapping the whole text.
+
+    Many models emit fenced diffs (``\`\`\`diff\n<diff>\n\`\`\``); we strip
+    those fences so the header-search below can find ``diff --git``
+    directly. Multi-fence responses (e.g. fences inside fences) are
+    left alone.
+    """
+    lines = text.splitlines()
+    if not lines:
+        return text
+    if not _FENCE_OPEN.match(lines[0]):
+        return text
+    # Find the matching close.
+    for index in range(1, len(lines)):
+        if _FENCE_CLOSE.match(lines[index]):
+            return "\n".join(lines[1:index])
+    return "\n".join(lines[1:])
+
+
+def _trim_to_diff_block(candidate: str) -> str:
+    """Trim trailing prose that is clearly not part of the diff.
+
+    We keep the diff block intact (including context lines) and stop at
+    the first blank or non-diff line that follows a substantive diff.
+    """
+    lines = candidate.splitlines()
+    kept: list[str] = []
+    for line in lines:
+        if not kept and not line.strip():
+            continue
+        if not _DIFF_OR_CONTEXT.match(line) and not line.startswith("diff --git "):
+            # We've hit prose. Stop here (but keep what we have).
+            if kept:
+                break
+            continue
+        kept.append(line)
+    return "\n".join(kept).rstrip() + "\n"
+
+
+def _looks_like_diff_line(line: str) -> bool:
+    if _DIFF_LINE.match(line):
+        return True
+    if line.startswith("diff --git "):
+        return True
+    return bool(line.startswith("index "))
+
+
+def is_file_requested(relative_path: str | Path, prompt: str, root_dir: Path) -> bool:
+    """Return True if the file was explicitly requested by the user or already exists."""
+    path_str = str(relative_path)
+    if path_str.startswith("b/") or path_str.startswith("a/"):
+        path_str = path_str[2:]
+
+    full_path = root_dir / path_str
+    # If the file already exists on disk, editing/modifying it is allowed.
+    if full_path.exists():
+        return True
+
+    # Clean the prompt: lower case, strip punctuation
+    cleaned_prompt = re.sub(r'[^\w\s-]', ' ', prompt.lower())
+    prompt_words = set(cleaned_prompt.split())
+
+    # Get path information
+    path_obj = Path(path_str)
+    filename = path_obj.name.lower()
+    stem = path_obj.stem.lower()
+
+    # Check if exact filename, stem, path, or part is in the cleaned prompt
+    cleaned_prompt_with_dashes = prompt.lower()
+    if filename in cleaned_prompt_with_dashes or stem in cleaned_prompt_with_dashes:
+        return True
+
+    for part in path_obj.parts:
+        part_lower = part.lower()
+        if part_lower in cleaned_prompt_with_dashes:
+            return True
+
+    # Tokenize the stem and check if any token matches
+    stem_tokens = re.split(r'[_.-]', stem)
+    for token in stem_tokens:
+        if len(token) >= 2 and token in prompt_words:
+            return True
+
+    # Special case: tests
+    has_test_in_prompt = any(w in prompt_words for w in ("test", "tests", "testing", "regression", "regressions", "assert"))
+    has_test_in_file = any(t in filename for t in ("test", "regression", "spec"))
+    if has_test_in_prompt and has_test_in_file:
+        return True
+
+    # Special case: python __init__.py package files
+    if filename == "__init__.py":
+        parent_name = path_obj.parent.name.lower()
+        if parent_name in cleaned_prompt_with_dashes or any(t in prompt_words for t in re.split(r'[_.-]', parent_name)):
+            return True
+
+    return False
+
+
+def filter_diff(diff_text: str, prompt: str, root_dir: Path) -> str:
+    """Filter the diff_text to only contain blocks for requested or existing files."""
+    if not diff_text:
+        return ""
+
+    blocks = []
+    lines = diff_text.splitlines(keepends=True)
+
+    current_block: list[str] = []
+    i = 0
+    while i < len(lines):
+        line = lines[i]
+        is_new_block = False
+        if line.startswith("diff --git ") or (line.startswith("--- ") and i + 1 < len(lines) and lines[i + 1].startswith("+++ ") and (
+            not current_block or not current_block[0].startswith("diff --git ")
+        )):
+            is_new_block = True
+
+        if is_new_block and current_block:
+            blocks.append(current_block)
+            current_block = []
+
+        current_block.append(line)
+        i += 1
+
+    if current_block:
+        blocks.append(current_block)
+
+    filtered_blocks = []
+    for block in blocks:
+        block_text = "".join(block)
+
+        # Try to find the new path
+        new_path = None
+        match = re.search(r"^\+\+\+\s+(?:b/)?(?P<path>[^\s\n]+)", block_text, re.MULTILINE)
+        if match:
+            new_path = match.group("path")
+            if new_path == "/dev/null":
+                # If deleted, check the old path
+                old_match = re.search(r"^---\s+(?:a/)?(?P<path>[^\s\n]+)", block_text, re.MULTILINE)
+                if old_match:
+                    new_path = old_match.group("path")
+        else:
+            # Fallback to diff --git header
+            git_match = re.search(r"^diff --git\s+(?:a/[^\s]+)\s+b/(?P<path>[^\s\n]+)", block_text, re.MULTILINE)
+            if git_match:
+                new_path = git_match.group("path")
+
+        if not new_path:
+            filtered_blocks.append(block_text)
+            continue
+
+        if new_path.startswith("b/") or new_path.startswith("a/"):
+            new_path = new_path[2:]
+
+        if is_file_requested(new_path, prompt, root_dir):
+            filtered_blocks.append(block_text)
+
+    return "".join(filtered_blocks)
+
+
+async def diff_extraction(context: BuildContext) -> BuildContext:
+    """Extract a unified diff from ``context.response`` and store in ``diff_text``."""
+    if context.response is None:
+        return context
+    extracted = extract_diff(context.response.message.content or "")
+    context.diff_text = filter_diff(extracted, context.prompt, context.root)
+    return context
+
+
+__all__ = ["diff_extraction", "extract_diff", "filter_diff", "is_file_requested"]