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

weco/__init__.py

@@ -5,11 +5,5 @@ import importlib.metadata
 __pkg_version__ = importlib.metadata.version("weco")
 __api_version__ = "v1"
 
-__base_url__ = f"https://api.weco.ai/{__api_version__}"
-# If user specifies a custom base URL, use that instead
-if os.environ.get("WECO_BASE_URL"):
-    __base_url__ = os.environ.get("WECO_BASE_URL")
-
-__dashboard_url__ = "https://dashboard.weco.ai"
-if os.environ.get("WECO_DASHBOARD_URL"):
-    __dashboard_url__ = os.environ.get("WECO_DASHBOARD_URL")
+__base_url__ = os.environ.get("WECO_BASE_URL", f"https://api.weco.ai/{__api_version__}")
+__dashboard_url__ = os.environ.get("WECO_DASHBOARD_URL", "https://dashboard.weco.ai")

weco/api.py

@@ -1,12 +1,30 @@
-from typing import Dict, Any, Optional
-import rich
-import requests
-from weco import __pkg_version__, __base_url__
 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__
 
-def handle_api_error(e: requests.exceptions.HTTPError, console: rich.console.Console) -> None:
+
+# --- 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
+
+
+def handle_api_error(e: requests.exceptions.HTTPError, console: Console) -> None:
     """Extract and display error messages from API responses in a structured format."""
     try:
         detail = e.response.json()["detail"]
@@ -18,7 +36,7 @@ def handle_api_error(e: requests.exceptions.HTTPError, console: rich.console.Con
 
 
 def start_optimization_run(
-    console: rich.console.Console,
+    console: Console,
     source_code: str,
     evaluation_command: str,
     metric_name: str,
@@ -30,12 +48,13 @@ def start_optimization_run(
     additional_instructions: str = None,
     api_keys: Dict[str, Any] = {},
     auth_headers: dict = {},
-    timeout: int = 800,
+    timeout: Union[int, Tuple[int, int]] = 800,
 ) -> Dict[str, Any]:
     """Start the optimization run."""
     with console.status("[bold green]Starting Optimization..."):
         try:
-            response = requests.post(
+            session = _get_weco_session()
+            response = session.post(
                 f"{__base_url__}/runs",
                 json={
                     "source_code": source_code,
@@ -68,11 +87,12 @@ def evaluate_feedback_then_suggest_next_solution(
     additional_instructions: str = None,
     api_keys: Dict[str, Any] = {},
     auth_headers: dict = {},
-    timeout: int = 800,
+    timeout: Union[int, Tuple[int, int]] = 800,
 ) -> Dict[str, Any]:
     """Evaluate the feedback and suggest the next solution."""
     try:
-        response = requests.post(
+        session = _get_weco_session()
+        response = session.post(
             f"{__base_url__}/runs/{run_id}/suggest",
             json={
                 "execution_output": execution_output,
@@ -94,11 +114,12 @@ def evaluate_feedback_then_suggest_next_solution(
 
 
 def get_optimization_run_status(
-    run_id: str, include_history: bool = False, auth_headers: dict = {}, timeout: int = 800
+    run_id: str, include_history: bool = False, auth_headers: dict = {}, timeout: Union[int, Tuple[int, int]] = 800
 ) -> Dict[str, Any]:
     """Get the current status of the optimization run."""
     try:
-        response = requests.get(
+        session = _get_weco_session()
+        response = session.get(
             f"{__base_url__}/runs/{run_id}", params={"include_history": include_history}, headers=auth_headers, timeout=timeout
         )
         response.raise_for_status()
@@ -111,10 +132,11 @@ def get_optimization_run_status(
         raise  # Re-raise
 
 
-def send_heartbeat(run_id: str, auth_headers: dict = {}, timeout: int = 10) -> bool:
+def send_heartbeat(run_id: str, auth_headers: dict = {}, timeout: Union[int, Tuple[int, int]] = 10) -> bool:
     """Send a heartbeat signal to the backend."""
     try:
-        response = requests.put(f"{__base_url__}/runs/{run_id}/heartbeat", headers=auth_headers, timeout=timeout)
+        session = _get_weco_session()
+        response = session.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:
@@ -129,11 +151,17 @@ def send_heartbeat(run_id: str, auth_headers: dict = {}, timeout: int = 10) -> b
 
 
 def report_termination(
-    run_id: str, status_update: str, reason: str, details: Optional[str] = None, auth_headers: dict = {}, timeout: int = 30
+    run_id: str,
+    status_update: str,
+    reason: str,
+    details: Optional[str] = None,
+    auth_headers: dict = {},
+    timeout: Union[int, Tuple[int, int]] = 30,
 ) -> bool:
     """Report the termination reason to the backend."""
     try:
-        response = requests.post(
+        session = _get_weco_session()
+        response = session.post(
             f"{__base_url__}/runs/{run_id}/terminate",
             json={"status_update": status_update, "termination_reason": reason, "termination_details": details},
             headers=auth_headers,
@@ -144,3 +172,188 @@ def report_termination(
     except requests.exceptions.RequestException as e:
         print(f"Warning: Failed to report termination to backend for run {run_id}: {e}", file=sys.stderr)
         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 model returned: {model}")
+
+    return model, api_key_dict
+
+
+def get_optimization_suggestions_from_codebase(
+    gitingest_summary: str,
+    gitingest_tree: str,
+    gitingest_content_str: str,
+    console: Console,
+    auth_headers: dict = {},
+    timeout: Union[int, Tuple[int, int]] = 800,
+) -> 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()
+        session = _get_weco_session()
+        response = session.post(
+            f"{__base_url__}/onboard/analyze-codebase",
+            json={
+                "gitingest_summary": gitingest_summary,
+                "gitingest_tree": gitingest_tree,
+                "gitingest_content": gitingest_content_str,
+                "model": model,
+                "metadata": api_key_dict,
+            },
+            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 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}[/]")
+        return None
+
+
+def generate_evaluation_script_and_metrics(
+    target_file: str,
+    description: str,
+    gitingest_content_str: str,
+    console: Console,
+    auth_headers: dict = {},
+    timeout: Union[int, Tuple[int, int]] = 800,
+) -> 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()
+        session = _get_weco_session()
+        response = session.post(
+            f"{__base_url__}/onboard/generate-script",
+            json={
+                "target_file": target_file,
+                "description": description,
+                "gitingest_content": gitingest_content_str,
+                "model": model,
+                "metadata": api_key_dict,
+            },
+            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 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}[/]")
+        return None, None, None, None
+
+
+def analyze_evaluation_environment(
+    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,
+) -> 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()
+        session = _get_weco_session()
+        response = session.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": api_key_dict,
+            },
+            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 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}[/]")
+        return None
+
+
+def analyze_script_execution_requirements(
+    script_content: str,
+    script_path: str,
+    target_file: str,
+    console: Console,
+    auth_headers: dict = {},
+    timeout: Union[int, Tuple[int, int]] = 800,
+) -> 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()
+        session = _get_weco_session()
+        response = session.post(
+            f"{__base_url__}/onboard/analyze-script",
+            json={
+                "script_content": script_content,
+                "script_path": script_path,
+                "target_file": target_file,
+                "model": model,
+                "metadata": api_key_dict,
+            },
+            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 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}[/]")
+        return f"python {script_path}"

weco/auth.py

@@ -3,6 +3,13 @@ import os
 import pathlib
 import json
 import stat
+import time
+import requests
+import webbrowser
+from rich.console import Console
+from rich.live import Live
+from rich.prompt import Prompt
+from . import __base_url__
 
 CONFIG_DIR = pathlib.Path.home() / ".config" / "weco"
 CREDENTIALS_FILE = CONFIG_DIR / "credentials.json"
@@ -27,10 +34,8 @@ def save_api_key(api_key: str):
             json.dump(credentials, f)
         # Set file permissions to read/write for owner only (600)
         os.chmod(CREDENTIALS_FILE, stat.S_IRUSR | stat.S_IWUSR)
-    except IOError as e:
-        print(f"Error: Could not write credentials file at {CREDENTIALS_FILE}: {e}")
     except OSError as e:
-        print(f"Warning: Could not set permissions on {CREDENTIALS_FILE}: {e}")
+        print(f"Error: Could not write credentials file or set permissions on {CREDENTIALS_FILE}: {e}")
 
 
 def load_weco_api_key() -> str | None:
@@ -62,3 +67,159 @@ def clear_api_key():
             print(f"Error: Could not remove credentials file at {CREDENTIALS_FILE}: {e}")
     else:
         print("Already logged out.")
+
+
+def perform_login(console: Console):
+    """Handles the device login flow."""
+    try:
+        # 1. Initiate device login
+        console.print("Initiating login...")
+        init_response = requests.post(f"{__base_url__}/auth/device/initiate")
+        init_response.raise_for_status()
+        init_data = init_response.json()
+
+        device_code = init_data["device_code"]
+        verification_uri = init_data["verification_uri"]
+        expires_in = init_data["expires_in"]
+        interval = init_data["interval"]
+
+        # 2. Display instructions
+        console.print("\n[bold yellow]Action Required:[/]")
+        console.print("Please open the following URL in your browser to authenticate:")
+        console.print(f"[link={verification_uri}]{verification_uri}[/link]")
+        console.print(f"This request will expire in {expires_in // 60} minutes.")
+        console.print("Attempting to open the authentication page in your default browser...")  # Notify user
+
+        # Automatically open the browser
+        try:
+            if not webbrowser.open(verification_uri):
+                console.print("[yellow]Could not automatically open the browser. Please open the link manually.[/]")
+        except Exception as browser_err:
+            console.print(
+                f"[yellow]Could not automatically open the browser ({browser_err}). Please open the link manually.[/]"
+            )
+
+        console.print("Waiting for authentication...", end="")
+
+        # 3. Poll for token
+        start_time = time.time()
+        # Use a simple text update instead of Spinner within Live for potentially better compatibility
+        polling_status = "Waiting..."
+        with Live(polling_status, refresh_per_second=1, transient=True, console=console) as live_status:
+            while True:
+                # Check for timeout
+                if time.time() - start_time > expires_in:
+                    console.print("\n[bold red]Error:[/] Login request timed out.")
+                    return False
+
+                time.sleep(interval)
+                live_status.update("Waiting... (checking status)")
+
+                try:
+                    token_response = requests.post(
+                        f"{__base_url__}/auth/device/token",
+                        json={"grant_type": "urn:ietf:params:oauth:grant-type:device_code", "device_code": device_code},
+                    )
+
+                    # Check for 202 Accepted - Authorization Pending
+                    if token_response.status_code == 202:
+                        token_data = token_response.json()
+                        if token_data.get("error") == "authorization_pending":
+                            live_status.update("Waiting... (authorization pending)")
+                            continue  # Continue polling
+                        else:
+                            # Unexpected 202 response format
+                            console.print(f"\n[bold red]Error:[/] Received unexpected 202 response: {token_data}")
+                            return False
+                    # Check for standard OAuth2 errors (often 400 Bad Request)
+                    elif token_response.status_code == 400:
+                        token_data = token_response.json()
+                        error_code = token_data.get("error", "unknown_error")
+                        if error_code == "slow_down":
+                            interval += 5  # Increase polling interval if instructed
+                            live_status.update(f"Waiting... (slowing down polling to {interval}s)")
+                            continue
+                        elif error_code == "expired_token":
+                            console.print("\n[bold red]Error:[/] Login request expired.")
+                            return False
+                        elif error_code == "access_denied":
+                            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.")
+                            console.print(f"\n[bold red]Error:[/] {error_desc} ({error_code})")
+                            return False
+
+                    # Check for other non-200/non-202/non-400 HTTP errors
+                    token_response.raise_for_status()
+                    # If successful (200 OK and no 'error' field)
+                    token_data = token_response.json()
+                    if "access_token" in token_data:
+                        api_key = token_data["access_token"]
+                        save_api_key(api_key)
+                        console.print("\n[bold green]Login successful![/]")
+                        return True
+                    else:
+                        # Unexpected successful response format
+                        console.print("\n[bold red]Error:[/] Received unexpected response from server during polling.")
+                        print(token_data)
+                        return False
+                except requests.exceptions.RequestException as e:
+                    # Handle network errors during polling gracefully
+                    live_status.update("Waiting... (network error, retrying)")
+                    console.print(f"\n[bold yellow]Warning:[/] Network error during polling: {e}. Retrying...")
+                    time.sleep(interval * 2)  # Simple backoff
+    except requests.exceptions.HTTPError as e:
+        from .api import handle_api_error  # Import here to avoid circular imports
+
+        handle_api_error(e, console)
+    except requests.exceptions.RequestException as e:
+        # Catch other request errors
+        console.print(f"\n[bold red]Network Error:[/] {e}")
+        return False
+    except Exception as e:
+        console.print(f"\n[bold red]An unexpected error occurred during login:[/] {e}")
+        return False
+
+
+def handle_authentication(console: Console, llm_api_keys: dict) -> tuple[str | None, dict]:
+    """
+    Handle the complete authentication flow.
+
+    Returns:
+        tuple: (weco_api_key, auth_headers)
+    """
+    weco_api_key = load_weco_api_key()
+
+    if not weco_api_key:
+        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",
+        ).lower()
+
+        if login_choice == "l":
+            console.print("[cyan]Starting login process...[/]")
+            if not perform_login(console):
+                console.print("[bold red]Login process failed or was cancelled.[/]")
+                return None, {}
+
+            weco_api_key = load_weco_api_key()
+            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, {}
+
+    # Build auth headers
+    auth_headers = {}
+    if weco_api_key:
+        auth_headers["Authorization"] = f"Bearer {weco_api_key}"
+
+    return weco_api_key, auth_headers