weco 0.3.4__py3-none-any.whl → 0.3.6__py3-none-any.whl

weco/api.py

@@ -1,10 +1,10 @@
 import sys
-from typing import Dict, Any, Optional, Union, Tuple, List
+from typing import Dict, Any, Optional, Union, Tuple
 import requests
 from rich.console import Console
 
 from weco import __pkg_version__, __base_url__
-from .utils import truncate_output, determine_model_for_onboarding
+from .utils import truncate_output
 
 
 def handle_api_error(e: requests.exceptions.HTTPError, console: Console) -> None:
@@ -109,31 +109,31 @@ def start_optimization_run(
     log_dir: str = ".runs",
     auth_headers: dict = {},
     timeout: Union[int, Tuple[int, int]] = (10, 3650),
+    api_keys: Optional[Dict[str, str]] = None,
 ) -> Optional[Dict[str, Any]]:
     """Start the optimization run."""
     with console.status("[bold green]Starting Optimization..."):
         try:
-            response = requests.post(
-                f"{__base_url__}/runs/",
-                json={
-                    "source_code": source_code,
-                    "source_path": source_path,
-                    "additional_instructions": additional_instructions,
-                    "objective": {"evaluation_command": evaluation_command, "metric_name": metric_name, "maximize": maximize},
-                    "optimizer": {
-                        "steps": steps,
-                        "code_generator": code_generator_config,
-                        "evaluator": evaluator_config,
-                        "search_policy": search_policy_config,
-                    },
-                    "eval_timeout": eval_timeout,
-                    "save_logs": save_logs,
-                    "log_dir": log_dir,
-                    "metadata": {"client_name": "cli", "client_version": __pkg_version__},
+            request_json = {
+                "source_code": source_code,
+                "source_path": source_path,
+                "additional_instructions": additional_instructions,
+                "objective": {"evaluation_command": evaluation_command, "metric_name": metric_name, "maximize": maximize},
+                "optimizer": {
+                    "steps": steps,
+                    "code_generator": code_generator_config,
+                    "evaluator": evaluator_config,
+                    "search_policy": search_policy_config,
                 },
-                headers=auth_headers,
-                timeout=timeout,
-            )
+                "eval_timeout": eval_timeout,
+                "save_logs": save_logs,
+                "log_dir": log_dir,
+                "metadata": {"client_name": "cli", "client_version": __pkg_version__},
+            }
+            if api_keys:
+                request_json["api_keys"] = api_keys
+
+            response = requests.post(f"{__base_url__}/runs/", json=request_json, headers=auth_headers, timeout=timeout)
             response.raise_for_status()
             result = response.json()
             # Handle None values for code and plan fields
@@ -156,11 +156,10 @@ def resume_optimization_run(
     """Request the backend to resume an interrupted run."""
     with console.status("[bold green]Resuming run..."):
         try:
+            request_json = {"metadata": {"client_name": "cli", "client_version": __pkg_version__}}
+
             response = requests.post(
-                f"{__base_url__}/runs/{run_id}/resume",
-                json={"metadata": {"client_name": "cli", "client_version": __pkg_version__}},
-                headers=auth_headers,
-                timeout=timeout,
+                f"{__base_url__}/runs/{run_id}/resume", json=request_json, headers=auth_headers, timeout=timeout
             )
             response.raise_for_status()
             result = response.json()
@@ -180,17 +179,19 @@ def evaluate_feedback_then_suggest_next_solution(
     execution_output: str,
     auth_headers: dict = {},
     timeout: Union[int, Tuple[int, int]] = (10, 3650),
+    api_keys: Optional[Dict[str, str]] = None,
 ) -> Dict[str, Any]:
     """Evaluate the feedback and suggest the next solution."""
     try:
         # Truncate the execution output before sending to backend
         truncated_output = truncate_output(execution_output)
 
+        request_json = {"execution_output": truncated_output, "metadata": {}}
+        if api_keys:
+            request_json["api_keys"] = api_keys
+
         response = requests.post(
-            f"{__base_url__}/runs/{run_id}/suggest",
-            json={"execution_output": truncated_output, "metadata": {}},
-            headers=auth_headers,
-            timeout=timeout,
+            f"{__base_url__}/runs/{run_id}/suggest", json=request_json, headers=auth_headers, timeout=timeout
         )
         response.raise_for_status()
         result = response.json()
@@ -314,145 +315,3 @@ def report_termination(
     except Exception as e:
         print(f"Warning: Failed to report termination to backend for run {run_id}: {e}", file=sys.stderr)
         return False
-
-
-def get_optimization_suggestions_from_codebase(
-    console: Console,
-    gitingest_summary: str,
-    gitingest_tree: str,
-    gitingest_content_str: str,
-    auth_headers: dict = {},
-    timeout: Union[int, Tuple[int, int]] = (10, 3650),
-) -> Optional[List[Dict[str, Any]]]:
-    """Analyze codebase and get optimization suggestions using the model-agnostic backend API."""
-    try:
-        model = determine_model_for_onboarding()
-        response = requests.post(
-            f"{__base_url__}/onboard/analyze-codebase",
-            json={
-                "gitingest_summary": gitingest_summary,
-                "gitingest_tree": gitingest_tree,
-                "gitingest_content": gitingest_content_str,
-                "model": model,
-                "metadata": {},
-            },
-            headers=auth_headers,
-            timeout=timeout,
-        )
-        response.raise_for_status()
-        result = response.json()
-        return [option for option in result.get("options", [])]
-
-    except requests.exceptions.HTTPError as e:
-        handle_api_error(e, console)
-        return None
-    except Exception as e:
-        console.print(f"[bold red]Error: {e}[/]")
-        return None
-
-
-def generate_evaluation_script_and_metrics(
-    console: Console,
-    target_file: str,
-    description: str,
-    gitingest_content_str: str,
-    auth_headers: dict = {},
-    timeout: Union[int, Tuple[int, int]] = (10, 3650),
-) -> Tuple[Optional[str], Optional[str], Optional[str], Optional[str]]:
-    """Generate evaluation script and determine metrics using the model-agnostic backend API."""
-    try:
-        model = determine_model_for_onboarding()
-        response = requests.post(
-            f"{__base_url__}/onboard/generate-script",
-            json={
-                "target_file": target_file,
-                "description": description,
-                "gitingest_content": gitingest_content_str,
-                "model": model,
-                "metadata": {},
-            },
-            headers=auth_headers,
-            timeout=timeout,
-        )
-        response.raise_for_status()
-        result = response.json()
-        return result.get("script_content"), result.get("metric_name"), result.get("goal"), result.get("reasoning")
-    except requests.exceptions.HTTPError as e:
-        handle_api_error(e, console)
-        return None, None, None, None
-    except Exception as e:
-        console.print(f"[bold red]Error: {e}[/]")
-        return None, None, None, None
-
-
-def analyze_evaluation_environment(
-    console: Console,
-    target_file: str,
-    description: str,
-    gitingest_summary: str,
-    gitingest_tree: str,
-    gitingest_content_str: str,
-    auth_headers: dict = {},
-    timeout: Union[int, Tuple[int, int]] = (10, 3650),
-) -> Optional[Dict[str, Any]]:
-    """Analyze existing evaluation scripts and environment using the model-agnostic backend API."""
-    try:
-        model = determine_model_for_onboarding()
-        response = requests.post(
-            f"{__base_url__}/onboard/analyze-environment",
-            json={
-                "target_file": target_file,
-                "description": description,
-                "gitingest_summary": gitingest_summary,
-                "gitingest_tree": gitingest_tree,
-                "gitingest_content": gitingest_content_str,
-                "model": model,
-                "metadata": {},
-            },
-            headers=auth_headers,
-            timeout=timeout,
-        )
-        response.raise_for_status()
-        return response.json()
-
-    except requests.exceptions.HTTPError as e:
-        handle_api_error(e, console)
-        return None
-    except Exception as e:
-        console.print(f"[bold red]Error: {e}[/]")
-        return None
-
-
-def analyze_script_execution_requirements(
-    console: Console,
-    script_content: str,
-    script_path: str,
-    target_file: str,
-    auth_headers: dict = {},
-    timeout: Union[int, Tuple[int, int]] = (10, 3650),
-) -> Optional[str]:
-    """Analyze script to determine proper execution command using the model-agnostic backend API."""
-    try:
-        model = determine_model_for_onboarding()
-        response = requests.post(
-            f"{__base_url__}/onboard/analyze-script",
-            json={
-                "script_content": script_content,
-                "script_path": script_path,
-                "target_file": target_file,
-                "model": model,
-                "metadata": {},
-            },
-            headers=auth_headers,
-            timeout=timeout,
-        )
-        response.raise_for_status()
-        result = response.json()
-        return result.get("command", f"python {script_path}")
-
-    except requests.exceptions.HTTPError as e:
-        handle_api_error(e, console)
-        return f"python {script_path}"
-    except Exception as e:
-        console.print(f"[bold red]Error: {e}[/]")
-        return f"python {script_path}"

weco/cli.py

@@ -1,16 +1,47 @@
 import argparse
 import sys
-import pathlib
 from rich.console import Console
 from rich.traceback import install
 
 from .auth import clear_api_key
-from .utils import check_for_cli_updates
+from .constants import DEFAULT_MODELS
+from .utils import check_for_cli_updates, get_default_model, UnrecognizedAPIKeysError, DefaultModelNotFoundError
+
 
 install(show_locals=True)
 console = Console()
 
 
+def parse_api_keys(api_key_args: list[str] | None) -> dict[str, str]:
+    """Parse API key arguments from CLI into a dictionary.
+
+    Args:
+        api_key_args: List of strings in format 'provider=key' (e.g., ['openai=sk-xxx', 'anthropic=sk-ant-yyy'])
+
+    Returns:
+        Dictionary mapping provider names to API keys. Returns empty dict if no keys provided.
+
+    Raises:
+        ValueError: If any argument is not in the correct format.
+    """
+    if not api_key_args:
+        return {}
+
+    api_keys = {}
+    for arg in api_key_args:
+        try:
+            provider, key = (s.strip() for s in arg.split("=", 1))
+        except Exception:
+            raise ValueError(f"Invalid API key format: '{arg}'. Expected format: 'provider=key'")
+
+        if not provider or not key:
+            raise ValueError(f"Invalid API key format: '{arg}'. Provider and key must be non-empty.")
+
+        api_keys[provider.lower()] = key
+
+    return api_keys
+
+
 # Function to define and return the run_parser (or configure it on a passed subparser object)
 # This helps keep main() cleaner and centralizes run command arg definitions.
 def configure_run_parser(run_parser: argparse.ArgumentParser) -> None:
@@ -78,6 +109,29 @@ def configure_run_parser(run_parser: argparse.ArgumentParser) -> None:
         help="Automatically apply the best solution to the source file without prompting",
     )
 
+    default_api_keys = " ".join([f"{provider}=xxx" for provider, _ in DEFAULT_MODELS])
+    supported_providers = ", ".join([provider for provider, _ in DEFAULT_MODELS])
+    default_models_for_providers = "\n".join([f"- {provider}: {model}" for provider, model in DEFAULT_MODELS])
+    run_parser.add_argument(
+        "--api-key",
+        nargs="+",
+        type=str,
+        default=None,
+        help=f"""Provide one or more API keys for supported LLM providers. Specify a model with the --model flag.
+Weco will use the default model for the provider if no model is specified.
+
+Use the format 'provider=KEY', separated by spaces to specify multiple keys.
+
+Example:
+    --api-key {default_api_keys}
+
+Supported provider names: {supported_providers}.
+
+Default models for providers:
+{default_models_for_providers}
+""",
+    )
+
 
 def configure_credits_parser(credits_parser: argparse.ArgumentParser) -> None:
     """Configure the credits command parser and all its subcommands."""
@@ -129,24 +183,62 @@ def configure_resume_parser(resume_parser: argparse.ArgumentParser) -> None:
         help="Automatically apply the best solution to the source file without prompting",
     )
 
+    default_api_keys = " ".join([f"{provider}=xxx" for provider, _ in DEFAULT_MODELS])
+    supported_providers = ", ".join([provider for provider, _ in DEFAULT_MODELS])
+
+    resume_parser.add_argument(
+        "--api-key",
+        nargs="+",
+        type=str,
+        default=None,
+        help=f"""Provide one or more API keys for supported LLM providers.
+Weco will use the model associated with the run you are resuming.
+
+Use the format 'provider=KEY', separated by spaces to specify multiple keys.
+
+Example:
+    --api-key {default_api_keys}
+
+Supported provider names: {supported_providers}.
+""",
+    )
+
 
 def execute_run_command(args: argparse.Namespace) -> None:
     """Execute the 'weco run' command with all its logic."""
     from .optimizer import execute_optimization
 
+    try:
+        api_keys = parse_api_keys(args.api_key)
+    except ValueError as e:
+        console.print(f"[bold red]Error parsing API keys: {e}[/]")
+        sys.exit(1)
+
+    model = args.model
+    if not model:
+        try:
+            model = get_default_model(api_keys=api_keys)
+        except (UnrecognizedAPIKeysError, DefaultModelNotFoundError) as e:
+            console.print(f"[bold red]Error: {e}[/]")
+            sys.exit(1)
+
+        if api_keys:
+            console.print(f"[bold yellow]Custom API keys provided. Using default model: {model} for the run.[/]")
+
     success = execute_optimization(
         source=args.source,
         eval_command=args.eval_command,
         metric=args.metric,
         goal=args.goal,
+        model=model,
         steps=args.steps,
-        model=args.model,
         log_dir=args.log_dir,
         additional_instructions=args.additional_instructions,
         console=console,
         eval_timeout=args.eval_timeout,
         save_logs=args.save_logs,
         apply_change=args.apply_change,
+        api_keys=api_keys,
     )
     exit_code = 0 if success else 1
     sys.exit(exit_code)
@@ -156,7 +248,13 @@ def execute_resume_command(args: argparse.Namespace) -> None:
     """Execute the 'weco resume' command with all its logic."""
     from .optimizer import resume_optimization
 
-    success = resume_optimization(run_id=args.run_id, console=console, apply_change=args.apply_change)
+    try:
+        api_keys = parse_api_keys(args.api_key)
+    except ValueError as e:
+        console.print(f"[bold red]Error parsing API keys: {e}[/]")
+        sys.exit(1)
+
+    success = resume_optimization(run_id=args.run_id, console=console, api_keys=api_keys, apply_change=args.apply_change)
     sys.exit(0 if success else 1)
 
 
@@ -169,22 +267,13 @@ def main() -> None:
         formatter_class=argparse.RawDescriptionHelpFormatter,
     )
 
-    # Add global model argument
-    parser.add_argument(
-        "-M",
-        "--model",
-        type=str,
-        default=None,
-        help="Model to use for optimization. Defaults to `o4-mini`. See full list at docs.weco.ai/cli/supported-models",
-    )
-
     subparsers = parser.add_subparsers(
         dest="command", help="Available commands"
     )  # Removed required=True for now to handle chatbot case easily
 
     # --- Run Command Parser Setup ---
     run_parser = subparsers.add_parser(
-        "run", help="Run code optimization", formatter_class=argparse.RawDescriptionHelpFormatter, allow_abbrev=False
+        "run", help="Run code optimization", formatter_class=argparse.RawTextHelpFormatter, allow_abbrev=False
     )
     configure_run_parser(run_parser)  # Use the helper to add arguments
 
@@ -199,86 +288,11 @@ def main() -> None:
     resume_parser = subparsers.add_parser(
         "resume",
         help="Resume an interrupted optimization run",
-        formatter_class=argparse.RawDescriptionHelpFormatter,
+        formatter_class=argparse.RawTextHelpFormatter,
         allow_abbrev=False,
     )
     configure_resume_parser(resume_parser)
 
-    # Check if we should run the chatbot
-    # This logic needs to be robust. If 'run' or 'logout' is present, or -h/--help, don't run chatbot.
-    # Otherwise, if it's just 'weco' or 'weco <path>' (with optional --model), run chatbot.
-
-    def should_run_chatbot(args_list):
-        """Determine if we should run chatbot by filtering out model arguments."""
-        filtered = []
-        i = 0
-        while i < len(args_list):
-            if args_list[i] in ["-M", "--model"]:
-                # Skip the model argument and its value (if it exists)
-                i += 1  # Skip the model flag
-                if i < len(args_list):  # Skip the model value if it exists
-                    i += 1
-            elif args_list[i].startswith("--model="):
-                i += 1  # Skip --model=value format
-            else:
-                filtered.append(args_list[i])
-                i += 1
-
-        # Apply existing chatbot detection logic to filtered args
-        return len(filtered) == 0 or (len(filtered) == 1 and not filtered[0].startswith("-"))
-
-    # Check for known commands by looking at the first non-option argument
-    def get_first_non_option_arg():
-        for arg in sys.argv[1:]:
-            if not arg.startswith("-"):
-                return arg
-        return None
-
-    first_non_option = get_first_non_option_arg()
-    is_known_command = first_non_option in ["run", "logout", "credits"]
-    is_help_command = len(sys.argv) > 1 and sys.argv[1] in ["-h", "--help"]  # Check for global help
-
-    should_run_chatbot_result = should_run_chatbot(sys.argv[1:])
-    should_run_chatbot_flag = not is_known_command and not is_help_command and should_run_chatbot_result
-
-    if should_run_chatbot_flag:
-        from .chatbot import run_onboarding_chatbot  # Moved import inside
-
-        # Create a simple parser just for extracting the model argument
-        model_parser = argparse.ArgumentParser(add_help=False)
-        model_parser.add_argument("-M", "--model", type=str, default=None)
-
-        # Parse args to extract model
-        args, unknown = model_parser.parse_known_args()
-
-        # Determine project path from remaining arguments
-        filtered_args = []
-        i = 1
-        while i < len(sys.argv):
-            if sys.argv[i] in ["-M", "--model"]:
-                # Skip the model argument and its value (if it exists)
-                i += 1  # Skip the model flag
-                if i < len(sys.argv):  # Skip the model value if it exists
-                    i += 1
-            elif sys.argv[i].startswith("--model="):
-                i += 1  # Skip --model=value format
-            else:
-                filtered_args.append(sys.argv[i])
-                i += 1
-
-        project_path = pathlib.Path(filtered_args[0]) if filtered_args else pathlib.Path.cwd()
-        if not project_path.is_dir():
-            console.print(
-                f"[bold red]Error:[/] The path '{project_path}' is not a valid directory. Please provide a valid directory path."
-            )
-            sys.exit(1)
-
-        # Pass the run_parser and model to the chatbot
-        run_onboarding_chatbot(project_path, console, run_parser, model=args.model)
-        sys.exit(0)
-
-    # If not running chatbot, proceed with normal arg parsing
-    # If we reached here, a command (run, logout) or help is expected.
     args = parser.parse_args()
 
     if args.command == "logout":

weco/constants.py

@@ -7,8 +7,8 @@ Constants for the Weco CLI package.
 TRUNCATION_THRESHOLD = 51000  # Maximum length before truncation
 TRUNCATION_KEEP_LENGTH = 25000  # Characters to keep from beginning and end
 
-# Default model configuration
-DEFAULT_MODEL = "o4-mini"
-
 # Supported file extensions for additional instructions
 SUPPORTED_FILE_EXTENSIONS = [".md", ".txt", ".rst"]
+
+# Default models for each provider in order of preference
+DEFAULT_MODELS = [("gemini", "gemini-3-pro-preview"), ("openai", "o4-mini"), ("vertex_ai", "claude-opus-4-5")]