cokit-cli 1.0.9 → 1.2.2

package/templates/repo/.github/skills/common/api_key_helper.py

@@ -0,0 +1,411 @@
+#!/usr/bin/env python3
+"""
+Common API Key Detection Helper for Gemini Skills
+
+Supports both Google AI Studio and Vertex AI endpoints.
+
+API Key Detection Order:
+1. Process environment variable
+2. Project root .env file
+3. ./.copilot/.env
+4. ./.copilot/skills/.env
+5. Skill directory .env file
+
+Vertex AI Configuration:
+- GEMINI_USE_VERTEX: Set to "true" to use Vertex AI endpoint
+- VERTEX_PROJECT_ID: GCP project ID (required for Vertex AI)
+- VERTEX_LOCATION: GCP region (default: us-central1)
+"""
+
+import os
+import re
+import sys
+from pathlib import Path
+from typing import Optional, Dict, Any, List
+
+
+def find_api_key(skill_dir: Optional[Path] = None) -> Optional[str]:
+    """
+    Find GEMINI_API_KEY using 5-step lookup:
+    1. Process environment
+    2. Project root .env
+    3. ./.copilot/.env
+    4. ./.copilot/skills/.env
+    5. Skill directory .env
+
+    Args:
+        skill_dir: Path to skill directory (optional, auto-detected if None)
+
+    Returns:
+        API key string or None if not found
+    """
+    # Step 1: Check process environment
+    api_key = os.getenv('GEMINI_API_KEY')
+    if api_key:
+        print("✓ Using API key from environment variable", file=sys.stderr)
+        return api_key
+
+    # Determine paths
+    if skill_dir is None:
+        skill_dir = Path(__file__).parent.parent
+    project_dir = skill_dir.parent.parent.parent  # 3 levels up from skill dir
+
+    # Step 2: Check project root .env
+    project_env = project_dir / '.env'
+    if project_env.exists():
+        api_key = load_env_file(project_env)
+        if api_key:
+            print(f"✓ Using API key from {project_env}", file=sys.stderr)
+            return api_key
+
+    # Step 3: Check ./.copilot/.env
+    config_env = project_dir / '.copilot' / '.env'
+    if config_env.exists():
+        api_key = load_env_file(config_env)
+        if api_key:
+            print(f"✓ Using API key from {config_env}", file=sys.stderr)
+            return api_key
+
+    # Step 4: Check ./.copilot/skills/.env
+    skills_config_env = project_dir / '.copilot' / 'skills' / '.env'
+    if skills_config_env.exists():
+        api_key = load_env_file(skills_config_env)
+        if api_key:
+            print(f"✓ Using API key from {skills_config_env}", file=sys.stderr)
+            return api_key
+
+    # Step 5: Check skill directory .env
+    skill_env = skill_dir / '.env'
+    if skill_env.exists():
+        api_key = load_env_file(skill_env)
+        if api_key:
+            print(f"✓ Using API key from {skill_env}", file=sys.stderr)
+            return api_key
+
+    return None
+
+
+def load_env_file(env_path: Path) -> Optional[str]:
+    """
+    Load GEMINI_API_KEY from .env file
+
+    Args:
+        env_path: Path to .env file
+
+    Returns:
+        API key or None
+    """
+    try:
+        with open(env_path, 'r') as f:
+            for line in f:
+                line = line.strip()
+                if line.startswith('GEMINI_API_KEY='):
+                    # Extract value, removing quotes if present
+                    value = line.split('=', 1)[1].strip()
+                    value = value.strip('"').strip("'")
+                    if value:
+                        return value
+    except Exception as e:
+        print(f"Warning: Error reading {env_path}: {e}", file=sys.stderr)
+
+    return None
+
+
+def load_env_var(env_path: Path, var_name: str) -> Optional[str]:
+    """
+    Load a specific environment variable from .env file
+
+    Args:
+        env_path: Path to .env file
+        var_name: Name of the environment variable
+
+    Returns:
+        Variable value or None
+    """
+    try:
+        with open(env_path, 'r') as f:
+            for line in f:
+                line = line.strip()
+                if line.startswith(f'{var_name}='):
+                    value = line.split('=', 1)[1].strip()
+                    value = value.strip('"').strip("'")
+                    if value:
+                        return value
+    except Exception as e:
+        print(f"Warning: Error reading {env_path}: {e}", file=sys.stderr)
+
+    return None
+
+
+def find_env_var(var_name: str, skill_dir: Optional[Path] = None) -> Optional[str]:
+    """
+    Find environment variable using 5-step lookup (same as API key)
+
+    Args:
+        var_name: Name of environment variable
+        skill_dir: Path to skill directory (optional)
+
+    Returns:
+        Variable value or None
+    """
+    # Step 1: Check process environment
+    value = os.getenv(var_name)
+    if value:
+        return value
+
+    # Determine paths
+    if skill_dir is None:
+        skill_dir = Path(__file__).parent.parent
+    project_dir = skill_dir.parent.parent.parent
+
+    # Step 2-5: Check .env files in order
+    env_files = [
+        project_dir / '.env',
+        project_dir / '.copilot' / '.env',
+        project_dir / '.copilot' / 'skills' / '.env',
+        skill_dir / '.env'
+    ]
+
+    for env_path in env_files:
+        if env_path.exists():
+            value = load_env_var(env_path, var_name)
+            if value:
+                return value
+
+    return None
+
+
+def find_all_api_keys(skill_dir: Optional[Path] = None) -> List[str]:
+    """
+    Find all Gemini API keys for rotation.
+
+    Searches for:
+    - GEMINI_API_KEY (primary, required)
+    - GEMINI_API_KEY_2, GEMINI_API_KEY_3, ... (additional, optional)
+
+    Uses the same 5-step lookup as find_api_key().
+
+    Args:
+        skill_dir: Path to skill directory (optional, auto-detected if None)
+
+    Returns:
+        List of API keys (may be empty if none found)
+    """
+    keys: List[str] = []
+    seen: set = set()  # Deduplicate keys
+
+    # Determine paths
+    if skill_dir is None:
+        skill_dir = Path(__file__).parent.parent
+    project_dir = skill_dir.parent.parent.parent
+
+    # Collect all .env file paths in priority order
+    env_files = [
+        project_dir / '.env',
+        project_dir / '.copilot' / '.env',
+        project_dir / '.copilot' / 'skills' / '.env',
+        skill_dir / '.env'
+    ]
+
+    def add_key(key: Optional[str]) -> None:
+        """Add key if valid and not duplicate."""
+        if key and key not in seen:
+            seen.add(key)
+            keys.append(key)
+
+    # Step 1: Check process environment for all GEMINI_API_KEY* vars
+    for env_key, value in os.environ.items():
+        if env_key == 'GEMINI_API_KEY' or re.match(r'^GEMINI_API_KEY_\d+$', env_key):
+            add_key(value)
+
+    # Step 2-5: Check .env files
+    for env_path in env_files:
+        if not env_path.exists():
+            continue
+
+        # Load all GEMINI_API_KEY* from this file
+        file_keys = _load_all_api_keys_from_file(env_path)
+        for key in file_keys:
+            add_key(key)
+
+    return keys
+
+
+def _load_all_api_keys_from_file(env_path: Path) -> List[str]:
+    """
+    Load all GEMINI_API_KEY* variables from an .env file.
+
+    Args:
+        env_path: Path to .env file
+
+    Returns:
+        List of API keys found in the file
+    """
+    keys: List[str] = []
+    pattern = re.compile(r'^GEMINI_API_KEY(_\d+)?=(.+)$')
+
+    try:
+        with open(env_path, 'r') as f:
+            for line in f:
+                line = line.strip()
+                if line.startswith('#') or not line:
+                    continue
+
+                match = pattern.match(line)
+                if match:
+                    value = match.group(2).strip()
+                    # Remove quotes if present
+                    value = value.strip('"').strip("'")
+                    if value and value != 'your_api_key_here':
+                        keys.append(value)
+    except Exception as e:
+        print(f"Warning: Error reading {env_path}: {e}", file=sys.stderr)
+
+    return keys
+
+
+def get_key_rotator(skill_dir: Optional[Path] = None, verbose: bool = False):
+    """
+    Get a KeyRotator instance with all available API keys.
+
+    Args:
+        skill_dir: Path to skill directory (optional)
+        verbose: Whether to enable verbose logging
+
+    Returns:
+        KeyRotator instance or None if no keys found
+    """
+    # Import here to avoid circular dependency
+    from api_key_rotator import KeyRotator
+
+    keys = find_all_api_keys(skill_dir)
+    if not keys:
+        return None
+
+    return KeyRotator(keys=keys, verbose=verbose)
+
+
+def get_vertex_config(skill_dir: Optional[Path] = None) -> Dict[str, Any]:
+    """
+    Get Vertex AI configuration from environment variables
+
+    Args:
+        skill_dir: Path to skill directory (optional)
+
+    Returns:
+        Dictionary with Vertex AI configuration:
+        {
+            'use_vertex': bool,
+            'project_id': str or None,
+            'location': str (default: 'us-central1')
+        }
+    """
+    use_vertex_str = find_env_var('GEMINI_USE_VERTEX', skill_dir)
+    use_vertex = use_vertex_str and use_vertex_str.lower() in ('true', '1', 'yes')
+
+    config = {
+        'use_vertex': use_vertex,
+        'project_id': find_env_var('VERTEX_PROJECT_ID', skill_dir) if use_vertex else None,
+        'location': find_env_var('VERTEX_LOCATION', skill_dir) or 'us-central1'
+    }
+
+    return config
+
+
+def get_api_key_or_exit(skill_dir: Optional[Path] = None) -> str:
+    """
+    Get API key or exit with helpful error message
+
+    Args:
+        skill_dir: Path to skill directory (optional, auto-detected if None)
+
+    Returns:
+        API key string
+    """
+    api_key = find_api_key(skill_dir)
+
+    if not api_key:
+        print("\n❌ Error: GEMINI_API_KEY not found!", file=sys.stderr)
+        print("\n📋 Please set your API key using one of these methods (in priority order):", file=sys.stderr)
+
+        if skill_dir is None:
+            skill_dir = Path(__file__).parent.parent
+        project_dir = skill_dir.parent.parent.parent
+
+        print("\n1️⃣  Environment variable (recommended for development):", file=sys.stderr)
+        print("   export GEMINI_API_KEY='your-api-key'", file=sys.stderr)
+
+        print("\n2️⃣  Project root .env file:", file=sys.stderr)
+        print(f"   echo 'GEMINI_API_KEY=your-api-key' > {project_dir}/.env", file=sys.stderr)
+
+        print("\n3️⃣  .copilot/.env file:", file=sys.stderr)
+        print(f"   echo 'GEMINI_API_KEY=your-api-key' > {project_dir}/.copilot/.env", file=sys.stderr)
+
+        print("\n4️⃣  .copilot/skills/.env file (shared across all Gemini skills):", file=sys.stderr)
+        print(f"   echo 'GEMINI_API_KEY=your-api-key' > {project_dir}/.copilot/skills/.env", file=sys.stderr)
+
+        print("\n5️⃣  Skill directory .env file:", file=sys.stderr)
+        print(f"   echo 'GEMINI_API_KEY=your-api-key' > {skill_dir}/.env", file=sys.stderr)
+
+        print("\n🔑 Get your API key at: https://aistudio.google.com/apikey", file=sys.stderr)
+        print("\n💡 Tip: Add .env files to .gitignore to avoid committing API keys", file=sys.stderr)
+        sys.exit(1)
+
+    return api_key
+
+
+def get_client(skill_dir: Optional[Path] = None):
+    """
+    Get appropriate Gemini client (AI Studio or Vertex AI)
+
+    Args:
+        skill_dir: Path to skill directory (optional)
+
+    Returns:
+        genai.Client or vertexai client
+    """
+    vertex_config = get_vertex_config(skill_dir)
+
+    if vertex_config['use_vertex']:
+        # Use Vertex AI
+        import vertexai
+        from vertexai.generative_models import GenerativeModel
+
+        if not vertex_config['project_id']:
+            print("\n❌ Error: VERTEX_PROJECT_ID required when GEMINI_USE_VERTEX=true!", file=sys.stderr)
+            print("\n📋 Set your GCP project ID:", file=sys.stderr)
+            print("   export VERTEX_PROJECT_ID='your-project-id'", file=sys.stderr)
+            print("   Or add to .env file: VERTEX_PROJECT_ID=your-project-id", file=sys.stderr)
+            sys.exit(1)
+
+        print(f"✓ Using Vertex AI endpoint", file=sys.stderr)
+        print(f"  Project: {vertex_config['project_id']}", file=sys.stderr)
+        print(f"  Location: {vertex_config['location']}", file=sys.stderr)
+
+        vertexai.init(
+            project=vertex_config['project_id'],
+            location=vertex_config['location']
+        )
+
+        return {'type': 'vertex', 'config': vertex_config}
+    else:
+        # Use AI Studio
+        from google import genai
+
+        api_key = get_api_key_or_exit(skill_dir)
+        client = genai.Client(api_key=api_key)
+
+        return {'type': 'aistudio', 'client': client}
+
+
+if __name__ == '__main__':
+    # Test the API key detection
+    api_key = get_api_key_or_exit()
+    print(f"✓ Found API key: {api_key[:8]}..." + "*" * (len(api_key) - 8))
+
+    # Test Vertex AI config
+    vertex_config = get_vertex_config()
+    if vertex_config['use_vertex']:
+        print(f"\n✓ Vertex AI enabled:")
+        print(f"  Project: {vertex_config['project_id']}")
+        print(f"  Location: {vertex_config['location']}")

package/templates/repo/.github/skills/common/api_key_rotator.py

@@ -0,0 +1,248 @@
+#!/usr/bin/env python3
+"""
+API Key Rotator for Gemini Skills
+
+Manages multiple API keys with automatic rotation on rate limit errors.
+Supports cooldown tracking per key to avoid hammering rate-limited keys.
+
+Usage:
+    from api_key_rotator import KeyRotator
+
+    rotator = KeyRotator(keys=['key1', 'key2', 'key3'])
+
+    # Get current key
+    key = rotator.get_key()
+
+    # Mark key as rate-limited (triggers rotation)
+    rotator.mark_rate_limited()
+
+    # Get next available key
+    key = rotator.get_key()
+"""
+
+import os
+import sys
+import time
+from pathlib import Path
+from typing import List, Optional, Dict
+
+
+# Default cooldown period in seconds after a key is rate-limited
+DEFAULT_COOLDOWN_SECONDS = 60
+
+
+class KeyRotator:
+    """
+    Manages API key rotation with cooldown tracking.
+
+    Keys are rotated when marked as rate-limited.
+    Rate-limited keys enter a cooldown period before being reused.
+    """
+
+    def __init__(
+        self,
+        keys: List[str],
+        cooldown_seconds: int = DEFAULT_COOLDOWN_SECONDS,
+        verbose: bool = False
+    ):
+        """
+        Initialize the key rotator.
+
+        Args:
+            keys: List of API keys to rotate through
+            cooldown_seconds: Seconds to wait before reusing a rate-limited key
+            verbose: Whether to print rotation events
+        """
+        if not keys:
+            raise ValueError("At least one API key is required")
+
+        self._keys = keys
+        self._cooldown_seconds = cooldown_seconds
+        self._verbose = verbose
+        self._current_index = 0
+
+        # Track cooldown end times per key (key index -> timestamp)
+        self._cooldowns: Dict[int, float] = {}
+
+        if verbose:
+            print(f"✓ KeyRotator initialized with {len(keys)} key(s)", file=sys.stderr)
+
+    @property
+    def key_count(self) -> int:
+        """Number of keys in the rotator."""
+        return len(self._keys)
+
+    @property
+    def current_key_masked(self) -> str:
+        """Current key with masking (first 8 chars + dots)."""
+        key = self._keys[self._current_index]
+        return f"{key[:8]}..." if len(key) > 8 else key
+
+    def _is_on_cooldown(self, index: int) -> bool:
+        """Check if a key is currently on cooldown."""
+        if index not in self._cooldowns:
+            return False
+        now = time.time()
+        if now >= self._cooldowns[index]:
+            # Cleanup expired cooldown entry to prevent memory leak
+            del self._cooldowns[index]
+            return False
+        return True
+
+    def _find_available_key(self) -> Optional[int]:
+        """Find the next available key that's not on cooldown."""
+        # Try all keys starting from current
+        for i in range(len(self._keys)):
+            index = (self._current_index + i) % len(self._keys)
+            if not self._is_on_cooldown(index):
+                return index
+        return None
+
+    def get_key(self) -> Optional[str]:
+        """
+        Get the current available API key.
+
+        Returns:
+            API key string or None if all keys are on cooldown
+        """
+        available = self._find_available_key()
+        if available is None:
+            if self._verbose:
+                print("⚠ All API keys are on cooldown", file=sys.stderr)
+            return None
+
+        if available != self._current_index:
+            self._current_index = available
+            if self._verbose:
+                print(f"→ Switched to key #{available + 1}: {self.current_key_masked}", file=sys.stderr)
+
+        return self._keys[self._current_index]
+
+    def mark_rate_limited(self, error_msg: str = "") -> bool:
+        """
+        Mark current key as rate-limited and rotate to next.
+
+        Args:
+            error_msg: Optional error message for logging
+
+        Returns:
+            True if rotation succeeded, False if all keys exhausted
+        """
+        # Put current key on cooldown
+        cooldown_until = time.time() + self._cooldown_seconds
+        self._cooldowns[self._current_index] = cooldown_until
+
+        if self._verbose:
+            print(
+                f"⚠ Key #{self._current_index + 1} rate-limited, "
+                f"cooldown {self._cooldown_seconds}s",
+                file=sys.stderr
+            )
+            if error_msg:
+                # Truncate long error messages
+                msg = error_msg[:100] + "..." if len(error_msg) > 100 else error_msg
+                print(f"  Error: {msg}", file=sys.stderr)
+
+        # Try to find next available key
+        available = self._find_available_key()
+        if available is None:
+            return False
+
+        self._current_index = available
+        if self._verbose:
+            print(f"→ Rotated to key #{available + 1}: {self.current_key_masked}", file=sys.stderr)
+
+        return True
+
+    def reset_cooldowns(self) -> None:
+        """Clear all cooldowns (useful for testing)."""
+        self._cooldowns.clear()
+        if self._verbose:
+            print("✓ All cooldowns cleared", file=sys.stderr)
+
+    def get_status(self) -> Dict:
+        """
+        Get current rotator status.
+
+        Returns:
+            Dict with keys: total, current_index, on_cooldown, available
+        """
+        now = time.time()
+        on_cooldown = sum(1 for cd in self._cooldowns.values() if now < cd)
+
+        return {
+            'total': len(self._keys),
+            'current_index': self._current_index,
+            'on_cooldown': on_cooldown,
+            'available': len(self._keys) - on_cooldown
+        }
+
+
+def is_rate_limit_error(error: Exception) -> bool:
+    """
+    Check if an exception indicates a rate limit error.
+
+    Detects:
+    - RESOURCE_EXHAUSTED errors
+    - 429 status codes
+    - Quota exceeded messages
+
+    Args:
+        error: The exception to check
+
+    Returns:
+        True if this is a rate limit error
+    """
+    error_str = str(error).lower()
+
+    rate_limit_indicators = [
+        'resource_exhausted',
+        'resourceexhausted',
+        '429',
+        'rate limit',
+        'ratelimit',
+        'quota exceeded',
+        'quota_exceeded',
+        'too many requests',
+        'limit: 0',
+    ]
+
+    return any(indicator in error_str for indicator in rate_limit_indicators)
+
+
+if __name__ == '__main__':
+    # Demo/test the rotator
+    print("=== KeyRotator Demo ===\n")
+
+    # Create rotator with test keys
+    test_keys = [
+        "AIzaTest1_XXXXXXXXXXXXXXXXXXXXXXXX",
+        "AIzaTest2_XXXXXXXXXXXXXXXXXXXXXXXX",
+        "AIzaTest3_XXXXXXXXXXXXXXXXXXXXXXXX",
+    ]
+
+    rotator = KeyRotator(keys=test_keys, cooldown_seconds=5, verbose=True)
+
+    print(f"\n1. Initial key: {rotator.current_key_masked}")
+    print(f"   Status: {rotator.get_status()}")
+
+    print("\n2. Simulating rate limit on key #1...")
+    rotator.mark_rate_limited("RESOURCE_EXHAUSTED: quota exceeded")
+    print(f"   Current key: {rotator.current_key_masked}")
+    print(f"   Status: {rotator.get_status()}")
+
+    print("\n3. Simulating rate limit on key #2...")
+    rotator.mark_rate_limited("429 Too Many Requests")
+    print(f"   Current key: {rotator.current_key_masked}")
+    print(f"   Status: {rotator.get_status()}")
+
+    print("\n4. Simulating rate limit on key #3...")
+    success = rotator.mark_rate_limited("Rate limit exceeded")
+    print(f"   Rotation succeeded: {success}")
+    print(f"   Status: {rotator.get_status()}")
+
+    print("\n5. Waiting for cooldown (5s)...")
+    time.sleep(6)
+    key = rotator.get_key()
+    print(f"   Key after cooldown: {rotator.current_key_masked}")
+    print(f"   Status: {rotator.get_status()}")