weco 0.2.22__py3-none-any.whl → 0.2.24__py3-none-any.whl

weco/api.py

@@ -1,27 +1,10 @@
 import sys
 from typing import Dict, Any, Optional, Union, Tuple, List
-
 import requests
-from requests.adapters import HTTPAdapter
-from urllib3.util.retry import Retry
 from rich.console import Console
 
 from weco import __pkg_version__, __base_url__
-
-
-# --- Session Configuration ---
-def _get_weco_session() -> requests.Session:
-    session = requests.Session()
-    retry_strategy = Retry(
-        total=3,
-        status_forcelist=[429, 500, 502, 503, 504],  # Retry on these server errors and rate limiting
-        allowed_methods=["HEAD", "GET", "PUT", "POST", "DELETE", "OPTIONS"],  # Case-insensitive
-        backoff_factor=1,  # e.g., sleep for 0s, 2s, 4s between retries (factor * (2 ** ({number of total retries} - 1)))
-    )
-    adapter = HTTPAdapter(max_retries=retry_strategy)
-    session.mount("http://", adapter)
-    session.mount("https://", adapter)
-    return session
+from .constants import DEFAULT_API_TIMEOUT
 
 
 def handle_api_error(e: requests.exceptions.HTTPError, console: Console) -> None:
@@ -48,13 +31,12 @@ def start_optimization_run(
     additional_instructions: str = None,
     api_keys: Dict[str, Any] = {},
     auth_headers: dict = {},
-    timeout: Union[int, Tuple[int, int]] = 800,
+    timeout: Union[int, Tuple[int, int]] = DEFAULT_API_TIMEOUT,
 ) -> Dict[str, Any]:
     """Start the optimization run."""
     with console.status("[bold green]Starting Optimization..."):
         try:
-            session = _get_weco_session()
-            response = session.post(
+            response = requests.post(
                 f"{__base_url__}/runs",
                 json={
                     "source_code": source_code,
@@ -72,27 +54,33 @@ def start_optimization_run(
                 timeout=timeout,
             )
             response.raise_for_status()
-            return response.json()
-        except requests.exceptions.HTTPError as e:
+            result = response.json()
+            # Handle None values for code and plan fields
+            if result.get("plan") is None:
+                result["plan"] = ""
+            if result.get("code") is None:
+                result["code"] = ""
+            return result
+        except Exception as e:
             handle_api_error(e, console)
             sys.exit(1)
-        except requests.exceptions.RequestException as e:
-            console.print(f"[bold red]Network Error starting run: {e}[/]")
+        except Exception as e:
+            console.print(f"[bold red]Error starting run: {e}[/]")
             sys.exit(1)
 
 
 def evaluate_feedback_then_suggest_next_solution(
+    console: Console,
     run_id: str,
     execution_output: str,
     additional_instructions: str = None,
     api_keys: Dict[str, Any] = {},
     auth_headers: dict = {},
-    timeout: Union[int, Tuple[int, int]] = 800,
+    timeout: Union[int, Tuple[int, int]] = DEFAULT_API_TIMEOUT,
 ) -> Dict[str, Any]:
     """Evaluate the feedback and suggest the next solution."""
     try:
-        session = _get_weco_session()
-        response = session.post(
+        response = requests.post(
             f"{__base_url__}/runs/{run_id}/suggest",
             json={
                 "execution_output": execution_output,
@@ -103,50 +91,73 @@ def evaluate_feedback_then_suggest_next_solution(
             timeout=timeout,
         )
         response.raise_for_status()
-        return response.json()
+        result = response.json()
+        # Handle None values for code and plan fields
+        if result.get("plan") is None:
+            result["plan"] = ""
+        if result.get("code") is None:
+            result["code"] = ""
+
+        return result
     except requests.exceptions.HTTPError as e:
         # Allow caller to handle suggest errors, maybe retry or terminate
-        handle_api_error(e, Console())  # Use default console if none passed
+        handle_api_error(e, console)  # Use default console if none passed
         raise  # Re-raise the exception
-    except requests.exceptions.RequestException as e:
-        print(f"Network Error during suggest: {e}")  # Use print as console might not be available
+    except Exception as e:
+        print(f"Error: {e}")  # Use print as console might not be available
         raise  # Re-raise the exception
 
 
 def get_optimization_run_status(
-    run_id: str, include_history: bool = False, auth_headers: dict = {}, timeout: Union[int, Tuple[int, int]] = 800
+    console: Console,
+    run_id: str,
+    include_history: bool = False,
+    auth_headers: dict = {},
+    timeout: Union[int, Tuple[int, int]] = DEFAULT_API_TIMEOUT,
 ) -> Dict[str, Any]:
     """Get the current status of the optimization run."""
     try:
-        session = _get_weco_session()
-        response = session.get(
+        response = requests.get(
             f"{__base_url__}/runs/{run_id}", params={"include_history": include_history}, headers=auth_headers, timeout=timeout
         )
         response.raise_for_status()
-        return response.json()
+        result = response.json()
+        # Handle None values for code and plan fields in best_result and nodes
+        if result.get("best_result"):
+            if result["best_result"].get("code") is None:
+                result["best_result"]["code"] = ""
+            if result["best_result"].get("plan") is None:
+                result["best_result"]["plan"] = ""
+        # Handle None values for code and plan fields in nodes array
+        if result.get("nodes"):
+            for i, node in enumerate(result["nodes"]):
+                if node.get("plan") is None:
+                    result["nodes"][i]["plan"] = ""
+                if node.get("code") is None:
+                    result["nodes"][i]["code"] = ""
+        return result
     except requests.exceptions.HTTPError as e:
-        handle_api_error(e, Console())  # Use default console
+        handle_api_error(e, console)  # Use default console
         raise  # Re-raise
-    except requests.exceptions.RequestException as e:
-        print(f"Network Error getting status: {e}")
+    except Exception as e:
+        print(f"Error getting run status: {e}")
         raise  # Re-raise
 
 
-def send_heartbeat(run_id: str, auth_headers: dict = {}, timeout: Union[int, Tuple[int, int]] = 10) -> bool:
+def send_heartbeat(run_id: str, auth_headers: dict = {}, timeout: Union[int, Tuple[int, int]] = (10, 10)) -> bool:
     """Send a heartbeat signal to the backend."""
     try:
-        session = _get_weco_session()
-        response = session.put(f"{__base_url__}/runs/{run_id}/heartbeat", headers=auth_headers, timeout=timeout)
+        response = requests.put(f"{__base_url__}/runs/{run_id}/heartbeat", headers=auth_headers, timeout=timeout)
         response.raise_for_status()
         return True
     except requests.exceptions.HTTPError as e:
         if e.response.status_code == 409:
-            print(f"Heartbeat ignored: Run {run_id} is not running.", file=sys.stderr)
+            print("Polling ignore: Run {run_id} is not running.", file=sys.stderr)
         else:
-            print(f"Heartbeat failed for run {run_id}: HTTP {e.response.status_code}", file=sys.stderr)
+            print(f"Polling failed for run {run_id}: HTTP {e.response.status_code}", file=sys.stderr)
         return False
-    except requests.exceptions.RequestException as e:
-        print(f"Heartbeat network error for run {run_id}: {e}", file=sys.stderr)
+    except Exception as e:
+        print(f"Error sending heartbeat for run {run_id}: {e}", file=sys.stderr)
         return False
 
 
@@ -156,12 +167,11 @@ def report_termination(
     reason: str,
     details: Optional[str] = None,
     auth_headers: dict = {},
-    timeout: Union[int, Tuple[int, int]] = 30,
+    timeout: Union[int, Tuple[int, int]] = (10, 30),
 ) -> bool:
     """Report the termination reason to the backend."""
     try:
-        session = _get_weco_session()
-        response = session.post(
+        response = requests.post(
             f"{__base_url__}/runs/{run_id}/terminate",
             json={"status_update": status_update, "termination_reason": reason, "termination_details": details},
             headers=auth_headers,
@@ -169,7 +179,7 @@ def report_termination(
         )
         response.raise_for_status()
         return True
-    except requests.exceptions.RequestException as e:
+    except Exception as e:
         print(f"Warning: Failed to report termination to backend for run {run_id}: {e}", file=sys.stderr)
         return False
 
@@ -195,24 +205,24 @@ def _determine_model_and_api_key() -> tuple[str, dict[str, str]]:
         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 model returned: {model}")
+        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,
-    console: Console,
     auth_headers: dict = {},
-    timeout: Union[int, Tuple[int, int]] = 800,
+    timeout: Union[int, Tuple[int, int]] = DEFAULT_API_TIMEOUT,
 ) -> Optional[List[Dict[str, Any]]]:
     """Analyze codebase and get optimization suggestions using the model-agnostic backend API."""
+    model, api_key_dict = _determine_model_and_api_key()
     try:
         model, api_key_dict = _determine_model_and_api_key()
-        session = _get_weco_session()
-        response = session.post(
+        response = requests.post(
             f"{__base_url__}/onboard/analyze-codebase",
             json={
                 "gitingest_summary": gitingest_summary,
@@ -228,30 +238,27 @@ def get_optimization_suggestions_from_codebase(
         result = response.json()
         return [option for option in result.get("options", [])]
 
-    except requests.exceptions.HTTPError as e:
+    except Exception as e:
         handle_api_error(e, console)
         return None
-    except requests.exceptions.RequestException as e:
-        console.print(f"[bold red]Network Error getting optimization suggestions: {e}[/]")
-        return None
     except Exception as e:
-        console.print(f"[bold red]Error calling backend API: {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,
-    console: Console,
     auth_headers: dict = {},
-    timeout: Union[int, Tuple[int, int]] = 800,
+    timeout: Union[int, Tuple[int, int]] = DEFAULT_API_TIMEOUT,
 ) -> Tuple[Optional[str], Optional[str], Optional[str], Optional[str]]:
     """Generate evaluation script and determine metrics using the model-agnostic backend API."""
+    model, api_key_dict = _determine_model_and_api_key()
     try:
         model, api_key_dict = _determine_model_and_api_key()
-        session = _get_weco_session()
-        response = session.post(
+        response = requests.post(
             f"{__base_url__}/onboard/generate-script",
             json={
                 "target_file": target_file,
@@ -266,33 +273,29 @@ def generate_evaluation_script_and_metrics(
         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 requests.exceptions.RequestException as e:
-        console.print(f"[bold red]Network Error generating evaluation script: {e}[/]")
-        return None, None, None, None
     except Exception as e:
-        console.print(f"[bold red]Error calling backend API: {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,
-    console: Console,
     auth_headers: dict = {},
-    timeout: Union[int, Tuple[int, int]] = 800,
+    timeout: Union[int, Tuple[int, int]] = DEFAULT_API_TIMEOUT,
 ) -> Optional[Dict[str, Any]]:
     """Analyze existing evaluation scripts and environment using the model-agnostic backend API."""
+    model, api_key_dict = _determine_model_and_api_key()
     try:
         model, api_key_dict = _determine_model_and_api_key()
-        session = _get_weco_session()
-        response = session.post(
+        response = requests.post(
             f"{__base_url__}/onboard/analyze-environment",
             json={
                 "target_file": target_file,
@@ -309,30 +312,27 @@ def analyze_evaluation_environment(
         response.raise_for_status()
         return response.json()
 
-    except requests.exceptions.HTTPError as e:
+    except Exception as e:
         handle_api_error(e, console)
         return None
-    except requests.exceptions.RequestException as e:
-        console.print(f"[bold red]Network Error analyzing evaluation environment: {e}[/]")
-        return None
     except Exception as e:
-        console.print(f"[bold red]Error calling backend API: {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,
-    console: Console,
     auth_headers: dict = {},
-    timeout: Union[int, Tuple[int, int]] = 800,
+    timeout: Union[int, Tuple[int, int]] = DEFAULT_API_TIMEOUT,
 ) -> Optional[str]:
     """Analyze script to determine proper execution command using the model-agnostic backend API."""
+    model, api_key_dict = _determine_model_and_api_key()
     try:
         model, api_key_dict = _determine_model_and_api_key()
-        session = _get_weco_session()
-        response = session.post(
+        response = requests.post(
             f"{__base_url__}/onboard/analyze-script",
             json={
                 "script_content": script_content,
@@ -348,12 +348,9 @@ def analyze_script_execution_requirements(
         result = response.json()
         return result.get("command", f"python {script_path}")
 
-    except requests.exceptions.HTTPError as e:
+    except Exception as e:
         handle_api_error(e, console)
         return f"python {script_path}"
-    except requests.exceptions.RequestException as e:
-        console.print(f"[bold red]Network Error analyzing script execution: {e}[/]")
-        return f"python {script_path}"
     except Exception as e:
-        console.print(f"[bold red]Error calling backend API: {e}[/]")
+        console.print(f"[bold red]Error: {e}[/]")
         return f"python {script_path}"

weco/auth.py

@@ -35,7 +35,7 @@ def save_api_key(api_key: str):
         # Set file permissions to read/write for owner only (600)
         os.chmod(CREDENTIALS_FILE, stat.S_IRUSR | stat.S_IWUSR)
     except OSError as e:
-        print(f"Error: Could not write credentials file or set permissions on {CREDENTIALS_FILE}: {e}")
+        print(f"Error: Unable to save credentials file or set permissions on {CREDENTIALS_FILE}: {e}")
 
 
 def load_weco_api_key() -> str | None:
@@ -53,7 +53,7 @@ def load_weco_api_key() -> str | None:
             credentials = json.load(f)
             return credentials.get("api_key")
     except (IOError, json.JSONDecodeError, OSError) as e:
-        print(f"Warning: Could not read or parse credentials file at {CREDENTIALS_FILE}: {e}")
+        print(f"Warning: Unable to read credentials file at {CREDENTIALS_FILE}: {e}")
         return None
 
 
@@ -64,7 +64,7 @@ def clear_api_key():
             os.remove(CREDENTIALS_FILE)
             print("Logged out successfully.")
         except OSError as e:
-            print(f"Error: Could not remove credentials file at {CREDENTIALS_FILE}: {e}")
+            print(f"Error: Unable to remove credentials file at {CREDENTIALS_FILE}: {e}")
     else:
         print("Already logged out.")
 
@@ -129,7 +129,9 @@ def perform_login(console: Console):
                             continue  # Continue polling
                         else:
                             # Unexpected 202 response format
-                            console.print(f"\n[bold red]Error:[/] Received unexpected 202 response: {token_data}")
+                            console.print(
+                                f"\n[bold red]Error:[/] Received unexpected response from authentication server: {token_data}"
+                            )
                             return False
                     # Check for standard OAuth2 errors (often 400 Bad Request)
                     elif token_response.status_code == 400:
@@ -146,7 +148,7 @@ def perform_login(console: Console):
                             console.print("\n[bold red]Error:[/] Authorization denied by user.")
                             return False
                         else:  # invalid_grant, etc.
-                            error_desc = token_data.get("error_description", "Unknown error during polling.")
+                            error_desc = token_data.get("error_description", "Unknown authentication error occurred.")
                             console.print(f"\n[bold red]Error:[/] {error_desc} ({error_code})")
                             return False
 

weco/chatbot.py

@@ -50,7 +50,7 @@ class UserInteractionHelper:
 
                 if attempts >= max_retries:
                     self.console.print(f"[red]Maximum retry attempts ({max_retries}) reached. Exiting.[/]")
-                    raise Exception("Maximum retry attempts exceeded")
+                    raise Exception("Maximum retry attempts exceeded. Please try again.")
 
                 # Show available options without the full prompt
                 if choices:
@@ -66,7 +66,7 @@ class UserInteractionHelper:
                 continue
 
         # This should never be reached due to the exception above, but just in case
-        raise Exception("Unexpected error in choice selection")
+        raise Exception("Unexpected error while selecting a choice")
 
     def get_choice_numeric(self, prompt: str, max_number: int, default: int = None, max_retries: int = 5) -> int:
         """Get numeric choice with validation and error handling."""
@@ -87,7 +87,7 @@ class UserInteractionHelper:
 
                 if attempts >= max_retries:
                     self.console.print(f"[red]Maximum retry attempts ({max_retries}) reached. Exiting.[/]")
-                    raise Exception("Maximum retry attempts exceeded")
+                    raise Exception("Maximum retry attempts exceeded. Please try again.")
 
                 # Show valid range
                 self.console.print(f"Please enter a number between [bold]1[/] and [bold]{max_number}[/]")
@@ -115,7 +115,7 @@ class UserInteractionHelper:
 
                 if attempts >= max_retries:
                     self.console.print(f"[red]Maximum retry attempts ({max_retries}) reached. Exiting.[/]")
-                    raise Exception("Maximum retry attempts exceeded")
+                    raise Exception("Maximum retry attempts exceeded. Please try again.")
 
                 self.console.print("Valid options: [bold]y[/] / [bold]n[/]")
                 if default:
@@ -123,7 +123,7 @@ class UserInteractionHelper:
 
                 continue
 
-        raise Exception("Unexpected error in yes/no selection")
+        raise Exception("Unexpected error while selecting an option")
 
     def display_optimization_options_table(self, options: List[Dict[str, str]]) -> None:
         """Display optimization options in a formatted table."""
@@ -215,7 +215,10 @@ class Chatbot:
 
             with self.console.status("[bold green]Generating optimization suggestions...[/]"):
                 result = get_optimization_suggestions_from_codebase(
-                    self.gitingest_summary, self.gitingest_tree, self.gitingest_content_str, self.console
+                    console=self.console,
+                    gitingest_summary=self.gitingest_summary,
+                    gitingest_tree=self.gitingest_tree,
+                    gitingest_content_str=self.gitingest_content_str,
                 )
 
                 if result and isinstance(result, list):
@@ -224,7 +227,7 @@ class Chatbot:
                     options = None
 
             if not options or not isinstance(options, list):
-                self.console.print("[red]Failed to get valid optimization options.[/]")
+                self.console.print("[red]Unable to retrieve valid optimization options from the backend.[/]")
                 return None
 
             if not options:
@@ -324,17 +327,19 @@ class Chatbot:
             elif action == "g" or action == "r":
                 with self.console.status("[bold green]Generating evaluation script and determining metrics...[/]"):
                     result = generate_evaluation_script_and_metrics(
-                        selected_option["target_file"],
-                        selected_option["description"],
-                        self.gitingest_content_str,
-                        self.console,
+                        console=self.console,
+                        target_file=selected_option["target_file"],
+                        description=selected_option["description"],
+                        gitingest_content_str=self.gitingest_content_str,
                     )
                 if result and result[0]:
                     eval_script_content, metric_name, goal, reasoning = result
                     if reasoning:
                         self.console.print(f"[dim]Reasoning: {reasoning}[/]")
                 else:
-                    self.console.print("[red]Failed to generate an evaluation script.[/]")
+                    self.console.print(
+                        "[red]Unable to generate an evaluation script. Please try providing a custom script path instead.[/]"
+                    )
                     eval_script_content = None
                     metric_name = None
                     goal = None
@@ -371,7 +376,10 @@ class Chatbot:
         # Analyze the script to determine the proper execution command
         with self.console.status("[bold green]Analyzing script execution requirements...[/]"):
             eval_command = analyze_script_execution_requirements(
-                eval_script_content, eval_script_path_str, selected_option["target_file"], self.console
+                console=self.console,
+                script_content=eval_script_content,
+                script_path=eval_script_path_str,
+                target_file=selected_option["target_file"],
             )
 
         return {
@@ -386,16 +394,16 @@ class Chatbot:
         """Get or create evaluation script configuration using intelligent conversation-guided approach."""
         with self.console.status("[bold green]Analyzing evaluation environment...[/]"):
             analysis = analyze_evaluation_environment(
-                selected_option["target_file"],
-                selected_option["description"],
-                self.gitingest_summary,
-                self.gitingest_tree,
-                self.gitingest_content_str,
-                self.console,
+                console=self.console,
+                target_file=selected_option["target_file"],
+                description=selected_option["description"],
+                gitingest_summary=self.gitingest_summary,
+                gitingest_tree=self.gitingest_tree,
+                gitingest_content_str=self.gitingest_content_str,
             )
 
         if not analysis:
-            self.console.print("[yellow]Failed to analyze evaluation environment. Falling back to generation.[/]")
+            self.console.print("[yellow]Unable to analyze evaluation environment. Falling back to script generation.[/]")
             return self.handle_script_generation_workflow(selected_option)
 
         self.evaluation_analysis = analysis
@@ -529,7 +537,10 @@ class Chatbot:
             if not eval_command or eval_command == f"python {script_path}":
                 with self.console.status("[bold green]Analyzing script execution requirements...[/]"):
                     eval_command = analyze_script_execution_requirements(
-                        script_content, script_path, selected_option["target_file"], self.console
+                        console=self.console,
+                        script_content=script_content,
+                        script_path=script_path,
+                        target_file=selected_option["target_file"],
                     )
 
             self.current_step = "confirmation"
@@ -711,7 +722,7 @@ class Chatbot:
         """Setup evaluation environment for the selected optimization."""
         eval_config = self.get_evaluation_configuration(selected_option)
         if not eval_config:
-            self.console.print("[red]Evaluation script setup failed.[/]")
+            self.console.print("[red]Evaluation script setup failed. Please check your script configuration and try again.[/]")
             return None
 
         eval_config = self.confirm_and_finalize_evaluation_config(eval_config)
@@ -791,7 +802,7 @@ def run_onboarding_chatbot(
         chatbot = Chatbot(project_path, console, run_parser, model)
         chatbot.start()
     except Exception as e:
-        console.print(f"[bold red]An unexpected error occurred in the chatbot: {e}[/]")
+        console.print(f"[bold red]An unexpected error occurred: {e}[/]")
         import traceback
 
         traceback.print_exc()

weco/cli.py

@@ -61,6 +61,12 @@ def configure_run_parser(run_parser: argparse.ArgumentParser) -> None:
         type=str,
         help="Description of additional instruction or path to a file containing additional instructions. Defaults to None.",
     )
+    run_parser.add_argument(
+        "--eval-timeout",
+        type=int,
+        default=None,
+        help="Timeout in seconds for each evaluation. No timeout by default. Example: --eval-timeout 3600",
+    )
 
 
 def execute_run_command(args: argparse.Namespace) -> None:
@@ -77,6 +83,7 @@ def execute_run_command(args: argparse.Namespace) -> None:
         log_dir=args.log_dir,
         additional_instructions=args.additional_instructions,
         console=console,
+        eval_timeout=args.eval_timeout,
     )
     exit_code = 0 if success else 1
     sys.exit(exit_code)
@@ -177,7 +184,9 @@ def main() -> None:
 
         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:[/] Path '{project_path}' is not a valid directory.")
+            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

weco/constants.py

@@ -0,0 +1,7 @@
+# weco/constants.py
+"""
+Constants for the Weco CLI package.
+"""
+
+# API timeout configuration (connect_timeout, read_timeout) in seconds
+DEFAULT_API_TIMEOUT = (10, 800)