feat(local): grammar-constrained structured output + thread-safe local batch

Author: ctrl-gaurav

effgen/cli/_main.py

@@ -2617,6 +2617,8 @@ def create_parser():
     eval_parser.add_argument('-o', '--output', help='Output file for results (JSON)')
     eval_parser.add_argument('--difficulty', choices=['easy', 'medium', 'hard'],
                               help='Filter test cases by difficulty')
+    eval_parser.add_argument('--max-cases', type=int, default=None,
+                              help='Only run the first N cases (quick subsample)')
     eval_parser.add_argument('--no-animation', action='store_true', default=argparse.SUPPRESS,
                               help='Disable the live progress bar (plain output)')
 
@@ -2628,11 +2630,18 @@ def create_parser():
                                       'for a bare id (e.g. '
                                       'groq:llama-3.1-8b-instant,gpt-5-nano).')
     compare_parser.add_argument('--suite', required=True,
-                                 help='Test suite name')
+                                 help='Built-in suite name (math, tool_use, '
+                                      'reasoning, safety, conversation) OR a path '
+                                      'to your own .jsonl/.json test cases')
     compare_parser.add_argument('--scoring', choices=['exact_match', 'contains', 'regex', 'semantic_similarity', 'llm_judge'],
                                  default='contains', help='Scoring mode (default: contains)')
     compare_parser.add_argument('--threshold', type=float, default=0.5,
                                  help='Pass threshold (default: 0.5)')
+    compare_parser.add_argument('--max-cases', type=int, default=None,
+                                 help='Only run the first N cases (quick bake-off '
+                                      'on a big suite)')
+    compare_parser.add_argument('--difficulty', choices=['easy', 'medium', 'hard'],
+                                 help='Filter test cases by difficulty')
     compare_parser.add_argument('-o', '--output', help='Output file for results (JSON or Markdown)')
     compare_parser.add_argument('--preset', choices=_preset_choices,
                                  help='Use a preset agent configuration')
@@ -3419,9 +3428,50 @@ def _handle_batch_command(args, cli) -> int:
                 logging.getLogger(__name__).debug("Batch agent close() failed", exc_info=True)
 
 
+def _resolve_eval_suite(suite_arg: str, difficulty=None, max_cases=None):
+    """Resolve a ``--suite`` argument to a ``TestSuite``.
+
+    Accepts a built-in suite name **or** a path to a ``.jsonl`` / ``.json`` file
+    of your own test cases (each an object with ``query``/``expected_output`` and
+    optional ``difficulty``/``tags``), so a bake-off can run on your own data —
+    not just the bundled suites. Optionally filters by ``difficulty`` and trims to
+    the first ``max_cases``. Raises ``KeyError`` (with the list of valid names)
+    for an unknown name, or ``FileNotFoundError`` / ``ValueError`` for a bad file.
+    """
+    from effgen.eval import get_suite
+    from effgen.eval.suites import TestSuite
+
+    p = Path(suite_arg)
+    if p.suffix.lower() in (".jsonl", ".json") or p.exists():
+        if not p.exists():
+            raise FileNotFoundError(f"Test-case file not found: {suite_arg}")
+        from effgen.eval.evaluator import TestCase
+        raw = p.read_text(encoding="utf-8")
+        records = []
+        if p.suffix.lower() == ".json":
+            data = json.loads(raw)
+            records = data if isinstance(data, list) else [data]
+        else:
+            records = [json.loads(line) for line in raw.splitlines() if line.strip()]
+        cases = [TestCase.from_dict(r) for r in records]
+        if not cases:
+            raise ValueError(f"No test cases found in {suite_arg}")
+        suite = TestSuite(test_cases=cases)
+        suite.name = p.stem
+    else:
+        suite = get_suite(suite_arg)
+
+    if difficulty:
+        from effgen.eval.evaluator import Difficulty
+        suite.test_cases = suite.filter(difficulty=Difficulty(difficulty))
+    if max_cases is not None and max_cases > 0:
+        suite.test_cases = suite.test_cases[:max_cases]
+    return suite
+
+
 def _handle_eval_command(args, cli) -> int:
     """Handle 'effgen eval' subcommand."""
-    from effgen.eval import AgentEvaluator, RegressionTracker, get_suite, list_suites
+    from effgen.eval import AgentEvaluator, RegressionTracker, list_suites
     from effgen.eval.evaluator import ScoringMode
 
     suite_name = args.suite
@@ -3430,6 +3480,7 @@ def _handle_eval_command(args, cli) -> int:
     scoring = ScoringMode(args.scoring)
     threshold = args.threshold
     difficulty = getattr(args, 'difficulty', None)
+    max_cases = getattr(args, 'max_cases', None)
 
     try:
         # List suites if requested
@@ -3439,13 +3490,13 @@ def _handle_eval_command(args, cli) -> int:
                 cli.print(f"  {name:16s} — {desc}")
             return 0
 
-        suite = get_suite(suite_name)
+        suite = _resolve_eval_suite(suite_name, difficulty=difficulty, max_cases=max_cases)
 
-        # Filter by difficulty if specified
+        # Report any narrowing applied
         if difficulty:
-            from effgen.eval.evaluator import Difficulty
-            suite.test_cases = suite.filter(difficulty=Difficulty(difficulty))
             cli.print(f"Filtered to {len(suite.test_cases)} {difficulty} test cases")
+        if max_cases:
+            cli.print(f"Limited to first {len(suite.test_cases)} cases")
 
         cli.print(f"Loading model {model_name}...")
 
@@ -3533,22 +3584,29 @@ def _handle_eval_command(args, cli) -> int:
 
 def _handle_compare_command(args, cli) -> int:
     """Handle 'effgen compare' subcommand."""
-    from effgen.eval import ModelComparison, get_suite
+    from effgen.eval import ModelComparison
     from effgen.eval.evaluator import ScoringMode
 
     model_names = [m.strip() for m in args.models.split(',')]
     suite_name = args.suite
     scoring = ScoringMode(args.scoring)
     threshold = args.threshold
     preset_name = getattr(args, 'preset', None)
+    difficulty = getattr(args, 'difficulty', None)
+    max_cases = getattr(args, 'max_cases', None)
 
     # Unknown suite is a user error, not a crash — report cleanly (no traceback)
-    # and exit 2 with the list of valid suites.
+    # and exit 2 with the list of valid suites. A bad data file is reported the
+    # same way.
     try:
-        suite = get_suite(suite_name)
-    except (KeyError, ValueError):
+        suite = _resolve_eval_suite(suite_name, difficulty=difficulty, max_cases=max_cases)
+    except (KeyError, ValueError, FileNotFoundError) as exc:
         from effgen.eval import list_suites
-        cli.print(f"Unknown suite '{suite_name}'. Available: {', '.join(list_suites())}.")
+        cli.print(
+            f"Could not load suite '{suite_name}' ({exc}). "
+            f"Use a built-in suite ({', '.join(list_suites())}) "
+            "or a path to a .jsonl/.json file of test cases."
+        )
         return 2
 
     agents: dict = {}

effgen/core/agent.py

@@ -13,6 +13,7 @@
 from __future__ import annotations
 
 import asyncio
+import functools
 import logging
 import time
 from collections.abc import Callable, Iterator
@@ -208,6 +209,38 @@ class AgentConfig:
     cache_tools: bool = True
 
 
+# Model-loading options belong to the engine (load_model), not the agent. Passing
+# one straight to AgentConfig otherwise raises a cryptic dataclass
+# "unexpected keyword argument 'engine'"; intercept it with an actionable hint.
+_MODEL_LOAD_KWARGS = frozenset({
+    "engine", "engine_config", "tensor_parallel_size", "gpu_memory_utilization",
+    "apply_chat_template", "quantization", "trust_remote_code",
+})
+
+
+def _agentconfig_init_guard(_dataclass_init):
+    """Wrap AgentConfig.__init__ to translate a stray model-loading kwarg into a
+    one-line "here's how to do it" instead of a bare dataclass TypeError."""
+    @functools.wraps(_dataclass_init)
+    def __init__(self, *args: Any, **kwargs: Any) -> None:
+        bad = _MODEL_LOAD_KWARGS.intersection(kwargs)
+        if bad:
+            opt = sorted(bad)
+            raise TypeError(
+                f"AgentConfig does not accept model-loading option(s) {opt}: "
+                "they configure the engine, not the agent. Either load the model "
+                "first — load_model(model_id, engine=\"transformers\") — and pass "
+                "the instance as model=, or use "
+                "create_agent(preset, model_id, engine=\"transformers\"), which "
+                "routes these to load_model for you."
+            )
+        _dataclass_init(self, *args, **kwargs)
+    return __init__
+
+
+AgentConfig.__init__ = _agentconfig_init_guard(AgentConfig.__init__)
+
+
 @dataclass
 class AgentResponse:
     """

effgen/core/agent_generation.py

@@ -447,17 +447,24 @@ def _accumulate_run_cost(self, result_metadata: dict[str, Any] | None) -> None:
         accum["calls"] = accum.get("calls", 0) + 1
 
     def _finalize_cost_metadata(self, response: AgentResponse) -> None:
-        """Fold the per-run cost accumulator onto ``response.metadata``.
+        """Fold per-run latency + the cost accumulator onto ``response.metadata``.
 
-        Surfaces ``cost_usd`` and ``prompt_tokens`` / ``completion_tokens`` /
-        ``total_tokens`` for the whole run (every model call summed, tool loops
-        included). Existing keys win, so a path that already set its own cost is
+        Surfaces ``latency_ms`` / ``duration_s`` (always available) and, when the
+        run made billable model calls, ``cost_usd`` plus ``prompt_tokens`` /
+        ``completion_tokens`` / ``total_tokens`` (every call summed, tool loops
+        included). Existing keys win, so a path that already set its own value is
         left untouched. Cost is rounded to whole microdollars.
         """
+        meta = response.metadata
+        # Per-run latency, mirrored from execution_time so callers can read it
+        # off metadata alongside cost/tokens — no need to wrap each call in a
+        # timer. (The eval layer already measures the same number.)
+        if response.execution_time:
+            meta.setdefault("latency_ms", round(response.execution_time * 1000.0, 1))
+            meta.setdefault("duration_s", round(response.execution_time, 4))
         accum = getattr(self, "_run_cost_accum", None)
         if not accum or not accum.get("calls"):
             return
-        meta = response.metadata
         if "cost_usd" in accum and "cost_usd" not in meta:
             meta["cost_usd"] = round(float(accum["cost_usd"]), 8)
         for key in ("prompt_tokens", "completion_tokens", "total_tokens"):

effgen/core/structured_output.py

@@ -334,9 +334,21 @@ def structured_generate(
         if raw is not None:
             last_raw = raw
 
+    final_error = last_error or "could not produce schema-valid output"
+    # Teachable failure: a local Transformers model that couldn't be coaxed into
+    # schema-valid JSON by prompting alone can be *guaranteed* to conform with
+    # grammar-constrained decoding — but only if ``outlines`` is installed. When
+    # it isn't, point the user at the one-line fix instead of leaving a dead end.
+    if _is_transformers_model(model) and not _outlines_available():
+        final_error += (
+            " — this local model did not follow the JSON schema by prompting. "
+            "Install grammar-constrained decoding (`pip install effgen[grammar]`) "
+            "for guaranteed schema-valid output, or use a stronger "
+            "instruction-following model."
+        )
     return StructuredOutcome(
         success=False, attempts=attempts, method="reprompt",
-        error=last_error or "could not produce schema-valid output",
+        error=final_error,
         raw_text=last_raw,
     )
 
@@ -387,7 +399,22 @@ def _outlines_available() -> bool:
     Checks the concrete symbols ``_try_grammar_constrained`` relies on (not just
     that the package exists), so an incompatible ``outlines`` build is never
     counted as a real grammar attempt — keeping the surfaced attempt count honest.
+    Supports both the current ``outlines`` 1.x API (``from_transformers`` +
+    ``Generator`` + ``types.JsonSchema``) and the legacy 0.0.x API
+    (``generate.json`` + ``models.Transformers``).
     """
+    try:
+        import outlines  # noqa: F401
+    except Exception:
+        return False
+    # outlines 1.x
+    if hasattr(outlines, "from_transformers") and hasattr(outlines, "Generator"):
+        try:
+            from outlines.types import JsonSchema  # noqa: F401
+        except Exception:
+            return False
+        return True
+    # legacy outlines 0.0.x
     try:
         from outlines import generate, models  # noqa: F401
     except Exception:
@@ -571,30 +598,61 @@ def _reprompt_once(
     return None, None, text, f"schema validation failed: {err}"
 
 
+def _grammar_max_new_tokens(gen_kwargs: dict[str, Any]) -> int:
+    """Token budget for one grammar-constrained call.
+
+    Honours an explicit ``max_new_tokens``/``max_tokens`` from the caller so a
+    long structured object isn't truncated mid-JSON; otherwise a sensible default
+    that comfortably fits a typical extraction record.
+    """
+    for key in ("max_new_tokens", "max_tokens"):
+        val = gen_kwargs.get(key)
+        if isinstance(val, int) and val > 0:
+            return val
+    return 512
+
+
 def _try_grammar_constrained(
     model: Any,
     prompt: str,
     schema: dict[str, Any],
     gen_kwargs: dict[str, Any],
 ) -> tuple[str, Any] | None:
-    """Try grammar-constrained decoding via outlines library."""
+    """Try grammar-constrained decoding via the ``outlines`` library.
+
+    The output is structurally guaranteed to match *schema* in one call, so even
+    weak instruction-followers (small Llama/Gemma/Phi locals) emit schema-valid
+    JSON. Supports both the current ``outlines`` 1.x API and the legacy 0.0.x API;
+    returns ``None`` (caller falls back to reprompting) if neither is usable.
+    """
+    # outlines drives the underlying HF model + tokenizer directly.
+    if not (hasattr(model, "model") and hasattr(model, "tokenizer")):
+        return None
     try:
-        import outlines  # noqa: F401
+        import outlines
+    except ImportError:
+        logger.debug("outlines not installed, skipping grammar-constrained decoding")
+        return None
+    try:
+        logger.debug("Using outlines grammar-constrained decoding for structured output")
+        if hasattr(outlines, "from_transformers") and hasattr(outlines, "Generator"):
+            # outlines 1.x
+            from outlines.types import JsonSchema
+
+            om = outlines.from_transformers(model.model, model.tokenizer)
+            generator = outlines.Generator(om, JsonSchema(schema))
+            raw = generator(prompt, max_new_tokens=_grammar_max_new_tokens(gen_kwargs))
+            parsed = json.loads(raw) if isinstance(raw, str) else raw
+            json_str = raw if isinstance(raw, str) else json.dumps(raw)
+            return json_str, parsed
+        # legacy outlines 0.0.x
         from outlines import generate
         from outlines import models as outlines_models
 
-        logger.debug("Using outlines grammar-constrained decoding for structured output")
-        # outlines requires access to the underlying HF model
-        if hasattr(model, 'model') and hasattr(model, 'tokenizer'):
-            hf_model = outlines_models.Transformers(model.model, model.tokenizer)
-            schema_str = json.dumps(schema)
-            generator = generate.json(hf_model, schema_str)
-            result = generator(prompt)
-            # result is already a parsed object
-            json_str = json.dumps(result)
-            return json_str, result
-    except ImportError:
-        logger.debug("outlines not installed, skipping grammar-constrained decoding")
+        hf_model = outlines_models.Transformers(model.model, model.tokenizer)
+        generator = generate.json(hf_model, json.dumps(schema))
+        result = generator(prompt)
+        return json.dumps(result), result
     except Exception as e:
         logger.warning(f"Grammar-constrained decoding failed: {e}")
     return None