weco 0.2.27__py3-none-any.whl → 0.3.0__py3-none-any.whl

weco/api.py

@@ -4,22 +4,54 @@ import requests
 from rich.console import Console
 
 from weco import __pkg_version__, __base_url__
-from .constants import DEFAULT_API_TIMEOUT
-from .utils import truncate_output
+from .constants import CODEGEN_API_TIMEOUT, STATUS_API_TIMEOUT
+from .utils import truncate_output, determine_model_for_onboarding
 
 
 def handle_api_error(e: requests.exceptions.HTTPError, console: Console) -> None:
     """Extract and display error messages from API responses in a structured format."""
+    status = getattr(e.response, "status_code", None)
     try:
-        detail = e.response.json()["detail"]
-    except (ValueError, KeyError):  # Handle cases where response is not JSON or detail key is missing
-        detail = f"HTTP {e.response.status_code} Error: {e.response.text}"
-    console.print(f"[bold red]{detail}[/]")
+        payload = e.response.json()
+        detail = payload.get("detail", payload)
+    except (ValueError, AttributeError):
+        detail = getattr(e.response, "text", "") or f"HTTP {status} Error"
+
+    def _render(detail_obj: Any) -> None:
+        if isinstance(detail_obj, str):
+            console.print(f"[bold red]{detail_obj}[/]")
+        elif isinstance(detail_obj, dict):
+            # Try common message keys in order of preference
+            message_keys = ("message", "error", "msg", "detail")
+            message = next((detail_obj.get(key) for key in message_keys if detail_obj.get(key)), None)
+            suggestion = detail_obj.get("suggestion")
+            if message:
+                console.print(f"[bold red]{message}[/]")
+            else:
+                console.print(f"[bold red]HTTP {status} Error[/]")
+            if suggestion:
+                console.print(f"[yellow]{suggestion}[/]")
+            extras = {
+                k: v
+                for k, v in detail_obj.items()
+                if k not in {"message", "error", "msg", "detail", "suggestion"} and v not in (None, "")
+            }
+            for key, value in extras.items():
+                console.print(f"[dim]{key}: {value}[/]")
+        elif isinstance(detail_obj, list) and detail_obj:
+            _render(detail_obj[0])
+            for extra in detail_obj[1:]:
+                console.print(f"[yellow]{extra}[/]")
+        else:
+            console.print(f"[bold red]{detail_obj or f'HTTP {status} Error'}[/]")
+
+    _render(detail)
 
 
 def start_optimization_run(
     console: Console,
     source_code: str,
+    source_path: str,
     evaluation_command: str,
     metric_name: str,
     maximize: bool,
@@ -28,9 +60,11 @@ def start_optimization_run(
     evaluator_config: Dict[str, Any],
     search_policy_config: Dict[str, Any],
     additional_instructions: str = None,
-    api_keys: Dict[str, Any] = {},
+    eval_timeout: Optional[int] = None,
+    save_logs: bool = False,
+    log_dir: str = ".runs",
     auth_headers: dict = {},
-    timeout: Union[int, Tuple[int, int]] = DEFAULT_API_TIMEOUT,
+    timeout: Union[int, Tuple[int, int]] = CODEGEN_API_TIMEOUT,
 ) -> Optional[Dict[str, Any]]:
     """Start the optimization run."""
     with console.status("[bold green]Starting Optimization..."):
@@ -39,6 +73,7 @@ def start_optimization_run(
                 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": {
@@ -47,7 +82,10 @@ def start_optimization_run(
                         "evaluator": evaluator_config,
                         "search_policy": search_policy_config,
                     },
-                    "metadata": {"client_name": "cli", "client_version": __pkg_version__, **api_keys},
+                    "eval_timeout": eval_timeout,
+                    "save_logs": save_logs,
+                    "log_dir": log_dir,
+                    "metadata": {"client_name": "cli", "client_version": __pkg_version__},
                 },
                 headers=auth_headers,
                 timeout=timeout,
@@ -68,14 +106,37 @@ def start_optimization_run(
             return None
 
 
+def resume_optimization_run(
+    console: Console, run_id: str, auth_headers: dict = {}, timeout: Union[int, Tuple[int, int]] = STATUS_API_TIMEOUT
+) -> Optional[Dict[str, Any]]:
+    """Request the backend to resume an interrupted run."""
+    with console.status("[bold green]Resuming run..."):
+        try:
+            response = requests.post(
+                f"{__base_url__}/runs/{run_id}/resume",
+                json={"metadata": {"client_name": "cli", "client_version": __pkg_version__}},
+                headers=auth_headers,
+                timeout=timeout,
+            )
+            response.raise_for_status()
+            result = response.json()
+            return result
+        except requests.exceptions.HTTPError as e:
+            handle_api_error(e, console)
+            return None
+        except Exception as e:
+            console.print(f"[bold red]Error resuming run: {e}[/]")
+            return None
+
+
 def evaluate_feedback_then_suggest_next_solution(
     console: Console,
     run_id: str,
+    step: int,
     execution_output: str,
     additional_instructions: str = None,
-    api_keys: Dict[str, Any] = {},
     auth_headers: dict = {},
-    timeout: Union[int, Tuple[int, int]] = DEFAULT_API_TIMEOUT,
+    timeout: Union[int, Tuple[int, int]] = CODEGEN_API_TIMEOUT,
 ) -> Dict[str, Any]:
     """Evaluate the feedback and suggest the next solution."""
     try:
@@ -84,11 +145,7 @@ def evaluate_feedback_then_suggest_next_solution(
 
         response = requests.post(
             f"{__base_url__}/runs/{run_id}/suggest",
-            json={
-                "execution_output": truncated_output,
-                "additional_instructions": additional_instructions,
-                "metadata": {**api_keys},
-            },
+            json={"execution_output": truncated_output, "additional_instructions": additional_instructions, "metadata": {}},
             headers=auth_headers,
             timeout=timeout,
         )
@@ -99,8 +156,44 @@ def evaluate_feedback_then_suggest_next_solution(
             result["plan"] = ""
         if result.get("code") is None:
             result["code"] = ""
-
         return result
+    except requests.exceptions.ReadTimeout as e:
+        # NOTE: This can occur when:
+        # 1. The server is busy and the request times out
+        # 2. When intermediaries drop the connection so even after the server responds,
+        # the client doesn't receive the response and times out
+
+        # Here we ONLY try to recover in the latter case i.e., server completed request but
+        # client didn't receive the response and timed out
+        run_status_recovery_response = get_optimization_run_status(
+            console=console, run_id=run_id, include_history=True, auth_headers=auth_headers
+        )
+        current_step = run_status_recovery_response.get("current_step")
+        current_status = run_status_recovery_response.get("status")
+        # The run should be "running" and the current step should correspond to the solution step we are attempting to generate
+        is_valid_run_state = current_status is not None and current_status == "running"
+        is_valid_step = current_step is not None and current_step == step
+        if is_valid_run_state and is_valid_step:
+            nodes = run_status_recovery_response.get("nodes") or []
+            # We need at least 2 nodes to reconstruct the expected response i.e., the last two nodes
+            if len(nodes) >= 2:
+                nodes_sorted_ascending = sorted(nodes, key=lambda n: n["step"])
+                latest_node = nodes_sorted_ascending[-1]
+                penultimate_node = nodes_sorted_ascending[-2]
+                # If the server finished generating the next candidate, it should be exactly this step
+                if latest_node and latest_node["step"] == step:
+                    # Try to reconstruct the expected response from the /suggest endpoint using the run status info
+                    reconstructed_expected_response = {
+                        "run_id": run_id,
+                        "previous_solution_metric_value": penultimate_node.get("metric_value"),
+                        "solution_id": latest_node.get("solution_id"),
+                        "code": latest_node.get("code"),
+                        "plan": latest_node.get("plan"),
+                        "is_done": False,
+                    }
+                    return reconstructed_expected_response
+        # If we couldn't recover, raise the timeout error so the run can be resumed by the user
+        raise requests.exceptions.ReadTimeout(e)
     except requests.exceptions.HTTPError as e:
         # Allow caller to handle suggest errors, maybe retry or terminate
         handle_api_error(e, console)
@@ -115,7 +208,7 @@ def get_optimization_run_status(
     run_id: str,
     include_history: bool = False,
     auth_headers: dict = {},
-    timeout: Union[int, Tuple[int, int]] = DEFAULT_API_TIMEOUT,
+    timeout: Union[int, Tuple[int, int]] = STATUS_API_TIMEOUT,
 ) -> Dict[str, Any]:
     """Get the current status of the optimization run."""
     try:
@@ -186,43 +279,17 @@ def report_termination(
         return False
 
 
-# --- Chatbot API Functions ---
-def _determine_model_and_api_key() -> tuple[str, dict[str, str]]:
-    """Determine the model and API key to use based on available environment variables.
-
-    Uses the shared model selection logic to maintain consistency.
-    Returns (model_name, api_key_dict)
-    """
-    from .utils import read_api_keys_from_env, determine_default_model
-
-    llm_api_keys = read_api_keys_from_env()
-    model = determine_default_model(llm_api_keys)
-
-    # Create API key dictionary with only the key for the selected model
-    if model == "o4-mini":
-        api_key_dict = {"OPENAI_API_KEY": llm_api_keys["OPENAI_API_KEY"]}
-    elif model == "claude-sonnet-4-0":
-        api_key_dict = {"ANTHROPIC_API_KEY": llm_api_keys["ANTHROPIC_API_KEY"]}
-    elif model == "gemini-2.5-pro":
-        api_key_dict = {"GEMINI_API_KEY": llm_api_keys["GEMINI_API_KEY"]}
-    else:
-        # This should never happen if determine_default_model works correctly
-        raise ValueError(f"Unknown default model choice: {model}")
-
-    return model, api_key_dict
-
-
 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]] = DEFAULT_API_TIMEOUT,
+    timeout: Union[int, Tuple[int, int]] = CODEGEN_API_TIMEOUT,
 ) -> Optional[List[Dict[str, Any]]]:
     """Analyze codebase and get optimization suggestions using the model-agnostic backend API."""
     try:
-        model, api_key_dict = _determine_model_and_api_key()
+        model = determine_model_for_onboarding()
         response = requests.post(
             f"{__base_url__}/onboard/analyze-codebase",
             json={
@@ -230,7 +297,7 @@ def get_optimization_suggestions_from_codebase(
                 "gitingest_tree": gitingest_tree,
                 "gitingest_content": gitingest_content_str,
                 "model": model,
-                "metadata": api_key_dict,
+                "metadata": {},
             },
             headers=auth_headers,
             timeout=timeout,
@@ -253,11 +320,11 @@ def generate_evaluation_script_and_metrics(
     description: str,
     gitingest_content_str: str,
     auth_headers: dict = {},
-    timeout: Union[int, Tuple[int, int]] = DEFAULT_API_TIMEOUT,
+    timeout: Union[int, Tuple[int, int]] = CODEGEN_API_TIMEOUT,
 ) -> Tuple[Optional[str], Optional[str], Optional[str], Optional[str]]:
     """Generate evaluation script and determine metrics using the model-agnostic backend API."""
     try:
-        model, api_key_dict = _determine_model_and_api_key()
+        model = determine_model_for_onboarding()
         response = requests.post(
             f"{__base_url__}/onboard/generate-script",
             json={
@@ -265,7 +332,7 @@ def generate_evaluation_script_and_metrics(
                 "description": description,
                 "gitingest_content": gitingest_content_str,
                 "model": model,
-                "metadata": api_key_dict,
+                "metadata": {},
             },
             headers=auth_headers,
             timeout=timeout,
@@ -289,11 +356,11 @@ def analyze_evaluation_environment(
     gitingest_tree: str,
     gitingest_content_str: str,
     auth_headers: dict = {},
-    timeout: Union[int, Tuple[int, int]] = DEFAULT_API_TIMEOUT,
+    timeout: Union[int, Tuple[int, int]] = CODEGEN_API_TIMEOUT,
 ) -> Optional[Dict[str, Any]]:
     """Analyze existing evaluation scripts and environment using the model-agnostic backend API."""
     try:
-        model, api_key_dict = _determine_model_and_api_key()
+        model = determine_model_for_onboarding()
         response = requests.post(
             f"{__base_url__}/onboard/analyze-environment",
             json={
@@ -303,7 +370,7 @@ def analyze_evaluation_environment(
                 "gitingest_tree": gitingest_tree,
                 "gitingest_content": gitingest_content_str,
                 "model": model,
-                "metadata": api_key_dict,
+                "metadata": {},
             },
             headers=auth_headers,
             timeout=timeout,
@@ -325,11 +392,11 @@ def analyze_script_execution_requirements(
     script_path: str,
     target_file: str,
     auth_headers: dict = {},
-    timeout: Union[int, Tuple[int, int]] = DEFAULT_API_TIMEOUT,
+    timeout: Union[int, Tuple[int, int]] = CODEGEN_API_TIMEOUT,
 ) -> Optional[str]:
     """Analyze script to determine proper execution command using the model-agnostic backend API."""
     try:
-        model, api_key_dict = _determine_model_and_api_key()
+        model = determine_model_for_onboarding()
         response = requests.post(
             f"{__base_url__}/onboard/analyze-script",
             json={
@@ -337,7 +404,7 @@ def analyze_script_execution_requirements(
                 "script_path": script_path,
                 "target_file": target_file,
                 "model": model,
-                "metadata": api_key_dict,
+                "metadata": {},
             },
             headers=auth_headers,
             timeout=timeout,

weco/auth.py

@@ -184,9 +184,9 @@ def perform_login(console: Console):
         return False
 
 
-def handle_authentication(console: Console, llm_api_keys: dict) -> tuple[str | None, dict]:
+def handle_authentication(console: Console) -> tuple[str | None, dict]:
     """
-    Handle the complete authentication flow.
+    Handle the complete authentication flow. Authentication is now mandatory.
 
     Returns:
         tuple: (weco_api_key, auth_headers)
@@ -194,13 +194,16 @@ def handle_authentication(console: Console, llm_api_keys: dict) -> tuple[str | N
     weco_api_key = load_weco_api_key()
 
     if not weco_api_key:
+        console.print("[bold yellow]Authentication Required[/]")
+        console.print("With our new credit-based billing system, authentication is required to use Weco.")
+        console.print("You'll receive free credits to get started!")
+        console.print("")
+
         login_choice = Prompt.ask(
-            "Log in to Weco to save run history or use anonymously? ([bold]L[/]ogin / [bold]S[/]kip)",
-            choices=["l", "s"],
-            default="s",
+            "Would you like to log in now? ([bold]Y[/]es / [bold]N[/]o)", choices=["y", "n"], default="y"
         ).lower()
 
-        if login_choice == "l":
+        if login_choice == "y":
             console.print("[cyan]Starting login process...[/]")
             if not perform_login(console):
                 console.print("[bold red]Login process failed or was cancelled.[/]")
@@ -210,14 +213,9 @@ def handle_authentication(console: Console, llm_api_keys: dict) -> tuple[str | N
             if not weco_api_key:
                 console.print("[bold red]Error: Login completed but failed to retrieve API key.[/]")
                 return None, {}
-
-        elif login_choice == "s":
-            console.print("[yellow]Proceeding anonymously. LLM API keys must be provided via environment variables.[/]")
-            if not llm_api_keys:
-                console.print(
-                    "[bold red]Error:[/] No LLM API keys found in environment (e.g., OPENAI_API_KEY). Cannot proceed anonymously."
-                )
-                return None, {}
+        else:
+            console.print("[yellow]Authentication is required to use Weco. Please run 'weco' again when ready to log in.[/]")
+            return None, {}
 
     # Build auth headers
     auth_headers = {}

weco/chatbot.py

@@ -220,6 +220,7 @@ class Chatbot:
                     gitingest_summary=self.gitingest_summary,
                     gitingest_tree=self.gitingest_tree,
                     gitingest_content_str=self.gitingest_content_str,
+                    auth_headers=getattr(self, "auth_headers", {}),
                 )
 
                 if result and isinstance(result, list):
@@ -332,6 +333,7 @@ class Chatbot:
                         target_file=selected_option["target_file"],
                         description=selected_option["description"],
                         gitingest_content_str=self.gitingest_content_str,
+                        auth_headers=getattr(self, "auth_headers", {}),
                     )
                 if result and result[0]:
                     eval_script_content, metric_name, goal, reasoning = result
@@ -381,6 +383,7 @@ class Chatbot:
                 script_content=eval_script_content,
                 script_path=eval_script_path_str,
                 target_file=selected_option["target_file"],
+                auth_headers=getattr(self, "auth_headers", {}),
             )
 
         return {
@@ -401,6 +404,7 @@ class Chatbot:
                 gitingest_summary=self.gitingest_summary,
                 gitingest_tree=self.gitingest_tree,
                 gitingest_content_str=self.gitingest_content_str,
+                auth_headers=getattr(self, "auth_headers", {}),
             )
 
         if not analysis:
@@ -542,6 +546,7 @@ class Chatbot:
                         script_content=script_content,
                         script_path=script_path,
                         target_file=selected_option["target_file"],
+                        auth_headers=getattr(self, "auth_headers", {}),
                     )
 
             self.current_step = "confirmation"
@@ -749,10 +754,9 @@ class Chatbot:
             self.resolved_model = self.user_specified_model
         else:
             # Use same default model selection as weco run
-            from .utils import determine_default_model, read_api_keys_from_env
+            from .utils import determine_model_for_onboarding
 
-            llm_api_keys = read_api_keys_from_env()
-            self.resolved_model = determine_default_model(llm_api_keys)
+            self.resolved_model = determine_model_for_onboarding()
 
         target_file = selected_option["target_file"]
         additional_instructions = selected_option["description"]
@@ -766,6 +770,17 @@ class Chatbot:
         self.console.print("[bold cyan]Welcome to Weco![/]")
         self.console.print(f"Let's optimize your codebase in: [cyan]{self.project_path}[/]\n")
 
+        # Mandatory authentication as per PLAN.md
+        from .auth import handle_authentication
+
+        weco_api_key, auth_headers = handle_authentication(self.console)
+        if not weco_api_key:
+            self.console.print("[yellow]Authentication is required to use Weco. Exiting...[/]")
+            return
+
+        # Store auth headers for API calls
+        self.auth_headers = auth_headers
+
         options = self.analyze_codebase_and_get_optimization_options()
         if not options:
             return

weco/cli.py

@@ -74,6 +74,52 @@ def configure_run_parser(run_parser: argparse.ArgumentParser) -> None:
     )
 
 
+def configure_credits_parser(credits_parser: argparse.ArgumentParser) -> None:
+    """Configure the credits command parser and all its subcommands."""
+    credits_subparsers = credits_parser.add_subparsers(dest="credits_command", help="Credit management commands")
+
+    # Credits balance command
+    _ = credits_subparsers.add_parser("balance", help="Check your current credit balance")
+
+    # Coerce CLI input into a float with two decimal precision for the API payload.
+    def _parse_credit_amount(value: str) -> float:
+        try:
+            amount = float(value)
+        except ValueError as exc:
+            raise argparse.ArgumentTypeError("Amount must be a number.") from exc
+
+        return round(amount, 2)
+
+    # Credits topup command
+    topup_parser = credits_subparsers.add_parser("topup", help="Purchase additional credits")
+    topup_parser.add_argument(
+        "amount",
+        nargs="?",
+        type=_parse_credit_amount,
+        default=_parse_credit_amount("10"),
+        metavar="CREDITS",
+        help="Amount of credits to purchase (minimum 2, defaults to 10)",
+    )
+
+    # Credits autotopup command
+    autotopup_parser = credits_subparsers.add_parser("autotopup", help="Configure automatic top-up")
+    autotopup_parser.add_argument("--enable", action="store_true", help="Enable automatic top-up")
+    autotopup_parser.add_argument("--disable", action="store_true", help="Disable automatic top-up")
+    autotopup_parser.add_argument(
+        "--threshold", type=float, default=4.0, help="Balance threshold to trigger auto top-up (default: 4.0 credits)"
+    )
+    autotopup_parser.add_argument(
+        "--amount", type=float, default=50.0, help="Amount to top up when threshold is reached (default: 50.0 credits)"
+    )
+
+
+def configure_resume_parser(resume_parser: argparse.ArgumentParser) -> None:
+    """Configure arguments for the resume command."""
+    resume_parser.add_argument(
+        "run_id", type=str, help="The UUID of the run to resume (e.g., '0002e071-1b67-411f-a514-36947f0c4b31')"
+    )
+
+
 def execute_run_command(args: argparse.Namespace) -> None:
     """Execute the 'weco run' command with all its logic."""
     from .optimizer import execute_optimization
@@ -95,6 +141,14 @@ def execute_run_command(args: argparse.Namespace) -> None:
     sys.exit(exit_code)
 
 
+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)
+    sys.exit(0 if success else 1)
+
+
 def main() -> None:
     """Main function for the Weco CLI."""
     check_for_cli_updates()
@@ -126,6 +180,19 @@ def main() -> None:
     # --- Logout Command Parser Setup ---
     _ = subparsers.add_parser("logout", help="Log out from Weco and clear saved API key.")
 
+    # --- Credits Command Parser Setup ---
+    credits_parser = subparsers.add_parser("credits", help="Manage your Weco credits")
+    configure_credits_parser(credits_parser)  # Use the helper to add subcommands and arguments
+
+    # --- Resume Command Parser Setup ---
+    resume_parser = subparsers.add_parser(
+        "resume",
+        help="Resume an interrupted optimization run",
+        formatter_class=argparse.RawDescriptionHelpFormatter,
+        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.
@@ -157,7 +224,7 @@ def main() -> None:
         return None
 
     first_non_option = get_first_non_option_arg()
-    is_known_command = first_non_option in ["run", "logout"]
+    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:])
@@ -208,6 +275,13 @@ def main() -> None:
         sys.exit(0)
     elif args.command == "run":
         execute_run_command(args)
+    elif args.command == "credits":
+        from .credits import handle_credits_command
+
+        handle_credits_command(args, console)
+        sys.exit(0)
+    elif args.command == "resume":
+        execute_resume_command(args)
     else:
         # This case should be hit if 'weco' is run alone and chatbot logic didn't catch it,
         # or if an invalid command is provided.

weco/constants.py

@@ -4,8 +4,15 @@ Constants for the Weco CLI package.
 """
 
 # API timeout configuration (connect_timeout, read_timeout) in seconds
-DEFAULT_API_TIMEOUT = (10, 800)
+CODEGEN_API_TIMEOUT = (30, 3650)
+STATUS_API_TIMEOUT = (10, 180)
 
 # Output truncation configuration
 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"]