""" Configuration module for AI Shell """ import os import shutil from pathlib import Path from typing import Dict, List, Optional try: from dotenv import load_dotenv except ImportError: load_dotenv = None try: import tomllib # Python 3.11+ except ImportError: try: import tomli as tomllib # fallback for older Python except ImportError: tomllib = None # Load .env file if it exists def load_env_file() -> None: """Load environment variables from .env file""" if load_dotenv is None: return # Try to find .env file in current directory or package directory env_paths = [ Path.cwd() / ".env", Path(__file__).parent.parent / ".env", Path.home() / ".ai-shell" / ".env", ] for env_path in env_paths: if env_path.exists(): load_dotenv(env_path) break # Load .env file on import load_env_file() # Default API configuration (fallback values) DEFAULT_API_KEY = "your_api_key_here" DEFAULT_BASE_URL = "https://api.openai.com/v1/" DEFAULT_MODEL = "gpt-3.5-turbo" def get_api_key() -> str: """Get API key from environment or use default""" api_key = os.getenv("AI_SHELL_API_KEY", DEFAULT_API_KEY) if api_key == DEFAULT_API_KEY: raise ValueError( "API key not configured. Please set AI_SHELL_API_KEY in .env file or environment variable." ) return api_key def get_base_url() -> str: """Get base URL from environment or use default""" return os.getenv("AI_SHELL_BASE_URL", DEFAULT_BASE_URL) def get_model() -> str: """Get model name from environment or use default""" return os.getenv("AI_SHELL_MODEL", DEFAULT_MODEL) def get_timeout() -> int: """Get request timeout from environment""" return int(os.getenv("AI_SHELL_TIMEOUT", "30")) def get_max_retries() -> int: """Get max retries from environment""" return int(os.getenv("AI_SHELL_MAX_RETRIES", "3")) def load_modern_commands_config() -> Dict[str, any]: """Load modern commands configuration from TOML file""" config_file = Path(__file__).parent / "modern_commands.toml" if not config_file.exists(): return {"commands": {}, "descriptions": {}, "categories": {}, "installation_hints": {}} if tomllib is None: # Fallback to basic parsing if tomllib not available return {"commands": {}, "descriptions": {}, "categories": {}, "installation_hints": {}} try: with open(config_file, "rb") as f: return tomllib.load(f) except Exception: return {"commands": {}, "descriptions": {}, "categories": {}, "installation_hints": {}} def get_modern_commands() -> Dict[str, str]: """Get modern command alternatives configuration""" # Load from TOML config file config = load_modern_commands_config() default_alternatives = config.get("commands", {}) # Try to load custom alternatives from environment custom_alternatives_str = os.getenv("AI_SHELL_MODERN_COMMANDS", "") custom_alternatives = {} if custom_alternatives_str: try: # Parse format: "old1:new1,old2:new2,old3:new3" pairs = custom_alternatives_str.split(",") for pair in pairs: if ":" in pair: old_cmd, new_cmd = pair.strip().split(":", 1) custom_alternatives[old_cmd.strip()] = new_cmd.strip() except Exception: pass # Ignore parsing errors # Merge default and custom alternatives alternatives = default_alternatives.copy() alternatives.update(custom_alternatives) return alternatives def get_available_modern_commands() -> Dict[str, str]: """Get only the modern commands that are actually installed on the system""" all_alternatives = get_modern_commands() available_alternatives = {} for old_cmd, new_cmd in all_alternatives.items(): # Check if the modern command is actually available if shutil.which(new_cmd): available_alternatives[old_cmd] = new_cmd # If modern command not available, keep the original (if it exists) elif shutil.which(old_cmd): available_alternatives[old_cmd] = old_cmd return available_alternatives def get_missing_modern_commands() -> Dict[str, str]: """Get modern commands that are configured but not installed""" all_alternatives = get_modern_commands() missing_commands = {} for old_cmd, new_cmd in all_alternatives.items(): # If modern command is not available but old command exists if not shutil.which(new_cmd) and shutil.which(old_cmd) and old_cmd != new_cmd: missing_commands[old_cmd] = new_cmd return missing_commands def get_installation_hint(command: str) -> Optional[str]: """Get installation hint for a modern command""" config = load_modern_commands_config() installation_hints = config.get("installation_hints", {}) return installation_hints.get(command) def get_command_description(command: str) -> Optional[str]: """Get description for a modern command""" config = load_modern_commands_config() descriptions = config.get("descriptions", {}) return descriptions.get(command) def generate_modern_commands_prompt() -> str: """Generate prompt text about available modern command alternatives""" available_commands = get_available_modern_commands() config = load_modern_commands_config() # Only include commands that have modern alternatives available active_alternatives = {k: v for k, v in available_commands.items() if k != v} if not active_alternatives: return "" prompt_parts = [ "\n\n**IMPORTANT: Modern Command Preferences**", "When generating shell commands, prefer these modern alternatives when available:", "" ] # Use categories from config file categories = config.get("categories", {}) descriptions = config.get("descriptions", {}) for category_name, commands in categories.items(): category_commands = [] for cmd in commands: if cmd in active_alternatives: new_cmd = active_alternatives[cmd] desc = descriptions.get(new_cmd, "") if desc: category_commands.append(f" - Use `{new_cmd}` instead of `{cmd}` ({desc})") else: category_commands.append(f" - Use `{new_cmd}` instead of `{cmd}`") if category_commands: # Convert category name to display format display_name = category_name.replace("_", " ").title() prompt_parts.append(f"**{display_name}:**") prompt_parts.extend(category_commands) prompt_parts.append("") # Add remaining commands not in categories categorized_commands = set() for commands in categories.values(): categorized_commands.update(commands) other_commands = [] for old_cmd, new_cmd in active_alternatives.items(): if old_cmd not in categorized_commands: desc = descriptions.get(new_cmd, "") if desc: other_commands.append(f" - Use `{new_cmd}` instead of `{old_cmd}` ({desc})") else: other_commands.append(f" - Use `{new_cmd}` instead of `{old_cmd}`") if other_commands: prompt_parts.append("**Other Tools:**") prompt_parts.extend(other_commands) prompt_parts.append("") prompt_parts.extend([ "These modern tools typically provide:", "- Better performance and speed", "- More intuitive syntax and output", "- Enhanced features and functionality", "- Better error messages and user experience", "", "Always prioritize these modern alternatives when they fulfill the user's requirements." ]) return "\n".join(prompt_parts) def setup_environment() -> None: """Setup environment variables for OpenAI client""" os.environ["OPENAI_API_KEY"] = get_api_key() os.environ["OPENAI_BASE_URL"] = get_base_url() def validate_config() -> bool: """Validate configuration""" try: get_api_key() get_base_url() get_model() return True except ValueError: return False