nimcode 0.1.21__tar.gz → 0.3.0__tar.gz

{nimcode-0.1.21 → nimcode-0.3.0}/PKG-INFO

@@ -1,14 +1,13 @@
 Metadata-Version: 2.4
 Name: nimcode
-Version: 0.1.21
+Version: 0.3.0
 Summary: A standalone, robust coding agent for NVIDIA NIM models.
 Author: Autonomous Agent
 Requires-Dist: httpx>=0.27.0
 Requires-Dist: rich>=13.7.0
 Requires-Dist: prompt_toolkit>=3.0.0
 Requires-Dist: mcp>=1.2.0
-Requires-Dist: keyring>=24.0.0
-Requires-Dist: watchdog>=4.0.0
+Requires-Dist: watchdog>=3.0.0
 Dynamic: author
 Dynamic: requires-dist
 Dynamic: summary

{nimcode-0.1.21 → nimcode-0.3.0}/setup.py

@@ -2,7 +2,7 @@ from setuptools import setup, find_packages
 
 setup(
     name="nimcode",
-    version="0.1.21",
+    version="0.3.0",
     packages=find_packages(where="src"),
     package_dir={"": "src"},
     install_requires=[
@@ -10,8 +10,7 @@ setup(
         "rich>=13.7.0",
         "prompt_toolkit>=3.0.0",
         "mcp>=1.2.0",
-        "keyring>=24.0.0",
-        "watchdog>=4.0.0"
+        "watchdog>=3.0.0"
     ],
     entry_points={
         "console_scripts": [

{nimcode-0.1.21 → nimcode-0.3.0}/src/nimcode/agent.py

@@ -8,7 +8,6 @@ from .permissions import PermissionEngine, PermissionMode
 from .config import load_settings, save_global_setting
 from .mcp_client import MCPManager
 from .memory import MemoryManager
-from .analytics import AnalyticsEngine
 
 logger = logging.getLogger(__name__)
 
@@ -28,7 +27,7 @@ Available Tools:
 - Bash: {"tool": "Bash", "args": {"command": "string"}}
 - Read: {"tool": "Read", "args": {"file_path": "string"}}
 - Write: {"tool": "Write", "args": {"file_path": "string", "content": "string"}}
-- Replace: {"tool": "Replace", "args": {"file_path": "string", "replacements": [{"old_string": "exact old", "new_string": "exact new"}]}}
+- Edit: {"tool": "Edit", "args": {"file_path": "string", "old_string": "string", "new_string": "string"}}
 - Glob: {"tool": "Glob", "args": {"pattern": "string"}}
 - Grep: {"tool": "Grep", "args": {"query": "string", "directory": "string"}}
 
@@ -49,42 +48,19 @@ When you have completely fulfilled the user's request and have no more tools to
 """
 
 class Agent:
-    def __init__(self, api_key: str, model: str = None, max_turns: int = 30, permission_mode: PermissionMode = PermissionMode.DEFAULT, max_tokens: int = 100000):
+    def __init__(self, api_key: str, model: str = None, max_turns: int = 30, permission_mode: PermissionMode = PermissionMode.DEFAULT, max_tokens: int = 4000):
         # Load global settings
         self.settings = load_settings()
         self.model = model or self.settings.get("model", "meta/llama-3.1-70b-instruct")
-        self.is_local = self.settings.get("is_local", False)
-        self.base_url = self.settings.get("base_url", None)
-        self.client = NimClient(api_key=api_key, base_url=self.base_url, model=self.model, is_local=self.is_local)
+        api_base_url = self.settings.get("api_base_url", "https://integrate.api.nvidia.com/v1")
+        self.client = NimClient(api_key=api_key, base_url=api_base_url, model=self.model)
         
         # Initialize MCP Manager
         self.mcp = MCPManager(self.settings)
-        self.analytics = AnalyticsEngine()
-        self.memory = MemoryManager(model_name=model, fallback_max_tokens=max_tokens)
         
         # Base system prompt
         final_prompt = SYSTEM_PROMPT + self.mcp.get_system_prompt_additions()
         
-        # Inject Repo Map
-        try:
-            from .repo_map import RepoMapper
-            mapper = RepoMapper(os.getcwd())
-            repo_map = mapper.generate_map()
-            final_prompt += f"\n\n--- REPOSITORY MAP ---\n{repo_map}\n----------------------\n"
-        except Exception as e:
-            logger.error(f"Failed to generate repo map: {e}")
-        
-        # Load .nimcoderules if present
-        rules_path = os.path.join(os.getcwd(), ".nimcoderules")
-        if os.path.exists(rules_path):
-            try:
-                with open(rules_path, "r", encoding="utf-8") as f:
-                    rules = f.read()
-                final_prompt += f"\n\nPROJECT-SPECIFIC RULES (.nimcoderules):\n{rules}\n"
-                logger.info("Loaded .nimcoderules")
-            except Exception as e:
-                logger.error(f"Failed to load .nimcoderules: {e}")
-        
         # Load skills if present
         skills_dir = os.path.join(os.getcwd(), ".nimcode", "skills")
         if os.path.exists(skills_dir) and os.path.isdir(skills_dir):
@@ -106,18 +82,13 @@ class Agent:
                 final_prompt += f"\n\nGIT CONTEXT:\nBranch: {branch}\nUncommitted changes:\n{status if status else 'None'}"
             except Exception as e:
                 logger.error(f"Failed to load git context: {e}")
-                
-        # Repo Map
-        repo_map = self._generate_repo_map(os.getcwd())
-        if repo_map:
-            final_prompt += f"\n\nREPOSITORY MAP:\n{repo_map}"
 
         self.messages: List[Dict[str, Any]] = [
             {"role": "system", "content": final_prompt}
         ]
         self.max_turns = max_turns
         self.permission_engine = PermissionEngine(mode=permission_mode)
-        self.memory = MemoryManager(model_name=model, fallback_max_tokens=max_tokens)
+        self.memory = MemoryManager(max_tokens=max_tokens)
 
     def save_history(self):
         """Saves current conversation to NIMCODE.md"""
@@ -138,28 +109,6 @@ class Agent:
         except Exception as e:
             logger.error(f"Failed to load history: {e}")
 
-    def _generate_repo_map(self, cwd: str) -> str:
-        """Generates a fast, lightweight file tree map."""
-        ignore_dirs = {'.git', 'node_modules', '__pycache__', '.venv', 'venv', 'env', '.nimcode'}
-        tree = []
-        for root, dirs, files in os.walk(cwd):
-            dirs[:] = [d for d in dirs if d not in ignore_dirs and not d.startswith('.')]
-            level = root.replace(cwd, '').count(os.sep)
-            indent = ' ' * 4 * level
-            basename = os.path.basename(root)
-            if basename:
-                tree.append(f"{indent}{basename}/")
-            subindent = ' ' * 4 * (level + 1)
-            for f in files:
-                if not f.endswith('.pyc') and not f.startswith('.'):
-                    tree.append(f"{subindent}{f}")
-                    
-        # Limit to 500 lines to save context
-        if len(tree) > 500:
-            tree = tree[:500] + ["... (truncated for context limit)"]
-            
-        return "\n".join(tree)
-
     async def _stream_response(self) -> str:
         from rich.live import Live
         from rich.markdown import Markdown
@@ -222,16 +171,23 @@ class Agent:
                 live.update(Markdown(response_text + "\n\n*[yellow]Stream interrupted by user.[/yellow]*", code_theme=code_theme))
                 c.print("\n[yellow]Generation interrupted.[/yellow]")
                 
-        # Analytics Token Tracker Update
-        est_prompt_tokens = self.client.count_tokens_approx(self.messages)
-        est_completion_tokens = len(response_text) // 4
+        # Approximate Token Tracker Update
+        est_tokens = len(response_text) // 4
+        if not hasattr(self, "session_tokens"):
+            self.session_tokens = 0
+        self.session_tokens += est_tokens
         
-        self.analytics.log_usage(self.model, est_prompt_tokens, est_completion_tokens)
+        # Approximate cost based on 70B typical rates ($3/1M tokens)
+        cost = (self.session_tokens / 1000000) * 3.0
+        c.print(f"[dim]Output est. tokens: {est_tokens} | Session Cost: ~${cost:.4f}[/dim]")
         
-        # Display today's cost
-        stats = self.analytics.get_summary()
-        today_cost = stats["today"]["cost_usd"]
-        c.print(f"[dim]Output est. tokens: {est_completion_tokens} | Today's Cost: ~${today_cost:.4f}[/dim]")
+        # Context usage warning
+        import json
+        total_context_chars = sum(len(str(m.get("content", ""))) for m in self.messages)
+        total_est_tokens = total_context_chars // 4
+        max_context = 128000  # Assume standard Llama-3.1 128k context for now
+        if total_est_tokens > max_context * 0.8:
+            c.print("[bold yellow]⚠️ Context window is over 80% full. Consider running /compact or /clear.[/bold yellow]")
             
         return response_text
 
@@ -242,10 +198,6 @@ class Agent:
         from .tools import ToolRegistry
         while turn < max_turns:
             turn += 1
-            if self.memory.count_messages_tokens(self.messages) > (self.memory.max_tokens * 0.8):
-                logger.info("Context full in headless mode. Distilling memory via LLM...")
-                self.messages = await self._distill_memory()
-                
             try:
                 response_text = await self.client.chat_one_shot(self.messages)
             except Exception as e:
@@ -305,7 +257,7 @@ class Agent:
             
             # Compact context before calling API
             from rich.console import Console
-            if self.memory.count_messages_tokens(self.messages) > (self.memory.max_tokens * 0.8):
+            if self.memory.count_messages_tokens(self.messages) > self.memory.max_tokens:
                 logger.info("Context full. Distilling memory via LLM...")
                 Console().print("[dim italic]🧠 Context full. Distilling memory into a summary to save tokens...[/dim italic]")
                 self.messages = await self._distill_memory()
@@ -338,6 +290,9 @@ class Agent:
                 
                 if not tool_calls:
                     # Model responded with plain text but didn't say TASK_COMPLETE.
+                    if "[Error: Model API returned" in full_content or "[Error communicating with" in full_content:
+                        break
+                        
                     self.messages.append({"role": "user", "content": "Please continue. Use a tool or output TASK_COMPLETE."})
                     continue
                     
@@ -346,7 +301,7 @@ class Agent:
                 tool_name = tool_call.get("tool", "Unknown")
                 
                 logger.info(f"Checking permissions for tool: {tool_name}")
-                if not await self.permission_engine.check_permission(tool_call):
+                if not self.permission_engine.check_permission(tool_call):
                     self.messages.append({
                         "role": "user",
                         "content": f"User explicitly denied permission to execute {tool_name}. Please choose another approach."

{nimcode-0.1.21 → nimcode-0.3.0}/src/nimcode/cli.py

@@ -14,20 +14,13 @@ def run_login():
     console.print("Get your API key from [bold underline blue]https://build.nvidia.com/[/bold underline blue]")
     
     import getpass
-    from .config import save_global_setting
-    
-    try:
-        api_key = getpass.getpass("Enter your NVIDIA NIM API Key: ")
-    except (EOFError, Exception):
-        console.print("[red]Could not prompt for API key. Please set NIM_API_KEY environment variable or run 'nimcode login' in a regular terminal.[/red]")
-        sys.exit(1)
-        
+    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 securely to settings[/green]")
+    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")
@@ -63,7 +56,19 @@ def install_hook():
     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 _silence_anyio_errors():
+    import sys
+    original_hook = sys.unraisablehook
+    def custom_unraisablehook(unraisable):
+        if unraisable.exc_type == RuntimeError and "exit cancel scope in a different task" in str(unraisable.exc_value):
+            return
+        if unraisable.exc_type == BaseExceptionGroup and "unhandled errors in a TaskGroup" in str(unraisable.exc_value):
+            return
+        original_hook(unraisable)
+    sys.unraisablehook = custom_unraisablehook
+
 def main():
+    _silence_anyio_errors()
     parser = argparse.ArgumentParser(description="NimCode: Autonomous Coding Agent for NVIDIA NIM APIs")
     
     # Check for doctor manually to avoid subparser conflict
@@ -74,25 +79,27 @@ def main():
         elif sys.argv[1] == "install-hook":
             install_hook()
             return
-        elif sys.argv[1] == "login" or sys.argv[1] == "/login":
+        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=None, help="Model ID to use from NIM (e.g. meta/llama-3.1-70b-instruct).")
+    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.")
-    parser.add_argument("--stdio", action="store_true", help="Run NimCode as a JSON-RPC server over stdio for IDE integration.")
     
     args = parser.parse_args()
     
     settings = load_settings()
+    api_base_url = settings.get("api_base_url", "https://integrate.api.nvidia.com/v1")
     
     final_key = args.api_key or os.environ.get("NIM_API_KEY") or settings.get("api_key")
-    if not final_key:
+    is_local = "localhost" in api_base_url or "127.0.0.1" in api_base_url
+    
+    if not final_key and not is_local:
         console.print("[yellow]No API Key found. Let's get you set up![/yellow]")
         run_login()
         settings = load_settings()
@@ -100,32 +107,12 @@ def main():
         if not final_key:
             console.print("[bold red]API Key is required to use NimCode. Exiting.[/bold red]")
             sys.exit(1)
-        
-    model = args.model or settings.get("model")
-    if not model and not args.stdio:
-        from rich.prompt import Prompt
-        from .nim_client import NimClient
-        client = NimClient(api_key=final_key)
-        console.print("[yellow]Fetching available models from NVIDIA NIM...[/yellow]")
-        models = asyncio.run(client.get_available_models())
-        console.print("\n[bold cyan]Available Models:[/bold cyan]")
-        for i, m in enumerate(models):
-            console.print(f"[{i+1}] {m}")
-        
-        choice = Prompt.ask("Select a model number", choices=[str(i) for i in range(1, len(models)+1)], default="1")
-        model = models[int(choice) - 1]
-        
-        from .config import save_global_setting
-        if Prompt.ask("Save this model as default?", choices=["y", "n"], default="y") == "y":
-            save_global_setting("model", model)
-            console.print(f"[green]Saved {model} as default model.[/green]")
             
-    if not model:
-        model = "meta/llama-3.1-70b-instruct" # ultimate fallback
-
+    final_key = final_key or "local-dummy-key"
+        
     agent = Agent(
         api_key=final_key,
-        model=model,
+        model=args.model,
         max_turns=args.max_turns,
         permission_mode=args.permission_mode
     )
@@ -134,15 +121,29 @@ def main():
         # Load from history if possible
         agent.load_history()
 
-    piped_input = None
-    if not sys.stdin.isatty():
-        import select
+    async def safe_start_repl(repl):
+        loop = asyncio.get_running_loop()
+        def custom_exception_handler(loop, context):
+            msg = context.get("message", "")
+            exception = context.get("exception", None)
+            if "unhandled errors in a TaskGroup" in str(msg) or "unhandled errors in a TaskGroup" in str(exception):
+                return
+            if exception and isinstance(exception, RuntimeError) and "exit cancel scope in a different task" in str(exception):
+                return
+            if "asynchronous generator" in str(msg):
+                return
+            loop.default_exception_handler(context)
+        loop.set_exception_handler(custom_exception_handler)
+        
         try:
-            # On Unix, we can check if stdin has data ready to read
-            if sys.platform != "win32" and select.select([sys.stdin], [], [], 0.0)[0]:
-                piped_input = sys.stdin.read().strip()
-        except Exception:
+            await repl.start_repl()
+        except asyncio.CancelledError:
             pass
+
+    piped_input = None
+    if not sys.stdin.isatty():
+        piped_input = sys.stdin.read().strip()
+
     if piped_input:
         console.print(f"[bold green]Starting NimCode[/bold green] with model [cyan]{args.model}[/cyan]")
         prompt = f"{piped_input}\n\n{args.prompt or ''}".strip()
@@ -152,14 +153,10 @@ def main():
         console.print(f"[bold green]Starting NimCode[/bold green] with model [cyan]{args.model}[/cyan]")
         console.print(f"Task: {args.prompt}")
         asyncio.run(agent.run(args.prompt))
-    elif args.stdio:
-        from .stdio_server import StdioServer
-        server = StdioServer(agent)
-        asyncio.run(server.start())
     else:
         from .repl import NimcodeREPL
         repl = NimcodeREPL(agent)
-        asyncio.run(repl.start_repl())
+        asyncio.run(safe_start_repl(repl))
     
     # We don't print "Done!" for REPL to keep it clean on exit
     if args.prompt or piped_input:

{nimcode-0.1.21 → nimcode-0.3.0}/src/nimcode/config.py

@@ -9,6 +9,7 @@ def load_settings() -> Dict[str, Any]:
     """Loads configuration from ~/.nimcode/settings.json and .nimcode/settings.json"""
     settings = {
         "model": "meta/llama-3.1-70b-instruct",
+        "api_base_url": "https://integrate.api.nvidia.com/v1",
         "mcp_servers": {}
     }
     
@@ -21,8 +22,7 @@ def load_settings() -> Dict[str, Any]:
                 settings.update(global_settings)
         except Exception as e:
             logger.error(f"Failed to load global settings: {e}")
-            
-    # Keyring caused 30-second hangs on Windows, so we no longer scrub api_key from settings.
+
     # Local settings
     local_path = os.path.join(os.getcwd(), ".nimcode", "settings.json")
     if os.path.exists(local_path):

nimcode-0.3.0/src/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-0.1.21 → nimcode-0.3.0}/src/nimcode/nim_client.py

@@ -8,58 +8,17 @@ from typing import List, Dict, Any, Optional, AsyncGenerator
 logger = logging.getLogger(__name__)
 
 class NimClient:
-    def __init__(self, api_key: str, base_url: str = None, model: str = None, is_local: bool = False):
+    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.is_local = is_local
-        self.base_url = (base_url or "https://integrate.api.nvidia.com/v1").rstrip("/")
-        self.model = model or "meta/llama-3.1-70b-instruct"
+        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"
         }
         
-    @staticmethod
-    def get_model_context_length(model_name: str) -> int:
-        """Heuristically determines the maximum token context length from the model name."""
-        model_lower = model_name.lower()
-        if "llama-3.1" in model_lower or "llama-3.2" in model_lower:
-            return 128000
-        elif "nemotron" in model_lower:
-            return 128000 # Nemotron-4 is usually high context
-        elif "mixtral" in model_lower:
-            return 64000
-        elif "phi" in model_lower:
-            return 128000
-        elif "gemma" in model_lower:
-            return 8192
-        # Default fallback
-        return 32000
-
     async def get_available_models(self) -> List[str]:
-        if self.is_local:
-            try:
-                async with httpx.AsyncClient() as client:
-                    response = await client.get(f"{self.base_url}/models", timeout=5.0)
-                    if response.status_code == 200:
-                        data = response.json()
-                        return [m["id"] for m in data.get("data", [])]
-            except Exception as e:
-                logger.error(f"Failed to fetch local models: {e}")
-                return ["local-model"]
-                
-        try:
-            async with httpx.AsyncClient() as client:
-                response = await client.get(f"{self.base_url}/models", headers={"Authorization": f"Bearer {self.api_key}"}, timeout=10.0)
-                if response.status_code == 200:
-                    data = response.json()
-                    return [m["id"] for m in data.get("data", [])]
-                else:
-                    logger.error(f"Failed to fetch models from API: {response.status_code}")
-        except Exception as e:
-            logger.error(f"Failed to fetch models: {e}")
-            
-        # Fallback list if API fails
         return [
             "meta/llama-3.1-70b-instruct",
             "meta/llama-3.1-8b-instruct",
@@ -131,31 +90,57 @@ class NimClient:
             "stream": stream
         }
         
-        async with httpx.AsyncClient() as client:
+        max_retries = 15
+        base_delay = 2.0
+        max_delay = 60.0
+        
+        for attempt in range(max_retries):
+            chunk_yielded = False
             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
+                async with httpx.AsyncClient() as client:
+                    async with client.stream("POST", f"{self.base_url}/chat/completions", headers=self.headers, json=payload, timeout=120.0) as response:
+                        if response.status_code in [408, 429, 500, 502, 503, 504, 529]:
+                            raise httpx.HTTPStatusError(f"Temporary server error {response.status_code}", request=response.request, response=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:
+                                        chunk_yielded = True
+                                        yield chunk
+                                except json.JSONDecodeError:
+                                    pass
+                        return  # Success
             except httpx.HTTPStatusError as e:
-                try:
-                    await e.response.aread()
-                    text = e.response.text
-                except Exception:
-                    text = "<unread stream>"
-                logger.error(f"API HTTP error: {e.response.status_code} - {text}")
-                yield f"\\n\\n[Error: Model API returned {e.response.status_code}. Please check your NVIDIA API key.]"
+                transient_codes = [408, 429, 500, 502, 503, 504, 529]
+                if e.response.status_code not in transient_codes or chunk_yielded or attempt == max_retries - 1:
+                    try:
+                        await e.response.aread()
+                        text = e.response.text
+                    except Exception:
+                        text = "<unread stream>"
+                    logger.error(f"API HTTP error: {e.response.status_code} - {text}")
+                    yield f"\n\n[Error: Model API returned {e.response.status_code}. Please check your API key if 401.]"
+                    return
+                    
+                delay = min(max_delay, base_delay * (2 ** attempt))
+                logger.warning(f"API HTTP error {e.response.status_code}. Retrying in {delay}s...")
+                await asyncio.sleep(delay)
+                
             except Exception as e:
-                logger.error(f"API connection error: {e}")
-                yield f"\\n\\n[Error communicating with NVIDIA API: {e}]"
+                if chunk_yielded or attempt == max_retries - 1:
+                    logger.error(f"API connection error: {type(e).__name__} - {e}")
+                    yield f"\n\n[Error communicating with NVIDIA API: {type(e).__name__} - {e}]"
+                    return
+                    
+                delay = min(max_delay, base_delay * (2 ** attempt))
+                logger.warning(f"API connection error: {type(e).__name__} - {e}. Retrying in {delay}s...")
+                await asyncio.sleep(delay)
 
     def count_tokens_approx(self, messages: List[Dict[str, Any]]) -> int:
         total_chars = 0