feat(sampling): forward seed/top_k/penalties through Agent, reject unknown run() kwargs

Author: ctrl-gaurav

effgen/cli/_main.py

@@ -777,6 +777,7 @@ def run_agent(self, args):
                     system_prompt=args.system_prompt or config.get("system_prompt"),
                     max_iterations=args.max_iterations,
                     temperature=args.temperature,
+                    max_tokens=args.max_tokens,
                     enable_streaming=args.stream,
                     session_id=getattr(args, 'session_id', None),
                     **_preset_overrides,
@@ -823,6 +824,7 @@ def run_agent(self, args):
                         "You are a helpful AI assistant."),
                     temperature=args.temperature or config.get("temperature", 0.7),
                     max_iterations=args.max_iterations or config.get("max_iterations", 10),
+                    max_tokens=args.max_tokens or config.get("max_tokens"),
                     enable_sub_agents=not args.no_sub_agents,
                     enable_streaming=args.stream
                 )
@@ -2572,8 +2574,17 @@ def create_parser():
     run_parser.add_argument('-n', '--name', help='Agent name')
     run_parser.add_argument('-t', '--tools', nargs='+', help='Tools to enable')
     run_parser.add_argument('-c', '--config', help='Configuration file')
-    run_parser.add_argument('--system-prompt', help='System prompt')
+    run_parser.add_argument(
+        '--system-prompt', '--persona', dest='system_prompt', metavar='TEXT',
+        help='Custom persona / system prompt for this run, e.g. '
+             '"You are a patient Socratic tutor who never gives the answer."',
+    )
     run_parser.add_argument('--temperature', type=float, help='Temperature')
+    run_parser.add_argument('--max-tokens', type=int,
+                            help='Max output tokens (raise for token-heavy or '
+                                 'reasoning models, e.g. gpt-5/o-series, which '
+                                 'spend part of the budget on hidden reasoning '
+                                 'before any visible text)')
     run_parser.add_argument('--max-iterations', type=int, help='Max iterations')
     run_parser.add_argument('--mode', choices=['auto', 'single', 'sub_agents'], help='Execution mode')
     run_parser.add_argument('--no-sub-agents', action='store_true', help='Disable sub-agents')
@@ -2641,6 +2652,11 @@ def create_parser():
              'Steers every reply (unlike --preset, which only labels the session).',
     )
     chat_parser.add_argument('--temperature', type=float, help='Temperature')
+    chat_parser.add_argument('--max-tokens', type=int,
+                             help='Max output tokens per reply (raise for token-heavy '
+                                  'or reasoning models, e.g. gpt-5/o-series, which '
+                                  'spend part of the budget on hidden reasoning '
+                                  'before any visible text)')
     chat_parser.add_argument('--no-sub-agents', action='store_true', help='Disable sub-agents')
     chat_parser.add_argument('-v', '--verbose', action='store_true', default=argparse.SUPPRESS,
                              help='Verbose output (show DEBUG/INFO logs)')

effgen/cli/chat.py

@@ -75,6 +75,7 @@ def __init__(self, cli: Any, args: Any):
         # and `effgen sessions`), so a customer's conversation can be continued.
         self.session_id = getattr(args, "session_id", None)
         self.temperature = getattr(args, "temperature", None)
+        self.max_tokens = getattr(args, "max_tokens", None)
         # Provider was validated by the caller; keep the canonical name.
         self.provider = getattr(args, "_provider", None) or getattr(args, "provider", None)
 
@@ -126,6 +127,7 @@ def _build_agent(self, carry_from: Any = None) -> Any:
             "provider": self.provider,
             "tools": tools,
             "temperature": self.temperature if self.temperature is not None else 0.7,
+            "max_tokens": self.max_tokens,
             "enable_sub_agents": not getattr(self.args, "no_sub_agents", False),
             "enable_streaming": True,
         }

effgen/core/agent.py

@@ -13,6 +13,7 @@
 from __future__ import annotations
 
 import asyncio
+import difflib
 import functools
 import logging
 import time
@@ -139,6 +140,14 @@ class AgentConfig:
         system_prompt: System-level instructions
         max_iterations: Maximum tool-use loop iterations
         temperature: Generation temperature
+        top_p: Nucleus-sampling threshold; overridden per call by run(top_p=...)
+        top_k: Top-k sampling cutoff (providers that don't support it ignore it)
+        seed: Sampling seed; a fixed seed plus temperature=0 reproduces a
+            generation exactly, on providers that support it
+        presence_penalty: Penalizes tokens already present anywhere in the text
+        frequency_penalty: Penalizes tokens proportionally to how often they
+            already appeared (the standard anti-repetition knob for long text)
+        repetition_penalty: Multiplicative repeat penalty used by local/HF engines
         enable_sub_agents: Enable sub-agent spawning
         enable_memory: Enable memory systems
         enable_streaming: Enable response streaming
@@ -168,6 +177,16 @@ class AgentConfig:
     # Default output-token budget for every run(). None lets the model pick a
     # size-aware default; run(max_tokens=...) overrides it for a single call.
     max_tokens: int | None = None
+    # Sampling controls. Pinned here they apply to every run(); a run(...)
+    # kwarg of the same name overrides them for a single call. seed and the
+    # penalties default to GenerationConfig's neutral values (no effect on
+    # generation) so existing agents are unaffected until a caller sets one.
+    top_p: float = 0.9
+    top_k: int = 50
+    seed: int | None = None
+    presence_penalty: float = 0.0
+    frequency_penalty: float = 0.0
+    repetition_penalty: float = 1.0
     enable_sub_agents: bool = True
     enable_memory: bool = True
     enable_streaming: bool = False
@@ -227,6 +246,18 @@ def __post_init__(self) -> None:
     "apply_chat_template", "quantization", "trust_remote_code",
 })
 
+# run()'s recognized **kwargs — generation controls plus the checkpoint/debug
+# knobs threaded through the tool loop. A name outside this set (and not
+# starting with "_", reserved for internal call-chain bookkeeping such as
+# resume()'s _resume_scratchpad) is almost always a typo, so run() rejects it
+# instead of silently ignoring it.
+_RUN_KWARGS = frozenset({
+    "debug", "max_tokens", "temperature", "top_p", "top_k", "seed",
+    "presence_penalty", "frequency_penalty", "repetition_penalty",
+    "stop_sequences", "reasoning_effort", "tools",
+    "checkpoint_dir", "checkpoint_interval", "max_iterations",
+})
+
 
 def _agentconfig_init_guard(_dataclass_init):
     """Wrap AgentConfig.__init__ to translate a stray model-loading kwarg into a
@@ -934,6 +965,16 @@ def run(self,
         Returns:
             AgentResponse with results
         """
+        unrecognized = {k for k in kwargs if k not in _RUN_KWARGS and not k.startswith("_")}
+        if unrecognized:
+            bad = sorted(unrecognized)[0]
+            close = difflib.get_close_matches(bad, sorted(_RUN_KWARGS), n=1, cutoff=0.5)
+            hint = f" Did you mean '{close[0]}'?" if close else ""
+            raise TypeError(
+                f"run() got an unexpected keyword argument '{bad}'.{hint} "
+                f"Recognized run() kwargs: {sorted(_RUN_KWARGS)}."
+            )
+
         start_time = time.time()
         context = context or {}
 
@@ -1030,6 +1071,8 @@ def run(self,
         if output_model is None and is_pydantic_model_class(raw_schema):
             output_model = raw_schema
 
+        self._warn_reasoning_budget(kwargs.get("max_tokens"), effective_schema is not None)
+
         # Track task start
         self.execution_tracker.track_event(ExecutionEvent(
             type=EventType.TASK_START,

effgen/core/agent_generation.py

@@ -13,7 +13,11 @@
 import time
 from typing import TYPE_CHECKING, Any
 
-from ..models._adapter_utils import FINISH_LENGTH, default_max_output_tokens
+from ..models._adapter_utils import (
+    FINISH_LENGTH,
+    default_max_output_tokens,
+    needs_reasoning_headroom,
+)
 from ..models.base import BaseModel, GenerationConfig
 from ..models.errors import (
     InvalidRequestError,
@@ -35,11 +39,20 @@
 _TRUNCATION_ESCALATION_FACTOR = 4
 _TRUNCATION_MAX_TOKENS_CEILING = 8192
 
+# A reasoning model asked for structured output needs headroom well beyond the
+# reasoning-family default (it spends budget on hidden reasoning *and* the
+# JSON/schema structure before any field value). Below this, warn at call time
+# instead of only after a billed, truncated failure.
+_REASONING_STRUCTURED_OUTPUT_MIN_TOKENS = 8192
+
 if TYPE_CHECKING:
     pass
 
 logger = logging.getLogger(__name__)
 _slog = get_structured_logger(__name__)
+# (model_name, kind) pairs already warned about in this process — a call-time
+# reasoning-budget heads-up fires once per model/kind, not on every call.
+_reasoning_budget_warned: set[tuple[str, str]] = set()
 # Canonical structured observability logger — emits redacted JSON lines with OTel context
 _obs_log = _get_obs_logger(__name__)
 
@@ -307,7 +320,12 @@ def _generate(self, prompt: Any, **kwargs) -> dict[str, Any]:
                     gen_config = GenerationConfig(
                         temperature=retry_temperature,
                         max_tokens=current_max_tokens,
-                        top_p=kwargs.get('top_p', 0.9),
+                        top_p=kwargs.get('top_p', self.config.top_p),
+                        top_k=kwargs.get('top_k', self.config.top_k),
+                        seed=kwargs.get('seed', self.config.seed),
+                        presence_penalty=kwargs.get('presence_penalty', self.config.presence_penalty),
+                        frequency_penalty=kwargs.get('frequency_penalty', self.config.frequency_penalty),
+                        repetition_penalty=kwargs.get('repetition_penalty', self.config.repetition_penalty),
                         stop_sequences=kwargs.get('stop_sequences', default_stop_sequences)
                     )
 
@@ -483,6 +501,47 @@ def _record_provider_metrics(
                 output_tokens=int(completion_tokens or 0),
             )
 
+    def _warn_reasoning_budget(self, max_tokens: int | None, structured_output: bool) -> None:
+        """Warn once per model when a reasoning model's budget looks too tight.
+
+        gpt-5 / o-series models spend part of ``max_tokens`` on hidden
+        reasoning before any visible token, so a small pinned budget — or a
+        generous one paired with structured output, which needs headroom for
+        the reasoning *and* the JSON/schema structure — can produce nothing
+        and still be billed. Logged before the call so the caller sees it
+        without first burning a failed, billed round-trip.
+        """
+        model = self.model
+        if model is None or not needs_reasoning_headroom(model):
+            return
+        name = getattr(model, "model_name", None) or self.model_name or "unknown"
+
+        if structured_output:
+            min_tokens = _REASONING_STRUCTURED_OUTPUT_MIN_TOKENS
+            if max_tokens is not None and max_tokens >= min_tokens:
+                return
+            kind, hint = "structured", (
+                f"'{name}' is a reasoning model asked for structured output with "
+                f"max_tokens={max_tokens if max_tokens is not None else 'default'} — "
+                f"reasoning plus the JSON structure can consume the whole budget "
+                f"before any field is filled. Consider max_tokens>={min_tokens}."
+            )
+        else:
+            min_tokens = default_max_output_tokens(model)
+            if max_tokens is None or max_tokens >= min_tokens:
+                return
+            kind, hint = "tight_budget", (
+                f"'{name}' is a reasoning model with max_tokens={max_tokens} — it can "
+                f"spend the whole budget on hidden reasoning and return no visible "
+                f"text. Consider max_tokens>={min_tokens}."
+            )
+
+        warn_key = (name, kind)
+        if warn_key in _reasoning_budget_warned:
+            return
+        _reasoning_budget_warned.add(warn_key)
+        logger.warning(hint)
+
     def _truncation_error_detail(self, model: Any, max_tokens: int) -> dict[str, Any]:
         """Structured error for a deterministic ``max_tokens`` truncation.
 
@@ -705,7 +764,12 @@ def _generate_speculative(self, prompt: str, **kwargs) -> dict[str, Any] | None:
         gen_config = GenerationConfig(
             temperature=base_temperature,
             max_tokens=kwargs.get('max_tokens', default_max_output_tokens(self.model)),
-            top_p=kwargs.get('top_p', 0.9),
+            top_p=kwargs.get('top_p', self.config.top_p),
+            top_k=kwargs.get('top_k', self.config.top_k),
+            seed=kwargs.get('seed', self.config.seed),
+            presence_penalty=kwargs.get('presence_penalty', self.config.presence_penalty),
+            frequency_penalty=kwargs.get('frequency_penalty', self.config.frequency_penalty),
+            repetition_penalty=kwargs.get('repetition_penalty', self.config.repetition_penalty),
             stop_sequences=kwargs.get('stop_sequences', default_stop_sequences),
         )
 

effgen/models/_cost.py

@@ -20,6 +20,7 @@
 
 import json
 import logging
+import os
 import threading
 import warnings
 from dataclasses import dataclass
@@ -31,11 +32,20 @@
 _BUDGET_CONFIG_PATH = Path.home() / ".effgen" / "budget.json"
 
 
+def _budget_config_path() -> Path:
+    """Resolve the budget-config file, honoring an ``EFFGEN_BUDGET_CONFIG``
+    override (mirrors ``EFFGEN_COST_DB``) so a sandbox or test run can point the
+    budget away from the user's real ``~/.effgen/budget.json``."""
+    env_path = os.environ.get("EFFGEN_BUDGET_CONFIG")
+    return Path(env_path) if env_path else _BUDGET_CONFIG_PATH
+
+
 def _load_budget() -> dict:
-    """Load budget config from ~/.effgen/budget.json (returns empty dict if absent)."""
+    """Load budget config (returns empty dict if absent)."""
     try:
-        if _BUDGET_CONFIG_PATH.exists():
-            return json.loads(_BUDGET_CONFIG_PATH.read_text())
+        path = _budget_config_path()
+        if path.exists():
+            return json.loads(path.read_text())
     except Exception:
         logger.debug("Failed to load budget config; treating as empty", exc_info=True)
     return {}

effgen/models/cerebras_adapter.py

@@ -315,6 +315,10 @@ def _do_generate(
             request_params["stop"] = config.stop_sequences
         if config.seed is not None:
             request_params["seed"] = config.seed
+        if config.presence_penalty:
+            request_params["presence_penalty"] = config.presence_penalty
+        if config.frequency_penalty:
+            request_params["frequency_penalty"] = config.frequency_penalty
 
         # Attach tools if provided and model supports them
         model_info_dict = CEREBRAS_MODELS.get(self.model_name, {})
@@ -516,6 +520,10 @@ def generate_stream(
             request_params["stop"] = config.stop_sequences
         if config.seed is not None:
             request_params["seed"] = config.seed
+        if config.presence_penalty:
+            request_params["presence_penalty"] = config.presence_penalty
+        if config.frequency_penalty:
+            request_params["frequency_penalty"] = config.frequency_penalty
 
         request_params.update(kwargs)
 

effgen/models/fireworks_adapter.py

@@ -313,6 +313,10 @@ def _do_generate(
             request_params["stop"] = config.stop_sequences
         if config.seed is not None:
             request_params["seed"] = config.seed
+        if config.presence_penalty:
+            request_params["presence_penalty"] = config.presence_penalty
+        if config.frequency_penalty:
+            request_params["frequency_penalty"] = config.frequency_penalty
 
         info = FIREWORKS_MODELS.get(self.model_name, {})
         if tools and info.get("supports_native_tools", False):
@@ -528,6 +532,10 @@ def generate_stream(
             request_params["stop"] = config.stop_sequences
         if config.seed is not None:
             request_params["seed"] = config.seed
+        if config.presence_penalty:
+            request_params["presence_penalty"] = config.presence_penalty
+        if config.frequency_penalty:
+            request_params["frequency_penalty"] = config.frequency_penalty
 
         request_params.update(kwargs)
 

effgen/models/gemini_adapter.py

@@ -240,6 +240,12 @@ def _build_config(self, config: GenerationConfig | None) -> Any:
         }
         if config.stop_sequences:
             kwargs["stop_sequences"] = config.stop_sequences
+        if config.seed is not None:
+            kwargs["seed"] = config.seed
+        if config.presence_penalty:
+            kwargs["presence_penalty"] = config.presence_penalty
+        if config.frequency_penalty:
+            kwargs["frequency_penalty"] = config.frequency_penalty
 
         # Thinking config — only for models that support it.
         if config.thinking_budget is not None:

effgen/models/groq_adapter.py

@@ -422,6 +422,10 @@ def _do_generate(
             request_params["stop"] = config.stop_sequences
         if config.seed is not None:
             request_params["seed"] = config.seed
+        if config.presence_penalty:
+            request_params["presence_penalty"] = config.presence_penalty
+        if config.frequency_penalty:
+            request_params["frequency_penalty"] = config.frequency_penalty
 
         info = GROQ_MODELS.get(self.model_name, {})
         if tools and info.get("supports_native_tools", False):
@@ -682,6 +686,10 @@ def generate_stream(
             request_params["stop"] = config.stop_sequences
         if config.seed is not None:
             request_params["seed"] = config.seed
+        if config.presence_penalty:
+            request_params["presence_penalty"] = config.presence_penalty
+        if config.frequency_penalty:
+            request_params["frequency_penalty"] = config.frequency_penalty
 
         request_params.update(kwargs)
 

effgen/models/together_adapter.py

@@ -373,6 +373,10 @@ def _do_generate(
             request_params["stop"] = config.stop_sequences
         if config.seed is not None:
             request_params["seed"] = config.seed
+        if config.presence_penalty:
+            request_params["presence_penalty"] = config.presence_penalty
+        if config.frequency_penalty:
+            request_params["frequency_penalty"] = config.frequency_penalty
 
         info = TOGETHER_MODELS.get(self.model_name, {})
         if tools and info.get("supports_native_tools", False):
@@ -623,6 +627,10 @@ def generate_stream(
             request_params["stop"] = config.stop_sequences
         if config.seed is not None:
             request_params["seed"] = config.seed
+        if config.presence_penalty:
+            request_params["presence_penalty"] = config.presence_penalty
+        if config.frequency_penalty:
+            request_params["frequency_penalty"] = config.frequency_penalty
 
         request_params.update(kwargs)