Merge pull request #136 from 73ai/cli-strengthen-types

refactor CLI: types, auth and cluster config

Author: priyanshujain

.gitignore

@@ -192,3 +192,5 @@ services/mcp/*
 
 .shoreman.pid
 .aider*
+
+.claude/*

cli/.gitignore

@@ -0,0 +1 @@
+.claude/*

cli/Makefile

@@ -0,0 +1,24 @@
+.PHONY: install lint format typecheck check test clean
+
+install:
+	uv sync
+
+lint:
+	uv run ruff check src/
+
+format:
+	uv run ruff format src/
+	uv run ruff check --fix src/
+
+typecheck:
+	uv run pyright src/
+
+check: lint typecheck
+
+test:
+	uv run pytest tests/
+
+clean:
+	rm -rf .pytest_cache .ruff_cache .mypy_cache __pycache__
+	find . -type d -name __pycache__ -exec rm -rf {} + 2>/dev/null || true
+	find . -type f -name "*.pyc" -delete 2>/dev/null || true

cli/pyproject.toml

@@ -43,5 +43,7 @@ infragpt = "infragpt.main:cli"
 
 [dependency-groups]
 dev = [
+    "pyright>=1.1.407",
     "pytest>=8.4.1",
+    "ruff>=0.14.10",
 ]

cli/src/infragpt/agent.py

@@ -50,11 +50,11 @@ def add_message(
         self,
         role: str,
         content: str,
-        tool_calls: Optional[List[Dict]] = None,
+        tool_calls: Optional[List[Dict[str, Any]]] = None,
         tool_call_id: Optional[str] = None,
-    ):
+    ) -> None:
         """Add a message to the conversation context."""
-        message_dict = {"role": role, "content": content}
+        message_dict: Dict[str, Any] = {"role": role, "content": content}
 
         if tool_calls:
             message_dict["tool_calls"] = tool_calls

cli/src/infragpt/api_client.py

@@ -1,7 +1,10 @@
 from dataclasses import dataclass
 from typing import Optional
+
 import httpx
 
+from infragpt.config import get_api_base_url
+
 
 @dataclass
 class DeviceFlowResponse:
@@ -52,8 +55,8 @@ def __init__(self, status_code: int, message: str):
 
 
 class InfraGPTClient:
-    def __init__(self, api_base_url: str, timeout: float = 30.0):
-        self.api_base_url = api_base_url.rstrip("/")
+    def __init__(self, timeout: float = 30.0):
+        self.api_base_url = get_api_base_url()
         self.timeout = timeout
 
     def _make_request(

cli/src/infragpt/auth.py

@@ -9,7 +9,7 @@
 import httpx
 from cryptography.fernet import InvalidToken
 
-from infragpt.config import CONFIG_DIR, console
+from infragpt.config import CONFIG_DIR, console, get_console_base_url
 from infragpt.encryption import (
     encrypt_data,
     decrypt_data,
@@ -35,20 +35,6 @@
 TOKEN_REFRESH_THRESHOLD_HOURS = 1
 
 
-def _get_api_base_url(api_base_url: Optional[str]) -> str:
-    DEFAULT = "https://api.infragpt.io"
-    if api_base_url:
-        return api_base_url.rstrip("/")
-    return DEFAULT
-
-
-def _get_console_base_url(console_base_url: Optional[str]) -> str:
-    DEFAULT = "https://app.infragpt.io"
-    if console_base_url:
-        return console_base_url.rstrip("/")
-    return DEFAULT
-
-
 @dataclass
 class AuthStatus:
     authenticated: bool
@@ -57,7 +43,6 @@ class AuthStatus:
     access_token: Optional[str] = None
     refresh_token: Optional[str] = None
     expires_at: Optional[str] = None
-    api_base_url: Optional[str] = None
 
 
 def _load_auth_data() -> Optional[dict]:
@@ -108,7 +93,6 @@ def get_auth_status() -> AuthStatus:
         access_token=data.get("access_token"),
         refresh_token=data.get("refresh_token"),
         expires_at=data.get("expires_at"),
-        api_base_url=data.get("api_base_url"),
     )
 
 
@@ -150,8 +134,7 @@ def refresh_token_if_needed() -> bool:
         return True
 
     try:
-        api_base_url = data.get("api_base_url")
-        client = InfraGPTClient(api_base_url=api_base_url)
+        client = InfraGPTClient()
         result = client.refresh_token(refresh_token)
 
         from datetime import timedelta
@@ -171,12 +154,9 @@ def refresh_token_if_needed() -> bool:
         return False
 
 
-def login(
-    api_base_url: Optional[str] = None, console_base_url: Optional[str] = None
-) -> None:
+def login() -> None:
     """Authenticate with InfraGPT platform using device flow."""
-    api_url = _get_api_base_url(api_base_url)
-    client = InfraGPTClient(api_base_url=api_url)
+    client = InfraGPTClient()
 
     console.print("\n[bold]Authenticating with InfraGPT...[/bold]\n")
 
@@ -190,7 +170,7 @@ def login(
     if flow.verification_url.startswith("http"):
         verification_url = flow.verification_url
     else:
-        console_url = _get_console_base_url(console_base_url)
+        console_url = get_console_base_url()
         path = flow.verification_url.lstrip("/")
         verification_url = f"{console_url}/{path}"
 
@@ -238,7 +218,6 @@ def login(
                 "organization_id": result.organization_id,
                 "user_id": result.user_id,
                 "expires_at": expires_at.isoformat(),
-                "api_base_url": api_base_url or InfraGPTClient.DEFAULT_SERVER_URL,
             }
             _save_auth_data(auth_data)
 
@@ -264,8 +243,7 @@ def logout() -> None:
 
     if data and data.get("access_token"):
         try:
-            api_base_url = data.get("api_base_url")
-            client = InfraGPTClient(api_base_url=api_base_url)
+            client = InfraGPTClient()
             client.revoke_token(data["access_token"])
         except (InfraGPTAPIError, httpx.RequestError):
             pass  # Token may already be invalid; don't fail logout
@@ -285,7 +263,7 @@ def fetch_gcp_credentials() -> Optional[GCPCredentials]:
         return None
 
     try:
-        client = InfraGPTClient(api_base_url=status.api_base_url)
+        client = InfraGPTClient()
         return client.get_gcp_credentials(status.access_token)
     except InfraGPTAPIError:
         return None
@@ -298,7 +276,7 @@ def fetch_gke_cluster_info() -> Optional[GKEClusterInfo]:
         return None
 
     try:
-        client = InfraGPTClient(api_base_url=status.api_base_url)
+        client = InfraGPTClient()
         return client.get_gke_cluster_info(status.access_token)
     except InfraGPTAPIError:
         return None
@@ -326,7 +304,7 @@ def validate_token_with_api() -> None:
         raise AuthValidationError("Not authenticated")
 
     try:
-        client = InfraGPTClient(api_base_url=status.api_base_url)
+        client = InfraGPTClient()
         client.validate_token(status.access_token)
     except InfraGPTAPIError as e:
         if e.status_code == 401:
@@ -365,8 +343,7 @@ def refresh_token_strict() -> None:
         return
 
     try:
-        api_base_url = data.get("api_base_url")
-        client = InfraGPTClient(api_base_url=api_base_url)
+        client = InfraGPTClient()
         result = client.refresh_token(refresh_token)
 
         from datetime import timedelta
@@ -394,7 +371,7 @@ def fetch_gcp_credentials_strict() -> GCPCredentials:
         raise GCPCredentialError("Not authenticated")
 
     try:
-        client = InfraGPTClient(api_base_url=status.api_base_url)
+        client = InfraGPTClient()
         return client.get_gcp_credentials(status.access_token)
     except InfraGPTAPIError as e:
         if e.status_code == 404:
@@ -413,7 +390,7 @@ def fetch_gke_cluster_info_strict() -> GKEClusterInfo:
         raise GKEClusterError("Not authenticated")
 
     try:
-        client = InfraGPTClient(api_base_url=status.api_base_url)
+        client = InfraGPTClient()
         return client.get_gke_cluster_info(status.access_token)
     except InfraGPTAPIError as e:
         if e.status_code == 404:

cli/src/infragpt/bump_version.py

@@ -6,7 +6,8 @@
 from pathlib import Path
 
 
-def get_current_version():
+def get_current_version() -> str:
+    """Get the current version from __init__.py."""
     init_file = Path("src/infragpt/__init__.py")
     if not init_file.exists():
         raise FileNotFoundError(f"Could not find {init_file}")
@@ -21,8 +22,8 @@ def get_current_version():
     return match.group(1)
 
 
-def update_version(new_version):
-    # Update version in __init__.py
+def update_version(new_version: str) -> None:
+    """Update version in __init__.py and pyproject.toml."""
     init_file = Path("src/infragpt/__init__.py")
     with open(init_file, "r") as f:
         content = f.read()
@@ -49,8 +50,8 @@ def update_version(new_version):
     print(f"Updated version to {new_version}")
 
 
-def commit_and_tag(version):
-    # Commit changes
+def commit_and_tag(version: str) -> None:
+    """Commit version bump and create a git tag."""
     subprocess.run(
         ["git", "add", "src/infragpt/__init__.py", "pyproject.toml"], check=True
     )
@@ -67,7 +68,8 @@ def commit_and_tag(version):
     print(f"  git push origin master && git push origin {tag_name}")
 
 
-def bump_version(part="patch"):
+def bump_version(part: str = "patch") -> str:
+    """Calculate the new version based on the bump type."""
     current = get_current_version()
     major, minor, patch = map(int, current.split("."))
 
@@ -83,7 +85,8 @@ def bump_version(part="patch"):
     return new_version
 
 
-def main():
+def main() -> None:
+    """Main entry point for the version bump CLI."""
     parser = argparse.ArgumentParser(description="Bump InfraGPT version")
     parser.add_argument(
         "part",

cli/src/infragpt/config.py

@@ -1,6 +1,7 @@
 import os
 import yaml
 import pathlib
+from typing import Any, Dict
 
 from rich.console import Console
 
@@ -17,7 +18,23 @@
 CONFIG_FILE = CONFIG_DIR / "config.yaml"
 
 
-def load_config():
+def is_dev_mode() -> bool:
+    return os.environ.get("INFRAGPT_DEV_MODE", "").lower() == "true"
+
+
+def get_api_base_url() -> str:
+    if is_dev_mode():
+        return "http://localhost:8080"
+    return "https://api.infragpt.io"
+
+
+def get_console_base_url() -> str:
+    if is_dev_mode():
+        return "http://localhost:5173"
+    return "https://app.infragpt.io"
+
+
+def load_config() -> Dict[str, Any]:
     """Load configuration from config file."""
     if not CONFIG_FILE.exists():
         return {}
@@ -30,7 +47,7 @@ def load_config():
         return {}
 
 
-def save_config(config):
+def save_config(config: Dict[str, Any]) -> None:
     """Save configuration to config file."""
     CONFIG_DIR.mkdir(parents=True, exist_ok=True)
 
@@ -41,7 +58,7 @@ def save_config(config):
         console.print(f"[yellow]Warning:[/yellow] Could not save config: {e}")
 
 
-def init_config():
+def init_config() -> None:
     """Initialize configuration file with environment variables if it doesn't exist."""
     if CONFIG_FILE.exists():
         return
@@ -52,24 +69,17 @@ def init_config():
 
     init_history_dir()
 
-    config = {}
-
-    from infragpt.llm import validate_env_api_keys
+    config: Dict[str, Any] = {}
 
     openai_key = os.getenv("OPENAI_API_KEY")
     anthropic_key = os.getenv("ANTHROPIC_API_KEY")
     env_model = os.getenv("INFRAGPT_MODEL")
 
-    model, api_key = validate_env_api_keys()
-
-    if model and api_key:
-        config["model"] = model
-        config["api_key"] = api_key
-    elif anthropic_key and (not env_model or env_model == "claude"):
-        config["model"] = "claude"
+    if anthropic_key and (not env_model or env_model == "claude"):
+        config["model"] = "anthropic:claude-sonnet-4-20250514"
         config["api_key"] = anthropic_key
     elif openai_key and (not env_model or env_model == "gpt4o"):
-        config["model"] = "gpt4o"
+        config["model"] = "openai:gpt-4o"
         config["api_key"] = openai_key
 
     if config: