aizen-ai-cli 2.4.0__tar.gz → 2.4.1__tar.gz

{aizen_ai_cli-2.4.0 → aizen_ai_cli-2.4.1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: aizen-ai-cli
-Version: 2.4.0
+Version: 2.4.1
 Summary: Aizen AI Agent — A professional-grade AI coding assistant for your terminal.
 Author: Irtaza Malik
 License: MIT

aizen_ai_cli-2.4.1/aizen/agent.py

@@ -0,0 +1,274 @@
+"""
+AgentRunner — Encapsulates the core agent turn loop.
+
+Extracted from main.py to enable:
+- Isolated testing with mocked clients
+- Cleaner separation between CLI plumbing and agent logic
+- Reuse in non-interactive contexts (e.g., scripted pipelines)
+"""
+
+import asyncio
+import json
+import random
+from typing import Any
+
+from rich.live import Live
+from rich.markdown import Markdown
+from rich.spinner import Spinner
+from rich.text import Text
+
+from .config import Theme, console, get_active_model
+from .logging_config import logger
+from .tools import execute_tool
+from .utils import Struct, truncate_output
+
+
+class AgentRunner:
+    """Handles a single conversational turn: stream → parse → execute tools → loop."""
+
+    def __init__(
+        self,
+        client,
+        active_tools: list[dict],
+        context_manager,
+        token_tracker,
+        mcp_manager=None,
+        auto_approve: bool = False,
+        is_auto_mode: bool = False,
+        auto_iteration_count: int = 0,
+        max_auto_iterations: int = 50,
+    ):
+        self.client = client
+        self.active_tools = active_tools
+        self.context_manager = context_manager
+        self.token_tracker = token_tracker
+        self.mcp_manager = mcp_manager
+        self.auto_approve = auto_approve
+        self.is_auto_mode = is_auto_mode
+        self.auto_iteration_count = auto_iteration_count
+        self.max_auto_iterations = max_auto_iterations
+
+    async def run_turn(self, messages: list[dict]) -> None:
+        """
+        Execute a full agent turn: stream the model's response, handle tool calls,
+        and loop until the model produces a final text response (no more tool calls).
+        """
+        while True:
+            if self.is_auto_mode:
+                self.auto_iteration_count += 1
+                if self.auto_iteration_count > self.max_auto_iterations:
+                    console.print(
+                        f"  [{Theme.WARNING}]⚠️  Autonomous mode reached iteration limit "
+                        f"({self.max_auto_iterations}). Exiting auto mode.[/{Theme.WARNING}]"
+                    )
+                    self.is_auto_mode = False
+                    messages.append({
+                        "role": "user",
+                        "content": (
+                            f"You have reached the maximum number of autonomous iterations "
+                            f"({self.max_auto_iterations}). Please provide a brief summary "
+                            f"of what you have accomplished and what remains."
+                        ),
+                    })
+                    self.auto_iteration_count = 0
+
+            # Stream the response
+            stream_result = await self._stream_response(messages)
+            if stream_result is None:
+                break  # Error occurred (already printed to console)
+
+            full_content, tool_calls_list, api_usage = stream_result
+
+            # Track tokens
+            self._track_tokens(messages, full_content, api_usage)
+
+            # Add assistant message to history
+            assistant_msg: dict[str, Any] = {
+                "role": "assistant",
+                "content": full_content or "",
+            }
+            if tool_calls_list:
+                assistant_msg["tool_calls"] = tool_calls_list
+            messages.append(assistant_msg)
+
+            # If no tool calls, we're done with this turn
+            if not tool_calls_list:
+                break
+
+            # Execute tool calls
+            tool_results = await self._execute_tools(tool_calls_list)
+            messages.extend(tool_results)
+
+            # Loop continues — model processes tool results
+
+    async def _stream_response(
+        self, messages: list[dict]
+    ) -> tuple[str, list[dict], Any] | None:
+        """
+        Stream a response from the model.
+
+        Returns (full_content, tool_calls_list, api_usage) or None on error.
+        """
+        full_content = ""
+        accumulated_tool_calls: dict[int, dict] = {}
+        api_usage = None
+
+        spinner_label = random.choice([
+            "Thinking...", "Analyzing...", "Reasoning...",
+            "Processing...", "Considering...", "Exploring...", "Synthesizing...",
+        ])
+
+        if self.is_auto_mode:
+            spinner_text = Text(
+                f" [Step {self.auto_iteration_count}/{self.max_auto_iterations}] {spinner_label}",
+                style=f"{Theme.MUTED} italic",
+            )
+        else:
+            spinner_text = Text(f" {spinner_label}", style=f"{Theme.MUTED} italic")
+
+        spinner_display = Spinner("dots2", text=spinner_text, style=f"{Theme.PRIMARY} bold")
+
+        try:
+            with Live(
+                spinner_display, console=console, refresh_per_second=8
+            ) as live:
+                from openai import AsyncStream
+
+                model = get_active_model()
+                api_params: dict[str, Any] = {
+                    "model": model,
+                    "messages": messages,
+                    "stream": True,
+                    "stream_options": {"include_usage": True},
+                }
+                if self.active_tools:
+                    api_params["tools"] = self.active_tools
+                    api_params["tool_choice"] = "auto"
+
+                stream: AsyncStream = await self.client.chat.completions.create(**api_params)
+
+                async for chunk in stream:
+                    if hasattr(chunk, "usage") and chunk.usage:
+                        api_usage = chunk.usage
+
+                    delta = chunk.choices[0].delta if chunk.choices else None
+                    if not delta:
+                        continue
+
+                    if delta.content:
+                        full_content += delta.content
+                        if full_content.strip():
+                            try:
+                                display_content = f"**◆ AIZEN:** {full_content}"
+                                rendered = Markdown(display_content)
+                                live.update(rendered)
+                            except Exception:
+                                display_text = Text.from_markup(
+                                    f"{Theme.BADGE} {full_content}"
+                                )
+                                live.update(display_text)
+
+                    if delta.tool_calls:
+                        for tc in delta.tool_calls:
+                            idx = tc.index
+                            if idx not in accumulated_tool_calls:
+                                accumulated_tool_calls[idx] = {
+                                    "id": "", "name": "", "arguments": "", "type": "function",
+                                }
+                            if tc.id:
+                                accumulated_tool_calls[idx]["id"] = tc.id
+                            if tc.function:
+                                if tc.function.name:
+                                    accumulated_tool_calls[idx]["name"] += tc.function.name
+                                if tc.function.arguments:
+                                    accumulated_tool_calls[idx]["arguments"] += tc.function.arguments
+
+                        names = [v["name"] for v in accumulated_tool_calls.values() if v["name"]]
+                        if names and not full_content.strip():
+                            tool_text = Text()
+                            tool_text.append("  ◆ ", style=f"bold {Theme.ACCENT}")
+                            tool_text.append("Invoking ", style=f"{Theme.TEXT}")
+                            tool_text.append(f"{', '.join(names)}", style=f"bold {Theme.ACCENT}")
+                            tool_text.append(" ...", style=f"{Theme.MUTED}")
+                            live.update(tool_text)
+
+        except Exception as e:
+            # Re-raise — let the caller (main_loop) handle specific exception types
+            raise
+
+        # Build tool calls list
+        tool_calls_list: list[dict[str, Any]] = []
+        for idx in sorted(accumulated_tool_calls.keys()):
+            tc = accumulated_tool_calls[idx]
+            tool_calls_list.append({
+                "id": tc["id"],
+                "type": "function",
+                "function": {
+                    "name": tc["name"],
+                    "arguments": tc["arguments"],
+                },
+            })
+
+        return full_content, tool_calls_list, api_usage
+
+    async def _execute_tools(self, tool_calls_list: list[dict]) -> list[dict]:
+        """Execute tool calls (in parallel where safe) and return tool result messages."""
+
+        async def _exec_tool(tc_dict: dict) -> dict:
+            func_name = tc_dict["function"]["name"]
+            if func_name.startswith("mcp_") and self.mcp_manager:
+                try:
+                    args = json.loads(tc_dict["function"]["arguments"])
+                    result = await self.mcp_manager.call_tool(func_name, args)
+                except json.JSONDecodeError:
+                    result = f"Error: Invalid JSON arguments for {func_name}."
+            else:
+                func_struct = Struct(**tc_dict["function"])
+                tc_struct = Struct(
+                    id=tc_dict["id"],
+                    type=tc_dict["type"],
+                    function=func_struct,
+                )
+                result = await asyncio.to_thread(execute_tool, tc_struct, self.auto_approve)
+
+            return {
+                "role": "tool",
+                "tool_call_id": tc_dict["id"],
+                "name": func_name,
+                "content": truncate_output(result),
+            }
+
+        tool_results = await asyncio.gather(
+            *[_exec_tool(tc) for tc in tool_calls_list],
+            return_exceptions=True,
+        )
+
+        # Handle individual tool failures gracefully
+        for i, result in enumerate(tool_results):
+            if isinstance(result, Exception):
+                logger.error("Tool execution failed: %s", result)
+                tool_results[i] = {
+                    "role": "tool",
+                    "tool_call_id": tool_calls_list[i]["id"],
+                    "name": tool_calls_list[i]["function"]["name"],
+                    "content": f"Error: Tool execution failed — {type(result).__name__}: {result}",
+                }
+
+        return list(tool_results)
+
+    def _track_tokens(self, messages, full_content, api_usage):
+        """Update token tracking from API usage or estimation."""
+        if api_usage and hasattr(api_usage, "prompt_tokens"):
+            self.token_tracker.add_api_usage(
+                api_usage.prompt_tokens or 0,
+                api_usage.completion_tokens or 0,
+            )
+            self.context_manager.update(
+                (api_usage.prompt_tokens or 0) + (api_usage.completion_tokens or 0)
+            )
+        elif full_content:
+            estimated_input = self.context_manager.estimate_messages_tokens(
+                messages, self.token_tracker.estimate_tokens
+            )
+            estimated_output = self.token_tracker.estimate_tokens(full_content)
+            self.token_tracker.add_usage(estimated_input, estimated_output)

{aizen_ai_cli-2.4.0 → aizen_ai_cli-2.4.1}/aizen/commands.py

@@ -5,8 +5,10 @@ import re
 import subprocess
 from datetime import datetime
 
+import questionary
+
 from prompt_toolkit.completion import Completer, Completion
-from prompt_toolkit.shortcuts import prompt
+from prompt_toolkit import PromptSession
 from rich.table import Table
 
 from .config import (
@@ -421,12 +423,8 @@ async def handle_slash_command(
                 # Attempt LLM-based summarization for much better context retention
                 console.print(f"  [{Theme.MUTED}]Summarizing conversation with AI...[/{Theme.MUTED}]")
                 try:
-                    from openai import AsyncOpenAI as _AsyncOpenAI
-
-                    _config = load_config()
-                    _api_key = _config.get("OPENROUTER_API_KEY", "")
-                    _api_base = _config.get("API_BASE_URL", "https://openrouter.ai/api/v1")
-                    _client = _AsyncOpenAI(base_url=_api_base, api_key=_api_key)
+                    # Use the client passed to handle_slash_command
+                    _client = client
 
                     # Build a summarization request from the middle messages
                     summary_messages = [
@@ -606,8 +604,8 @@ async def handle_slash_command(
                     console.print(f"  [{Theme.WARNING}]No changes found to commit.[/{Theme.WARNING}]\n")
                     return False
 
-                answer = prompt("No staged changes. Stage all current changes? [Y/n] ")
-                if answer.lower() not in ("y", "yes", ""):
+                answer = await questionary.confirm("No staged changes. Stage all current changes?").ask_async()
+                if not answer:
                     console.print(f"  [{Theme.WARNING}]Commit aborted.[/{Theme.WARNING}]\n")
                     return False
 
@@ -631,24 +629,40 @@ async def handle_slash_command(
                 messages=commit_messages,
                 max_tokens=200,
             )
-            commit_msg = response.choices[0].message.content.strip()
+            commit_content = response.choices[0].message.content
+            commit_msg = commit_content.strip() if commit_content else ""
             # Remove any markdown codeblocks if model didn't listen
             commit_msg = commit_msg.replace("```text", "").replace("```", "").strip()
 
-            console.print(f"\n  [bold {Theme.TEXT}]Generated Commit Message:[/bold {Theme.TEXT}]")
-            console.print(f"  [{Theme.ACCENT}]{commit_msg}[/{Theme.ACCENT}]\n")
-
-            action = prompt("Commit with this message? [Y/n/e(dit)] ")
-            action = action.lower().strip()
+            if not commit_msg:
+                console.print(f"\n  [{Theme.WARNING}]⚠️ The model failed to generate a commit message.[/{Theme.WARNING}]")
+                action = "Edit message"
+            else:
+                console.print(f"\n  [bold {Theme.TEXT}]Generated Commit Message:[/bold {Theme.TEXT}]")
+                console.print(f"  [{Theme.ACCENT}]{commit_msg}[/{Theme.ACCENT}]\n")
+
+                action = await questionary.select(
+                    "Commit with this message?",
+                    choices=[
+                        "Yes, commit this",
+                        "Edit message",
+                        "Cancel"
+                    ]
+                ).ask_async()
 
-            if action in ("y", "yes", ""):
+            if action == "Yes, commit this":
                 final_msg = commit_msg
-            elif action in ("e", "edit"):
-                final_msg = prompt("Edit message: ", default=commit_msg)
+            elif action == "Edit message":
+                final_msg = await questionary.text("Edit message:", default=commit_msg).ask_async()
             else:
                 console.print("[yellow]Commit aborted.[/yellow]\n")
                 return False
 
+            final_msg = final_msg.strip()
+            if not final_msg:
+                console.print(f"  [{Theme.ERROR}]Error: Commit message cannot be empty. Aborted.[/{Theme.ERROR}]\n")
+                return False
+
             subprocess.run(["git", "commit", "-m", final_msg], check=True)
             console.print(f"  [{Theme.SUCCESS}]✓ Committed successfully.[/{Theme.SUCCESS}]\n")
 

{aizen_ai_cli-2.4.0 → aizen_ai_cli-2.4.1}/aizen/config.py

@@ -20,7 +20,7 @@ logger = logging.getLogger("aizen")
 
 # 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.4.0"
+_FALLBACK_VERSION = "2.4.1"
 try:
     VERSION = _pkg_version("aizen-ai-cli")
 except PackageNotFoundError:
@@ -169,12 +169,14 @@ def build_system_prompt(config: dict | None = None) -> str:
 
     return "\n".join(parts)
 
-# Global state for active model
+# Global state for active model (protected by lock for thread safety)
 active_model = DEFAULT_MODEL
+_model_lock = threading.Lock()
 
 def set_active_model(model_name: str, save: bool = False):
     global active_model
-    active_model = model_name
+    with _model_lock:
+        active_model = model_name
     if save:
         try:
             config = load_config()
@@ -185,7 +187,8 @@ def set_active_model(model_name: str, save: bool = False):
             logger.error("Failed to save default model: %s", e)
 
 def get_active_model() -> str:
-    return active_model
+    with _model_lock:
+        return active_model
 
 # ─── Configuration ──────────────────────────────────────────────────────────────
 

{aizen_ai_cli-2.4.0 → aizen_ai_cli-2.4.1}/aizen/context.py

@@ -171,3 +171,64 @@ class ContextManager:
     def get_footer_text(self) -> str:
         """Get a compact footer string showing context usage."""
         return f"[{Theme.MUTED}]ctx:[/{Theme.MUTED}] {self.get_usage_bar(10)}"
+
+
+class ContextPruner:
+    """Handles smart pruning and summarization of old conversation context."""
+
+    @staticmethod
+    def prune_attached_contexts(messages: list[dict]) -> int:
+        """
+        Removes <file_context>, <url_context>, etc. blocks from older user messages.
+        Returns the number of messages modified.
+        """
+        import re
+        dropped_count = 0
+        
+        # Keep the system prompt and the last couple of turns intact
+        if len(messages) <= 3:
+            return 0
+
+        for msg in messages[1:-2]:
+            if msg.get("role") == "user" and msg.get("content"):
+                old_content = msg["content"]
+                new_content = re.sub(r'<file_context path="[^"]+">.*?</file_context>', '[File context dropped]', old_content, flags=re.DOTALL)
+                new_content = re.sub(r'<url_context url="[^"]+">.*?</url_context>', '[URL context dropped]', new_content, flags=re.DOTALL)
+                new_content = re.sub(r'<directory_context path="[^"]+">.*?</directory_context>', '[Directory context dropped]', new_content, flags=re.DOTALL)
+                new_content = re.sub(r'<command_context cmd="[^"]+">.*?</command_context>', '[Command context dropped]', new_content, flags=re.DOTALL)
+
+                if old_content != new_content:
+                    msg["content"] = new_content
+                    dropped_count += 1
+                    
+        return dropped_count
+
+    @staticmethod
+    def summarize_old_messages(messages: list[dict], recent_count: int = 4) -> list[dict]:
+        """
+        Condenses older messages into a naive summary to save tokens.
+        Modifies the `messages` list in place and returns the summary message text.
+        """
+        if len(messages) <= recent_count + 2:
+            return ""
+
+        system_msg = messages[0]
+        recent = messages[-recent_count:]
+        middle = messages[1:-recent_count]
+        
+        user_topics = [
+            m["content"][:100].replace('\n', ' ') 
+            for m in middle 
+            if m.get("role") == "user" and m.get("content")
+        ]
+        
+        summary = "Previous conversation summary: The user and assistant discussed " + "; ".join(user_topics[:5]) + ". The assistant helped with these requests."
+        
+        messages[:] = [
+            system_msg,
+            {"role": "user", "content": f"Previous conversation summary:\n{summary}"},
+            {"role": "assistant", "content": "Understood. I have the context. How can I continue helping?"},
+        ] + recent
+        
+        return summary
+