simplification

scripts/get_setup_env.py

@@ -47,24 +47,75 @@ def get_setup_env_vars(setup_name, suite_name=None):
     return setup.env
 
 
+def get_setup_config(setup_name, suite_name=None):
+    """
+    Get full configuration (env vars + defaults) for a setup.
+
+    Args:
+        setup_name: Name of the setup (e.g., 'ollama', 'gpt')
+        suite_name: Optional suite name to get default setup if setup_name is None
+
+    Returns:
+        Dictionary with 'env' and 'defaults' keys
+    """
+    # If no setup specified, try to get default from suite
+    if not setup_name and suite_name:
+        suite = SUITE_DEFINITIONS.get(suite_name)
+        if suite and suite.default_setup:
+            setup_name = suite.default_setup
+
+    if not setup_name:
+        return {"env": {}, "defaults": {}}
+
+    setup = SETUP_DEFINITIONS.get(setup_name)
+    if not setup:
+        print(
+            f"Error: Unknown setup '{setup_name}'. Available: {', '.join(sorted(SETUP_DEFINITIONS.keys()))}",
+            file=sys.stderr,
+        )
+        sys.exit(1)
+
+    return {"env": setup.env, "defaults": setup.defaults}
+
+
 def main():
-    parser = argparse.ArgumentParser(description="Extract environment variables from a test setup")
+    parser = argparse.ArgumentParser(description="Extract environment variables and defaults from a test setup")
     parser.add_argument("--setup", help="Setup name (e.g., ollama, gpt)")
     parser.add_argument("--suite", help="Suite name to get default setup from if --setup not provided")
     parser.add_argument("--format", choices=["bash", "json"], default="bash", help="Output format (default: bash)")
+    parser.add_argument(
+        "--include-defaults",
+        action="store_true",
+        help="Include setup defaults (text_model, embedding_model, etc.) in addition to env vars",
+    )
 
     args = parser.parse_args()
 
-    env_vars = get_setup_env_vars(args.setup, args.suite)
+    if args.include_defaults:
+        config = get_setup_config(args.setup, args.suite)
+        if args.format == "bash":
+            # Output env vars as bash export statements
+            for key, value in config["env"].items():
+                print(f"export {key}='{value}'")
+            # Output defaults as bash export statements with LLAMA_STACK_TEST_ prefix
+            for key, value in config["defaults"].items():
+                env_key = f"LLAMA_STACK_TEST_{key.upper()}"
+                print(f"export {env_key}='{value}'")
+        elif args.format == "json":
+            import json
 
-    if args.format == "bash":
-        # Output as bash export statements
-        for key, value in env_vars.items():
-            print(f"export {key}='{value}'")
-    elif args.format == "json":
-        import json
+            print(json.dumps(config))
+    else:
+        env_vars = get_setup_env_vars(args.setup, args.suite)
 
-        print(json.dumps(env_vars))
+        if args.format == "bash":
+            # Output as bash export statements
+            for key, value in env_vars.items():
+                print(f"export {key}='{value}'")
+        elif args.format == "json":
+            import json
+
+            print(json.dumps(env_vars))
 
 
 if __name__ == "__main__":