axons 0.1.0__py3-none-any.whl

axon/__init__.py

@@ -0,0 +1,6 @@
+"""
+Axon - AI-powered coding agent.
+CLI tool and MCP server for AI-assisted code generation and editing.
+"""
+
+__version__ = "0.1.0"

axon/agent.py

@@ -0,0 +1,436 @@
+"""Agent components — memory manager, tool executor, verification, and orchestration loop."""
+
+import asyncio
+import json
+from collections import Counter
+from datetime import datetime
+from pathlib import Path
+from typing import Optional
+
+from rich.console import Console
+from rich.markdown import Markdown
+from rich.panel import Panel
+from rich.syntax import Syntax
+from rich import box
+
+from axon.config import AxonConfig
+from axon.llm import LLMMessage, ToolCall, create_provider, LLMProvider
+from axon.tools import TOOL_DEFINITIONS, PLANNING_TOOL_DEFINITIONS, execute_tool, ToolResult
+
+console = Console()
+
+CONTEXT_WINDOW_KEEP = 20
+MAX_IDENTICAL_RESULTS = 3
+# Detect alternating-loop patterns over this many recent calls
+LOOP_WINDOW = 8
+LOOP_REPEAT_THRESHOLD = 3  # if any single signature appears this many times in the window
+
+# Approval levels
+MODE_AUTO = "auto"    # auto-run all tools
+MODE_NORMAL = "normal"  # ask before write_file, edit_file, apply_patch
+MODE_SAFE = "safe"      # ask before write, edit, run_command, git_commit
+
+# Which tools need approval per mode
+RISKY_TOOLS = {"write_file", "apply_patch"}
+VERY_RISKY_TOOLS = {"write_file", "apply_patch", "run_command", "git_commit"}
+
+
+# ─── Memory Manager ──────────────────────────────────────────────────────────────
+
+class ConversationMemory:
+    """Manages conversation history with safe compression.
+
+    The first user message (the original task) is pinned and survives all
+    compression so the model never forgets what it was asked to do.
+    """
+
+    def __init__(self, system_prompt: str, persist: bool = False):
+        self.messages: list[LLMMessage] = []
+        self.persist = persist
+        self._pinned_task: Optional[LLMMessage] = None  # #1: pinned original task
+        self.messages.append(LLMMessage(role="system", content=system_prompt))
+
+    def add_user(self, content: str):
+        msg = LLMMessage(role="user", content=content)
+        # Pin the very first user message as the original task
+        if self._pinned_task is None:
+            self._pinned_task = msg
+        self.messages.append(msg)
+
+    def add_assistant(self, content: str, tool_calls: Optional[list[ToolCall]] = None):
+        self.messages.append(LLMMessage(role="assistant", content=content, tool_calls=tool_calls or []))
+
+    def add_tool_result(self, tc: ToolCall, result: ToolResult):
+        self.messages.append(
+            LLMMessage(role="tool", content=result.to_str(), tool_call_id=tc.id)
+        )
+
+    def compress(self):
+        """Trim old messages while preserving valid role structure.
+
+        Rules:
+        - Always keep messages[0] (system prompt).
+        - Always re-inject the pinned task message after trimming.
+        - Keep the most recent CONTEXT_WINDOW_KEEP non-system messages.
+        - After trimming, drop orphaned tool/assistant messages at the boundary.
+        - Prepend a trim notice only when the kept slice doesn't start with a user msg.
+        """
+        non_system = self.messages[1:]
+        if len(non_system) <= CONTEXT_WINDOW_KEEP:
+            return
+
+        kept = non_system[-(CONTEXT_WINDOW_KEEP):]
+
+        # Drop orphaned tool results at the front (no paired assistant above them)
+        while kept and kept[0].role == "tool":
+            kept = kept[1:]
+
+        # Drop an assistant-with-tool-calls at the front if its results were trimmed
+        if kept and kept[0].role == "assistant" and kept[0].tool_calls:
+            call_ids = {tc.id for tc in kept[0].tool_calls}
+            result_ids = {m.tool_call_id for m in kept if m.role == "tool"}
+            if not call_ids.issubset(result_ids):
+                kept = kept[1:]
+                while kept and kept[0].role == "tool":
+                    kept = kept[1:]
+
+        if not kept:
+            return  # nothing safe to keep — don't corrupt history
+
+        # #4 FIX: only insert trim notice when the first kept msg is NOT already a user msg
+        # (previously both branches did the same insert, causing double user messages)
+        if kept[0].role != "user":
+            trim_notice = LLMMessage(
+                role="user",
+                content="[Earlier conversation was trimmed. Continue working on the current task.]"
+            )
+            kept.insert(0, trim_notice)
+
+        # #1: Re-inject pinned task at the front if it's not already there
+        if self._pinned_task is not None and kept[0] is not self._pinned_task:
+            # Avoid duplicating if the pinned task is already in kept
+            if self._pinned_task not in kept:
+                kept.insert(0, self._pinned_task)
+
+        self.messages = [self.messages[0]] + kept
+
+    def get_messages(self) -> list[LLMMessage]:
+        return self.messages
+
+
+# ─── Verification Manager ──────────────────────────────────────────────────────
+
+class VerificationManager:
+    """Tracks verification checkpoints and ensures changes are verified."""
+
+    def __init__(self):
+        self.changes_made: list[dict] = []  # list of {tool, path, time}
+        self.verified = False
+
+    def record_change(self, tool_name: str, path: str):
+        self.changes_made.append({
+            "tool": tool_name,
+            "path": path,
+            "time": datetime.now().isoformat(),
+        })
+        self.verified = False
+
+    def mark_verified(self):
+        self.verified = True
+
+    def has_unverified_changes(self) -> bool:
+        return len(self.changes_made) > 0 and not self.verified
+
+    def summary(self) -> str:
+        if not self.changes_made:
+            return ""
+        lines = [f"[Changes made this session ({len(self.changes_made)}):]"]
+        for c in self.changes_made:
+            lines.append(f"  - {c['tool']}: {c['path']}")
+        if self.verified:
+            lines.append("  ✓ Verified")
+        else:
+            lines.append("  ⚠ Not yet verified")
+        return "\n".join(lines)
+
+
+# ─── Tool Executor ─────────────────────────────────────────────────────────────
+
+class ToolExecutor:
+    """Handles tool call execution and result tracking."""
+
+    def __init__(self, workspace: Path, approval_mode: str = MODE_AUTO):
+        self.workspace = workspace
+        self.last_signatures: list[str] = []
+        self.approval_mode = approval_mode
+
+    def set_approval_mode(self, mode: str):
+        self.approval_mode = mode
+
+    def needs_approval(self, tool_name: str) -> bool:
+        if self.approval_mode == MODE_AUTO:
+            return False
+        if self.approval_mode == MODE_SAFE:
+            return tool_name in VERY_RISKY_TOOLS
+        if self.approval_mode == MODE_NORMAL:
+            return tool_name in RISKY_TOOLS
+        return False
+
+    async def execute(self, tc: ToolCall) -> ToolResult:
+        # Ask for approval if needed
+        if self.needs_approval(tc.name):
+            args_str = json.dumps(tc.arguments, indent=2)[:500]
+            print(f"\n[Axon needs approval]")
+            print(f"Tool: {tc.name}")
+            print(f"Args:\n{args_str}")
+            try:
+                loop = asyncio.get_event_loop()
+                answer = await loop.run_in_executor(None, input, "Run this? (Y/n): ")
+                if answer.strip().lower() == "n":
+                    return ToolResult(
+                        success=False,
+                        error="Skipped — user declined approval.",
+                        tool_name=tc.name,
+                    )
+            except (EOFError, KeyboardInterrupt):
+                return ToolResult(
+                    success=False,
+                    error="Skipped — interrupted.",
+                    tool_name=tc.name,
+                )
+
+        result = await execute_tool(tc.name, tc.arguments, self.workspace)
+        sig = f"{tc.name}:{json.dumps(tc.arguments, sort_keys=True)}"
+        self.last_signatures.append(sig)
+        return result
+
+    def is_looping(self) -> bool:
+        """Detect both identical-repeat and alternating-loop patterns."""
+        if len(self.last_signatures) < MAX_IDENTICAL_RESULTS:
+            return False
+
+        # Classic: last N calls all identical
+        recent = self.last_signatures[-MAX_IDENTICAL_RESULTS:]
+        if len(set(recent)) == 1:
+            return True
+
+        # Alternating: within the last LOOP_WINDOW calls, any signature
+        # appears LOOP_REPEAT_THRESHOLD or more times
+        if len(self.last_signatures) >= LOOP_WINDOW:
+            window = self.last_signatures[-LOOP_WINDOW:]
+            counts = Counter(window)
+            if counts.most_common(1)[0][1] >= LOOP_REPEAT_THRESHOLD:
+                return True
+
+        return False
+
+    def reset_loop_detection(self):
+        self.last_signatures = []
+
+
+# ─── Orchestrator ──────────────────────────────────────────────────────────────
+
+class Agent:
+    """High-level agent that orchestrates the AI coding loop."""
+
+    def __init__(
+        self,
+        config: AxonConfig,
+        workspace: Path,
+        task: str,
+        interactive: bool = True,
+        persist: bool = False,
+        approval_mode: str = MODE_AUTO,
+    ):
+        self.config = config
+        self.workspace = workspace
+        self.interactive = interactive
+        self.approval_mode = approval_mode
+        self.memory = ConversationMemory(config.system_prompt, persist=persist)
+        self.executor = ToolExecutor(workspace, approval_mode)
+        self.verifier = VerificationManager()
+        self.provider: LLMProvider = create_provider(config.llm)
+        self.persist = persist
+        self.active_plan: Optional[str] = None
+
+        self.memory.add_user(task)
+
+    def set_approval_mode(self, mode: str):
+        self.approval_mode = mode
+        self.executor.set_approval_mode(mode)
+
+    def add_task(self, new_task: str):
+        """Add a follow-up task to the existing conversation (reuse provider)."""
+        self.memory.add_user(new_task)
+
+    def set_plan(self, plan: str):
+        """Set a structured plan that the agent can reference."""
+        self.active_plan = plan
+        self.memory.add_user(
+            f"[PLAN FOR THIS TASK]\n{plan}\n\n"
+            "Follow this plan step by step. Update progress after each step."
+        )
+
+    async def _run_planning_phase(self) -> str:
+        """#5: Dedicated planning phase — one call with read-only tools, returns the plan.
+
+        The model is given only safe/read tools so it can't accidentally start
+        making changes while planning. It must respond with a numbered plan.
+        """
+        console.print("[dim]Planning...[/]")
+
+        # Inject a planning-specific instruction
+        planning_prompt = LLMMessage(
+            role="user",
+            content=(
+                "[PLANNING PHASE]\n"
+                "Before doing anything, produce a numbered step-by-step plan for this task.\n"
+                "- Use the available tools to explore the codebase if needed.\n"
+                "- Your response MUST end with a section starting with '## Plan' listing numbered steps.\n"
+                "- Do NOT make any file edits or run any commands yet — this is planning only.\n"
+                "- Keep steps concrete: specify which files and what changes."
+            ),
+        )
+
+        planning_messages = self.memory.get_messages() + [planning_prompt]
+
+        content, _ = await self.provider.chat(planning_messages, PLANNING_TOOL_DEFINITIONS)
+
+        if content:
+            _print_assistant(content, title="Axon · Planning")
+
+        return content or ""
+
+    async def run(self) -> str:
+        iteration = 0
+        final_result = ""
+
+        # #5: Structural planning phase (unless plan was manually set via /plan)
+        if self.active_plan is None:
+            plan_text = await self._run_planning_phase()
+            if plan_text:
+                self.active_plan = plan_text
+                self.memory.add_assistant(plan_text)
+                self.memory.add_user(
+                    "[Planning complete. Now execute the plan step by step. "
+                    "After each step update progress. Verify changes after edits. "
+                    "Call finish_task when done.]"
+                )
+
+        while iteration < self.config.max_iterations:
+            iteration += 1
+
+            content, tool_calls = await self.provider.chat(
+                self.memory.get_messages(), TOOL_DEFINITIONS
+            )
+
+            # #8: Buffer all content — print once per turn as one panel
+            turn_content_parts: list[str] = []
+            if content:
+                turn_content_parts.append(content)
+
+            if tool_calls:
+                self.memory.add_assistant(content, tool_calls)
+
+                for tc in tool_calls:
+                    _print_tool_call(tc)
+                    result = await self.executor.execute(tc)
+                    _print_tool_result(tc.name, result.to_str())
+                    self.memory.add_tool_result(tc, result)
+
+                    # Track changes for verification
+                    if tc.name in ("write_file", "apply_patch"):
+                        path = tc.arguments.get("path", "unknown")
+                        self.verifier.record_change(tc.name, path)
+
+                    # After a change tool, prompt verification
+                    if tc.name in ("write_file", "apply_patch"):
+                        self.memory.add_user(
+                            "[Change made. Now verify: check the diff with git_diff, run relevant tests, "
+                            "and confirm everything works before continuing. If you wrote a new file, check it compiles/runs correctly.]"
+                        )
+
+                    # Mark verified when the agent checks the diff
+                    if tc.name == "git_diff" and result.success:
+                        self.verifier.mark_verified()
+
+                    # #6: Handle ask_question with optional numbered options
+                    if tc.name == "ask_question":
+                        answer = result.to_str()
+                        # answer is already captured by _ask_question in tools.py
+                        # just make sure it feeds back into memory (already done via add_tool_result)
+
+                    if tc.name == "finish_task":
+                        # #2: Guard finish_task — block if there are unverified changes
+                        if self.verifier.has_unverified_changes():
+                            self.memory.add_user(
+                                "[finish_task blocked: you have unverified changes. "
+                                "Run git_diff to review your edits, then run tests or the linter "
+                                "to confirm everything works. Only call finish_task after verifying.]"
+                            )
+                            # Don't break — let the loop continue so the model can verify
+                        else:
+                            final_result = tc.arguments.get("result", "")
+                            break
+
+                if final_result:
+                    break
+
+                # Loop detection
+                if self.executor.is_looping():
+                    self.memory.add_user(
+                        "[You appear to be stuck in a loop — the same actions keep repeating. "
+                        "Step back, think about why the approach isn't working, and try a fundamentally different strategy. "
+                        "If you are truly blocked, call ask_question to get help from the user.]"
+                    )
+                    self.executor.reset_loop_detection()
+            else:
+                self.memory.add_assistant(content)
+
+            # #8: Print all buffered content for this turn as one panel
+            if turn_content_parts:
+                _print_assistant("\n\n".join(turn_content_parts))
+
+            self.memory.compress()
+
+            if self.interactive and content:
+                try:
+                    resp = input("\n[Axon] Continue? (Y/n): ").strip().lower()
+                    if resp == "n":
+                        final_result = "Session ended by user."
+                        break
+                except (EOFError, KeyboardInterrupt):
+                    final_result = "Session ended by user."
+                    break
+
+        if self.persist and final_result:
+            summary = final_result
+            if self.verifier.has_unverified_changes():
+                summary += "\n⚠ Some changes may not have been verified."
+            self.memory.add_user(f"[Previous task completed] {summary}")
+
+        return final_result or "Task completed."
+
+
+# ─── UI helpers ────────────────────────────────────────────────────────────────
+
+def _print_assistant(content: str, title: str = "Axon"):
+    """#8: Print assistant content as a single panel — one box per call."""
+    console.print()
+    console.print(Panel(Markdown(content), title=title, border_style="blue", title_align="left"))
+
+
+def _print_tool_call(tc: ToolCall):
+    args_str = json.dumps(tc.arguments, indent=2)
+    console.print(Panel(
+        f"[bold yellow]Tool:[/] [cyan]{tc.name}[/]\n[bold yellow]Args:[/]\n{args_str}",
+        title="🔧 Tool Call", border_style="yellow", title_align="left",
+    ))
+
+
+def _print_tool_result(tool_name: str, result: str):
+    display = result[:1000] + "\n... (truncated)" if len(result) > 1000 else result
+    console.print(Panel(
+        Syntax(display, "text", theme="monokai"),
+        title=f"📦 Result: {tool_name}", border_style="green", title_align="left",
+    ))