Upload 7 files

README.md

@@ -75,19 +75,29 @@ regressions."
 
 ## Models
 
-This shipped version is intentionally **deterministic and model-free** so the
-demo works with **no model and no GPU**.
+The core pipeline is **deterministic and model-free** so the demo works with
+**no model and no GPU**: every test comes from executing the real code on
+type-hint-driven inputs.
 
 | Job | Current implementation | Fallback |
 |---|---|---|
 | Input generation | Type-hint-driven deterministic cases | Fixed mixed literals |
 | Test expectation generation | Real code execution | None needed |
 | Quality proof | Hand-rolled mutation scoring | None needed |
+| Edge-case assist (optional) | `Qwen/Qwen2.5-Coder-0.5B-Instruct` via `@spaces.GPU` proposes extra argument tuples | Skipped if the model/GPU is unavailable; deterministic cases still run |
+
+Checking **"Use small coder model"** sends each function's signature, docstring,
+and already-tried inputs to a small (0.5B parameter) coder model running on
+ZeroGPU, which proposes a couple of extra edge-case argument tuples (empty,
+negative, boundary values, ...). Those tuples are **not trusted directly** --
+they go through the same `capture`-then-assert path as the deterministic cases
+(`generator.capture`), so the recorded expectation always comes from running
+the real code. A bad or redundant suggestion just becomes a redundant (still
+correct) test; it can never make a test pass on model guesswork alone.
 
 This keeps the demo fast, inspectable, and stable under hackathon time
-pressure. A future enhancement path is to let a small coder model propose extra
-edge-case inputs, but only accept those cases if they improve coverage or the
-mutation score.
+pressure while satisfying the **ZeroGPU** requirement via
+`model_suggest.suggest_extra_inputs`.
 
 ---
 
@@ -131,6 +141,7 @@ The project's own test suite currently covers:
 - PR patch export
 - forge pipeline smoke path
 - injected regression causing exactly one failure
+- model-suggestion parsing (arity-matched argument tuples only)
 
 ## Project layout
 
@@ -143,6 +154,7 @@ runner.py              subprocess pytest execution and parsing
 mutator.py             hand-rolled mutation catalog and mutation score
 patch.py               PR-style unified diff export
 inject.py              controlled one-click demo regression
+model_suggest.py       optional @spaces.GPU small-model edge-case assist
 samples/legacy_repo/   bundled untested Python target repo
 tests/test_testforge.py
 ```

agent.py

@@ -8,8 +8,9 @@ from dataclasses import dataclass
 from pathlib import Path
 
 from analyzer import Analysis, analyze
-from generator import GeneratedSuite, generate_suite, write_suite
+from generator import GeneratedSuite, deterministic_inputs, generate_suite, write_suite
 from inject import reset_sample
+from model_suggest import suggest_extra_inputs
 from mutator import MUTATIONS, Mutation, MutationScore, mutation_score
 from patch import write_patch
 from runner import RunResult, run_pytest
@@ -35,8 +36,15 @@ class ForgeArtifacts:
 def forge_legacy_repo(
     max_cases_per_function: int = 4,
     mutations: tuple[Mutation, ...] = MUTATIONS,
+    use_model: bool = False,
 ) -> ForgeArtifacts:
-    """Run the deterministic TestForge pipeline for the bundled sample."""
+    """Run the TestForge pipeline for the bundled sample.
+
+    The deterministic path (`use_model=False`) is the demo backbone and needs
+    no model or GPU. When `use_model=True`, a small coder model proposes extra
+    edge-case argument tuples per function via `model_suggest.suggest_extra_inputs`;
+    every suggestion is still captured by real execution before it becomes a test.
+    """
     run_dir = _fresh_run_dir()
     logs: list[str] = []
 
@@ -48,7 +56,25 @@ def forge_legacy_repo(
     for fn in analysis.functions:
         logs.append(f"  - {fn.module}.{fn.qualname}({', '.join(fn.parameters)})")
 
-    suite = generate_suite(analysis, max_cases_per_function=max_cases_per_function)
+    extra_inputs: dict[str, list[tuple]] = {}
+    if use_model:
+        logs.append("Model assist: asking small coder model for extra edge cases")
+        for fn in analysis.functions:
+            if fn.arity == 0:
+                continue
+            deterministic = deterministic_inputs(fn, max_cases=max_cases_per_function)
+            suggestions = suggest_extra_inputs(fn, deterministic, max_new=2)
+            if suggestions:
+                extra_inputs[f"{fn.module}.{fn.qualname}"] = suggestions
+                logs.append(f"  +{len(suggestions)} model-proposed case(s) for {fn.qualname}")
+        if not extra_inputs:
+            logs.append("  no extra cases proposed (model/GPU unavailable or nothing new found)")
+
+    suite = generate_suite(
+        analysis,
+        max_cases_per_function=max_cases_per_function,
+        extra_inputs=extra_inputs,
+    )
     write_suite(suite, run_dir)
     logs.append(f"Captured {suite.assertion_count} behaviors into {len(suite.files)} test files")
 

app.py

@@ -44,8 +44,7 @@ CSS = """
 
 
 def forge(use_model: bool, max_cases: int):
-    del use_model  # deterministic backbone is the demo path.
-    artifacts = forge_legacy_repo(max_cases_per_function=max_cases)
+    artifacts = forge_legacy_repo(max_cases_per_function=max_cases, use_model=use_model)
     LAST_RUN["run_dir"] = artifacts.run_dir
     log = "\n".join(artifacts.logs)
     preview = _suite_preview(artifacts.suite.files)

generator.py

@@ -33,13 +33,27 @@ class GeneratedSuite:
     assertion_count: int
 
 
-def generate_suite(analysis: Analysis, max_cases_per_function: int = 4) -> GeneratedSuite:
-    """Generate a green characterization suite by executing current behavior."""
+def generate_suite(
+    analysis: Analysis,
+    max_cases_per_function: int = 4,
+    extra_inputs: dict[str, list[tuple[Any, ...]]] | None = None,
+) -> GeneratedSuite:
+    """Generate a green characterization suite by executing current behavior.
+
+    `extra_inputs` maps "<module>.<qualname>" to additional argument tuples
+    (e.g. model-suggested edge cases) to capture alongside the deterministic
+    ones. Every entry still goes through `capture`, so the resulting assertion
+    is always grounded in real execution.
+    """
     _ensure_import_root(analysis.root.parent)
+    extra_inputs = extra_inputs or {}
     files: dict[str, str] = {}
     cases_by_function: dict[str, list[CapturedCase]] = {}
     for fn in analysis.functions:
         candidates = deterministic_inputs(fn, max_cases=max_cases_per_function)
+        for case in extra_inputs.get(f"{fn.module}.{fn.qualname}", []):
+            if case not in candidates:
+                candidates.append(case)
         cases = capture(fn, candidates)
         if not cases:
             continue

model_suggest.py

@@ -0,0 +1,121 @@
+"""Optional small-model assist: propose extra characterization inputs.
+
+This never weakens the "capture-then-assert" guarantee in generator.py -- any
+input tuple suggested here is executed by `generator.capture` exactly like the
+deterministic inputs, so the recorded expectation always comes from running the
+real code, never from the model. A model that proposes a redundant or useless
+input just yields a redundant (still-correct) test; it can never make a test
+green by assertion alone.
+"""
+
+from __future__ import annotations
+
+import ast
+import re
+
+from analyzer import FunctionInfo
+
+try:
+    import spaces
+except ImportError:  # local/dev environments without the `spaces` package
+    class _SpacesShim:
+        @staticmethod
+        def GPU(*args, **kwargs):
+            if args and callable(args[0]):
+                return args[0]
+
+            def decorator(fn):
+                return fn
+
+            return decorator
+
+    spaces = _SpacesShim()
+
+
+MODEL_ID = "Qwen/Qwen2.5-Coder-0.5B-Instruct"
+MAX_NEW_TOKENS = 96
+
+_model = None
+_tokenizer = None
+
+
+def _load_model():
+    global _model, _tokenizer
+    if _model is None:
+        import torch
+        from transformers import AutoModelForCausalLM, AutoTokenizer
+
+        _tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
+        _model = AutoModelForCausalLM.from_pretrained(MODEL_ID, torch_dtype=torch.bfloat16)
+    return _model, _tokenizer
+
+
+@spaces.GPU(duration=120)
+def suggest_extra_inputs(fn: FunctionInfo, existing: list[tuple], max_new: int = 2) -> list[tuple]:
+    """Ask a small coder model for extra argument tuples to try for `fn`.
+
+    Returns [] on any failure (no GPU, model unavailable, bad output, ...) so the
+    deterministic pipeline never depends on this succeeding.
+    """
+    try:
+        import torch
+
+        model, tokenizer = _load_model()
+        model.to("cuda" if torch.cuda.is_available() else "cpu")
+    except Exception:
+        return []
+
+    prompt = _build_prompt(fn, existing, max_new)
+    try:
+        messages = [{"role": "user", "content": prompt}]
+        input_ids = tokenizer.apply_chat_template(
+            messages, add_generation_prompt=True, return_tensors="pt"
+        ).to(model.device)
+        output_ids = model.generate(
+            input_ids,
+            max_new_tokens=MAX_NEW_TOKENS,
+            do_sample=False,
+            pad_token_id=tokenizer.eos_token_id,
+        )
+        text = tokenizer.decode(output_ids[0, input_ids.shape[1] :], skip_special_tokens=True)
+    except Exception:
+        return []
+
+    return _parse_tuples(text, arity=len(fn.parameters), limit=max_new)
+
+
+def _build_prompt(fn: FunctionInfo, existing: list[tuple], max_new: int) -> str:
+    signature = ", ".join(fn.parameters) or "(no arguments)"
+    existing_repr = ", ".join(repr(case) for case in existing[:4])
+    return (
+        f"Function under test: {fn.qualname}({signature})\n"
+        f"Docstring: {fn.docstring or '(none)'}\n"
+        f"Argument tuples already tried: [{existing_repr}]\n\n"
+        f"Suggest {max_new} new argument tuples that probe edge cases "
+        "(empty, zero, negative, boundary, or unusual values) and differ from "
+        "the ones already tried. Respond with ONLY a Python list of tuples, "
+        "e.g. [(0, 'x'), (-1, '')]. No explanation, no code block."
+    )
+
+
+def _parse_tuples(text: str, arity: int, limit: int) -> list[tuple]:
+    """Pull up to `limit` arity-matched argument tuples out of free-form model text."""
+    match = re.search(r"\[.*\]", text, re.DOTALL)
+    if not match:
+        return []
+    try:
+        parsed = ast.literal_eval(match.group(0))
+    except (ValueError, SyntaxError):
+        return []
+    if not isinstance(parsed, list):
+        return []
+
+    cases: list[tuple] = []
+    for item in parsed:
+        if isinstance(item, tuple) and len(item) == arity:
+            cases.append(item)
+        elif arity == 1 and isinstance(item, (int, float, str, bool)):
+            cases.append((item,))
+        if len(cases) >= limit:
+            break
+    return cases

requirements.txt

@@ -1,3 +1,7 @@
 gradio==6.18.0
+spaces
+torch
+transformers>=4.45
+accelerate
 pytest>=8.0
 pytest-cov>=5.0

tests/test_testforge.py

@@ -6,6 +6,7 @@ from analyzer import analyze
 from agent import forge_legacy_repo
 from generator import capture, generate_suite, write_suite
 from inject import reset_sample, run_injected_suite
+from model_suggest import _parse_tuples
 from mutator import MUTATIONS, mutation_score
 from patch import make_pr_patch
 from runner import run_pytest
@@ -122,6 +123,15 @@ def test_forge_pipeline_produces_green_suite_and_patch():
     assert artifacts.patch_path.exists()
 
 
+def test_model_suggest_parses_only_arity_matched_tuples():
+    text = "Here you go: [(0, 'x'), (-1, ''), (1, 2, 3)]"
+
+    assert _parse_tuples(text, arity=2, limit=5) == [(0, "x"), (-1, "")]
+    assert _parse_tuples("[(0, 1, 2)]", arity=2, limit=2) == []
+    assert _parse_tuples("no list here", arity=2, limit=2) == []
+    assert _parse_tuples("[5, -1, 0]", arity=1, limit=2) == [(5,), (-1,)]
+
+
 def test_inject_regression_causes_exactly_one_failure(tmp_path):
     analysis = analyze(LEGACY_ROOT, package="legacy_repo")
     suite = generate_suite(analysis)