ai-shell/ai_shell/main.py

"""
Main entry point for AI Shell
"""

import os
import sys
import argparse
from typing import List, Optional

from .agent import create_agent
from . import __version__

def execute_command(command: str) -> None:
    """Execute a shell command"""
    os.system(command)

def get_user_confirmation(command: str) -> bool:
    """Get user confirmation before executing command"""
    print(f"(AI Answer): {command}")
    response = input("Execute? [Y/n]: ").strip().lower()
    return response in ["", "y", "yes"]

def process_prompt(prompt: str) -> None:
    """Process user prompt and generate shell command"""
    if not prompt.strip():
        print("Error: No prompt provided")
        sys.exit(1)
    
    # Create AI agent
    agent = create_agent()
    
    try:
        # Generate response
        resp = agent.run_sync(prompt)
        answer = resp.output
        
        if answer.success and answer.cmd is not None:
            if get_user_confirmation(answer.cmd):
                execute_command(answer.cmd)
        else:
            print(f"(AI Answer): {answer.failure}")
            print("Command generation failed")
            sys.exit(1)
            
    except Exception as e:
        print(f"Error: {e}")
        sys.exit(1)

def create_parser() -> argparse.ArgumentParser:
    """Create command line argument parser"""
    parser = argparse.ArgumentParser(
        prog="ai",
        description="AI-powered shell command generator using DeepSeek V3",
        epilog="Examples:\n"
               "  ai list files\n"
               "  ai \"show disk usage\"\n"
               "  ai 显示当前目录\n",
        formatter_class=argparse.RawDescriptionHelpFormatter
    )
    
    parser.add_argument(
        "prompt",
        nargs="*",
        help="Natural language description of the command you want"
    )
    
    parser.add_argument(
        "--version",
        action="version",
        version=f"ai-shell {__version__}"
    )
    
    parser.add_argument(
        "--config",
        action="store_true",
        help="Show current configuration"
    )
    
    return parser

def show_config() -> None:
    """Show current configuration"""
    from .config import get_api_key, get_base_url, get_model, get_timeout, get_max_retries, validate_config

    print("AI Shell Configuration:")
    print(f"  Model: {get_model()}")
    print(f"  Base URL: {get_base_url()}")

    try:
        api_key = get_api_key()
        print(f"  API Key: {api_key[:8]}...{api_key[-4:]}")
    except ValueError as e:
        print(f"  API Key: ❌ {e}")

    print(f"  Timeout: {get_timeout()}s")
    print(f"  Max Retries: {get_max_retries()}")

    print(f"\nConfiguration Status: {'✅ Valid' if validate_config() else '❌ Invalid'}")

    print("\nConfiguration Sources (in priority order):")
    print("  1. Environment variables")
    print("  2. .env file in current directory")
    print("  3. .env file in package directory")
    print("  4. ~/.ai-shell/.env file")
    print("  5. Default values")

    print("\nEnvironment Variables:")
    print("  AI_SHELL_API_KEY - API key")
    print("  AI_SHELL_BASE_URL - API base URL")
    print("  AI_SHELL_MODEL - Model name")
    print("  AI_SHELL_TIMEOUT - Request timeout (seconds)")
    print("  AI_SHELL_MAX_RETRIES - Maximum retry attempts")

def main() -> None:
    """Main entry point"""
    parser = create_parser()
    args = parser.parse_args()

    if args.config:
        show_config()
        return

    if not args.prompt:
        parser.print_help()
        sys.exit(1)

    # Validate configuration before processing
    from .config import validate_config
    if not validate_config():
        print("❌ Configuration error: API key not configured.")
        print("Please set AI_SHELL_API_KEY in .env file or environment variable.")
        print("Run 'ai --config' to see current configuration.")
        sys.exit(1)

    # Join all prompt arguments into a single string
    prompt = " ".join(args.prompt)
    process_prompt(prompt)

if __name__ == "__main__":
    main()