z3

ARCHITECTURE.md

@@ -47,6 +47,7 @@ The repository adopts responsibility-based modular organization:
 | [`prompts.py`](file:///Users/shreshthdhimole/AutoSetter/autosetter/prompts.py) | Template loader and safe `{JSON}` substitution without string format hazards. |
 | [`extractor.py`](file:///Users/shreshthdhimole/AutoSetter/autosetter/extractor.py) | Extraction of `problem.json` via VLM, markdown code fence stripping, schema validation. |
 | [`generator.py`](file:///Users/shreshthdhimole/AutoSetter/autosetter/generator.py) | Downstream generation orchestrator for statement markdown and testlib C++ code. |
+| `testgen/` | Z3 test generation from `test_spec.json`: spec validation, Z3 integer solving per seed strategy, array/string/tree/graph builders, sample checking. |
 | [`sandbox.py`](file:///Users/shreshthdhimole/AutoSetter/autosetter/sandbox.py) | C++ compilation and execution backends (host `g++` and Docker+NsJail HTTP). |
 | [`pipeline.py`](file:///Users/shreshthdhimole/AutoSetter/autosetter/pipeline.py) | Test generation, validation, jury answer solving, attribution, and checker probing. |
 | [`packager.py`](file:///Users/shreshthdhimole/AutoSetter/autosetter/packager.py) | Assembly of release bundle, test pairing, and manifest creation. |
@@ -70,7 +71,7 @@ The repository adopts responsibility-based modular organization:
 - Serializes `problem.json` and renders five specialized prompt templates:
   1. `statement.txt` ➔ `generated/statement.md`
   2. `validator.txt` ➔ `generated/validator.cpp` (uses `testlib.h`)
-  3. `generator.txt` ➔ `generated/generator.cpp` (uses `testlib.h`)
+  3. `test_spec.txt` ➔ `generated/test_spec.json` (input spec for the Z3 test generator in `autosetter/testgen/`; checked against the official samples and retried with the errors if it rejects them)
   4. `solution.txt` ➔ `generated/solution.cpp` (optimal C++17 solution)
   5. `checker.txt` ➔ `generated/checker.cpp` (uses `testlib.h`)
 - Runs text inference against a local coding model (`qwen2.5-coder:7b`).

README.md

@@ -7,7 +7,8 @@ statement image / PDF
          │  Qwen-VL
          ▼
     problem.json  ──►  statement.md, solution.cpp, validator.cpp,
-         │             generator.cpp, checker.cpp     (one Ollama call each)
+         │             test_spec.json, checker.cpp    (one Ollama call each)
+         │             test_spec.json ──► Z3 ──► test inputs
          ▼
       validate  (compile, generate, validate, solve, check, probe)
          ▼
@@ -198,6 +199,28 @@ When all C++ files are generated by a single model from the same JSON, they can
 
 ---
 
+## Z3 Test Generation
+
+Instead of writing a generator program, the text model writes `test_spec.json`: a declarative description of the input (integers with bounds, arrays, strings, permutations, matrices, rows of queries, trees, graphs), per-test constraints such as `k <= n`, file-wide constraints such as `sum(n) <= 200000`, and the line layout. The prompt is `autosetter/prompts/test_spec.txt`; the engine is `autosetter/testgen/`.
+
+- **Z3 solves the integers.** Every integer of every test case in the file is a Z3 variable with its bounds and constraints. They are fixed one at a time: Z3 reports the feasible range, and the seed's strategy picks within it, so relations and sums always stay satisfiable.
+- **Builders fill in the bulk.** Arrays, strings, trees and graphs are built directly at the sizes Z3 chose, so max-size tests (e.g. n = 2·10⁵) take about a second.
+- **Each seed has a strategy**: 1 = all minimum, 2 = all maximum, then random, small, near-maximum, log-scaled, and so on (`testgen.engine.PLAN`), with matching shapes (sorted arrays, path/star trees, ...).
+- **The spec is checked against the official samples.** Samples are parsed with the spec's layout; a spec that rejects a sample is sent back to the model with the exact problem, like a C++ compile error.
+- **Every generated input is parsed back and re-checked** against the spec before it reaches the validator.
+
+Try a spec on its own:
+```bash
+python -m autosetter.testgen out/generated/test_spec.json 1 10 -o /tmp/tests
+python -m autosetter.testgen out/generated/test_spec.json --check-samples out/problem.json
+```
+
+Settings: `AUTOSETTER_Z3_TIMEOUT_MS` (per solver call, default 10000) and `AUTOSETTER_Z3_MAX_CASES` (most test cases per multi-test file, default 30).
+
+Limits: the spec describes value ranges and structure, not properties that relate elements to each other ("exactly one pair sums to target", "the answer exists"). The validator still rejects tests that break such guarantees, and the pipeline reports the generator as at fault. Z3-generated tests ship as files in `package/tests/`; Polygon cannot run Z3, so no generator script is emitted. A `generator.py` or `generator.cpp` in `out/generated/` is still used if there is no `test_spec.json`.
+
+---
+
 ## Testing
 
 ```bash

autosetter/cli.py

@@ -228,7 +228,15 @@ def generate_from_image(
                     # Analyze failure for next iteration
                     targets = []
                     feedback_context = {}
-
+
+                    # Files that failed to build (including an unusable test spec)
+                    for name, error in test_report.compilation.errors.items():
+                        if name in ("validator", "generator", "solution", "checker"):
+                            targets.append(name)
+                            feedback_context[name] = (
+                                f"Your file could not be used by the pipeline:\n{error[:2000]}"
+                            )
+
                     if "validator rejects official samples" in test_report.diagnosis:
                         targets.append("validator")
                         feedback_context["validator"] = "The validator you generated rejected the official problem samples provided in the problem description."

autosetter/config.py

@@ -53,6 +53,12 @@
 DEFAULT_EXECUTION_TIMEOUT = int(os.environ.get("AUTOSETTER_TIMEOUT", "5"))
 DEFAULT_COMPILE_TIMEOUT = int(os.environ.get("AUTOSETTER_COMPILE_TIMEOUT", "60"))
 
+# Z3 test generation (autosetter.testgen)
+# Per solver call; one test makes a few calls per integer variable.
+Z3_TIMEOUT_MS = int(os.environ.get("AUTOSETTER_Z3_TIMEOUT_MS", "10000"))
+# Most test cases packed into one multi-test input file (each adds solver variables).
+Z3_MAX_CASES = int(os.environ.get("AUTOSETTER_Z3_MAX_CASES", "30"))
+
 # Vision / Image Processing
 PDF_RENDER_DPI = int(os.environ.get("AUTOSETTER_PDF_DPI", "200"))
 SUPPORTED_RASTER_EXTENSIONS = {".png", ".jpg", ".jpeg"}

autosetter/generator.py

@@ -37,6 +37,15 @@
 # ─────────────────────────────────────────────────────────────────────────────
 from autosetter.prompts import PromptError, load_and_render_prompt
 
+from autosetter.extractor import JSONExtractionError, parse_model_json
+from autosetter.testgen import (
+    SpecError,
+    TestGenError,
+    check_samples,
+    generate_test,
+    load_spec,
+)
+
 
 # =============================================================================
 # Exception Classes
@@ -71,6 +80,7 @@ class ArtifactSpec:
     is_cpp: bool = False           # True if the artifact is C++ source code
     is_testlib: bool = False       # True if the artifact requires testlib.h
     strip_code_fence: bool = True  # True to strip ```...``` code fences from LLM output
+    is_test_spec: bool = False     # True for the Z3 test spec (JSON, checked against samples)
 
 
 ARTIFACTS: List[ArtifactSpec] = [
@@ -92,15 +102,17 @@ class ArtifactSpec:
         is_testlib=True,
         strip_code_fence=True,
     ),
-    # ── Ollama-routed artifact: Test case generator (UNCHANGED) ──
-    # This artifact continues to use the existing Ollama backend.
+    # ── Test generator: a declarative input spec that autosetter.testgen
+    # turns into tests with Z3. Named "generator" so validation feedback and
+    # self-healing target it like any other generator. ──
     ArtifactSpec(
         name="generator",
-        prompt_template="generator.txt",
-        output_filename="generator.py",
+        prompt_template="test_spec.txt",
+        output_filename="test_spec.json",
         is_cpp=False,
         is_testlib=False,
-        strip_code_fence=True,
+        strip_code_fence=False,
+        is_test_spec=True,
     ),
     ArtifactSpec(
         name="solution",
@@ -269,6 +281,42 @@ def check_cpp_syntax(code: str, include_dir: Path) -> Tuple[bool, str]:
         return (False, str(exc))
 
 
+def prepare_test_spec(raw_reply: str, json_payload: str) -> Tuple[str, str]:
+    """
+    Parse and verify a model-written test spec.
+
+    Returns (content_to_write, error). The spec must be valid, accept every
+    official sample input, and generate the smallest and largest tests
+    (seeds 1 and 2) without error. `error` is empty on success.
+    """
+    try:
+        data = parse_model_json(raw_reply)
+    except JSONExtractionError as exc:
+        return raw_reply.strip() + "\n", f"The reply is not a JSON object: {str(exc)[:500]}"
+
+    content = json.dumps(data, indent=2, ensure_ascii=False) + "\n"
+    try:
+        test_spec = load_spec(data)
+    except SpecError as exc:
+        return content, f"The spec is invalid:\n{exc}"
+
+    samples = (json.loads(json_payload) or {}).get("samples") or []
+    problems = check_samples(test_spec, samples)
+    if problems:
+        return content, (
+            "The spec rejects the problem's official sample input(s), so it does not "
+            "describe the input correctly:\n" + "\n".join(f"- {p}" for p in problems)
+        )
+
+    for seed in (1, 2):
+        try:
+            generate_test(test_spec, seed)
+        except TestGenError as exc:
+            return content, f"Generating a test from the spec failed (seed {seed}):\n{exc}"
+
+    return content, ""
+
+
 # =============================================================================
 # Single Artifact Generation — Ollama Backend (UNCHANGED LOGIC)
 # =============================================================================
@@ -316,6 +364,7 @@ def generate_single_artifact(
         )
 
     last_error = ""
+    base_prompt = current_prompt
 
     # ── Retry loop: generate → check syntax → repair if needed ──
     for attempt in range(max_retries + 1):
@@ -331,6 +380,19 @@ def generate_single_artifact(
                 f"Ollama text inference failed while generating '{spec.name}': {exc}"
             ) from exc
 
+        # ── Test spec: validate against the schema and the official samples ──
+        if spec.is_test_spec:
+            content, last_error = prepare_test_spec(raw_reply, json_payload)
+            if not last_error:
+                break
+            current_prompt = (
+                f"{base_prompt}\n\n=========================================\n"
+                f"Your previous test spec was rejected:\n{last_error[:2000]}\n\n"
+                f"--- PREVIOUS SPEC ---\n{content.strip()[:4000]}\n\n"
+                "Fix every problem listed above. Output ONLY the corrected JSON spec."
+            )
+            continue
+
         # ── Post-process: strip code fences and sanitize C++ headers ──
         content = strip_code_fence(raw_reply) if spec.strip_code_fence else raw_reply
         if spec.is_cpp:

autosetter/packager.py

@@ -11,7 +11,7 @@
 │   └── solution.cpp      # Reference solution
 ├── files/
 │   ├── validator.cpp     # Input validator (testlib.h)
-│   ├── generator.cpp     # Test generator (testlib.h)
+│   ├── test_spec.json    # Z3 test spec (or generator.py / generator.cpp)
 │   ├── checker.cpp       # Output checker (testlib.h)
 │   └── testlib.h         # Bundled testlib header
 ├── tests/
@@ -108,15 +108,20 @@ def build(
         _log("Packaging testlib files...")
         files_dir = self.package_dir / "files"
         files_dir.mkdir(exist_ok=True)
-        gen_found = False
-        for gen_name in ("generator.py", "generator.cpp"):
+        gen_found = ""
+        for gen_name in ("test_spec.json", "generator.py", "generator.cpp"):
             gen_src = self.generated_dir / gen_name
             if gen_src.exists():
                 shutil.copy2(gen_src, files_dir / gen_name)
-                gen_found = True
+                gen_found = gen_name
                 break
         if not gen_found:
-            _log("  ⚠️  generator.py / generator.cpp not found, skipping")
+            _log("  ⚠️  test_spec.json / generator.py / generator.cpp not found, skipping")
+
+        # Only a testlib C++ generator can be re-run by Polygon from a script.
+        # Z3 (test_spec.json) and Python generators need libraries Polygon does
+        # not have, so their tests ship as files.
+        uses_script = gen_found == "generator.cpp"
 
         for name in ("validator.cpp", "checker.cpp"):
             src = self.generated_dir / name
@@ -143,7 +148,7 @@ def build(
 
         generated_indices = set()
         report_src = self.tests_dir / "validation_report.json"
-        if report_src.exists():
+        if uses_script and report_src.exists():
             try:
                 report_data = json.loads(report_src.read_text(encoding="utf-8"))
                 for tc in report_data.get("test_cases", []):
@@ -182,10 +187,12 @@ def build(
             _log("  ⚠️  Package contains no tests")
 
         # 7. Generate script file for Polygon (if tests were generated)
-        _log("Generating test script...")
         script_content = ""
         report_src = self.tests_dir / "validation_report.json"
-        if report_src.exists():
+        if not uses_script:
+            _log("Tests are shipped as files (no Polygon generator script).")
+        elif report_src.exists():
+            _log("Generating test script...")
             try:
                 report_data = json.loads(report_src.read_text(encoding="utf-8"))
                 for tc in report_data.get("test_cases", []):
@@ -196,7 +203,7 @@ def build(
 
         if script_content:
             (self.package_dir / "script").write_text(script_content, encoding="utf-8")
-        else:
+        elif uses_script:
             _log("  ⚠️  Could not generate script, missing validation report or test_cases")
 
         # 8. Generate manifest.json