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.
- {nimcode-0.1.12 → nimcode-0.1.13}/PKG-INFO +3 -1
- {nimcode-0.1.12 → nimcode-0.1.13}/setup.py +4 -2
- {nimcode-0.1.12 → nimcode-0.1.13}/src/nimcode/agent.py +70 -21
- nimcode-0.1.13/src/nimcode/analytics.py +68 -0
- {nimcode-0.1.12 → nimcode-0.1.13}/src/nimcode/cli.py +37 -6
- {nimcode-0.1.12 → nimcode-0.1.13}/src/nimcode/config.py +14 -0
- nimcode-0.1.13/src/nimcode/memory.py +140 -0
- {nimcode-0.1.12 → nimcode-0.1.13}/src/nimcode/nim_client.py +44 -3
- {nimcode-0.1.12 → nimcode-0.1.13}/src/nimcode/permissions.py +34 -8
- nimcode-0.1.13/src/nimcode/plugin_manager.py +53 -0
- {nimcode-0.1.12 → nimcode-0.1.13}/src/nimcode/repl.py +247 -87
- nimcode-0.1.13/src/nimcode/repo_map.py +53 -0
- nimcode-0.1.13/src/nimcode/secret_scanner.py +48 -0
- nimcode-0.1.13/src/nimcode/stdio_server.py +134 -0
- nimcode-0.1.13/src/nimcode/swarm.py +42 -0
- {nimcode-0.1.12 → nimcode-0.1.13}/src/nimcode/tools.py +151 -93
- nimcode-0.1.13/src/nimcode/updater.py +34 -0
- nimcode-0.1.13/src/nimcode/watcher.py +81 -0
- {nimcode-0.1.12 → nimcode-0.1.13}/src/nimcode.egg-info/PKG-INFO +3 -1
- {nimcode-0.1.12 → nimcode-0.1.13}/src/nimcode.egg-info/SOURCES.txt +15 -1
- {nimcode-0.1.12 → nimcode-0.1.13}/src/nimcode.egg-info/requires.txt +2 -0
- {nimcode-0.1.12 → nimcode-0.1.13}/tests/test_agent.py +7 -6
- nimcode-0.1.13/tests/test_analytics.py +22 -0
- {nimcode-0.1.12 → nimcode-0.1.13}/tests/test_cli.py +19 -11
- {nimcode-0.1.12 → nimcode-0.1.13}/tests/test_memory.py +4 -4
- nimcode-0.1.13/tests/test_nimcoderules.py +19 -0
- {nimcode-0.1.12 → nimcode-0.1.13}/tests/test_permissions.py +31 -21
- nimcode-0.1.13/tests/test_plugin_manager.py +35 -0
- {nimcode-0.1.12 → nimcode-0.1.13}/tests/test_repl.py +1 -0
- nimcode-0.1.13/tests/test_repo_map.py +22 -0
- nimcode-0.1.13/tests/test_secret_scanner.py +31 -0
- {nimcode-0.1.12 → nimcode-0.1.13}/tests/test_tools.py +48 -38
- nimcode-0.1.13/tests/test_updater.py +39 -0
- nimcode-0.1.12/src/nimcode/memory.py +0 -69
- {nimcode-0.1.12 → nimcode-0.1.13}/README.md +0 -0
- {nimcode-0.1.12 → nimcode-0.1.13}/setup.cfg +0 -0
- {nimcode-0.1.12 → nimcode-0.1.13}/src/nimcode/__init__.py +0 -0
- {nimcode-0.1.12 → nimcode-0.1.13}/src/nimcode/lenient_parser.py +0 -0
- {nimcode-0.1.12 → nimcode-0.1.13}/src/nimcode/mcp_client.py +0 -0
- {nimcode-0.1.12 → nimcode-0.1.13}/src/nimcode.egg-info/dependency_links.txt +0 -0
- {nimcode-0.1.12 → nimcode-0.1.13}/src/nimcode.egg-info/entry_points.txt +0 -0
- {nimcode-0.1.12 → nimcode-0.1.13}/src/nimcode.egg-info/top_level.txt +0 -0
- {nimcode-0.1.12 → nimcode-0.1.13}/tests/test_config.py +0 -0
- {nimcode-0.1.12 → nimcode-0.1.13}/tests/test_lenient_parser.py +0 -0
- {nimcode-0.1.12 → nimcode-0.1.13}/tests/test_mcp_client.py +0 -0
- {nimcode-0.1.12 → nimcode-0.1.13}/tests/test_nim_client.py +0 -0
- {nimcode-0.1.12 → nimcode-0.1.13}/tests/test_repl_extra.py +0 -0
- {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.
|
|
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.
|
|
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
|
-
-
|
|
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 =
|
|
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.
|
|
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(
|
|
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
|
-
#
|
|
174
|
-
|
|
175
|
-
|
|
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
|
-
|
|
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
|
-
#
|
|
184
|
-
|
|
185
|
-
|
|
186
|
-
|
|
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
|
-
|
|
23
|
-
console.print("[green][OK] API Key saved
|
|
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=
|
|
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
|
-
|
|
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=
|
|
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 =
|
|
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.
|
|
14
|
-
self.
|
|
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
|
|
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
|
-
|
|
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))
|