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

aizen/mcp.py

@@ -0,0 +1,110 @@
+import asyncio
+import contextlib
+import os
+import threading
+
+from mcp import ClientSession, StdioServerParameters
+from mcp.client.stdio import stdio_client
+
+from .config import console
+from .logging_config import logger
+
+
+class MCPManager:
+    def __init__(self, mcp_servers_config: dict):
+        self.config = mcp_servers_config
+        self.sessions: dict[str, ClientSession] = {}
+        self.exit_stack: contextlib.AsyncExitStack | None = None
+        self._loop: asyncio.AbstractEventLoop | None = None
+        self._thread: threading.Thread | None = None
+        self._ready_event = threading.Event()
+        self.tools_cache: list[dict] = []
+
+    async def start(self):
+        if not self.config:
+            return
+
+        logger.info("Starting MCP servers: %s", list(self.config.keys()))
+
+        try:
+            await self._init_all_servers()
+        except Exception as e:
+            console.print(f"[dim yellow]⚠️  Error initializing MCP servers: {e}[/dim yellow]")
+
+    async def _init_all_servers(self):
+        self.exit_stack = contextlib.AsyncExitStack()
+
+        for name, server_config in self.config.items():
+            command = server_config.get("command")
+            args = server_config.get("args", [])
+            env = server_config.get("env")
+
+            if not command:
+                continue
+
+            try:
+                # Inherit environment, merging any custom env
+                merged_env = os.environ.copy()
+                if env:
+                    merged_env.update(env)
+
+                server_params = StdioServerParameters(command=command, args=args, env=merged_env)
+                stdio_transport = await self.exit_stack.enter_async_context(stdio_client(server_params))
+                read, write = stdio_transport
+
+                session = await self.exit_stack.enter_async_context(ClientSession(read, write))
+                await session.initialize()
+
+                self.sessions[name] = session
+
+                # Fetch tools and format them for OpenAI
+                response = await session.list_tools()
+                for tool in response.tools:
+                    self.tools_cache.append({
+                        "type": "function",
+                        "function": {
+                            # Prefix with mcp_serverName_ to avoid collisions
+                            "name": f"mcp_{name}_{tool.name}",
+                            "description": tool.description or f"MCP tool {tool.name} from {name}",
+                            "parameters": tool.inputSchema,
+                        }
+                    })
+                logger.info("MCP server '%s' connected with %d tools", name, len(response.tools))
+                console.print(f"  [dim green]✓ Connected to MCP server: {name}[/dim green]")
+            except Exception as e:
+                logger.exception("Failed to connect to MCP server '%s'", name)
+                console.print(f"  [dim yellow]⚠️  Failed to connect to MCP server {name}: {e}[/dim yellow]")
+
+    async def stop(self):
+        if self.exit_stack:
+            await self.exit_stack.aclose()
+        self.sessions.clear()
+
+    def get_tools(self) -> list[dict]:
+        return self.tools_cache
+
+    async def call_tool(self, full_tool_name: str, arguments: dict) -> str:
+        logger.debug("MCP call_tool: %s args=%s", full_tool_name, list(arguments.keys()))
+        for server_name in self.sessions:
+            prefix = f"mcp_{server_name}_"
+            if full_tool_name.startswith(prefix):
+                tool_name = full_tool_name[len(prefix):]
+                session = self.sessions[server_name]
+                try:
+                    result = await session.call_tool(tool_name, arguments=arguments)
+
+                    if not result.content:
+                        return "Tool executed successfully but returned no content."
+
+                    output = []
+                    for item in result.content:
+                        if item.type == "text":
+                            output.append(item.text)
+                        else:
+                            output.append(f"[Non-text content: {item.type}]")
+
+                    return "\n".join(output)
+                except Exception as e:
+                    return f"Error executing MCP tool {tool_name} on {server_name}: {e}"
+
+        return f"Error: MCP server for tool '{full_tool_name}' not found."

aizen/plugins.py

@@ -0,0 +1,63 @@
+import importlib.util
+import os
+import sys
+from collections.abc import Callable
+
+from .logging_config import logger
+
+PLUGINS_DIR = os.path.expanduser("~/.aizen/plugins")
+
+class PluginManager:
+    """Manages loading and executing tools from user-provided Python scripts."""
+    def __init__(self):
+        self.plugins = {}
+        self.tools = []
+        self.handlers: dict[str, Callable] = {}
+        self._load_plugins()
+
+    def _load_plugins(self):
+        if not os.path.exists(PLUGINS_DIR):
+            try:
+                os.makedirs(PLUGINS_DIR, exist_ok=True)
+            except Exception as e:
+                logger.debug("Failed to create plugins directory: %s", e)
+            return
+
+        for filename in os.listdir(PLUGINS_DIR):
+            if filename.endswith(".py") and not filename.startswith("_"):
+                name = filename[:-3]
+                path = os.path.join(PLUGINS_DIR, filename)
+                try:
+                    spec = importlib.util.spec_from_file_location(name, path)
+                    if spec and spec.loader:
+                        module = importlib.util.module_from_spec(spec)
+                        # Add to sys.modules so plugins can import each other if needed
+                        sys.modules[f"aizen_plugin_{name}"] = module
+                        spec.loader.exec_module(module)
+
+                        if hasattr(module, "get_tools") and hasattr(module, "execute_tool"):
+                            plugin_tools = module.get_tools()
+                            self.plugins[name] = module
+                            self.tools.extend(plugin_tools)
+                            for t in plugin_tools:
+                                self.handlers[t["function"]["name"]] = module.execute_tool
+                            logger.info("Loaded plugin '%s' with %d tools", name, len(plugin_tools))
+                except Exception as e:
+                    logger.error("Failed to load plugin '%s': %s", filename, e)
+
+    def get_tools(self) -> list[dict]:
+        return self.tools
+
+    def execute_tool(self, tool_call, auto_approve: bool = False) -> str | None:
+        """Executes a plugin tool. Returns None if tool is not handled by plugins."""
+        func_name = tool_call.function.name
+        if func_name in self.handlers:
+            try:
+                return self.handlers[func_name](tool_call, auto_approve)
+            except Exception as e:
+                logger.error("Plugin tool error: %s", e)
+                return f"Error executing plugin tool {func_name}: {e}"
+        return None
+
+# Global instance
+plugin_manager = PluginManager()

aizen/retry.py

@@ -0,0 +1,133 @@
+"""
+Retry logic with exponential backoff + jitter for transient API errors.
+
+Supports both synchronous and asynchronous functions — the decorator
+auto-detects coroutine functions and uses asyncio.sleep accordingly.
+"""
+
+import asyncio
+import functools
+import inspect
+import random
+import time
+
+from rich.text import Text
+
+from .config import console
+
+
+def _compute_delay(backoff_base: float, attempt: int, jitter: bool) -> float:
+    """Calculate retry delay with optional jitter."""
+    delay = backoff_base ** attempt
+    if jitter:
+        delay *= 1.0 + random.uniform(-0.25, 0.25)
+    return delay
+
+
+def _print_retry_message(exception: BaseException, delay: float, attempt: int, max_retries: int) -> None:
+    """Print a formatted retry notice to the console."""
+    retry_msg = Text()
+    retry_msg.append("  ⏳ ", style="yellow")
+    retry_msg.append(f"{type(exception).__name__}. ", style="dim")
+    retry_msg.append(
+        f"Retrying in {delay:.1f}s... ({attempt + 1}/{max_retries})",
+        style="dim italic",
+    )
+    console.print(retry_msg)
+
+
+def _is_retryable_503(e: BaseException) -> bool:
+    """Check if an exception represents a 503 Service Unavailable."""
+    return hasattr(e, "status_code") and e.status_code == 503
+
+
+def retry_with_backoff(
+    max_retries: int = 3,
+    backoff_base: float = 2.0,
+    retryable_exceptions: tuple | None = None,
+    jitter: bool = True,
+):
+    """
+    Decorator that retries a function on transient failures with exponential backoff.
+
+    Automatically detects async functions and uses ``asyncio.sleep`` instead of
+    ``time.sleep`` so that the event loop is never blocked.
+
+    Args:
+        max_retries: Maximum number of retry attempts.
+        backoff_base: Base for exponential backoff (delay = base ** attempt).
+        retryable_exceptions: Tuple of exception types to retry on.
+        jitter: If True, adds random jitter (±25%) to prevent thundering herd.
+    """
+    if retryable_exceptions is None:
+        # Import here to avoid circular imports — these are the standard transient errors
+        from openai import (
+            APIConnectionError as OpenAIConnectionError,
+        )
+        from openai import (
+            APITimeoutError,
+        )
+        from openai import (
+            RateLimitError as OpenAIRateLimitError,
+        )
+        retryable_exceptions = (
+            OpenAIRateLimitError,
+            APITimeoutError,
+            OpenAIConnectionError,
+        )
+
+    def decorator(func):
+        # ── Async wrapper ──
+        if inspect.iscoroutinefunction(func):
+            @functools.wraps(func)
+            async def async_wrapper(*args, **kwargs):
+                last_exception: BaseException = RuntimeError("Retry exhausted")
+                for attempt in range(max_retries + 1):
+                    try:
+                        return await func(*args, **kwargs)
+                    except retryable_exceptions as e:
+                        last_exception = e
+                        if attempt < max_retries:
+                            delay = _compute_delay(backoff_base, attempt, jitter)
+                            _print_retry_message(e, delay, attempt, max_retries)
+                            await asyncio.sleep(delay)
+                    except Exception as e:
+                        if _is_retryable_503(e):
+                            last_exception = e
+                            if attempt < max_retries:
+                                delay = _compute_delay(backoff_base, attempt, jitter)
+                                _print_retry_message(e, delay, attempt, max_retries)
+                                await asyncio.sleep(delay)
+                                continue
+                        raise
+                raise last_exception
+
+            return async_wrapper
+
+        # ── Sync wrapper ──
+        @functools.wraps(func)
+        def sync_wrapper(*args, **kwargs):
+            last_exception: BaseException = RuntimeError("Retry exhausted")
+            for attempt in range(max_retries + 1):
+                try:
+                    return func(*args, **kwargs)
+                except retryable_exceptions as e:
+                    last_exception = e
+                    if attempt < max_retries:
+                        delay = _compute_delay(backoff_base, attempt, jitter)
+                        _print_retry_message(e, delay, attempt, max_retries)
+                        time.sleep(delay)
+                except Exception as e:
+                    if _is_retryable_503(e):
+                        last_exception = e
+                        if attempt < max_retries:
+                            delay = _compute_delay(backoff_base, attempt, jitter)
+                            _print_retry_message(e, delay, attempt, max_retries)
+                            time.sleep(delay)
+                            continue
+                    raise
+            raise last_exception
+
+        return sync_wrapper
+
+    return decorator

aizen/session.py

@@ -0,0 +1,137 @@
+import json
+import os
+import re
+import sqlite3
+from datetime import datetime
+
+from .config import SESSIONS_DIR
+from .logging_config import logger
+from .utils import TokenTracker
+
+# ─── Singleton DB connection ────────────────────────────────────────────────────
+
+_db_connection: sqlite3.Connection | None = None
+_db_path_cached: str | None = None
+
+
+def _get_db() -> sqlite3.Connection:
+    """Return a singleton SQLite connection, creating the schema if needed.
+
+    Automatically reconnects if SESSIONS_DIR has changed (e.g. during testing).
+    """
+    global _db_connection, _db_path_cached
+
+    os.makedirs(SESSIONS_DIR, exist_ok=True)
+    db_path = os.path.join(SESSIONS_DIR, "aizen.db")
+
+    # Reconnect if the path changed (supports monkeypatch in tests)
+    if _db_connection is not None and _db_path_cached == db_path:
+        return _db_connection
+
+    if _db_connection is not None:
+        try:
+            _db_connection.close()
+        except Exception:
+            pass
+
+    _db_connection = sqlite3.connect(db_path)
+    _db_path_cached = db_path
+    _db_connection.execute('''
+        CREATE TABLE IF NOT EXISTS sessions (
+            name TEXT PRIMARY KEY,
+            saved_at TEXT,
+            message_count INTEGER,
+            messages TEXT,
+            input_tokens INTEGER,
+            output_tokens INTEGER
+        )
+    ''')
+    _db_connection.commit()
+    return _db_connection
+
+
+def _migrate_legacy_sessions():
+    """Migrate old .json files into the SQLite DB once."""
+    if not os.path.exists(SESSIONS_DIR):
+        return
+
+    conn = _get_db()
+    migrated_any = False
+    for f in os.listdir(SESSIONS_DIR):
+        if f.endswith(".json"):
+            filepath = os.path.join(SESSIONS_DIR, f)
+            try:
+                with open(filepath) as fh:
+                    data = json.load(fh)
+                    name = data.get("name", f[:-5])
+                    msgs = data.get("messages", [])
+                    saved_at = data.get("saved_at", datetime.now().isoformat())
+
+                    conn.execute(
+                        "INSERT OR IGNORE INTO sessions (name, saved_at, message_count, messages, input_tokens, output_tokens) VALUES (?, ?, ?, ?, ?, ?)",
+                        (name, saved_at, len(msgs), json.dumps(msgs), 0, 0)
+                    )
+                # Mark as migrated
+                os.rename(filepath, filepath + ".migrated")
+                migrated_any = True
+            except Exception as e:
+                logger.debug("Failed to migrate session file %s: %s", filepath, e)
+
+    if migrated_any:
+        conn.commit()
+
+# Run migration on import
+_migrate_legacy_sessions()
+
+def save_session(
+    messages: list, name: str | None = None, token_tracker: TokenTracker | None = None
+) -> str:
+    if not name:
+        name = datetime.now().strftime("session_%Y%m%d_%H%M%S")
+
+    # Sanitize
+    name = re.sub(r"[^\w\-]", "_", name)
+
+    input_toks = token_tracker.input_tokens if token_tracker else 0
+    output_toks = token_tracker.output_tokens if token_tracker else 0
+    saved_at = datetime.now().isoformat()
+
+    conn = _get_db()
+    conn.execute(
+        "REPLACE INTO sessions (name, saved_at, message_count, messages, input_tokens, output_tokens) VALUES (?, ?, ?, ?, ?, ?)",
+        (name, saved_at, len(messages), json.dumps(messages), input_toks, output_toks)
+    )
+    conn.commit()
+    return f"sqlite://{name}"
+
+
+def load_session(name: str) -> list | None:
+    # If the user passed the legacy filename by accident
+    if name.endswith(".json"):
+        name = name[:-5]
+
+    conn = _get_db()
+    cur = conn.execute("SELECT messages FROM sessions WHERE name = ?", (name,))
+    row = cur.fetchone()
+    if row:
+        try:
+            return json.loads(row[0])
+        except json.JSONDecodeError:
+            return None
+    return None
+
+
+def list_sessions() -> list:
+    if not os.path.exists(SESSIONS_DIR):
+        return []
+
+    conn = _get_db()
+    cur = conn.execute("SELECT name, saved_at, message_count FROM sessions ORDER BY saved_at DESC")
+    sessions = []
+    for row in cur.fetchall():
+        sessions.append({
+            "name": row[0],
+            "saved_at": row[1],
+            "messages": row[2]
+        })
+    return sessions