synth-ai 0.2.9.dev4__py3-none-any.whl → 0.2.9.dev6__py3-none-any.whl

examples/warming_up_to_rl/run_local_rollout_modal.py

@@ -7,14 +7,12 @@ import argparse
 import asyncio
 import json
 import os
+import sys
 from pathlib import Path
 from typing import Any
 
-import sys
-
 import httpx
 from dotenv import load_dotenv
-
 from synth_ai.task import (
     RolloutEnvSpec,
     RolloutPolicySpec,
@@ -25,7 +23,9 @@ from synth_ai.task import (
 )
 
 
-def build_rollout_request(seed: int, run_id: str, *, model: str, inference_url: str, ops: list[str], api_key: str) -> RolloutRequest:
+def build_rollout_request(
+    seed: int, run_id: str, *, model: str, inference_url: str, ops: list[str], api_key: str
+) -> RolloutRequest:
     policy_config = {
         "model": model,
         "inference_url": inference_url,
@@ -45,7 +45,11 @@ def build_rollout_request(seed: int, run_id: str, *, model: str, inference_url:
 
 
 def summarise_response(data: Any) -> dict[str, Any]:
-    metrics = data.metrics.model_dump() if hasattr(data.metrics, "model_dump") else data.get("metrics", {})
+    metrics = (
+        data.metrics.model_dump()
+        if hasattr(data.metrics, "model_dump")
+        else data.get("metrics", {})
+    )
     return {
         "run_id": getattr(data, "run_id", None) or data.get("run_id"),
         "num_episodes": metrics.get("num_episodes"),
@@ -57,21 +61,54 @@ def summarise_response(data: Any) -> dict[str, Any]:
 
 
 async def main() -> None:
+    # Load .env file from current directory first if it exists
+    default_env = Path.cwd() / ".env"
+    if default_env.exists():
+        load_dotenv(default_env, override=False)
+
     parser = argparse.ArgumentParser(description=__doc__)
     parser.add_argument("--base-url", default="http://localhost:8010", help="Task app base URL")
     parser.add_argument("--env-file", type=str, default=None, help="Path to .env file with keys")
     parser.add_argument("--seed", type=int, default=42, help="Env seed to rollout")
     parser.add_argument("--run-id", default="modal-eval", help="Run identifier")
-    parser.add_argument("--model", required=True, help="Model identifier for the Crafter policy")
-    parser.add_argument("--inference-url", required=True, help="Modal backend inference base URL (e.g., http://localhost:8000/api)")
-    parser.add_argument("--task-app-key", default=None, help="Environment API key for the task app (fallback ENVIRONMENT_API_KEY)")
-    parser.add_argument("--modal-key", default=None, help="Synth/Modal API key for inference (fallback SYNTH_API_KEY)")
-    parser.add_argument("--max-llm-calls", type=int, default=20, help="Number of policy inference calls")
-    parser.add_argument("--ops", default=None, help="Comma-separated rollout ops (advanced override)")
-    parser.add_argument("--max-policy-tokens", type=int, default=None, help="Optional per-call token limit forwarded to the policy config")
-    parser.add_argument("--verbose", action="store_true", help="Print resolved configuration and headers")
+    parser.add_argument(
+        "--model",
+        required=False,
+        help="Model identifier for the Crafter policy (e.g., fft:Qwen/Qwen3-4B:job_xxx)",
+    )
+    parser.add_argument(
+        "--inference-url",
+        required=False,
+        help="Modal backend inference base URL (e.g., http://localhost:8000/api)",
+    )
+    parser.add_argument(
+        "--task-app-key",
+        default=None,
+        help="Environment API key for the task app (fallback ENVIRONMENT_API_KEY)",
+    )
+    parser.add_argument(
+        "--modal-key",
+        default=None,
+        help="Synth/Modal API key for inference (fallback SYNTH_API_KEY)",
+    )
+    parser.add_argument(
+        "--max-llm-calls", type=int, default=20, help="Number of policy inference calls"
+    )
+    parser.add_argument(
+        "--ops", default=None, help="Comma-separated rollout ops (advanced override)"
+    )
+    parser.add_argument(
+        "--max-policy-tokens",
+        type=int,
+        default=None,
+        help="Optional per-call token limit forwarded to the policy config",
+    )
+    parser.add_argument(
+        "--verbose", action="store_true", help="Print resolved configuration and headers"
+    )
     args = parser.parse_args()
 
+    # Also load from explicit --env-file if provided
     if args.env_file:
         env_path = Path(args.env_file).expanduser()
         if not env_path.exists():
@@ -79,16 +116,51 @@ async def main() -> None:
         else:
             load_dotenv(env_path, override=False)
 
+    # Prompt for required parameters if not provided
+    base_url = args.base_url
+    if args.base_url == "http://localhost:8010":
+        print("\nTask app configuration:")
+        base_url_input = input("Task app base URL [http://localhost:8001]: ").strip()
+        base_url = base_url_input if base_url_input else "http://localhost:8001"
+
+    model = args.model
+    if not model:
+        print("\nFine-tuned model configuration:")
+        print(
+            "Note: This should be the model ID returned from training (e.g., fft:Qwen/Qwen3-4B:job_abc123)"
+        )
+        model_input = input("Fine-tuned model ID: ").strip()
+        if not model_input:
+            parser.error("Model identifier is required")
+        model = model_input
+
+    inference_url = args.inference_url
+    if not inference_url:
+        inference_url_input = input("Inference URL [http://localhost:8000/api]: ").strip()
+        inference_url = inference_url_input if inference_url_input else "http://localhost:8000/api"
+
+    # Override args
+    args.base_url = base_url
+    args.model = model
+    args.inference_url = inference_url
+
+    # Check environment variables first (loaded from .env)
     task_app_key = args.task_app_key or os.getenv("ENVIRONMENT_API_KEY")
     if not task_app_key:
-        parser.error("Missing task app API key (set ENVIRONMENT_API_KEY or pass --task-app-key)")
+        print("\n[INFO] ENVIRONMENT_API_KEY not found in environment or .env file")
+        task_app_key = input("RL Environment API key: ").strip()
+        if not task_app_key:
+            parser.error("Missing task app API key")
 
     modal_key = args.modal_key or os.getenv("SYNTH_API_KEY")
     if not modal_key:
-        parser.error("Missing Synth/Modal API key (set SYNTH_API_KEY or pass --modal-key)")
+        print("[INFO] SYNTH_API_KEY not found in environment or .env file")
+        modal_key = input("Synth API key: ").strip()
+        if not modal_key:
+            parser.error("Missing Synth/Modal API key")
 
-    if synth_key and "openai.com" not in args.inference_url.lower():
-        os.environ["OPENAI_API_KEY"] = synth_key
+    if modal_key and "openai.com" not in args.inference_url.lower():
+        os.environ["OPENAI_API_KEY"] = modal_key
 
     if args.ops:
         ops = [op.strip() for op in args.ops.split(",") if op.strip()]
@@ -103,6 +175,7 @@ async def main() -> None:
             ops.extend(["agent", "env"])
 
     if args.verbose:
+
         def _mask(val: str | None) -> str:
             if not val:
                 return "<unset>"
@@ -115,11 +188,15 @@ async def main() -> None:
         print(f"  Modal API key      : {_mask(modal_key)}")
         print(f"  Ops (count={len(ops)}) : {ops}")
 
-    inf_url_norm = args.inference_url.rstrip('/')
-    if '/api' not in inf_url_norm:
-        print('[WARN] Inference URL is missing /api prefix; proxy endpoints usually live at /api/inference/v1/chat/completions.')
-    elif not inf_url_norm.lower().endswith('/api'):
-        print('[INFO] Using inference base URL; policy will append /v1/chat/completions automatically.')
+    inf_url_norm = args.inference_url.rstrip("/")
+    if "/api" not in inf_url_norm:
+        print(
+            "[WARN] Inference URL is missing /api prefix; proxy endpoints usually live at /api/inference/v1/chat/completions."
+        )
+    elif not inf_url_norm.lower().endswith("/api"):
+        print(
+            "[INFO] Using inference base URL; policy will append /v1/chat/completions automatically."
+        )
 
     async with TaskAppClient(args.base_url, api_key=task_app_key) as client:
         try:
@@ -139,20 +216,29 @@ async def main() -> None:
             if args.verbose:
                 print(f"Request headers: {request.policy.config.get('extra_headers', {})}")
             if args.max_policy_tokens is not None:
-                request.policy.config.update({
-                    "max_completion_tokens": args.max_policy_tokens,
-                    "max_tokens": args.max_policy_tokens,
-                })
+                request.policy.config.update(
+                    {
+                        "max_completion_tokens": args.max_policy_tokens,
+                        "max_tokens": args.max_policy_tokens,
+                    }
+                )
             print("Requesting rollout…")
             response = await client.rollout(request)
             summary = summarise_response(response)
             print(json.dumps(summary, indent=2))
             print(f"Ops executed: {ops}")
         except httpx.HTTPStatusError as exc:
-            detail = exc.response.json() if exc.response.headers.get("content-type", "").startswith("application/json") else exc.response.text
+            detail = (
+                exc.response.json()
+                if exc.response.headers.get("content-type", "").startswith("application/json")
+                else exc.response.text
+            )
             print(f"HTTP error {exc.response.status_code}: {detail}", file=sys.stderr)
             if exc.response.status_code in (401, 503):
-                print("Hint: ensure ENVIRONMENT_API_KEY and SYNTH_API_KEY are correctly set.", file=sys.stderr)
+                print(
+                    "Hint: ensure ENVIRONMENT_API_KEY and SYNTH_API_KEY are correctly set.",
+                    file=sys.stderr,
+                )
             raise
 
 

examples/warming_up_to_rl/run_local_rollout_parallel.py

@@ -5,7 +5,6 @@ from __future__ import annotations
 
 import argparse
 import asyncio
-import json
 import os
 from collections import Counter
 from pathlib import Path
@@ -13,15 +12,13 @@ from statistics import mean, median
 from typing import Any
 
 from dotenv import load_dotenv
-
-from synth_ai.task import TaskAppClient
-
 from synth_ai.task import (
     RolloutEnvSpec,
     RolloutPolicySpec,
     RolloutRecordConfig,
     RolloutRequest,
     RolloutSafetyConfig,
+    TaskAppClient,
 )
 
 
@@ -31,12 +28,17 @@ def build_rollout_request(
     run_id: str,
     model: str,
     inference_url: str,
+    inference_api_key: str,
     ops: list[str],
     extra_headers: dict[str, str] | None = None,
     trace_format: str = "compact",
     return_trace: bool = False,
 ) -> RolloutRequest:
-    policy_config = {"model": model, "inference_url": inference_url}
+    policy_config = {
+        "model": model,
+        "inference_url": inference_url,
+        "api_key": inference_api_key,
+    }
     if extra_headers:
         policy_config["extra_headers"] = extra_headers
     record_cfg = RolloutRecordConfig(
@@ -123,7 +125,9 @@ def analyse_rollout_response(response: Any) -> dict[str, Any]:
         if isinstance(final_list, list):
             final_achievements = [str(item) for item in final_list]
 
-    decision_rewards = trace_payload.get("decision_rewards") if isinstance(trace_payload, dict) else []
+    decision_rewards = (
+        trace_payload.get("decision_rewards") if isinstance(trace_payload, dict) else []
+    )
     trace_all: list[str] = []
     if isinstance(decision_rewards, list):
         for item in decision_rewards:
@@ -180,7 +184,9 @@ def summarise_runs(run_summaries: list[dict[str, Any]]) -> dict[str, Any]:
     return stats
 
 
-def print_summary(stats: dict[str, Any], *, run_details: list[dict[str, Any]], total_runs: int) -> None:
+def print_summary(
+    stats: dict[str, Any], *, run_details: list[dict[str, Any]], total_runs: int
+) -> None:
     if not stats:
         print("No successful rollouts to summarise.")
         return
@@ -234,7 +240,22 @@ async def execute_rollouts(args: argparse.Namespace) -> None:
 
     api_key = args.api_key or os.getenv("ENVIRONMENT_API_KEY")
     if not api_key:
-        raise RuntimeError("Missing --api-key or ENVIRONMENT_API_KEY")
+        import sys
+
+        print("Please enter your RL Environment API key:", file=sys.stderr, flush=True)
+        api_key = input("> ").strip()
+        if not api_key:
+            raise RuntimeError("RL Environment API key is required")
+
+    # Prompt for Groq API key if not set
+    groq_api_key = os.getenv("GROQ_API_KEY")
+    if not groq_api_key:
+        import sys
+
+        print("Please enter your Groq API key:", file=sys.stderr, flush=True)
+        groq_api_key = input("> ").strip()
+        if not groq_api_key:
+            raise RuntimeError("Groq API key is required")
 
     synth_key = os.getenv("SYNTH_API_KEY")
     extra_headers: dict[str, str] | None = None
@@ -252,29 +273,41 @@ async def execute_rollouts(args: argparse.Namespace) -> None:
 
     ops = build_ops(args.max_llm_calls, args.ops)
 
+    print(f"\n🚀 Starting {args.count} rollouts with {args.parallel} parallel workers...")
+    print(f"📊 Each rollout: {len(ops)} ops ({args.max_llm_calls} LLM calls)\n")
+
     async with TaskAppClient(args.base_url, api_key=api_key, timeout=args.timeout) as client:
+
         async def run_single(index: int) -> dict[str, Any]:
             run_id = f"{args.run_id}-{index:03d}"
             seed = args.seed + index * args.seed_stride
+            print(f"\n▶️  [{index + 1}/{args.count}] Starting rollout {run_id} (seed={seed})...")
+
             request = build_rollout_request(
                 seed=seed,
                 run_id=run_id,
                 model=args.model,
                 inference_url=args.inference_url,
+                inference_api_key=groq_api_key,
                 ops=ops,
                 extra_headers=extra_headers,
                 trace_format=args.trace_format,
                 return_trace=True,
             )
             if args.max_policy_tokens is not None:
-                request.policy.config.update({
-                    "max_completion_tokens": args.max_policy_tokens,
-                    "max_tokens": args.max_policy_tokens,
-                })
+                request.policy.config.update(
+                    {
+                        "max_completion_tokens": args.max_policy_tokens,
+                        "max_tokens": args.max_policy_tokens,
+                    }
+                )
 
             try:
                 response = await client.rollout(request)
                 summary = analyse_rollout_response(response)
+                print(
+                    f"\n✅ [{index + 1}/{args.count}] Completed {run_id} (outcome={summary.get('outcome_score', 'N/A')})"
+                )
                 return {
                     "ok": True,
                     "run_id": run_id,
@@ -283,6 +316,7 @@ async def execute_rollouts(args: argparse.Namespace) -> None:
                     "summary": summary,
                 }
             except Exception as exc:  # pragma: no cover - surface errors
+                print(f"\n❌ [{index + 1}/{args.count}] Failed {run_id}: {exc}")
                 return {
                     "ok": False,
                     "run_id": run_id,
@@ -302,6 +336,7 @@ async def execute_rollouts(args: argparse.Namespace) -> None:
     successes = [item for item in results if item.get("ok")]
     failures = [item for item in results if not item.get("ok")]
 
+    print(f"\n{'=' * 100}\n")
     stats = summarise_runs([item["summary"] for item in successes])
     print_summary(stats, run_details=successes, total_runs=args.count)
 
@@ -317,17 +352,43 @@ def parse_args() -> argparse.Namespace:
     parser.add_argument("--base-url", default="http://localhost:8001", help="Task app base URL")
     parser.add_argument("--api-key", help="Environment API key (or set via --env-file)")
     parser.add_argument("--env-file", help="Path to .env file providing API keys")
-    parser.add_argument("--model", default="gpt-4o-mini", help="Model identifier for the Crafter policy")
-    parser.add_argument("--inference-url", default="https://api.openai.com", help="Inference base URL for the policy")
+    parser.add_argument(
+        "--model", default="gpt-4o-mini", help="Model identifier for the Crafter policy"
+    )
+    parser.add_argument(
+        "--inference-url",
+        default="https://api.openai.com",
+        help="Inference base URL for the policy",
+    )
     parser.add_argument("--seed", type=int, default=42, help="Base seed for the first rollout")
-    parser.add_argument("--seed-stride", type=int, default=1, help="Increment applied to the seed for each rollout")
-    parser.add_argument("--count", type=int, default=20, help="Number of rollout trajectories to execute")
+    parser.add_argument(
+        "--seed-stride", type=int, default=1, help="Increment applied to the seed for each rollout"
+    )
+    parser.add_argument(
+        "--count", type=int, default=20, help="Number of rollout trajectories to execute"
+    )
     parser.add_argument("--parallel", type=int, default=4, help="Maximum concurrent rollouts")
     parser.add_argument("--ops", help="Comma-separated rollout ops (advanced override)")
-    parser.add_argument("--max-llm-calls", type=int, default=20, help="Number of agent/env pairs per rollout when --ops not provided")
-    parser.add_argument("--max-policy-tokens", type=int, help="Optional per-call token limit forwarded to the policy config")
-    parser.add_argument("--timeout", type=float, default=600.0, help="HTTP timeout (seconds) for task app requests")
-    parser.add_argument("--trace-format", default="compact", choices=["compact", "full"], help="Trace format requested from the task app")
+    parser.add_argument(
+        "--max-llm-calls",
+        type=int,
+        default=20,
+        help="Number of agent/env pairs per rollout when --ops not provided",
+    )
+    parser.add_argument(
+        "--max-policy-tokens",
+        type=int,
+        help="Optional per-call token limit forwarded to the policy config",
+    )
+    parser.add_argument(
+        "--timeout", type=float, default=600.0, help="HTTP timeout (seconds) for task app requests"
+    )
+    parser.add_argument(
+        "--trace-format",
+        default="compact",
+        choices=["compact", "full"],
+        help="Trace format requested from the task app",
+    )
     parser.add_argument("--run-id", default="batch-demo", help="Run ID prefix for rollouts")
     parser.add_argument("--verbose", action="store_true", help="Print resolved configuration")
     return parser.parse_args()

examples/warming_up_to_rl/run_local_rollout_traced.py

@@ -6,13 +6,12 @@ from __future__ import annotations
 import argparse
 import asyncio
 import json
+import os
+import sys
 from pathlib import Path
 from typing import Any
 
-import sys
-
 import httpx
-
 from synth_ai.task import (
     RolloutEnvSpec,
     RolloutPolicySpec,
@@ -29,6 +28,7 @@ def build_rollout_request(
     run_id: str,
     model: str,
     inference_url: str,
+    inference_api_key: str,
     ops: list[str],
     return_trace: bool,
     trace_format: str,
@@ -37,6 +37,7 @@ def build_rollout_request(
     policy_config = {
         "model": model,
         "inference_url": inference_url,
+        "api_key": inference_api_key,
     }
     if max_policy_tokens is not None:
         policy_config.update(
@@ -64,7 +65,11 @@ def build_rollout_request(
 
 
 def summarise_rollout(response: Any) -> dict[str, Any]:
-    metrics = response.metrics.model_dump() if hasattr(response, "metrics") else response.get("metrics", {})
+    metrics = (
+        response.metrics.model_dump()
+        if hasattr(response, "metrics")
+        else response.get("metrics", {})
+    )
     return {
         "run_id": getattr(response, "run_id", None) or response.get("run_id"),
         "num_episodes": metrics.get("num_episodes"),
@@ -83,17 +88,25 @@ def summarise_trace(trace: Any) -> dict[str, Any]:
 
     format_hint = "compact" if "events_count" in trace or "lm_calls" in trace else "full"
     events_count = trace.get("events_count")
-    if events_count is None and "event_history" in trace and isinstance(trace["event_history"], list):
+    if (
+        events_count is None
+        and "event_history" in trace
+        and isinstance(trace["event_history"], list)
+    ):
         events_count = len(trace["event_history"])
     messages_count = trace.get("messages_count")
-    if messages_count is None and "markov_blanket_message_history" in trace and isinstance(
-        trace["markov_blanket_message_history"], list
+    if (
+        messages_count is None
+        and "markov_blanket_message_history" in trace
+        and isinstance(trace["markov_blanket_message_history"], list)
     ):
         messages_count = len(trace["markov_blanket_message_history"])
 
     metadata = trace.get("metadata") if isinstance(trace.get("metadata"), dict) else {}
     lm_calls = trace.get("lm_calls") if isinstance(trace.get("lm_calls"), list) else []
-    decision_rewards = trace.get("decision_rewards") if isinstance(trace.get("decision_rewards"), list) else []
+    decision_rewards = (
+        trace.get("decision_rewards") if isinstance(trace.get("decision_rewards"), list) else []
+    )
 
     return {
         "session_id": trace.get("session_id"),
@@ -215,11 +228,13 @@ def print_reward_summary(
     if decision_rewards:
         print("  Decision rewards:")
         for entry in decision_rewards:
-            turn = entry.get('turn')
-            ach_delta = entry.get('ach_delta')
-            unique_delta = entry.get('unique_delta')
-            achievements = entry.get('achievements') or []
-            print(f"    turn={turn}, ach_delta={ach_delta}, unique_delta={unique_delta}, achievements={achievements}")
+            turn = entry.get("turn")
+            ach_delta = entry.get("ach_delta")
+            unique_delta = entry.get("unique_delta")
+            achievements = entry.get("achievements") or []
+            print(
+                f"    turn={turn}, ach_delta={ach_delta}, unique_delta={unique_delta}, achievements={achievements}"
+            )
     else:
         print("  Decision rewards: none recorded")
 
@@ -242,16 +257,40 @@ def print_reward_summary(
 
 
 async def main() -> None:
+    # Load .env file from current directory if it exists
+    env_file = Path.cwd() / ".env"
+    if env_file.exists():
+        from dotenv import load_dotenv
+
+        load_dotenv(env_file)
+
     parser = argparse.ArgumentParser(description=__doc__)
-    parser.add_argument("--base-url", default="http://localhost:8010", help="Task app base URL")
-    parser.add_argument("--api-key", required=True, help="Environment API key")
+    parser.add_argument("--base-url", default="http://localhost:8001", help="Task app base URL")
+    parser.add_argument("--api-key", help="RL Environment API key (will prompt if not provided)")
+    parser.add_argument(
+        "--inference-api-key", help="Inference provider API key (will prompt if not provided)"
+    )
     parser.add_argument("--seed", type=int, default=42, help="Environment seed")
     parser.add_argument("--run-id", default="local-trace", help="Run identifier")
     parser.add_argument("--model", default="gpt-4o-mini", help="OpenAI-compatible model id")
-    parser.add_argument("--inference-url", default="https://api.openai.com", help="Inference base URL (OpenAI/Groq)")
-    parser.add_argument("--ops", help="Comma-separated rollout ops (fallback: alternating agent/env)")
-    parser.add_argument("--max-llm-calls", type=int, default=1, help="Number of agent/env pairs when --ops not supplied")
-    parser.add_argument("--max-policy-tokens", type=int, default=None, help="Optional max token budget forwarded to policy")
+    parser.add_argument(
+        "--inference-url", default="https://api.openai.com", help="Inference base URL (OpenAI/Groq)"
+    )
+    parser.add_argument(
+        "--ops", help="Comma-separated rollout ops (fallback: alternating agent/env)"
+    )
+    parser.add_argument(
+        "--max-llm-calls",
+        type=int,
+        default=1,
+        help="Number of agent/env pairs when --ops not supplied",
+    )
+    parser.add_argument(
+        "--max-policy-tokens",
+        type=int,
+        default=None,
+        help="Optional max token budget forwarded to policy",
+    )
     parser.add_argument(
         "--trace-format",
         choices=["compact", "full"],
@@ -286,10 +325,69 @@ async def main() -> None:
     )
     args = parser.parse_args()
 
+    # Prompt for required parameters if not provided
+    base_url = args.base_url
+    if args.base_url == "http://localhost:8001":
+        print("\nTask app configuration:")
+        base_url_input = input("Task app base URL [http://localhost:8001]: ").strip()
+        base_url = base_url_input if base_url_input else "http://localhost:8001"
+
+    api_key = args.api_key or os.getenv("ENVIRONMENT_API_KEY")
+    if not api_key:
+        api_key = input("RL Environment API key (from ENVIRONMENT_API_KEY): ").strip()
+        if not api_key:
+            parser.error("RL Environment API key is required")
+
+    # Use Groq by default
+    model = "llama-3.3-70b-versatile"
+    inference_url = "https://api.groq.com/openai"
+
+    print("\nInference configuration (Groq):")
+    inference_api_key = args.inference_api_key or os.getenv("GROQ_API_KEY")
+    if not inference_api_key:
+        inference_api_key = input("Groq API key: ").strip()
+        if not inference_api_key:
+            parser.error("Groq API key is required")
+
+        # Save to .env for future use
+        env_path = Path.cwd() / ".env"
+        try:
+            # Read existing .env
+            existing_lines = []
+            if env_path.exists():
+                existing_lines = env_path.read_text().splitlines()
+
+            # Check if GROQ_API_KEY already exists
+            key_exists = any(line.strip().startswith("GROQ_API_KEY=") for line in existing_lines)
+
+            if not key_exists:
+                # Append to .env
+                with open(env_path, "a") as f:
+                    if existing_lines and not existing_lines[-1].strip():
+                        # File exists and last line is not empty
+                        pass
+                    elif existing_lines:
+                        # Add newline before appending
+                        f.write("\n")
+                    f.write(f"GROQ_API_KEY={inference_api_key}\n")
+                print(f"[INFO] Saved GROQ_API_KEY to {env_path}")
+        except Exception as e:
+            print(f"[WARN] Could not save GROQ_API_KEY to .env: {e}")
+
+    print("\nRollout configuration:")
+    max_llm_calls = args.max_llm_calls
+    if args.max_llm_calls == 1:
+        max_llm_calls_input = input("Max LLM calls [10]: ").strip()
+        max_llm_calls = int(max_llm_calls_input) if max_llm_calls_input else 10
+
+    # Override args with prompted values
+    args.base_url = base_url
+    args.max_llm_calls = max_llm_calls
+
     ops = ensure_ops(args.ops, args.max_llm_calls)
     return_trace = not args.no_trace
 
-    async with TaskAppClient(args.base_url, api_key=args.api_key, timeout=args.timeout) as client:
+    async with TaskAppClient(args.base_url, api_key=api_key, timeout=args.timeout) as client:
         try:
             print(f"Fetching task_info for seed {args.seed}…")
             task_info = await client.task_info(seeds=[args.seed])
@@ -302,8 +400,9 @@ async def main() -> None:
             request = build_rollout_request(
                 seed=args.seed,
                 run_id=args.run_id,
-                model=args.model,
-                inference_url=args.inference_url,
+                model=model,
+                inference_url=inference_url,
+                inference_api_key=inference_api_key,
                 ops=ops,
                 return_trace=return_trace,
                 trace_format=args.trace_format,
@@ -350,7 +449,11 @@ async def main() -> None:
                 "Tip: export TASKAPP_TRACING_ENABLED=1 and optionally TASKAPP_SFT_OUTPUT_DIR before running `uvx synth-ai serve …` to persist traces/SFT."
             )
         except httpx.HTTPStatusError as exc:
-            detail = exc.response.json() if exc.response.headers.get("content-type", "").startswith("application/json") else exc.response.text
+            detail = (
+                exc.response.json()
+                if exc.response.headers.get("content-type", "").startswith("application/json")
+                else exc.response.text
+            )
             print(f"HTTP error {exc.response.status_code}: {detail}", file=sys.stderr)
             if exc.response.status_code in (401, 503):
                 print(