nimcode 0.1.0__py3-none-any.whl

nimcode/cli.py

@@ -0,0 +1,127 @@
+import argparse
+import asyncio
+import os
+import sys
+from rich.console import Console
+from .agent import Agent
+from .config import load_settings, save_global_setting
+from .permissions import PermissionMode
+
+console = Console()
+
+def run_login():
+    console.print("[bold cyan]NimCode Login[/bold cyan]")
+    console.print("Get your API key from [bold underline blue]https://build.nvidia.com/[/bold underline blue]")
+    
+    import getpass
+    api_key = getpass.getpass("Enter your NVIDIA NIM API Key: ")
+    if not api_key.strip():
+        console.print("[red]API Key cannot be empty.[/red]")
+        return
+        
+    save_global_setting("api_key", api_key.strip())
+    console.print("[green][OK] API Key saved successfully to ~/.nimcode/settings.json[/green]")
+
+def run_doctor():
+    console.print("[bold cyan]NimCode Doctor[/bold cyan] - Diagnostics")
+    key = os.environ.get("NIM_API_KEY")
+    if key:
+        console.print("[green][OK][/green] NIM_API_KEY environment variable is set.")
+    else:
+        console.print("[red][X][/red] NIM_API_KEY environment variable is missing.")
+    
+    # Check .nimcode existence
+    if os.path.exists(".nimcode"):
+        console.print("[green][OK][/green] .nimcode directory found in current project.")
+    else:
+        console.print("[yellow][!][/yellow] No .nimcode directory found. Standard settings apply.")
+    
+    console.print("[green][OK][/green] MCP SDK installed.")
+    console.print("Diagnostics complete.")
+
+def install_hook():
+    if not os.path.exists(".git"):
+        console.print("[red][X][/red] Not a git repository.")
+        return
+        
+    hook_path = os.path.join(".git", "hooks", "prepare-commit-msg")
+    with open(hook_path, "w", encoding="utf-8") as f:
+        f.write("#!/bin/sh\n")
+        f.write("# NimCode auto-commit hook\n")
+        f.write("if [ -z \"$(cat $1)\" ]; then\n")
+        f.write("  nimcode /commit > $1\n")
+        f.write("fi\n")
+    
+    import stat
+    os.chmod(hook_path, os.stat(hook_path).st_mode | stat.S_IEXEC)
+    console.print(f"[green][OK][/green] Git hook installed to {hook_path}")
+
+def main():
+    parser = argparse.ArgumentParser(description="NimCode: Autonomous Coding Agent for NVIDIA NIM APIs")
+    
+    # Check for doctor manually to avoid subparser conflict
+    if len(sys.argv) > 1:
+        if sys.argv[1] == "doctor":
+            run_doctor()
+            return
+        elif sys.argv[1] == "install-hook":
+            install_hook()
+            return
+        elif sys.argv[1] == "login":
+            run_login()
+            return
+        
+    # Main CLI arguments
+    parser.add_argument("prompt", nargs="?", default=None, help="The task you want NimCode to accomplish. If omitted, starts interactive REPL.")
+    parser.add_argument("--api-key", "-k", default=None, help="NVIDIA NIM API Key. Can also be set via NIM_API_KEY environment variable.")
+    parser.add_argument("--model", "-m", default="meta/llama-3.1-70b-instruct", help="Model ID to use from NIM.")
+    parser.add_argument("--max-turns", "-t", type=int, default=30, help="Maximum number of turns the agent is allowed to run.")
+    parser.add_argument("--permission-mode", "-p", type=PermissionMode, choices=list(PermissionMode), default=PermissionMode.DEFAULT, help="Permission mode for mutating tools.")
+    parser.add_argument("--resume", "-r", action="store_true", help="Resume from the last session stored in NIMCODE.md.")
+    
+    args = parser.parse_args()
+    
+    settings = load_settings()
+    
+    final_key = args.api_key or os.environ.get("NIM_API_KEY") or settings.get("api_key")
+    if not final_key:
+        console.print("[yellow]No API Key found. Let's get you set up![/yellow]")
+        run_login()
+        settings = load_settings()
+        final_key = settings.get("api_key")
+        if not final_key:
+            console.print("[bold red]API Key is required to use NimCode. Exiting.[/bold red]")
+            sys.exit(1)
+        
+    console.print(f"[bold green]Starting NimCode[/bold green] with model [cyan]{args.model}[/cyan]")
+    
+    agent = Agent(
+        api_key=final_key,
+        model=args.model,
+        max_turns=args.max_turns,
+        permission_mode=args.permission_mode
+    )
+    
+    if args.resume:
+        # Load from history if possible
+        agent.load_history()
+
+    piped_input = None
+    if not sys.stdin.isatty():
+        piped_input = sys.stdin.read().strip()
+        
+    if piped_input:
+        prompt = f"{piped_input}\n\n{args.prompt or ''}".strip()
+        console.print(f"Task (with piped input): {prompt}")
+        asyncio.run(agent.run(prompt))
+    elif args.prompt:
+        console.print(f"Task: {args.prompt}")
+        asyncio.run(agent.run(args.prompt))
+    else:
+        console.print("[bold yellow]Entering Interactive REPL Mode[/bold yellow]. Type /exit to quit, /plan for planning mode, /code for coding mode.")
+        asyncio.run(agent.start_repl())
+    
+    console.print("[bold green]Done![/bold green]")
+
+if __name__ == "__main__":
+    main()

nimcode/config.py

@@ -0,0 +1,54 @@
+import os
+import json
+import logging
+from typing import Dict, Any
+
+logger = logging.getLogger(__name__)
+
+def load_settings() -> Dict[str, Any]:
+    """Loads configuration from ~/.nimcode/settings.json and .nimcode/settings.json"""
+    settings = {
+        "model": "meta/llama-3.1-70b-instruct",
+        "mcp_servers": {}
+    }
+    
+    # Global settings
+    global_path = os.path.expanduser("~/.nimcode/settings.json")
+    if os.path.exists(global_path):
+        try:
+            with open(global_path, "r", encoding="utf-8") as f:
+                global_settings = json.load(f)
+                settings.update(global_settings)
+        except Exception as e:
+            logger.error(f"Failed to load global settings: {e}")
+
+    # Local settings
+    local_path = os.path.join(os.getcwd(), ".nimcode", "settings.json")
+    if os.path.exists(local_path):
+        try:
+            with open(local_path, "r", encoding="utf-8") as f:
+                local_settings = json.load(f)
+                settings.update(local_settings)
+        except Exception as e:
+            logger.error(f"Failed to load local settings: {e}")
+
+    return settings
+
+def save_global_setting(key: str, value: Any) -> None:
+    """Saves a setting to ~/.nimcode/settings.json"""
+    global_dir = os.path.expanduser("~/.nimcode")
+    os.makedirs(global_dir, exist_ok=True)
+    global_path = os.path.join(global_dir, "settings.json")
+    
+    settings = {}
+    if os.path.exists(global_path):
+        try:
+            with open(global_path, "r", encoding="utf-8") as f:
+                settings = json.load(f)
+        except Exception:
+            pass
+            
+    settings[key] = value
+    
+    with open(global_path, "w", encoding="utf-8") as f:
+        json.dump(settings, f, indent=4)

nimcode/lenient_parser.py

@@ -0,0 +1,101 @@
+import re
+import json
+import logging
+from typing import Dict, Any, List, Tuple, Optional
+
+logger = logging.getLogger(__name__)
+
+class LenientParser:
+    @staticmethod
+    def extract_tool_calls(text: str) -> List[str]:
+        """Extracts JSON strings from <tool_call> fenced blocks."""
+        # The model might forget the closing tag, or put markdown fences around it.
+        # We look for <tool_call> and extract everything until </tool_call> or EOF.
+        pattern = re.compile(r"<tool_call>\s*(.*?)(?:</tool_call>|\Z)", re.DOTALL | re.IGNORECASE)
+        matches = pattern.findall(text)
+        
+        # Clean up markdown code blocks if the model wrapped the JSON in them
+        cleaned_matches = []
+        for match in matches:
+            match = match.strip()
+            if match.startswith("```json"):
+                match = match[7:]
+            elif match.startswith("```"):
+                match = match[3:]
+            
+            if match.endswith("```"):
+                match = match[:-3]
+                
+            cleaned_matches.append(match.strip())
+            
+        return cleaned_matches
+
+    @staticmethod
+    def repair_json(json_str: str) -> str:
+        """Attempts to fix common model JSON formatting errors."""
+        repaired = json_str.strip()
+        
+        # 1. Remove trailing commas in objects and arrays
+        repaired = re.sub(r",\s*}", "}", repaired)
+        repaired = re.sub(r",\s*]", "]", repaired)
+        
+        # 2. Fix unescaped newlines within strings.
+        # A simple approach: we find things that look like string literals and escape newlines.
+        # We'll use a state machine or regex for simple cases. 
+        # For a robust approach without a full parser, we replace literal newlines that are inside quotes.
+        # Actually, python's json.loads is quite strict about unescaped newlines.
+        def escape_newlines_in_strings(match):
+            return match.group(0).replace("\n", "\\n")
+            
+        # Match string literals: " followed by anything except unescaped quote, followed by "
+        # We need to be careful with escaped quotes inside the string.
+        string_regex = re.compile(r'"(?:[^"\\]|\\.)*"', re.DOTALL)
+        repaired = string_regex.sub(escape_newlines_in_strings, repaired)
+        
+        # 3. Very rudimentary single quote to double quote conversion for keys/values
+        # (Only if it fails standard parsing, but we can do it proactively for top-level if needed).
+        # We will try standard json.loads first, and if it fails, maybe do more aggressive repair.
+        return repaired
+
+    @classmethod
+    def parse_tool_call(cls, tool_call_text: str) -> Dict[str, Any]:
+        """Parses a tool call string into a dictionary."""
+        try:
+            return json.loads(tool_call_text)
+        except json.JSONDecodeError:
+            pass
+            
+        # Try repairing
+        repaired = cls.repair_json(tool_call_text)
+        try:
+            return json.loads(repaired)
+        except json.JSONDecodeError as e:
+            logger.error(f"Failed to parse tool call even after repair: {repaired}")
+            raise ValueError(f"Malformed tool call JSON: {e}") from e
+
+    @classmethod
+    def process_model_response(cls, text: str) -> Tuple[str, List[Dict[str, Any]]]:
+        """
+        Parses the raw text response from the model.
+        Returns (plain_text_message, list_of_tool_calls).
+        """
+        # First, extract the tool calls
+        tool_call_strings = cls.extract_tool_calls(text)
+        
+        tool_calls = []
+        for tc_str in tool_call_strings:
+            if not tc_str:
+                continue
+            try:
+                tc = cls.parse_tool_call(tc_str)
+                tool_calls.append(tc)
+            except ValueError:
+                # We could append an error message or raise. We'll raise to let the agent loop handle it and re-prompt.
+                raise
+                
+        # The plain text is whatever is not in the tool_call blocks.
+        # We remove the tool call blocks from the original text to get the assistant's prose.
+        plain_text = re.sub(r"<tool_call>\s*(.*?)(?:</tool_call>|\Z)", "", text, flags=re.DOTALL | re.IGNORECASE)
+        plain_text = plain_text.strip()
+        
+        return plain_text, tool_calls

nimcode/mcp_client.py

@@ -0,0 +1,77 @@
+import logging
+from typing import Dict, Any, List
+from contextlib import AsyncExitStack
+
+logger = logging.getLogger(__name__)
+
+class MCPManager:
+    """Manages connections to MCP servers."""
+    
+    def __init__(self, mcp_config: Dict[str, Any]):
+        self.config = mcp_config
+        self.servers = mcp_config.get("mcp_servers", {})
+        self.exit_stack = AsyncExitStack()
+        self.sessions = {}  # name -> ClientSession
+        self.server_tools = {} # name -> list of tools
+        
+    async def connect_all(self):
+        """Initialize all MCP servers."""
+        if not self.servers:
+            return
+            
+        try:
+            from mcp.client.stdio import stdio_client, StdioServerParameters
+            from mcp.client.session import ClientSession
+        except ImportError:
+            logger.warning("MCP SDK not installed. Skipping MCP initialization.")
+            return
+
+        for name, details in self.servers.items():
+            cmd = details.get("command")
+            args = details.get("args", [])
+            env = details.get("env", None)
+            
+            logger.info(f"Connecting to MCP server '{name}'...")
+            try:
+                params = StdioServerParameters(command=cmd, args=args, env=env)
+                read, write = await self.exit_stack.enter_async_context(stdio_client(params))
+                session = await self.exit_stack.enter_async_context(ClientSession(read, write))
+                await session.initialize()
+                self.sessions[name] = session
+                
+                # Fetch tools
+                tools_response = await session.list_tools()
+                self.server_tools[name] = tools_response.tools if hasattr(tools_response, "tools") else []
+                logger.info(f"MCP Server '{name}' initialized with {len(self.server_tools[name])} tools.")
+            except Exception as e:
+                logger.error(f"Failed to connect to MCP server '{name}': {e}")
+
+    async def close(self):
+        await self.exit_stack.aclose()
+        
+    def get_system_prompt_additions(self) -> str:
+        """Returns extra text to add to the system prompt about MCPs."""
+        if not self.sessions:
+            return ""
+            
+        lines = ["\n\nAvailable MCP Servers & Tools:"]
+        for name, tools in self.server_tools.items():
+            lines.append(f"- Server '{name}':")
+            for t in tools:
+                lines.append(f"  * {t.name}: {t.description}")
+        return "\n".join(lines)
+        
+        
+    async def call_tool_by_name(self, tool_name: str, arguments: dict) -> Any:
+        for server_name, tools in self.server_tools.items():
+            for t in tools:
+                if t.name == tool_name:
+                    return await self.sessions[server_name].call_tool(tool_name, arguments)
+        raise ValueError(f"MCP Tool {tool_name} not found.")
+
+    async def call_tool(self, server_name: str, tool_name: str, arguments: dict) -> Any:
+        if server_name not in self.sessions:
+            raise ValueError(f"Server {server_name} not found or not connected.")
+        session = self.sessions[server_name]
+        result = await session.call_tool(tool_name, arguments)
+        return result

nimcode/memory.py

@@ -0,0 +1,69 @@
+import os
+from typing import List, Dict, Any
+
+class MemoryManager:
+    def __init__(self, max_tokens: int = 4000):
+        # We assume 1 token ~= 4 chars roughly
+        self.max_tokens = max_tokens
+
+    @staticmethod
+    def count_tokens(text: str) -> int:
+        """Roughly count tokens in a string."""
+        if not text:
+            return 0
+        return len(text) // 4 + 1
+
+    @classmethod
+    def count_messages_tokens(cls, messages: List[Dict[str, Any]]) -> int:
+        total = 0
+        for msg in messages:
+            content = msg.get("content", "")
+            total += cls.count_tokens(content)
+        return total
+
+    def compact_context(self, messages: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
+        """
+        If messages exceed max_tokens, drops the oldest messages.
+        Always keeps the System prompt (first message).
+        Always keeps the most recent user prompt.
+        """
+        total_tokens = self.count_messages_tokens(messages)
+        if total_tokens <= self.max_tokens:
+            return messages
+
+        # We need to compact
+        if not messages:
+            return messages
+
+        compacted = [messages[0]] # System prompt
+        remaining_messages = messages[1:]
+        
+        # We start from the end and add backwards until we hit the limit
+        # Reserved tokens for system prompt
+        current_tokens = self.count_tokens(messages[0].get("content", ""))
+        
+        kept_messages = []
+        for msg in reversed(remaining_messages):
+            msg_tokens = self.count_tokens(msg.get("content", ""))
+            if current_tokens + msg_tokens > self.max_tokens:
+                # If we haven't even kept the most recent message, we MUST keep it and just truncate its content
+                if not kept_messages:
+                    truncated_content = msg.get("content", "")[:(self.max_tokens - current_tokens) * 4]
+                    kept_messages.insert(0, {"role": msg["role"], "content": truncated_content + "...[TRUNCATED]"})
+                break
+            
+            kept_messages.insert(0, msg)
+            current_tokens += msg_tokens
+
+        compacted.extend(kept_messages)
+        return compacted
+
+    @staticmethod
+    def log_to_nimcode_md(turn: int, prompt: str, response: str, cwd: str = ".") -> None:
+        """Appends the interaction to NIMCODE.md for persistent session history."""
+        file_path = os.path.join(cwd, "NIMCODE.md")
+        with open(file_path, "a", encoding="utf-8") as f:
+            f.write(f"## Turn {turn}\n\n")
+            f.write(f"**User**: {prompt}\n\n")
+            f.write(f"**Agent**: {response}\n\n")
+            f.write("---\n\n")

nimcode/nim_client.py

@@ -0,0 +1,120 @@
+import httpx
+import asyncio
+import logging
+import json
+import random
+from typing import List, Dict, Any, Optional, AsyncGenerator
+
+logger = logging.getLogger(__name__)
+
+class NimClient:
+    def __init__(self, api_key: str, base_url: str = "https://integrate.api.nvidia.com/v1", model: str = "meta/llama-3.1-70b-instruct"):
+        self.api_key = api_key
+        self.base_url = base_url.rstrip("/")
+        self.model = model
+        self.headers = {
+            "Authorization": f"Bearer {self.api_key}",
+            "Content-Type": "application/json",
+            "Accept": "text/event-stream"
+        }
+        
+    async def get_available_models(self) -> List[str]:
+        return [
+            "meta/llama-3.1-70b-instruct",
+            "meta/llama-3.1-8b-instruct",
+            "meta/llama-3.1-405b-instruct",
+            "nvidia/nemotron-4-340b-instruct",
+            "mistralai/mixtral-8x22b-instruct-v0.1"
+        ]
+
+    async def chat_one_shot(self, prompt: str) -> str:
+        payload = {
+            "model": self.model,
+            "messages": [{"role": "user", "content": prompt}],
+            "max_tokens": 1024,
+            "temperature": 0.2
+        }
+        async with httpx.AsyncClient() as client:
+            response = await client.post(
+                f"{self.base_url}/chat/completions",
+                headers=self.headers,
+                json=payload,
+                timeout=60.0
+            )
+            response.raise_for_status()
+            data = response.json()
+            return data["choices"][0]["message"]["content"]
+
+    async def chat_vision(self, base64_image: str, prompt: str) -> str:
+        # For vision, we might need a specific vision model, but we'll try with the default if it supports multimodal
+        # Or switch to a known vision model like nv-llama-3.2-90b-vision-instruct
+        vision_model = "meta/llama-3.2-90b-vision-instruct"
+        payload = {
+            "model": vision_model,
+            "messages": [
+                {
+                    "role": "user",
+                    "content": [
+                        {"type": "text", "text": prompt},
+                        {"type": "image_url", "image_url": {"url": f"data:image/png;base64,{base64_image}"}}
+                    ]
+                }
+            ],
+            "max_tokens": 1024,
+            "temperature": 0.2
+        }
+        
+        async with httpx.AsyncClient() as client:
+            try:
+                response = await client.post(
+                    f"{self.base_url}/chat/completions",
+                    headers=self.headers,
+                    json=payload,
+                    timeout=60.0
+                )
+                response.raise_for_status()
+                data = response.json()
+                return data["choices"][0]["message"]["content"]
+            except httpx.HTTPStatusError as e:
+                logger.error(f"Vision API error: {e.response.text}")
+                return f"Vision processing failed: {e.response.text}"
+            except Exception as e:
+                return f"Vision error: {str(e)}"
+
+    async def chat(self, messages: List[Dict[str, Any]], tools: Optional[List[Dict[str, Any]]] = None, stream: bool = True) -> AsyncGenerator[str, None]:
+        payload = {
+            "model": self.model,
+            "messages": messages,
+            "max_tokens": 4096,
+            "temperature": 0.2,
+            "stream": stream
+        }
+        
+        async with httpx.AsyncClient() as client:
+            try:
+                async with client.stream("POST", f"{self.base_url}/chat/completions", headers=self.headers, json=payload, timeout=60.0) as response:
+                    response.raise_for_status()
+                    async for line in response.aiter_lines():
+                        if line.startswith("data: ") and line != "data: [DONE]":
+                            data_str = line[6:]
+                            try:
+                                data_json = json.loads(data_str)
+                                chunk = data_json["choices"][0]["delta"].get("content", "")
+                                if chunk:
+                                    yield chunk
+                            except json.JSONDecodeError:
+                                pass
+            except httpx.HTTPStatusError as e:
+                await e.response.aread()
+                logger.error(f"API HTTP error: {e.response.status_code} - {e.response.text}")
+                yield f"\\n\\n[Error: Model API returned {e.response.status_code}. Please check your NVIDIA API key.]"
+            except Exception as e:
+                logger.error(f"API connection error: {e}")
+                yield f"\\n\\n[Error communicating with NVIDIA API: {e}]"
+
+    def count_tokens_approx(self, messages: List[Dict[str, Any]]) -> int:
+        total_chars = 0
+        for msg in messages:
+            total_chars += len(msg.get("role", ""))
+            total_chars += len(str(msg.get("content", "")))
+        return total_chars // 4 + (len(messages) * 4)

nimcode/permissions.py

@@ -0,0 +1,78 @@
+import typer
+import logging
+from typing import Dict, Any
+from enum import Enum
+from rich.console import Console
+import sys
+
+logger = logging.getLogger(__name__)
+console = Console()
+
+class PermissionMode(str, Enum):
+    DEFAULT = "default"  # Prompts for everything (simulated in headless)
+    BYPASS = "bypass"    # Allows everything
+    AUTO = "auto"        # Allows safe reads, prompts for dangerous actions
+
+class PermissionEngine:
+    def __init__(self, mode: PermissionMode = PermissionMode.DEFAULT):
+        self.mode = mode
+        self.safe_tools = {"Read", "Glob", "Grep"}
+
+    def check_permission(self, tool_call: dict) -> bool:
+        """Returns True if permitted, False otherwise."""
+        if self.mode == PermissionMode.BYPASS:
+            return True
+            
+        tool_name = tool_call.get("tool")
+        
+        if self.mode == PermissionMode.AUTO and tool_name in self.safe_tools:
+            return True
+            
+        return self._prompt_user(tool_call)
+            
+    def _prompt_user(self, tool_call: Dict[str, Any]) -> bool:
+        from rich.panel import Panel
+        from rich.prompt import Prompt
+        import json
+        
+        tool_name = tool_call.get("tool")
+        args = tool_call.get("args", {})
+        
+        if not sys.stdin.isatty():
+            console.print(f"[yellow]Running non-interactive. Approving {tool_name} for tests.[/yellow]")
+            return True
+            
+        while True:
+            content = ""
+            if tool_name == "Bash":
+                content = f"[bold]Command:[/bold]\n{args.get('command')}"
+            elif tool_name == "Write":
+                content = f"[bold]File:[/bold] {args.get('file_path')}\n[bold]Action:[/bold] Overwrite with new content"
+            elif tool_name == "Edit":
+                content = f"[bold]File:[/bold] {args.get('file_path')}\n[bold]Replacing:[/bold] '{args.get('old_string')}' -> '{args.get('new_string')}'"
+            else:
+                content = f"[bold]Args:[/bold] {json.dumps(args, indent=2)}"
+                
+            panel = Panel(content, title=f"NimCode Wants to Run: {tool_name}", border_style="yellow")
+            console.print(panel)
+            
+            choice = Prompt.ask("[bold cyan]Action[/bold cyan] [green](a)ccept[/green] / [red](r)eject[/red] / [blue](e)dit[/blue]", choices=["a", "r", "e"], default="a", show_choices=False)
+            
+            if choice == "a":
+                return True
+            elif choice == "r":
+                return False
+            elif choice == "e":
+                from prompt_toolkit import prompt as pt_prompt
+                if tool_name == "Bash":
+                    new_cmd = pt_prompt("Edit Command: ", default=args.get("command", ""))
+                    tool_call["args"]["command"] = new_cmd
+                else:
+                    new_args_str = pt_prompt("Edit Args (JSON): ", default=json.dumps(args))
+                    try:
+                        tool_call["args"] = json.loads(new_args_str)
+                    except json.JSONDecodeError:
+                        console.print("[red]Invalid JSON. Please try again.[/red]")
+                        continue
+                args = tool_call.get("args", {})
+                # Loop repeats and shows the updated panel for final confirmation