stackfix 0.1.0__py3-none-any.whl

cloudgym/fixer/lambda_handler.py

@@ -0,0 +1,126 @@
+"""AWS Lambda handler for IaC repair using GGUF model.
+
+Deploy with the 0.5B model (~300MB) for sub-$0.001/invocation cost.
+
+Model loading:
+  - Lambda layer: Place model at /opt/model.gguf
+  - S3: Set MODEL_S3_BUCKET and MODEL_S3_KEY env vars
+  - Bundled: Set MODEL_PATH env var
+
+Example event:
+    {
+        "config": "resource aws_s3_bucket my_bucket {\\n  acl = \\"public\\"\\n}",
+        "errors": ["Error: 'acl' argument is deprecated"],
+        "format": "terraform"
+    }
+
+Response:
+    {
+        "repaired": "resource aws_s3_bucket my_bucket {\\n  ...\\n}",
+        "original_errors": 1,
+        "verified": true,
+        "remaining_errors": 0
+    }
+"""
+
+from __future__ import annotations
+
+import json
+import logging
+import os
+import tempfile
+from pathlib import Path
+
+logger = logging.getLogger(__name__)
+
+# Lazy-loaded repairer (persists across warm invocations)
+_repairer = None
+
+
+def _get_model_path() -> str:
+    """Resolve the GGUF model path from environment or known locations."""
+    # Explicit path
+    if path := os.environ.get("MODEL_PATH"):
+        return path
+
+    # Lambda layer
+    layer_path = "/opt/model.gguf"
+    if Path(layer_path).exists():
+        return layer_path
+
+    # S3 download
+    bucket = os.environ.get("MODEL_S3_BUCKET")
+    key = os.environ.get("MODEL_S3_KEY")
+    if bucket and key:
+        local_path = "/tmp/model.gguf"
+        if not Path(local_path).exists():
+            import boto3
+
+            logger.info("Downloading model from s3://%s/%s", bucket, key)
+            s3 = boto3.client("s3")
+            s3.download_file(bucket, key, local_path)
+            logger.info("Model downloaded to %s", local_path)
+        return local_path
+
+    raise RuntimeError(
+        "No model found. Set MODEL_PATH, place model at /opt/model.gguf, "
+        "or set MODEL_S3_BUCKET + MODEL_S3_KEY."
+    )
+
+
+def _get_repairer():
+    """Get or create the GGUFRepairer (cached across warm invocations)."""
+    global _repairer
+    if _repairer is None:
+        from cloudgym.fixer.repairer import GGUFRepairer
+
+        _repairer = GGUFRepairer(
+            model_path=_get_model_path(),
+            n_gpu_layers=0,
+            max_tokens=4096,
+        )
+    return _repairer
+
+
+def handler(event, context=None):
+    """AWS Lambda entry point for IaC repair.
+
+    Args:
+        event: Dict with 'config' (str), 'errors' (list[str]), optional 'format' (str).
+        context: Lambda context (unused).
+
+    Returns:
+        Dict with 'repaired' (str), 'verified' (bool), etc.
+    """
+    config = event.get("config", "")
+    errors = event.get("errors", [])
+    iac_format = event.get("format")
+
+    if not config:
+        return {"error": "Missing 'config' field", "statusCode": 400}
+
+    if not errors:
+        return {"repaired": config, "original_errors": 0, "verified": True, "remaining_errors": 0}
+
+    repairer = _get_repairer()
+    repaired = repairer.repair(config, errors)
+
+    # Optional: verify the repair
+    verified = False
+    remaining_errors = -1
+    try:
+        from cloudgym.fixer.detector import IaCFormat, validate_content_sync
+
+        fmt = IaCFormat(iac_format) if iac_format else None
+        _, result = validate_content_sync(repaired, fmt)
+        verified = result.valid
+        remaining_errors = len(result.errors)
+    except Exception:
+        logger.exception("Verification failed")
+
+    return {
+        "repaired": repaired,
+        "original_errors": len(errors),
+        "verified": verified,
+        "remaining_errors": remaining_errors,
+    }

cloudgym/fixer/repairer.py

@@ -0,0 +1,237 @@
+"""Model-based IaC repair engine."""
+
+from __future__ import annotations
+
+import logging
+from pathlib import Path
+
+logger = logging.getLogger(__name__)
+
+REPAIR_SYSTEM_PROMPT = (
+    "You are an expert Infrastructure-as-Code engineer. "
+    "Fix the broken configuration below. Return ONLY the fixed configuration "
+    "with no explanation, no markdown fences, and no comments about the fix."
+)
+
+DISCUSS_SYSTEM_PROMPT = (
+    "You are an expert Infrastructure-as-Code engineer and teacher. "
+    "Analyze the configuration and validation errors below. Explain:\n"
+    "1. What each error means in plain language\n"
+    "2. Why it's a problem (what would happen in production)\n"
+    "3. How to fix it\n"
+    "4. Best practices to avoid this in future\n"
+    "Be concise but thorough."
+)
+
+# Default model/adapter paths (relative to package or absolute)
+DEFAULT_BASE_MODEL = "mlx-community/Qwen2.5-Coder-3B-Instruct-4bit"
+DEFAULT_ADAPTER_PATH = "data/models/iac-repair-adapter"
+
+
+def _strip_markdown_fences(text: str) -> str:
+    """Strip markdown code fences and chat tokens from model output."""
+    for stop in ("<|im_end|>", "<|endoftext|>", "<|end|>"):
+        if stop in text:
+            text = text[:text.index(stop)]
+
+    text = text.strip()
+    lines = text.splitlines()
+    if lines and lines[0].startswith("```"):
+        lines = lines[1:]
+    if lines and lines[-1].strip() == "```":
+        lines = lines[:-1]
+    return "\n".join(lines)
+
+
+def _build_prompt(broken_config: str, errors: list[str]) -> str:
+    """Build the repair prompt from config content and error messages."""
+    error_text = "\n".join(errors) if errors else "Unknown error"
+    return (
+        f"This IaC configuration has validation errors:\n\n"
+        f"Errors:\n{error_text}\n\n"
+        f"Broken config:\n```\n{broken_config}\n```\n\n"
+        f"Return the fixed configuration:"
+    )
+
+
+def _build_discuss_prompt(config: str, errors: list[str]) -> str:
+    """Build the discussion prompt for explaining errors."""
+    error_text = "\n".join(errors) if errors else "No validation errors"
+    return (
+        f"Analyze this IaC configuration:\n\n"
+        f"Validation errors:\n{error_text}\n\n"
+        f"Configuration:\n```\n{config}\n```\n\n"
+        f"Explain what's wrong and how to fix it:"
+    )
+
+
+class MLXRepairer:
+    """Repair IaC configs using a local MLX fine-tuned model."""
+
+    def __init__(
+        self,
+        base_model: str = DEFAULT_BASE_MODEL,
+        adapter_path: str = DEFAULT_ADAPTER_PATH,
+        temp: float = 0.3,
+        max_tokens: int = 4096,
+    ):
+        self.base_model = base_model
+        self.adapter_path = adapter_path
+        self.temp = temp
+        self.max_tokens = max_tokens
+        self._model = None
+        self._tokenizer = None
+
+    def _ensure_loaded(self):
+        """Lazy-load model and tokenizer."""
+        if self._model is not None:
+            return
+
+        from mlx_lm import load
+
+        adapter = self.adapter_path if Path(self.adapter_path).exists() else None
+        logger.info("Loading model %s (adapter: %s)", self.base_model, adapter)
+        self._model, self._tokenizer = load(
+            self.base_model, adapter_path=adapter
+        )
+        logger.info("Model loaded.")
+
+    def _generate(self, system_prompt: str, user_prompt: str) -> str:
+        """Generate a response with the given prompts."""
+        self._ensure_loaded()
+
+        from mlx_lm import generate
+        from mlx_lm.sample_utils import make_sampler
+
+        messages = [
+            {"role": "system", "content": system_prompt},
+            {"role": "user", "content": user_prompt},
+        ]
+        prompt_text = self._tokenizer.apply_chat_template(
+            messages, tokenize=False, add_generation_prompt=True
+        )
+        sampler = make_sampler(temp=self.temp)
+        return generate(
+            self._model, self._tokenizer, prompt=prompt_text,
+            max_tokens=self.max_tokens, sampler=sampler,
+        )
+
+    def repair(self, broken_config: str, errors: list[str]) -> str:
+        """Generate a repaired config from broken config + errors."""
+        return _strip_markdown_fences(
+            self._generate(REPAIR_SYSTEM_PROMPT, _build_prompt(broken_config, errors))
+        )
+
+    def discuss(self, config: str, errors: list[str]) -> str:
+        """Explain errors and suggest fixes in natural language."""
+        return self._generate(DISCUSS_SYSTEM_PROMPT, _build_discuss_prompt(config, errors))
+
+
+class GGUFRepairer:
+    """Repair IaC configs using a GGUF model via llama-cpp-python (cross-platform CPU)."""
+
+    def __init__(
+        self,
+        model_path: str,
+        n_gpu_layers: int = 0,
+        temp: float = 0.3,
+        max_tokens: int = 4096,
+        n_ctx: int = 8192,
+    ):
+        self.model_path = model_path
+        self.n_gpu_layers = n_gpu_layers
+        self.temp = temp
+        self.max_tokens = max_tokens
+        self.n_ctx = n_ctx
+        self._llm = None
+
+    def _ensure_loaded(self):
+        """Lazy-load the GGUF model."""
+        if self._llm is not None:
+            return
+
+        from llama_cpp import Llama
+
+        logger.info("Loading GGUF model from %s", self.model_path)
+        self._llm = Llama(
+            model_path=self.model_path,
+            n_gpu_layers=self.n_gpu_layers,
+            n_ctx=self.n_ctx,
+            verbose=False,
+            chat_format="chatml",
+        )
+        logger.info("GGUF model loaded.")
+
+    def _chat(self, system_prompt: str, user_prompt: str) -> str:
+        """Generate a response via llama.cpp chat completion."""
+        self._ensure_loaded()
+
+        response = self._llm.create_chat_completion(
+            messages=[
+                {"role": "system", "content": system_prompt},
+                {"role": "user", "content": user_prompt},
+            ],
+            temperature=self.temp,
+            max_tokens=self.max_tokens,
+        )
+        return response["choices"][0]["message"]["content"]
+
+    def repair(self, broken_config: str, errors: list[str]) -> str:
+        """Generate a repaired config from broken config + errors."""
+        return _strip_markdown_fences(
+            self._chat(REPAIR_SYSTEM_PROMPT, _build_prompt(broken_config, errors))
+        )
+
+    def discuss(self, config: str, errors: list[str]) -> str:
+        """Explain errors and suggest fixes in natural language."""
+        return self._chat(DISCUSS_SYSTEM_PROMPT, _build_discuss_prompt(config, errors))
+
+    def unload(self):
+        """Explicitly release model memory."""
+        if self._llm is not None:
+            del self._llm
+            self._llm = None
+            logger.info("GGUF model unloaded.")
+
+    def __enter__(self):
+        return self
+
+    def __exit__(self, *args):
+        self.unload()
+
+
+class OllamaRepairer:
+    """Repair IaC configs using an Ollama-served model."""
+
+    def __init__(self, model: str = "qwen2.5-coder:3b", base_url: str = "http://localhost:11434"):
+        self.model = model
+        self.base_url = base_url
+
+    def _chat(self, system_prompt: str, user_prompt: str) -> str:
+        """Send a chat request to Ollama."""
+        import httpx
+
+        resp = httpx.post(
+            f"{self.base_url}/api/chat",
+            json={
+                "model": self.model,
+                "messages": [
+                    {"role": "system", "content": system_prompt},
+                    {"role": "user", "content": user_prompt},
+                ],
+                "stream": False,
+            },
+            timeout=120.0,
+        )
+        resp.raise_for_status()
+        return resp.json()["message"]["content"]
+
+    def repair(self, broken_config: str, errors: list[str]) -> str:
+        """Generate a repaired config via Ollama API."""
+        return _strip_markdown_fences(
+            self._chat(REPAIR_SYSTEM_PROMPT, _build_prompt(broken_config, errors))
+        )
+
+    def discuss(self, config: str, errors: list[str]) -> str:
+        """Explain errors and suggest fixes in natural language."""
+        return self._chat(DISCUSS_SYSTEM_PROMPT, _build_discuss_prompt(config, errors))

cloudgym/generator/formatter.py

@@ -0,0 +1,142 @@
+"""Output formatter for training data (JSONL with train/val/test splits)."""
+
+from __future__ import annotations
+
+import hashlib
+import json
+import logging
+from dataclasses import asdict, dataclass, field
+from pathlib import Path
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass
+class TrainingRecord:
+    """A single training record (gold + broken pair)."""
+
+    id: str
+    format: str  # "terraform" or "cloudformation"
+    gold_config: str
+    broken_config: str
+    errors: list[str]
+    warnings: list[str]
+    fault_types: list[str]
+    fault_description: str
+    difficulty: str  # "low", "medium", "high"
+    source: str  # "programmatic" or "agentic"
+    split: str = ""  # "train", "val", or "test" — assigned by format_and_split
+    gold_hash: str = ""  # Hash of gold config for dedup/split
+
+    def __post_init__(self):
+        if not self.gold_hash:
+            self.gold_hash = hashlib.md5(self.gold_config.encode()).hexdigest()
+
+
+def format_and_split(
+    records: list[TrainingRecord],
+    output_dir: str | Path,
+    ratios: tuple[float, float, float] = (0.8, 0.1, 0.1),
+) -> dict:
+    """Split records by gold config and write JSONL files.
+
+    Splitting is done by gold config hash (not by record) to prevent
+    data leakage — all variants of the same gold config go to the same split.
+
+    Args:
+        records: List of TrainingRecord objects.
+        output_dir: Directory to write train.jsonl, val.jsonl, test.jsonl.
+        ratios: (train, val, test) split ratios.
+
+    Returns:
+        Metadata dict with counts and split info.
+    """
+    output_path = Path(output_dir)
+    output_path.mkdir(parents=True, exist_ok=True)
+
+    # Group records by gold_hash
+    gold_groups: dict[str, list[TrainingRecord]] = {}
+    for record in records:
+        gold_groups.setdefault(record.gold_hash, []).append(record)
+
+    # Deterministic ordering by hash
+    sorted_hashes = sorted(gold_groups.keys())
+    total_golds = len(sorted_hashes)
+
+    train_end = int(total_golds * ratios[0])
+    val_end = train_end + int(total_golds * ratios[1])
+
+    train_hashes = set(sorted_hashes[:train_end])
+    val_hashes = set(sorted_hashes[train_end:val_end])
+    test_hashes = set(sorted_hashes[val_end:])
+
+    # Assign splits
+    splits: dict[str, list[TrainingRecord]] = {"train": [], "val": [], "test": []}
+    for h in sorted_hashes:
+        if h in train_hashes:
+            split = "train"
+        elif h in val_hashes:
+            split = "val"
+        else:
+            split = "test"
+        for record in gold_groups[h]:
+            record.split = split
+            splits[split].append(record)
+
+    # Write JSONL files
+    counts = {}
+    for split_name, split_records in splits.items():
+        filepath = output_path / f"{split_name}.jsonl"
+        with open(filepath, "w") as f:
+            for record in split_records:
+                f.write(json.dumps(asdict(record), ensure_ascii=False) + "\n")
+        counts[split_name] = len(split_records)
+        logger.info("Wrote %d records to %s", len(split_records), filepath)
+
+    # Write metadata
+    metadata = {
+        "total_records": len(records),
+        "total_gold_configs": total_golds,
+        "splits": counts,
+        "ratios": list(ratios),
+        "fault_type_distribution": _count_fault_types(records),
+        "format_distribution": _count_formats(records),
+        "source_distribution": _count_sources(records),
+        "difficulty_distribution": _count_difficulties(records),
+    }
+
+    meta_path = output_path / "metadata.json"
+    with open(meta_path, "w") as f:
+        json.dump(metadata, f, indent=2)
+    logger.info("Wrote metadata to %s", meta_path)
+
+    return metadata
+
+
+def _count_fault_types(records: list[TrainingRecord]) -> dict[str, int]:
+    counts: dict[str, int] = {}
+    for r in records:
+        for ft in r.fault_types:
+            counts[ft] = counts.get(ft, 0) + 1
+    return dict(sorted(counts.items()))
+
+
+def _count_formats(records: list[TrainingRecord]) -> dict[str, int]:
+    counts: dict[str, int] = {}
+    for r in records:
+        counts[r.format] = counts.get(r.format, 0) + 1
+    return counts
+
+
+def _count_sources(records: list[TrainingRecord]) -> dict[str, int]:
+    counts: dict[str, int] = {}
+    for r in records:
+        counts[r.source] = counts.get(r.source, 0) + 1
+    return counts
+
+
+def _count_difficulties(records: list[TrainingRecord]) -> dict[str, int]:
+    counts: dict[str, int] = {}
+    for r in records:
+        counts[r.difficulty] = counts.get(r.difficulty, 0) + 1
+    return counts