nimcode 0.1.12__tar.gz → 0.1.13__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (48) hide show
  1. {nimcode-0.1.12 → nimcode-0.1.13}/PKG-INFO +3 -1
  2. {nimcode-0.1.12 → nimcode-0.1.13}/setup.py +4 -2
  3. {nimcode-0.1.12 → nimcode-0.1.13}/src/nimcode/agent.py +70 -21
  4. nimcode-0.1.13/src/nimcode/analytics.py +68 -0
  5. {nimcode-0.1.12 → nimcode-0.1.13}/src/nimcode/cli.py +37 -6
  6. {nimcode-0.1.12 → nimcode-0.1.13}/src/nimcode/config.py +14 -0
  7. nimcode-0.1.13/src/nimcode/memory.py +140 -0
  8. {nimcode-0.1.12 → nimcode-0.1.13}/src/nimcode/nim_client.py +44 -3
  9. {nimcode-0.1.12 → nimcode-0.1.13}/src/nimcode/permissions.py +34 -8
  10. nimcode-0.1.13/src/nimcode/plugin_manager.py +53 -0
  11. {nimcode-0.1.12 → nimcode-0.1.13}/src/nimcode/repl.py +247 -87
  12. nimcode-0.1.13/src/nimcode/repo_map.py +53 -0
  13. nimcode-0.1.13/src/nimcode/secret_scanner.py +48 -0
  14. nimcode-0.1.13/src/nimcode/stdio_server.py +134 -0
  15. nimcode-0.1.13/src/nimcode/swarm.py +42 -0
  16. {nimcode-0.1.12 → nimcode-0.1.13}/src/nimcode/tools.py +151 -93
  17. nimcode-0.1.13/src/nimcode/updater.py +34 -0
  18. nimcode-0.1.13/src/nimcode/watcher.py +81 -0
  19. {nimcode-0.1.12 → nimcode-0.1.13}/src/nimcode.egg-info/PKG-INFO +3 -1
  20. {nimcode-0.1.12 → nimcode-0.1.13}/src/nimcode.egg-info/SOURCES.txt +15 -1
  21. {nimcode-0.1.12 → nimcode-0.1.13}/src/nimcode.egg-info/requires.txt +2 -0
  22. {nimcode-0.1.12 → nimcode-0.1.13}/tests/test_agent.py +7 -6
  23. nimcode-0.1.13/tests/test_analytics.py +22 -0
  24. {nimcode-0.1.12 → nimcode-0.1.13}/tests/test_cli.py +19 -11
  25. {nimcode-0.1.12 → nimcode-0.1.13}/tests/test_memory.py +4 -4
  26. nimcode-0.1.13/tests/test_nimcoderules.py +19 -0
  27. {nimcode-0.1.12 → nimcode-0.1.13}/tests/test_permissions.py +31 -21
  28. nimcode-0.1.13/tests/test_plugin_manager.py +35 -0
  29. {nimcode-0.1.12 → nimcode-0.1.13}/tests/test_repl.py +1 -0
  30. nimcode-0.1.13/tests/test_repo_map.py +22 -0
  31. nimcode-0.1.13/tests/test_secret_scanner.py +31 -0
  32. {nimcode-0.1.12 → nimcode-0.1.13}/tests/test_tools.py +48 -38
  33. nimcode-0.1.13/tests/test_updater.py +39 -0
  34. nimcode-0.1.12/src/nimcode/memory.py +0 -69
  35. {nimcode-0.1.12 → nimcode-0.1.13}/README.md +0 -0
  36. {nimcode-0.1.12 → nimcode-0.1.13}/setup.cfg +0 -0
  37. {nimcode-0.1.12 → nimcode-0.1.13}/src/nimcode/__init__.py +0 -0
  38. {nimcode-0.1.12 → nimcode-0.1.13}/src/nimcode/lenient_parser.py +0 -0
  39. {nimcode-0.1.12 → nimcode-0.1.13}/src/nimcode/mcp_client.py +0 -0
  40. {nimcode-0.1.12 → nimcode-0.1.13}/src/nimcode.egg-info/dependency_links.txt +0 -0
  41. {nimcode-0.1.12 → nimcode-0.1.13}/src/nimcode.egg-info/entry_points.txt +0 -0
  42. {nimcode-0.1.12 → nimcode-0.1.13}/src/nimcode.egg-info/top_level.txt +0 -0
  43. {nimcode-0.1.12 → nimcode-0.1.13}/tests/test_config.py +0 -0
  44. {nimcode-0.1.12 → nimcode-0.1.13}/tests/test_lenient_parser.py +0 -0
  45. {nimcode-0.1.12 → nimcode-0.1.13}/tests/test_mcp_client.py +0 -0
  46. {nimcode-0.1.12 → nimcode-0.1.13}/tests/test_nim_client.py +0 -0
  47. {nimcode-0.1.12 → nimcode-0.1.13}/tests/test_repl_extra.py +0 -0
  48. {nimcode-0.1.12 → nimcode-0.1.13}/tests/test_repl_trust.py +0 -0
@@ -1,12 +1,14 @@
1
1
  Metadata-Version: 2.4
2
2
  Name: nimcode
3
- Version: 0.1.12
3
+ Version: 0.1.13
4
4
  Summary: A standalone, robust coding agent for NVIDIA NIM models.
5
5
  Author: Autonomous Agent
6
6
  Requires-Dist: httpx>=0.27.0
7
7
  Requires-Dist: rich>=13.7.0
8
8
  Requires-Dist: prompt_toolkit>=3.0.0
9
9
  Requires-Dist: mcp>=1.2.0
10
+ Requires-Dist: keyring>=24.0.0
11
+ Requires-Dist: watchdog>=4.0.0
10
12
  Dynamic: author
11
13
  Dynamic: requires-dist
12
14
  Dynamic: summary
@@ -2,14 +2,16 @@ from setuptools import setup, find_packages
2
2
 
3
3
  setup(
4
4
  name="nimcode",
5
- version="0.1.12",
5
+ version="0.1.13",
6
6
  packages=find_packages(where="src"),
7
7
  package_dir={"": "src"},
8
8
  install_requires=[
9
9
  "httpx>=0.27.0",
10
10
  "rich>=13.7.0",
11
11
  "prompt_toolkit>=3.0.0",
12
- "mcp>=1.2.0"
12
+ "mcp>=1.2.0",
13
+ "keyring>=24.0.0",
14
+ "watchdog>=4.0.0"
13
15
  ],
14
16
  entry_points={
15
17
  "console_scripts": [
@@ -8,6 +8,7 @@ from .permissions import PermissionEngine, PermissionMode
8
8
  from .config import load_settings, save_global_setting
9
9
  from .mcp_client import MCPManager
10
10
  from .memory import MemoryManager
11
+ from .analytics import AnalyticsEngine
11
12
 
12
13
  logger = logging.getLogger(__name__)
13
14
 
@@ -27,7 +28,7 @@ Available Tools:
27
28
  - Bash: {"tool": "Bash", "args": {"command": "string"}}
28
29
  - Read: {"tool": "Read", "args": {"file_path": "string"}}
29
30
  - Write: {"tool": "Write", "args": {"file_path": "string", "content": "string"}}
30
- - Edit: {"tool": "Edit", "args": {"file_path": "string", "old_string": "string", "new_string": "string"}}
31
+ - Replace: {"tool": "Replace", "args": {"file_path": "string", "replacements": [{"old_string": "exact old", "new_string": "exact new"}]}}
31
32
  - Glob: {"tool": "Glob", "args": {"pattern": "string"}}
32
33
  - Grep: {"tool": "Grep", "args": {"query": "string", "directory": "string"}}
33
34
 
@@ -48,18 +49,42 @@ When you have completely fulfilled the user's request and have no more tools to
48
49
  """
49
50
 
50
51
  class Agent:
51
- def __init__(self, api_key: str, model: str = None, max_turns: int = 30, permission_mode: PermissionMode = PermissionMode.DEFAULT, max_tokens: int = 4000):
52
+ def __init__(self, api_key: str, model: str = None, max_turns: int = 30, permission_mode: PermissionMode = PermissionMode.DEFAULT, max_tokens: int = 100000):
52
53
  # Load global settings
53
54
  self.settings = load_settings()
54
55
  self.model = model or self.settings.get("model", "meta/llama-3.1-70b-instruct")
55
- self.client = NimClient(api_key=api_key, model=self.model)
56
+ self.is_local = self.settings.get("is_local", False)
57
+ self.base_url = self.settings.get("base_url", None)
58
+ self.client = NimClient(api_key=api_key, base_url=self.base_url, model=self.model, is_local=self.is_local)
56
59
 
57
60
  # Initialize MCP Manager
58
61
  self.mcp = MCPManager(self.settings)
62
+ self.analytics = AnalyticsEngine()
63
+ self.memory = MemoryManager(model_name=model, fallback_max_tokens=max_tokens)
59
64
 
60
65
  # Base system prompt
61
66
  final_prompt = SYSTEM_PROMPT + self.mcp.get_system_prompt_additions()
62
67
 
68
+ # Inject Repo Map
69
+ try:
70
+ from .repo_map import RepoMapper
71
+ mapper = RepoMapper(os.getcwd())
72
+ repo_map = mapper.generate_map()
73
+ final_prompt += f"\n\n--- REPOSITORY MAP ---\n{repo_map}\n----------------------\n"
74
+ except Exception as e:
75
+ logger.error(f"Failed to generate repo map: {e}")
76
+
77
+ # Load .nimcoderules if present
78
+ rules_path = os.path.join(os.getcwd(), ".nimcoderules")
79
+ if os.path.exists(rules_path):
80
+ try:
81
+ with open(rules_path, "r", encoding="utf-8") as f:
82
+ rules = f.read()
83
+ final_prompt += f"\n\nPROJECT-SPECIFIC RULES (.nimcoderules):\n{rules}\n"
84
+ logger.info("Loaded .nimcoderules")
85
+ except Exception as e:
86
+ logger.error(f"Failed to load .nimcoderules: {e}")
87
+
63
88
  # Load skills if present
64
89
  skills_dir = os.path.join(os.getcwd(), ".nimcode", "skills")
65
90
  if os.path.exists(skills_dir) and os.path.isdir(skills_dir):
@@ -81,13 +106,18 @@ class Agent:
81
106
  final_prompt += f"\n\nGIT CONTEXT:\nBranch: {branch}\nUncommitted changes:\n{status if status else 'None'}"
82
107
  except Exception as e:
83
108
  logger.error(f"Failed to load git context: {e}")
109
+
110
+ # Repo Map
111
+ repo_map = self._generate_repo_map(os.getcwd())
112
+ if repo_map:
113
+ final_prompt += f"\n\nREPOSITORY MAP:\n{repo_map}"
84
114
 
85
115
  self.messages: List[Dict[str, Any]] = [
86
116
  {"role": "system", "content": final_prompt}
87
117
  ]
88
118
  self.max_turns = max_turns
89
119
  self.permission_engine = PermissionEngine(mode=permission_mode)
90
- self.memory = MemoryManager(max_tokens=max_tokens)
120
+ self.memory = MemoryManager(model_name=model, fallback_max_tokens=max_tokens)
91
121
 
92
122
  def save_history(self):
93
123
  """Saves current conversation to NIMCODE.md"""
@@ -108,6 +138,28 @@ class Agent:
108
138
  except Exception as e:
109
139
  logger.error(f"Failed to load history: {e}")
110
140
 
141
+ def _generate_repo_map(self, cwd: str) -> str:
142
+ """Generates a fast, lightweight file tree map."""
143
+ ignore_dirs = {'.git', 'node_modules', '__pycache__', '.venv', 'venv', 'env', '.nimcode'}
144
+ tree = []
145
+ for root, dirs, files in os.walk(cwd):
146
+ dirs[:] = [d for d in dirs if d not in ignore_dirs and not d.startswith('.')]
147
+ level = root.replace(cwd, '').count(os.sep)
148
+ indent = ' ' * 4 * level
149
+ basename = os.path.basename(root)
150
+ if basename:
151
+ tree.append(f"{indent}{basename}/")
152
+ subindent = ' ' * 4 * (level + 1)
153
+ for f in files:
154
+ if not f.endswith('.pyc') and not f.startswith('.'):
155
+ tree.append(f"{subindent}{f}")
156
+
157
+ # Limit to 500 lines to save context
158
+ if len(tree) > 500:
159
+ tree = tree[:500] + ["... (truncated for context limit)"]
160
+
161
+ return "\n".join(tree)
162
+
111
163
  async def _stream_response(self) -> str:
112
164
  from rich.live import Live
113
165
  from rich.markdown import Markdown
@@ -170,23 +222,16 @@ class Agent:
170
222
  live.update(Markdown(response_text + "\n\n*[yellow]Stream interrupted by user.[/yellow]*", code_theme=code_theme))
171
223
  c.print("\n[yellow]Generation interrupted.[/yellow]")
172
224
 
173
- # Approximate Token Tracker Update
174
- est_tokens = len(response_text) // 4
175
- if not hasattr(self, "session_tokens"):
176
- self.session_tokens = 0
177
- self.session_tokens += est_tokens
225
+ # Analytics Token Tracker Update
226
+ est_prompt_tokens = self.client.count_tokens_approx(self.messages)
227
+ est_completion_tokens = len(response_text) // 4
178
228
 
179
- # Approximate cost based on 70B typical rates ($3/1M tokens)
180
- cost = (self.session_tokens / 1000000) * 3.0
181
- c.print(f"[dim]Output est. tokens: {est_tokens} | Session Cost: ~${cost:.4f}[/dim]")
229
+ self.analytics.log_usage(self.model, est_prompt_tokens, est_completion_tokens)
182
230
 
183
- # Context usage warning
184
- import json
185
- total_context_chars = sum(len(str(m.get("content", ""))) for m in self.messages)
186
- total_est_tokens = total_context_chars // 4
187
- max_context = 128000 # Assume standard Llama-3.1 128k context for now
188
- if total_est_tokens > max_context * 0.8:
189
- c.print("[bold yellow]⚠️ Context window is over 80% full. Consider running /compact or /clear.[/bold yellow]")
231
+ # Display today's cost
232
+ stats = self.analytics.get_summary()
233
+ today_cost = stats["today"]["cost_usd"]
234
+ c.print(f"[dim]Output est. tokens: {est_completion_tokens} | Today's Cost: ~${today_cost:.4f}[/dim]")
190
235
 
191
236
  return response_text
192
237
 
@@ -197,6 +242,10 @@ class Agent:
197
242
  from .tools import ToolRegistry
198
243
  while turn < max_turns:
199
244
  turn += 1
245
+ if self.memory.count_messages_tokens(self.messages) > (self.memory.max_tokens * 0.8):
246
+ logger.info("Context full in headless mode. Distilling memory via LLM...")
247
+ self.messages = await self._distill_memory()
248
+
200
249
  try:
201
250
  response_text = await self.client.chat_one_shot(self.messages)
202
251
  except Exception as e:
@@ -256,7 +305,7 @@ class Agent:
256
305
 
257
306
  # Compact context before calling API
258
307
  from rich.console import Console
259
- if self.memory.count_messages_tokens(self.messages) > self.memory.max_tokens:
308
+ if self.memory.count_messages_tokens(self.messages) > (self.memory.max_tokens * 0.8):
260
309
  logger.info("Context full. Distilling memory via LLM...")
261
310
  Console().print("[dim italic]🧠 Context full. Distilling memory into a summary to save tokens...[/dim italic]")
262
311
  self.messages = await self._distill_memory()
@@ -297,7 +346,7 @@ class Agent:
297
346
  tool_name = tool_call.get("tool", "Unknown")
298
347
 
299
348
  logger.info(f"Checking permissions for tool: {tool_name}")
300
- if not self.permission_engine.check_permission(tool_call):
349
+ if not await self.permission_engine.check_permission(tool_call):
301
350
  self.messages.append({
302
351
  "role": "user",
303
352
  "content": f"User explicitly denied permission to execute {tool_name}. Please choose another approach."
@@ -0,0 +1,68 @@
1
+ import sqlite3
2
+ import os
3
+ import time
4
+
5
+ class AnalyticsEngine:
6
+ def __init__(self, db_path=None):
7
+ if not db_path:
8
+ global_dir = os.path.expanduser("~/.nimcode")
9
+ os.makedirs(global_dir, exist_ok=True)
10
+ self.db_path = os.path.join(global_dir, "analytics.db")
11
+ else:
12
+ self.db_path = db_path
13
+
14
+ self._init_db()
15
+
16
+ def _init_db(self):
17
+ with sqlite3.connect(self.db_path) as conn:
18
+ cursor = conn.cursor()
19
+ cursor.execute('''
20
+ CREATE TABLE IF NOT EXISTS token_usage (
21
+ id INTEGER PRIMARY KEY AUTOINCREMENT,
22
+ timestamp REAL,
23
+ model TEXT,
24
+ prompt_tokens INTEGER,
25
+ completion_tokens INTEGER,
26
+ cost_usd REAL
27
+ )
28
+ ''')
29
+ conn.commit()
30
+
31
+ def log_usage(self, model: str, prompt_tokens: int, completion_tokens: int):
32
+ # Calculate cost based on current NIM pricing (estimated)
33
+ # e.g., $0.50 per 1M prompt, $1.50 per 1M completion for llama3-70b
34
+ cost_usd = (prompt_tokens / 1_000_000) * 0.50 + (completion_tokens / 1_000_000) * 1.50
35
+
36
+ with sqlite3.connect(self.db_path) as conn:
37
+ cursor = conn.cursor()
38
+ cursor.execute(
39
+ "INSERT INTO token_usage (timestamp, model, prompt_tokens, completion_tokens, cost_usd) VALUES (?, ?, ?, ?, ?)",
40
+ (time.time(), model, prompt_tokens, completion_tokens, cost_usd)
41
+ )
42
+ conn.commit()
43
+
44
+ def get_summary(self):
45
+ with sqlite3.connect(self.db_path) as conn:
46
+ cursor = conn.cursor()
47
+
48
+ # Total stats
49
+ cursor.execute("SELECT SUM(prompt_tokens), SUM(completion_tokens), SUM(cost_usd) FROM token_usage")
50
+ total_stats = cursor.fetchone()
51
+
52
+ # Today's stats
53
+ today_start = time.time() - (time.time() % 86400) # Start of UTC day approx
54
+ cursor.execute("SELECT SUM(prompt_tokens), SUM(completion_tokens), SUM(cost_usd) FROM token_usage WHERE timestamp >= ?", (today_start,))
55
+ today_stats = cursor.fetchone()
56
+
57
+ return {
58
+ "total": {
59
+ "prompt_tokens": total_stats[0] or 0,
60
+ "completion_tokens": total_stats[1] or 0,
61
+ "cost_usd": total_stats[2] or 0.0
62
+ },
63
+ "today": {
64
+ "prompt_tokens": today_stats[0] or 0,
65
+ "completion_tokens": today_stats[1] or 0,
66
+ "cost_usd": today_stats[2] or 0.0
67
+ }
68
+ }
@@ -14,13 +14,15 @@ def run_login():
14
14
  console.print("Get your API key from [bold underline blue]https://build.nvidia.com/[/bold underline blue]")
15
15
 
16
16
  import getpass
17
+ import keyring
18
+
17
19
  api_key = getpass.getpass("Enter your NVIDIA NIM API Key: ")
18
20
  if not api_key.strip():
19
21
  console.print("[red]API Key cannot be empty.[/red]")
20
22
  return
21
23
 
22
- save_global_setting("api_key", api_key.strip())
23
- console.print("[green][OK] API Key saved successfully to ~/.nimcode/settings.json[/green]")
24
+ keyring.set_password("nimcode", "api_key", api_key.strip())
25
+ console.print("[green][OK] API Key saved securely to OS Keyring[/green]")
24
26
 
25
27
  def run_doctor():
26
28
  console.print("[bold cyan]NimCode Doctor[/bold cyan] - Diagnostics")
@@ -74,28 +76,53 @@ def main():
74
76
  # Main CLI arguments
75
77
  parser.add_argument("prompt", nargs="?", default=None, help="The task you want NimCode to accomplish. If omitted, starts interactive REPL.")
76
78
  parser.add_argument("--api-key", "-k", default=None, help="NVIDIA NIM API Key. Can also be set via NIM_API_KEY environment variable.")
77
- parser.add_argument("--model", "-m", default="meta/llama-3.1-70b-instruct", help="Model ID to use from NIM.")
79
+ parser.add_argument("--model", "-m", default=None, help="Model ID to use from NIM (e.g. meta/llama-3.1-70b-instruct).")
78
80
  parser.add_argument("--max-turns", "-t", type=int, default=30, help="Maximum number of turns the agent is allowed to run.")
79
81
  parser.add_argument("--permission-mode", "-p", type=PermissionMode, choices=list(PermissionMode), default=PermissionMode.DEFAULT, help="Permission mode for mutating tools.")
80
82
  parser.add_argument("--resume", "-r", action="store_true", help="Resume from the last session stored in NIMCODE.md.")
83
+ parser.add_argument("--stdio", action="store_true", help="Run NimCode as a JSON-RPC server over stdio for IDE integration.")
81
84
 
82
85
  args = parser.parse_args()
83
86
 
84
87
  settings = load_settings()
85
88
 
86
- final_key = args.api_key or os.environ.get("NIM_API_KEY") or settings.get("api_key")
89
+ import keyring
90
+ final_key = args.api_key or os.environ.get("NIM_API_KEY") or keyring.get_password("nimcode", "api_key") or settings.get("api_key")
87
91
  if not final_key:
88
92
  console.print("[yellow]No API Key found. Let's get you set up![/yellow]")
89
93
  run_login()
90
94
  settings = load_settings()
91
- final_key = settings.get("api_key")
95
+ final_key = keyring.get_password("nimcode", "api_key") or settings.get("api_key")
92
96
  if not final_key:
93
97
  console.print("[bold red]API Key is required to use NimCode. Exiting.[/bold red]")
94
98
  sys.exit(1)
95
99
 
100
+ model = args.model or settings.get("model")
101
+ if not model and not args.stdio:
102
+ from rich.prompt import Prompt
103
+ from .nim_client import NimClient
104
+ import asyncio
105
+ client = NimClient(api_key=final_key)
106
+ console.print("[yellow]Fetching available models from NVIDIA NIM...[/yellow]")
107
+ models = asyncio.run(client.get_available_models())
108
+ console.print("\n[bold cyan]Available Models:[/bold cyan]")
109
+ for i, m in enumerate(models):
110
+ console.print(f"[{i+1}] {m}")
111
+
112
+ choice = Prompt.ask("Select a model number", choices=[str(i) for i in range(1, len(models)+1)], default="1")
113
+ model = models[int(choice) - 1]
114
+
115
+ from .config import save_global_setting
116
+ if Prompt.ask("Save this model as default?", choices=["y", "n"], default="y") == "y":
117
+ save_global_setting("model", model)
118
+ console.print(f"[green]Saved {model} as default model.[/green]")
119
+
120
+ if not model:
121
+ model = "meta/llama-3.1-70b-instruct" # ultimate fallback
122
+
96
123
  agent = Agent(
97
124
  api_key=final_key,
98
- model=args.model,
125
+ model=model,
99
126
  max_turns=args.max_turns,
100
127
  permission_mode=args.permission_mode
101
128
  )
@@ -117,6 +144,10 @@ def main():
117
144
  console.print(f"[bold green]Starting NimCode[/bold green] with model [cyan]{args.model}[/cyan]")
118
145
  console.print(f"Task: {args.prompt}")
119
146
  asyncio.run(agent.run(args.prompt))
147
+ elif args.stdio:
148
+ from .stdio_server import StdioServer
149
+ server = StdioServer(agent)
150
+ asyncio.run(server.start())
120
151
  else:
121
152
  from .repl import NimcodeREPL
122
153
  repl = NimcodeREPL(agent)
@@ -21,6 +21,20 @@ def load_settings() -> Dict[str, Any]:
21
21
  settings.update(global_settings)
22
22
  except Exception as e:
23
23
  logger.error(f"Failed to load global settings: {e}")
24
+
25
+ # Scrub legacy api_key if found to ensure it uses keyring
26
+ if "api_key" in settings:
27
+ import keyring
28
+ try:
29
+ if not keyring.get_password("nimcode", "api_key"):
30
+ keyring.set_password("nimcode", "api_key", settings["api_key"])
31
+
32
+ # Remove from file
33
+ del settings["api_key"]
34
+ with open(global_path, "w", encoding="utf-8") as f:
35
+ json.dump(settings, f, indent=4)
36
+ except Exception:
37
+ pass
24
38
 
25
39
  # Local settings
26
40
  local_path = os.path.join(os.getcwd(), ".nimcode", "settings.json")
@@ -0,0 +1,140 @@
1
+ import os
2
+ import json
3
+ import time
4
+ import urllib.request
5
+ import logging
6
+ from typing import List, Dict, Any, Optional
7
+
8
+ logger = logging.getLogger(__name__)
9
+
10
+ class DynamicContextFetcher:
11
+ CACHE_FILE = os.path.expanduser('~/.nimcode/context_map.json')
12
+ CACHE_EXPIRY = 24 * 3600 # 24 hours
13
+ URL = 'https://raw.githubusercontent.com/BerriAI/litellm/main/model_prices_and_context_window.json'
14
+
15
+ @classmethod
16
+ def get_max_tokens(cls, model_name: str, default_limit: int = 8000) -> int:
17
+ if not model_name:
18
+ return default_limit
19
+
20
+ data = cls._load_map()
21
+ if not data:
22
+ return default_limit
23
+
24
+ # Exact match
25
+ for k, v in data.items():
26
+ if model_name.lower() == k.lower():
27
+ return v.get('max_tokens', default_limit) or default_limit
28
+
29
+ # Heuristic matching
30
+ clean_model = model_name.split('/')[-1].lower()
31
+
32
+ matches = []
33
+ for k, v in data.items():
34
+ if clean_model in k.lower():
35
+ tokens = v.get('max_tokens')
36
+ if tokens:
37
+ matches.append(tokens)
38
+
39
+ if matches:
40
+ return max(matches)
41
+
42
+ from .nim_client import NimClient
43
+ return NimClient.get_model_context_length(model_name)
44
+
45
+ @classmethod
46
+ def _load_map(cls) -> Optional[dict]:
47
+ if os.path.exists(cls.CACHE_FILE):
48
+ if time.time() - os.path.getmtime(cls.CACHE_FILE) < cls.CACHE_EXPIRY:
49
+ try:
50
+ with open(cls.CACHE_FILE, 'r', encoding='utf-8') as f:
51
+ return json.load(f)
52
+ except:
53
+ pass
54
+
55
+ # Fetch live
56
+ try:
57
+ req = urllib.request.Request(cls.URL, headers={'User-Agent': 'Mozilla/5.0'})
58
+ with urllib.request.urlopen(req, timeout=5.0) as res:
59
+ data = json.loads(res.read().decode('utf-8'))
60
+ os.makedirs(os.path.dirname(cls.CACHE_FILE), exist_ok=True)
61
+ with open(cls.CACHE_FILE, 'w', encoding='utf-8') as f:
62
+ json.dump(data, f)
63
+ return data
64
+ except Exception as e:
65
+ logger.debug(f"Failed to fetch dynamic context map: {e}")
66
+ return None
67
+
68
+ class MemoryManager:
69
+ def __init__(self, model_name: str = None, fallback_max_tokens: int = 8000):
70
+ # We assume 1 token ~= 4 chars roughly
71
+ self.model_name = model_name
72
+ self.fallback_max_tokens = fallback_max_tokens
73
+
74
+ @property
75
+ def max_tokens(self) -> int:
76
+ if not self.model_name:
77
+ return self.fallback_max_tokens
78
+ return DynamicContextFetcher.get_max_tokens(self.model_name, self.fallback_max_tokens)
79
+
80
+ @staticmethod
81
+ def count_tokens(text: str) -> int:
82
+ """Roughly count tokens in a string."""
83
+ if not text:
84
+ return 0
85
+ return len(text) // 4 + 1
86
+
87
+ @classmethod
88
+ def count_messages_tokens(cls, messages: List[Dict[str, Any]]) -> int:
89
+ total = 0
90
+ for msg in messages:
91
+ content = msg.get("content", "")
92
+ total += cls.count_tokens(content)
93
+ return total
94
+
95
+ def compact_context(self, messages: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
96
+ """
97
+ If messages exceed max_tokens, drops the oldest messages.
98
+ Always keeps the System prompt (first message).
99
+ Always keeps the most recent user prompt.
100
+ """
101
+ total_tokens = self.count_messages_tokens(messages)
102
+ if total_tokens <= self.max_tokens:
103
+ return messages
104
+
105
+ # We need to compact
106
+ if not messages:
107
+ return messages
108
+
109
+ compacted = [messages[0]] # System prompt
110
+ remaining_messages = messages[1:]
111
+
112
+ # We start from the end and add backwards until we hit the limit
113
+ # Reserved tokens for system prompt
114
+ current_tokens = self.count_tokens(messages[0].get("content", ""))
115
+
116
+ kept_messages = []
117
+ for msg in reversed(remaining_messages):
118
+ msg_tokens = self.count_tokens(msg.get("content", ""))
119
+ if current_tokens + msg_tokens > self.max_tokens:
120
+ # If we haven't even kept the most recent message, we MUST keep it and just truncate its content
121
+ if not kept_messages:
122
+ truncated_content = msg.get("content", "")[:(self.max_tokens - current_tokens) * 4]
123
+ kept_messages.insert(0, {"role": msg["role"], "content": truncated_content + "...[TRUNCATED]"})
124
+ break
125
+
126
+ kept_messages.insert(0, msg)
127
+ current_tokens += msg_tokens
128
+
129
+ compacted.extend(kept_messages)
130
+ return compacted
131
+
132
+ @staticmethod
133
+ def log_to_nimcode_md(turn: int, prompt: str, response: str, cwd: str = ".") -> None:
134
+ """Appends the interaction to NIMCODE.md for persistent session history."""
135
+ file_path = os.path.join(cwd, "NIMCODE.md")
136
+ with open(file_path, "a", encoding="utf-8") as f:
137
+ f.write(f"## Turn {turn}\n\n")
138
+ f.write(f"**User**: {prompt}\n\n")
139
+ f.write(f"**Agent**: {response}\n\n")
140
+ f.write("---\n\n")
@@ -8,17 +8,58 @@ from typing import List, Dict, Any, Optional, AsyncGenerator
8
8
  logger = logging.getLogger(__name__)
9
9
 
10
10
  class NimClient:
11
- def __init__(self, api_key: str, base_url: str = "https://integrate.api.nvidia.com/v1", model: str = "meta/llama-3.1-70b-instruct"):
11
+ def __init__(self, api_key: str, base_url: str = None, model: str = None, is_local: bool = False):
12
12
  self.api_key = api_key
13
- self.base_url = base_url.rstrip("/")
14
- self.model = model
13
+ self.is_local = is_local
14
+ self.base_url = (base_url or "https://integrate.api.nvidia.com/v1").rstrip("/")
15
+ self.model = model or "meta/llama-3.1-70b-instruct"
15
16
  self.headers = {
16
17
  "Authorization": f"Bearer {self.api_key}",
17
18
  "Content-Type": "application/json",
18
19
  "Accept": "text/event-stream"
19
20
  }
20
21
 
22
+ @staticmethod
23
+ def get_model_context_length(model_name: str) -> int:
24
+ """Heuristically determines the maximum token context length from the model name."""
25
+ model_lower = model_name.lower()
26
+ if "llama-3.1" in model_lower or "llama-3.2" in model_lower:
27
+ return 128000
28
+ elif "nemotron" in model_lower:
29
+ return 128000 # Nemotron-4 is usually high context
30
+ elif "mixtral" in model_lower:
31
+ return 64000
32
+ elif "phi" in model_lower:
33
+ return 128000
34
+ elif "gemma" in model_lower:
35
+ return 8192
36
+ # Default fallback
37
+ return 32000
38
+
21
39
  async def get_available_models(self) -> List[str]:
40
+ if self.is_local:
41
+ try:
42
+ async with httpx.AsyncClient() as client:
43
+ response = await client.get(f"{self.base_url}/models", timeout=5.0)
44
+ if response.status_code == 200:
45
+ data = response.json()
46
+ return [m["id"] for m in data.get("data", [])]
47
+ except Exception as e:
48
+ logger.error(f"Failed to fetch local models: {e}")
49
+ return ["local-model"]
50
+
51
+ try:
52
+ async with httpx.AsyncClient() as client:
53
+ response = await client.get(f"{self.base_url}/models", headers={"Authorization": f"Bearer {self.api_key}"}, timeout=10.0)
54
+ if response.status_code == 200:
55
+ data = response.json()
56
+ return [m["id"] for m in data.get("data", [])]
57
+ else:
58
+ logger.error(f"Failed to fetch models from API: {response.status_code}")
59
+ except Exception as e:
60
+ logger.error(f"Failed to fetch models: {e}")
61
+
62
+ # Fallback list if API fails
22
63
  return [
23
64
  "meta/llama-3.1-70b-instruct",
24
65
  "meta/llama-3.1-8b-instruct",
@@ -18,7 +18,7 @@ class PermissionEngine:
18
18
  self.mode = mode
19
19
  self.safe_tools = {"Read", "Glob", "Grep"}
20
20
 
21
- def check_permission(self, tool_call: dict) -> bool:
21
+ async def check_permission(self, tool_call: dict) -> bool:
22
22
  """Returns True if permitted, False otherwise."""
23
23
  if self.mode == PermissionMode.BYPASS:
24
24
  return True
@@ -28,9 +28,9 @@ class PermissionEngine:
28
28
  if self.mode == PermissionMode.AUTO and tool_name in self.safe_tools:
29
29
  return True
30
30
 
31
- return self._prompt_user(tool_call)
31
+ return await self._prompt_user(tool_call)
32
32
 
33
- def _prompt_user(self, tool_call: Dict[str, Any]) -> bool:
33
+ async def _prompt_user(self, tool_call: Dict[str, Any]) -> bool:
34
34
  from rich.panel import Panel
35
35
  from rich.prompt import Prompt
36
36
  import json
@@ -46,18 +46,44 @@ class PermissionEngine:
46
46
  content = ""
47
47
  if tool_name == "Bash":
48
48
  content = f"Command:\n{args.get('command')}"
49
- elif tool_name == "Write":
50
- content_lines = len(args.get('content', '').split('\n'))
51
- content = f"File: [bold cyan]{args.get('file_path')}[/bold cyan]\nAction: Overwrite with [bold green]{content_lines} new lines[/bold green]"
52
- elif tool_name in ["Edit", "ASTReplace"]:
49
+ elif tool_name in ["Edit", "ASTReplace", "Replace", "Write"]:
53
50
  import difflib
54
51
  import os
55
52
  file_path = args.get('file_path', '')
53
+ cwd = os.getcwd() # Assumption: running in cwd
54
+ full_path = os.path.join(cwd, file_path)
56
55
 
57
56
  # Default to basic diff summary
58
57
  diff_str = ""
59
58
 
60
- if tool_name == "Edit":
59
+ old_lines = []
60
+ if os.path.exists(full_path):
61
+ try:
62
+ with open(full_path, "r", encoding="utf-8") as f:
63
+ old_lines = f.readlines()
64
+ except:
65
+ pass
66
+
67
+ if tool_name == "Write":
68
+ new_content = args.get('content', '')
69
+ new_lines = new_content.splitlines(keepends=True)
70
+ diff = list(difflib.unified_diff(old_lines, new_lines, fromfile=file_path, tofile=file_path))
71
+ diff_str = "".join(diff)
72
+ elif tool_name == "Replace":
73
+ replacements = args.get("replacements", [])
74
+ try:
75
+ with open(full_path, "r", encoding="utf-8") as f:
76
+ new_content = f.read()
77
+ for rep in replacements:
78
+ old_s = rep.get("old_string", "")
79
+ new_s = rep.get("new_string", "")
80
+ new_content = new_content.replace(old_s, new_s, 1)
81
+ new_lines = new_content.splitlines(keepends=True)
82
+ diff = list(difflib.unified_diff(old_lines, new_lines, fromfile=file_path, tofile=file_path))
83
+ diff_str = "".join(diff)
84
+ except Exception:
85
+ diff_str = "Could not compute diff for Replace."
86
+ elif tool_name == "Edit":
61
87
  old_str = args.get('old_string', '')
62
88
  new_str = args.get('new_string', '')
63
89
  diff = list(difflib.unified_diff(old_str.splitlines(keepends=True), new_str.splitlines(keepends=True), fromfile=file_path, tofile=file_path))