aizen-ai-cli 2.2.2__py3-none-any.whl

aizen/config.py

@@ -0,0 +1,363 @@
+import getpass
+import json
+import logging
+import os
+import shutil
+import sys
+import threading
+import time
+import urllib.request
+import urllib.error
+import ssl
+from importlib.metadata import PackageNotFoundError
+from importlib.metadata import version as _pkg_version
+
+from dotenv import load_dotenv
+from rich.console import Console
+
+logger = logging.getLogger("aizen")
+
+# ─── Constants ──────────────────────────────────────────────────────────────────
+
+# Read version from installed package metadata (stays in sync with pyproject.toml).
+# Falls back to a hardcoded value only when running from source without installing.
+_FALLBACK_VERSION = "2.2.2"
+try:
+    VERSION = _pkg_version("aizen-ai-cli")
+except PackageNotFoundError:
+    VERSION = _FALLBACK_VERSION
+CONFIG_PATH = os.path.expanduser("~/.aizen_config.json")
+SESSIONS_DIR = os.path.expanduser("~/.aizen_sessions")
+BACKUPS_DIR = os.path.expanduser("~/.aizen_backups")
+DEFAULT_MODEL = "nvidia/nemotron-3-super-120b-a12b:free"
+
+AIZEN_ASCII = r"""[bold magenta]
+    _    _
+   / \  (_)_______ _ __
+  / _ \ | |_  / _ \ '_ \
+ / ___ \| |/ /  __/ | | |
+/_/   \_\_/___\___|_| |_|
+[/bold magenta]
+    [dim]by Irtaza Malik[/dim]
+"""
+
+# Safe commands that auto-execute without confirmation
+SAFE_COMMAND_PREFIXES = [
+    "ls", "cat", "head", "tail", "wc", "file",
+    "git status", "git log", "git diff", "git branch", "git show", "git rev-parse",
+    "pwd", "echo", "which", "type", "tree", "du", "df",
+    "python --version", "python3 --version", "node --version",
+    "npm --version", "pip --version", "pip list", "pip show",
+    "cargo --version", "rustc --version", "go version",
+    "date", "whoami", "uname", "printenv",
+]
+
+# Dangerous patterns that always require confirmation
+DANGEROUS_PATTERNS = [
+    r"\brm\s", r"\bsudo\b", r"\bchmod\b", r"\bchown\b", r"\bmkfs\b",
+    r"\bdd\b", r":\(\)\{", r"\bkill\b", r"\bpkill\b", r"\bshutdown\b",
+    r"\breboot\b", r">\s*/dev/", r"\bcurl\b.*\|\s*(ba)?sh",
+    # Shell injection patterns
+    r"`[^`]+`",           # Backtick command substitution
+    r"\$\([^)]+\)",       # $() command substitution
+    r"\beval\b",          # eval execution
+    r"\bexec\b",          # exec execution
+    r"\bsource\b",        # source execution
+    r"\|\s*(ba)?sh\b",    # Pipe to shell
+    r"\|\s*zsh\b",        # Pipe to zsh
+    r"\|\s*python",       # Pipe to python
+    r"\bwget\b.*\|\s*",   # wget piped to anything
+    r"\bnohup\b",         # Background with nohup
+]
+
+SYSTEM_PROMPT = """\
+You are Aizen, an expert AI coding assistant running in a user's terminal. \
+You help users write, debug, understand, and refactor code with precision and care.
+
+## Your Workflow
+1. **Understand**: Always read relevant files first. Don't guess at file contents or structure.
+2. **Plan**: Briefly explain your approach before making changes.
+3. **Implement**: Make precise, targeted changes. Use `edit_file` for modifying existing files \
+(surgical edits). Use `write_file` only for creating new files.
+4. **Verify**: After making changes, read the modified file or run tests to confirm correctness.
+
+## Guidelines
+- Be concise but thorough in explanations.
+- Use tools iteratively to explore and understand the codebase.
+- Prefer small, focused changes over large rewrites.
+- When modifying existing files, ALWAYS use `edit_file` with the exact `old_content` to replace. \
+Never use `write_file` to modify existing files unless a full rewrite is truly needed.
+- Run tests or linting commands after changes when applicable.
+- If unsure about something, ask the user rather than guessing.
+- Use fenced code blocks with language identifiers when showing code.
+
+## Tool Preferences
+- `edit_file` > `write_file` for modifications (surgical precision)
+- `grep_search` for finding patterns across the codebase
+- `find_files` for locating files by name
+- `list_directory` for understanding project structure
+- `run_command` for running tests, builds, and verification"""
+
+console = Console()
+
+
+# Project-level rules files (checked in order, first found wins)
+_PROJECT_RULES_FILES = [".aizen_rules", ".cursorrules"]
+
+
+def build_system_prompt(config: dict | None = None) -> str:
+    """
+    Build the final system prompt by merging:
+    1. Default SYSTEM_PROMPT
+    2. User override from ~/.aizen_config.json ("system_prompt" key)
+    3. Project-specific rules from .aizen_rules or .cursorrules in CWD
+
+    This allows per-project customization without modifying source code.
+    """
+    parts = []
+
+    # Start with default or config override
+    if config and config.get("system_prompt"):
+        parts.append(config["system_prompt"])
+    else:
+        parts.append(SYSTEM_PROMPT)
+
+    # Append project-specific rules if present
+    for rules_file in _PROJECT_RULES_FILES:
+        if os.path.isfile(rules_file):
+            try:
+                with open(rules_file, encoding="utf-8", errors="ignore") as f:
+                    project_rules = f.read().strip()
+                if project_rules:
+                    parts.append(
+                        f"\n\n## Project-Specific Rules\n"
+                        f"The following rules are defined by the project maintainers "
+                        f"(from {rules_file}):\n\n{project_rules}"
+                    )
+                    console.print(f"  [dim]📋 Loaded project rules from {rules_file}[/dim]")
+                break  # Only use the first rules file found
+            except Exception as e:
+                logger.debug("Failed to load project rules from %s: %s", rules_file, e)
+
+    return "\n".join(parts)
+
+# Global state for active model
+active_model = DEFAULT_MODEL
+
+def set_active_model(model_name: str, save: bool = False):
+    global active_model
+    active_model = model_name
+    if save:
+        try:
+            config = load_config()
+            config["DEFAULT_MODEL"] = model_name
+            with open(CONFIG_PATH, "w") as f:
+                json.dump(config, f, indent=4)
+        except Exception as e:
+            logger.error("Failed to save default model: %s", e)
+
+def get_active_model() -> str:
+    return active_model
+
+# ─── Configuration ──────────────────────────────────────────────────────────────
+
+
+def migrate_legacy_data():
+    """Migrate legacy Aether config/sessions to Aizen."""
+    legacy_config = os.path.expanduser("~/.aether_config.json")
+    if os.path.exists(legacy_config) and not os.path.exists(CONFIG_PATH):
+        try:
+            shutil.copy2(legacy_config, CONFIG_PATH)
+            console.print("[dim]Migrated legacy config to ~/.aizen_config.json[/dim]")
+        except Exception as e:
+            logger.debug(f"Failed to migrate config: {e}")
+
+    legacy_sessions = os.path.expanduser("~/.aether_sessions")
+    if os.path.exists(legacy_sessions) and not os.path.exists(SESSIONS_DIR):
+        try:
+            shutil.copytree(legacy_sessions, SESSIONS_DIR)
+            console.print("[dim]Migrated legacy sessions to ~/.aizen_sessions[/dim]")
+        except Exception as e:
+            logger.debug(f"Failed to migrate sessions: {e}")
+
+def load_config() -> dict:
+    migrate_legacy_data()
+    if os.path.exists(CONFIG_PATH):
+        try:
+            with open(CONFIG_PATH) as f:
+                return json.load(f)
+        except Exception as e:
+            logger.debug("Failed to load config file: %s", e)
+    return {}
+
+
+def get_mcp_servers(config: dict) -> dict:
+    """Returns the configured MCP servers."""
+    return config.get("mcp_servers", {})
+
+
+def save_config(config: dict):
+    try:
+        with open(CONFIG_PATH, 'w') as f:
+            json.dump(config, f, indent=2)
+    except Exception as e:
+        console.print(f"[yellow]⚠️  Could not save config: {e}[/yellow]\n")
+
+
+def get_api_key(config: dict, reset: bool = False) -> str:
+    if reset:
+        config.pop("OPENROUTER_API_KEY", None)
+        save_config(config)
+
+    key = config.get("OPENROUTER_API_KEY")
+    if key:
+        return key
+
+    load_dotenv()
+    env_key = os.getenv("OPENROUTER_API_KEY")
+    if env_key and env_key != "your_api_key_here":
+        return env_key
+
+    console.print(AIZEN_ASCII)
+    console.print("[bold]Welcome to Aizen![/bold]\n")
+    console.print("To get started, enter your OpenRouter API key.")
+    console.print("[dim](Get one free at https://openrouter.ai/keys)[/dim]\n")
+
+    key = getpass.getpass("API Key: ").strip()
+    if not key:
+        console.print("[bold red]Error:[/bold red] API Key cannot be empty.")
+        sys.exit(1)
+
+    config["OPENROUTER_API_KEY"] = key
+    save_config(config)
+    console.print(f"[green]✓ API key saved to {CONFIG_PATH}[/green]\n")
+    return key
+
+
+# ─── Update Checker (Truly Non-Blocking) ────────────────────────────────────────
+
+# Cache TTL: only check PyPI once every 24 hours
+_UPDATE_CHECK_INTERVAL = 86400  # 24 hours in seconds
+
+
+def _should_check_updates(config: dict) -> bool:
+    """Determine if enough time has passed since the last update check."""
+    last_check = config.get("_last_update_check", 0)
+    return (time.time() - last_check) > _UPDATE_CHECK_INTERVAL
+
+
+def _do_update_check(config: dict):
+    """
+    Background thread target: fetch latest version from PyPI
+    and print a notice if an update is available.
+    """
+    try:
+        ctx = ssl.create_default_context()
+        try:
+            import certifi
+            ctx.load_verify_locations(cafile=certifi.where())
+        except ImportError:
+            ctx = ssl._create_unverified_context()
+            
+        url = "https://pypi.org/pypi/aizen-ai-cli/json"
+        req = urllib.request.Request(url, headers={"User-Agent": "aizen-ai-cli"})
+        with urllib.request.urlopen(req, timeout=3, context=ctx) as response:
+            data = json.loads(response.read().decode())
+            latest = data["info"]["version"]
+
+            # Update the last-check timestamp
+            config["_last_update_check"] = time.time()
+            config["_latest_version"] = latest
+            try:
+                save_config(config)
+            except Exception as e:
+                logger.debug("Failed to save config after update check: %s", e)
+
+            if latest != VERSION:
+                console.print(
+                    f"\n[bold magenta]🔔 Update available:[/bold magenta] v{VERSION} → v{latest}"
+                )
+                console.print("[dim]Run: pip install -U aizen-ai-cli (or brew upgrade aizen)[/dim]")
+                console.print("[dim]Then restart Aizen to use the new version![/dim]\n")
+    except Exception as e:
+        logger.debug("Update check failed (network/parsing): %s", e)
+
+
+def check_for_updates(config: dict | None = None):
+    """
+    Launch a non-blocking background thread to check for updates.
+    Respects a 24-hour cache to avoid repeated network calls.
+    """
+    if config is None:
+        config = load_config()
+
+    if not _should_check_updates(config):
+        # Check if we have a cached latest version that's newer
+        cached = config.get("_latest_version")
+        if cached and cached != VERSION:
+            console.print(
+                f"\n[bold magenta]🔔 Update available:[/bold magenta] v{VERSION} → v{cached}"
+            )
+            console.print("[dim]Run: pip install -U aizen-ai-cli (or brew upgrade aizen)[/dim]")
+            console.print("[dim]Then restart Aizen to use the new version![/dim]\n")
+        return
+
+    thread = threading.Thread(target=_do_update_check, args=(config,), daemon=True)
+    thread.start()
+
+# ─── OpenRouter Models Cache ────────────────────────────────────────────────────
+
+MODELS_CACHE_PATH = os.path.expanduser("~/.aizen_models.json")
+_MODELS_CACHE_TTL = 86400  # 24 hours
+
+def get_cached_models() -> list[dict]:
+    if os.path.exists(MODELS_CACHE_PATH):
+        try:
+            with open(MODELS_CACHE_PATH, encoding="utf-8") as f:
+                data = json.load(f)
+                if time.time() - data.get("timestamp", 0) < _MODELS_CACHE_TTL:
+                    return data.get("models", [])
+        except Exception as e:
+            logger.debug("Failed to load models cache: %s", e)
+    return []
+
+def _do_fetch_models():
+    try:
+        ctx = ssl.create_default_context()
+        try:
+            import certifi
+            ctx.load_verify_locations(cafile=certifi.where())
+        except ImportError:
+            ctx = ssl._create_unverified_context()
+            
+        req = urllib.request.Request("https://openrouter.ai/api/v1/models")
+        with urllib.request.urlopen(req, timeout=5, context=ctx) as response:
+            data = json.loads(response.read().decode())
+            models = data.get("data", [])
+            simplified_models = []
+            for m in models:
+                simplified_models.append({
+                    "id": m.get("id"),
+                    "name": m.get("name"),
+                    "context_length": m.get("context_length", "Unknown"),
+                    "pricing": m.get("pricing", {})
+                })
+
+            with open(MODELS_CACHE_PATH, "w", encoding="utf-8") as f:
+                json.dump({"timestamp": time.time(), "models": simplified_models}, f)
+    except Exception as e:
+        logger.debug("Failed to fetch OpenRouter models: %s", e)
+
+def fetch_openrouter_models_bg():
+    """Fetches OpenRouter models in the background if the cache is stale."""
+    if os.path.exists(MODELS_CACHE_PATH):
+        try:
+            with open(MODELS_CACHE_PATH, encoding="utf-8") as f:
+                data = json.load(f)
+                if time.time() - data.get("timestamp", 0) < _MODELS_CACHE_TTL:
+                    return
+        except Exception:
+            pass
+
+    thread = threading.Thread(target=_do_fetch_models, daemon=True)
+    thread.start()

aizen/context.py

@@ -0,0 +1,171 @@
+"""
+Context window management for Aizen.
+
+Tracks token usage against model context limits and auto-compacts
+conversations when approaching the boundary.
+"""
+
+import json
+
+# Known context window sizes for popular models (in tokens).
+# Users can override via config.
+MODEL_CONTEXT_WINDOWS: dict[str, int] = {
+    # Anthropic
+    "anthropic/claude-sonnet-4": 200_000,
+    "anthropic/claude-3.5-sonnet": 200_000,
+    "anthropic/claude-3.7-sonnet": 200_000,
+    "anthropic/claude-3-opus": 200_000,
+    "anthropic/claude-3-haiku": 200_000,
+    "anthropic/claude-3.5-haiku": 200_000,
+    "anthropic/claude-4-opus": 200_000,
+    # OpenAI
+    "openai/gpt-4o": 128_000,
+    "openai/gpt-4o-mini": 128_000,
+    "openai/gpt-4-turbo": 128_000,
+    "openai/gpt-4": 8_192,
+    "openai/gpt-4.1": 1_047_576,
+    "openai/gpt-4.1-mini": 1_047_576,
+    "openai/gpt-4.1-nano": 1_047_576,
+    "openai/o1": 200_000,
+    "openai/o1-mini": 128_000,
+    "openai/o3": 200_000,
+    "openai/o3-mini": 200_000,
+    "openai/o4-mini": 200_000,
+    # Google
+    "google/gemini-2.5-pro": 1_048_576,
+    "google/gemini-2.5-flash": 1_048_576,
+    "google/gemini-2.0-flash": 1_048_576,
+    "google/gemini-2.0-flash-001": 1_048_576,
+    "google/gemini-pro-1.5": 1_048_576,
+    # Meta
+    "meta-llama/llama-4-maverick": 1_048_576,
+    "meta-llama/llama-3.3-70b-instruct": 131_072,
+    "meta-llama/llama-3.1-405b-instruct": 131_072,
+    "meta-llama/llama-3.1-70b-instruct": 131_072,
+    "meta-llama/llama-3.1-8b-instruct": 131_072,
+    # Nvidia
+    "nvidia/nemotron-3-super-120b-a12b:free": 32_768,
+    # DeepSeek
+    "deepseek/deepseek-chat-v3": 128_000,
+    "deepseek/deepseek-chat": 64_000,
+    "deepseek/deepseek-coder": 64_000,
+    "deepseek/deepseek-r1": 128_000,
+    # Mistral
+    "mistralai/mistral-large": 128_000,
+    "mistralai/mixtral-8x7b-instruct": 32_768,
+    # Qwen
+    "qwen/qwen-2.5-72b-instruct": 131_072,
+    "qwen/qwen3-235b-a22b": 131_072,
+}
+
+# Default context window when model is unknown
+DEFAULT_CONTEXT_WINDOW = 32_768
+
+# Warn when usage exceeds this fraction of the context window
+WARNING_THRESHOLD = 0.75
+
+# Auto-compact when usage exceeds this fraction
+AUTO_COMPACT_THRESHOLD = 0.85
+
+
+class ContextManager:
+    """Tracks token usage against model context limits."""
+
+    def __init__(self, model: str, custom_limit: int | None = None):
+        self.model = model
+        self._custom_limit = custom_limit
+        self._total_tokens = 0
+        self._warned = False
+
+    @property
+    def context_limit(self) -> int:
+        """Get the context window size for the current model."""
+        if self._custom_limit:
+            return self._custom_limit
+        return MODEL_CONTEXT_WINDOWS.get(self.model, DEFAULT_CONTEXT_WINDOW)
+
+    @property
+    def usage_fraction(self) -> float:
+        """Current usage as a fraction of the context window (0.0 to 1.0+)."""
+        if self.context_limit == 0:
+            return 0.0
+        return self._total_tokens / self.context_limit
+
+    @property
+    def usage_percent(self) -> int:
+        """Current usage as a percentage."""
+        return int(self.usage_fraction * 100)
+
+    def update(self, total_tokens: int) -> None:
+        """Update the tracked token count."""
+        self._total_tokens = total_tokens
+
+    def estimate_messages_tokens(self, messages: list, estimator) -> int:
+        """Estimate total tokens across all messages using the provided estimator function."""
+        total = 0
+        for msg in messages:
+            content = msg.get("content", "") or ""
+            total += estimator(content)
+            # Account for tool calls in the message
+            if msg.get("tool_calls"):
+                total += estimator(json.dumps(msg["tool_calls"]))
+        return total
+
+    def set_model(self, model: str) -> None:
+        """Update the model (resets warning state)."""
+        self.model = model
+        self._warned = False
+
+    def check_and_warn(self) -> str | None:
+        """
+        Check usage against thresholds.
+        Returns a warning message if threshold exceeded, None otherwise.
+        """
+        fraction = self.usage_fraction
+
+        if fraction >= AUTO_COMPACT_THRESHOLD:
+            return (
+                f"⚠️  Context window is {self.usage_percent}% full "
+                f"({self._total_tokens:,}/{self.context_limit:,} tokens). "
+                f"Consider using /compact to free up space."
+            )
+        elif fraction >= WARNING_THRESHOLD and not self._warned:
+            self._warned = True
+            return (
+                f"💡 Context window is {self.usage_percent}% full "
+                f"({self._total_tokens:,}/{self.context_limit:,} tokens). "
+                f"Use /compact if the conversation gets long."
+            )
+        return None
+
+    def needs_auto_compact(self) -> bool:
+        """Returns True if the conversation should be auto-compacted."""
+        return self.usage_fraction >= AUTO_COMPACT_THRESHOLD
+
+    def get_usage_bar(self, width: int = 20) -> str:
+        """
+        Generate a visual usage bar for the footer.
+
+        Example: [████████░░░░░░░░░░░░] 42%
+        """
+        fraction = min(self.usage_fraction, 1.0)
+        filled = int(width * fraction)
+        empty = width - filled
+
+        # Color coding based on usage
+        if fraction >= AUTO_COMPACT_THRESHOLD:
+            bar_char = "█"
+            style = "bold red"
+        elif fraction >= WARNING_THRESHOLD:
+            bar_char = "█"
+            style = "yellow"
+        else:
+            bar_char = "█"
+            style = "green"
+
+        bar = f"[{style}]{bar_char * filled}[/{style}][dim]{'░' * empty}[/dim]"
+        return f"[{bar}] {self.usage_percent}%"
+
+    def get_footer_text(self) -> str:
+        """Get a compact footer string showing context usage."""
+        return f"ctx: {self.get_usage_bar(10)}"

aizen/exceptions.py

@@ -0,0 +1,46 @@
+"""
+Aizen custom exception hierarchy.
+
+All Aizen-specific errors inherit from AizenError, enabling
+typed error handling and user-friendly error messages.
+"""
+
+
+class AizenError(Exception):
+    """Base exception for all Aizen errors."""
+    pass
+
+
+class APIKeyError(AizenError):
+    """Raised when the API key is missing, invalid, or rejected (HTTP 401)."""
+    pass
+
+
+class APIConnectionError(AizenError):
+    """Raised when the API is unreachable or the request times out."""
+    pass
+
+
+class RateLimitError(AizenError):
+    """Raised when the API returns HTTP 429 (rate limited)."""
+    pass
+
+
+class ToolExecutionError(AizenError):
+    """Raised when a tool fails to execute properly."""
+    pass
+
+
+class FileOperationError(AizenError):
+    """Raised when a file read/write/edit operation fails."""
+    pass
+
+
+class SessionCorruptedError(AizenError):
+    """Raised when a session file cannot be loaded or is malformed."""
+    pass
+
+
+class ContextWindowExceededError(AizenError):
+    """Raised when the conversation exceeds the model's context window."""
+    pass

aizen/logging_config.py

@@ -0,0 +1,65 @@
+"""
+Structured logging configuration for Aizen.
+
+Provides a rotating file logger at ~/.aizen_logs/aizen.log plus
+an optional console handler controlled by --verbose.
+"""
+
+import logging
+import os
+from logging.handlers import RotatingFileHandler
+
+# ─── Constants ──────────────────────────────────────────────────────────────────
+
+LOG_DIR = os.path.expanduser("~/.aizen_logs")
+LOG_FILE = os.path.join(LOG_DIR, "aizen.log")
+MAX_LOG_BYTES = 5 * 1024 * 1024  # 5 MB per file
+BACKUP_COUNT = 3  # Keep 3 rotated log files
+
+# Module-level logger used throughout the application
+logger = logging.getLogger("aizen")
+
+
+def setup_logging(verbose: bool = False) -> logging.Logger:
+    """
+    Configure logging for the application.
+
+    - Always logs to ~/.aizen_logs/aizen.log (DEBUG level, rotating).
+    - When verbose=True, also logs DEBUG to stderr.
+    - When verbose=False, only WARNING+ goes to stderr.
+
+    Returns the configured root "aizen" logger.
+    """
+    os.makedirs(LOG_DIR, exist_ok=True)
+
+    logger.setLevel(logging.DEBUG)
+
+    # Clear any existing handlers (e.g. on re-init)
+    logger.handlers.clear()
+
+    # ── File handler (always DEBUG) ──
+    file_fmt = logging.Formatter(
+        fmt="%(asctime)s │ %(levelname)-8s │ %(name)s │ %(message)s",
+        datefmt="%Y-%m-%d %H:%M:%S",
+    )
+    file_handler = RotatingFileHandler(
+        LOG_FILE,
+        maxBytes=MAX_LOG_BYTES,
+        backupCount=BACKUP_COUNT,
+        encoding="utf-8",
+    )
+    file_handler.setLevel(logging.DEBUG)
+    file_handler.setFormatter(file_fmt)
+    logger.addHandler(file_handler)
+
+    # ── Console handler (level depends on --verbose) ──
+    console_fmt = logging.Formatter(
+        fmt="%(levelname)-8s │ %(message)s",
+    )
+    console_handler = logging.StreamHandler()
+    console_handler.setLevel(logging.DEBUG if verbose else logging.WARNING)
+    console_handler.setFormatter(console_fmt)
+    logger.addHandler(console_handler)
+
+    logger.debug("Aizen logging initialized (verbose=%s)", verbose)
+    return logger