nimcode 0.8.6__tar.gz → 0.8.8__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.8.6 → nimcode-0.8.8}/PKG-INFO +1 -1
- nimcode-0.8.8/src/nimcode/__version__.py +1 -0
- {nimcode-0.8.6 → nimcode-0.8.8}/src/nimcode/agent.py +3 -1
- nimcode-0.8.8/src/nimcode/model_registry.py +129 -0
- {nimcode-0.8.6 → nimcode-0.8.8}/src/nimcode/tools.py +4 -0
- {nimcode-0.8.6 → nimcode-0.8.8}/src/nimcode.egg-info/PKG-INFO +1 -1
- nimcode-0.8.8/tests/test_model_registry.py +90 -0
- nimcode-0.8.6/src/nimcode/__version__.py +0 -1
- nimcode-0.8.6/src/nimcode/model_registry.py +0 -54
- nimcode-0.8.6/tests/test_model_registry.py +0 -38
- {nimcode-0.8.6 → nimcode-0.8.8}/README.md +0 -0
- {nimcode-0.8.6 → nimcode-0.8.8}/setup.cfg +0 -0
- {nimcode-0.8.6 → nimcode-0.8.8}/setup.py +0 -0
- {nimcode-0.8.6 → nimcode-0.8.8}/src/nimcode/__init__.py +0 -0
- {nimcode-0.8.6 → nimcode-0.8.8}/src/nimcode/agents/__init__.py +0 -0
- {nimcode-0.8.6 → nimcode-0.8.8}/src/nimcode/agents/qa_agent.py +0 -0
- {nimcode-0.8.6 → nimcode-0.8.8}/src/nimcode/cli.py +0 -0
- {nimcode-0.8.6 → nimcode-0.8.8}/src/nimcode/config.py +0 -0
- {nimcode-0.8.6 → nimcode-0.8.8}/src/nimcode/lenient_parser.py +0 -0
- {nimcode-0.8.6 → nimcode-0.8.8}/src/nimcode/mcp_client.py +0 -0
- {nimcode-0.8.6 → nimcode-0.8.8}/src/nimcode/memory.py +0 -0
- {nimcode-0.8.6 → nimcode-0.8.8}/src/nimcode/nim_client.py +0 -0
- {nimcode-0.8.6 → nimcode-0.8.8}/src/nimcode/permissions.py +0 -0
- {nimcode-0.8.6 → nimcode-0.8.8}/src/nimcode/plugin_manager.py +0 -0
- {nimcode-0.8.6 → nimcode-0.8.8}/src/nimcode/rag.py +0 -0
- {nimcode-0.8.6 → nimcode-0.8.8}/src/nimcode/repl.py +0 -0
- {nimcode-0.8.6 → nimcode-0.8.8}/src/nimcode/repo_map.py +0 -0
- {nimcode-0.8.6 → nimcode-0.8.8}/src/nimcode/task_manager.py +0 -0
- {nimcode-0.8.6 → nimcode-0.8.8}/src/nimcode/updater.py +0 -0
- {nimcode-0.8.6 → nimcode-0.8.8}/src/nimcode/watcher.py +0 -0
- {nimcode-0.8.6 → nimcode-0.8.8}/src/nimcode.egg-info/SOURCES.txt +0 -0
- {nimcode-0.8.6 → nimcode-0.8.8}/src/nimcode.egg-info/dependency_links.txt +0 -0
- {nimcode-0.8.6 → nimcode-0.8.8}/src/nimcode.egg-info/entry_points.txt +0 -0
- {nimcode-0.8.6 → nimcode-0.8.8}/src/nimcode.egg-info/requires.txt +0 -0
- {nimcode-0.8.6 → nimcode-0.8.8}/src/nimcode.egg-info/top_level.txt +0 -0
- {nimcode-0.8.6 → nimcode-0.8.8}/tests/test_advanced_blockers.py +0 -0
- {nimcode-0.8.6 → nimcode-0.8.8}/tests/test_agent.py +0 -0
- {nimcode-0.8.6 → nimcode-0.8.8}/tests/test_bash_blockers.py +0 -0
- {nimcode-0.8.6 → nimcode-0.8.8}/tests/test_cli.py +0 -0
- {nimcode-0.8.6 → nimcode-0.8.8}/tests/test_config.py +0 -0
- {nimcode-0.8.6 → nimcode-0.8.8}/tests/test_lazy_blockers.py +0 -0
- {nimcode-0.8.6 → nimcode-0.8.8}/tests/test_lenient_parser.py +0 -0
- {nimcode-0.8.6 → nimcode-0.8.8}/tests/test_mcp_client.py +0 -0
- {nimcode-0.8.6 → nimcode-0.8.8}/tests/test_memory.py +0 -0
- {nimcode-0.8.6 → nimcode-0.8.8}/tests/test_nim_client.py +0 -0
- {nimcode-0.8.6 → nimcode-0.8.8}/tests/test_permissions.py +0 -0
- {nimcode-0.8.6 → nimcode-0.8.8}/tests/test_plan_injection.py +0 -0
- {nimcode-0.8.6 → nimcode-0.8.8}/tests/test_plugin_manager.py +0 -0
- {nimcode-0.8.6 → nimcode-0.8.8}/tests/test_qa_agent.py +0 -0
- {nimcode-0.8.6 → nimcode-0.8.8}/tests/test_repl.py +0 -0
- {nimcode-0.8.6 → nimcode-0.8.8}/tests/test_repl_extra.py +0 -0
- {nimcode-0.8.6 → nimcode-0.8.8}/tests/test_repl_fixes.py +0 -0
- {nimcode-0.8.6 → nimcode-0.8.8}/tests/test_repl_trust.py +0 -0
- {nimcode-0.8.6 → nimcode-0.8.8}/tests/test_repo_map.py +0 -0
- {nimcode-0.8.6 → nimcode-0.8.8}/tests/test_task_manager.py +0 -0
- {nimcode-0.8.6 → nimcode-0.8.8}/tests/test_tools.py +0 -0
- {nimcode-0.8.6 → nimcode-0.8.8}/tests/test_updater.py +0 -0
- {nimcode-0.8.6 → nimcode-0.8.8}/tests/test_version.py +0 -0
|
@@ -0,0 +1 @@
|
|
|
1
|
+
__version__ = "0.8.8"
|
|
@@ -277,7 +277,9 @@ class Agent:
|
|
|
277
277
|
self.max_turns = max_turns if max_turns is not None else (0 if cfg_max_turns == 0 else cfg_max_turns)
|
|
278
278
|
|
|
279
279
|
cfg_max_tokens = self.settings.get("max_tokens", 120000)
|
|
280
|
-
|
|
280
|
+
from .model_registry import get_context_window
|
|
281
|
+
model_limit = get_context_window(self.model)
|
|
282
|
+
effective_max_tokens = max_tokens if max_tokens is not None else min(cfg_max_tokens, model_limit)
|
|
281
283
|
self.memory = MemoryManager(max_tokens=effective_max_tokens)
|
|
282
284
|
|
|
283
285
|
self.permission_engine = PermissionEngine(mode=permission_mode)
|
|
@@ -0,0 +1,129 @@
|
|
|
1
|
+
import logging
|
|
2
|
+
import os
|
|
3
|
+
import json
|
|
4
|
+
import urllib.request
|
|
5
|
+
import threading
|
|
6
|
+
import time
|
|
7
|
+
|
|
8
|
+
logger = logging.getLogger(__name__)
|
|
9
|
+
|
|
10
|
+
# Known NVIDIA NIM models and their context window sizes (in tokens)
|
|
11
|
+
# This acts as a fallback + enrichment layer when API is unavailable
|
|
12
|
+
MODEL_CONTEXT_WINDOWS = {
|
|
13
|
+
# Llama 3.1 family
|
|
14
|
+
"meta/llama-3.1-8b-instruct": 128000,
|
|
15
|
+
"meta/llama-3.3-70b-instruct": 128000,
|
|
16
|
+
"meta/llama-3.1-405b-instruct": 128000,
|
|
17
|
+
# Llama 3.2 family
|
|
18
|
+
"meta/llama-3.2-1b-instruct": 128000,
|
|
19
|
+
"meta/llama-3.2-3b-instruct": 128000,
|
|
20
|
+
"meta/llama-3.2-11b-vision-instruct": 128000,
|
|
21
|
+
"meta/llama-3.2-90b-vision-instruct": 128000,
|
|
22
|
+
# Llama 3.3 family
|
|
23
|
+
"meta/llama-3.3-70b-instruct": 128000,
|
|
24
|
+
# Nemotron
|
|
25
|
+
"nvidia/nemotron-4-340b-instruct": 4096,
|
|
26
|
+
"nvidia/llama-3.1-nemotron-70b-instruct": 128000,
|
|
27
|
+
# Mistral
|
|
28
|
+
"mistralai/mixtral-8x22b-instruct-v0.1": 65536,
|
|
29
|
+
"mistralai/mistral-large-2-instruct": 128000,
|
|
30
|
+
# DeepSeek
|
|
31
|
+
"deepseek-ai/deepseek-r1": 128000,
|
|
32
|
+
"deepseek-ai/deepseek-v4-pro": 128000,
|
|
33
|
+
# Qwen
|
|
34
|
+
"qwen/qwen2.5-72b-instruct": 128000,
|
|
35
|
+
"qwen/qwq-32b": 128000,
|
|
36
|
+
}
|
|
37
|
+
|
|
38
|
+
# Default context window for unknown models
|
|
39
|
+
DEFAULT_CONTEXT_WINDOW = 8192
|
|
40
|
+
|
|
41
|
+
# Hardcoded fallback model list (used when API is unreachable)
|
|
42
|
+
FALLBACK_MODELS = [
|
|
43
|
+
"meta/llama-3.3-70b-instruct",
|
|
44
|
+
"meta/llama-3.1-8b-instruct",
|
|
45
|
+
"meta/llama-3.1-405b-instruct",
|
|
46
|
+
"meta/llama-3.3-70b-instruct",
|
|
47
|
+
"nvidia/llama-3.1-nemotron-70b-instruct",
|
|
48
|
+
"deepseek-ai/deepseek-r1",
|
|
49
|
+
"mistralai/mixtral-8x22b-instruct-v0.1",
|
|
50
|
+
"qwen/qwen2.5-72b-instruct",
|
|
51
|
+
]
|
|
52
|
+
|
|
53
|
+
_CACHE_FILE = os.path.expanduser("~/.nimcode/model_contexts.json")
|
|
54
|
+
_CACHE_MAX_AGE_DAYS = 7
|
|
55
|
+
_DYNAMIC_CACHE = {}
|
|
56
|
+
_CACHE_LOADED = False
|
|
57
|
+
_UPDATE_THREAD_STARTED = False
|
|
58
|
+
|
|
59
|
+
def _fetch_and_cache_models():
|
|
60
|
+
"""Background task to fetch latest model info from LiteLLM."""
|
|
61
|
+
url = "https://raw.githubusercontent.com/BerriAI/litellm/main/model_prices_and_context_window.json"
|
|
62
|
+
try:
|
|
63
|
+
req = urllib.request.Request(url, headers={'User-Agent': 'Mozilla/5.0 Nimcode/1.0'})
|
|
64
|
+
with urllib.request.urlopen(req, timeout=15) as response:
|
|
65
|
+
data = json.loads(response.read().decode('utf-8'))
|
|
66
|
+
|
|
67
|
+
processed = {}
|
|
68
|
+
for k, v in data.items():
|
|
69
|
+
if isinstance(v, dict) and "max_tokens" in v and isinstance(v["max_tokens"], int):
|
|
70
|
+
processed[k] = v["max_tokens"]
|
|
71
|
+
basename = k.split("/")[-1]
|
|
72
|
+
if basename not in processed or processed[basename] < v["max_tokens"]:
|
|
73
|
+
processed[basename] = v["max_tokens"]
|
|
74
|
+
|
|
75
|
+
os.makedirs(os.path.dirname(_CACHE_FILE), exist_ok=True)
|
|
76
|
+
with open(_CACHE_FILE, "w", encoding="utf-8") as f:
|
|
77
|
+
json.dump(processed, f, indent=4)
|
|
78
|
+
|
|
79
|
+
global _DYNAMIC_CACHE
|
|
80
|
+
_DYNAMIC_CACHE.update(processed)
|
|
81
|
+
logger.debug(f"Successfully cached {len(processed)} models from LiteLLM.")
|
|
82
|
+
except Exception as e:
|
|
83
|
+
logger.debug(f"Failed to fetch model contexts: {e}")
|
|
84
|
+
|
|
85
|
+
def _load_or_update_cache():
|
|
86
|
+
"""Loads cache from disk, spawns background update if old/missing."""
|
|
87
|
+
global _CACHE_LOADED, _DYNAMIC_CACHE, _UPDATE_THREAD_STARTED
|
|
88
|
+
if _CACHE_LOADED:
|
|
89
|
+
return
|
|
90
|
+
|
|
91
|
+
_CACHE_LOADED = True
|
|
92
|
+
needs_update = True
|
|
93
|
+
|
|
94
|
+
if os.path.exists(_CACHE_FILE):
|
|
95
|
+
try:
|
|
96
|
+
with open(_CACHE_FILE, "r", encoding="utf-8") as f:
|
|
97
|
+
_DYNAMIC_CACHE = json.load(f)
|
|
98
|
+
|
|
99
|
+
# Check age
|
|
100
|
+
mtime = os.path.getmtime(_CACHE_FILE)
|
|
101
|
+
age_days = (time.time() - mtime) / (24 * 3600)
|
|
102
|
+
if age_days < _CACHE_MAX_AGE_DAYS:
|
|
103
|
+
needs_update = False
|
|
104
|
+
except Exception as e:
|
|
105
|
+
logger.debug(f"Error reading cache file: {e}")
|
|
106
|
+
|
|
107
|
+
if needs_update and not _UPDATE_THREAD_STARTED:
|
|
108
|
+
_UPDATE_THREAD_STARTED = True
|
|
109
|
+
t = threading.Thread(target=_fetch_and_cache_models, daemon=True)
|
|
110
|
+
t.start()
|
|
111
|
+
|
|
112
|
+
def get_context_window(model_id: str) -> int:
|
|
113
|
+
"""Returns the context window size for a given model ID."""
|
|
114
|
+
_load_or_update_cache()
|
|
115
|
+
|
|
116
|
+
# 1. Check hardcoded exact match
|
|
117
|
+
if model_id in MODEL_CONTEXT_WINDOWS:
|
|
118
|
+
return MODEL_CONTEXT_WINDOWS[model_id]
|
|
119
|
+
|
|
120
|
+
# 2. Check dynamic cache exact match
|
|
121
|
+
if model_id in _DYNAMIC_CACHE:
|
|
122
|
+
return _DYNAMIC_CACHE[model_id]
|
|
123
|
+
|
|
124
|
+
# 3. Check dynamic cache basename match
|
|
125
|
+
basename = model_id.split("/")[-1]
|
|
126
|
+
if basename in _DYNAMIC_CACHE:
|
|
127
|
+
return _DYNAMIC_CACHE[basename]
|
|
128
|
+
|
|
129
|
+
return DEFAULT_CONTEXT_WINDOW
|
|
@@ -306,6 +306,8 @@ class ToolRegistry:
|
|
|
306
306
|
|
|
307
307
|
@staticmethod
|
|
308
308
|
def _check_lazy_code(content: str, file_path: str):
|
|
309
|
+
if file_path.lower().endswith(('.md', '.txt', '.json', '.yaml', '.yml', '.csv', '.ini', '.toml')):
|
|
310
|
+
return
|
|
309
311
|
import re
|
|
310
312
|
# Broad regex for // TO DO, // TODO, # TODO, // FIXME, etc. with flexible spaces
|
|
311
313
|
if re.search(r'(?i)(?://|#)\s*(?:to\s*do|fix\s*me)', content):
|
|
@@ -680,6 +682,8 @@ class ToolRegistry:
|
|
|
680
682
|
|
|
681
683
|
@staticmethod
|
|
682
684
|
def _check_dependency_hallucination(content: str, file_path: str, cwd: str):
|
|
685
|
+
if file_path.lower().endswith(('.md', '.txt', '.json', '.yaml', '.yml', '.csv', '.ini', '.toml')):
|
|
686
|
+
return
|
|
683
687
|
import re
|
|
684
688
|
pass
|
|
685
689
|
import json
|
|
@@ -0,0 +1,90 @@
|
|
|
1
|
+
import os
|
|
2
|
+
import json
|
|
3
|
+
import time
|
|
4
|
+
import pytest
|
|
5
|
+
from unittest.mock import patch, mock_open, MagicMock
|
|
6
|
+
|
|
7
|
+
import nimcode.model_registry as mr
|
|
8
|
+
|
|
9
|
+
@pytest.fixture(autouse=True)
|
|
10
|
+
def reset_registry_state():
|
|
11
|
+
"""Reset the module-level state before each test."""
|
|
12
|
+
mr._CACHE_LOADED = False
|
|
13
|
+
mr._UPDATE_THREAD_STARTED = False
|
|
14
|
+
mr._DYNAMIC_CACHE = {}
|
|
15
|
+
yield
|
|
16
|
+
|
|
17
|
+
def test_get_context_window_hardcoded():
|
|
18
|
+
# Known model in hardcoded dictionary
|
|
19
|
+
assert mr.get_context_window("meta/llama-3.1-8b-instruct") == 128000
|
|
20
|
+
|
|
21
|
+
def test_get_context_window_unknown():
|
|
22
|
+
# Completely unknown model
|
|
23
|
+
assert mr.get_context_window("unknown/model-123") == mr.DEFAULT_CONTEXT_WINDOW
|
|
24
|
+
|
|
25
|
+
@patch("os.path.exists")
|
|
26
|
+
@patch("os.path.getmtime")
|
|
27
|
+
@patch("builtins.open", new_callable=mock_open, read_data='{"provider/test-model": 64000, "test-model": 64000}')
|
|
28
|
+
def test_get_context_window_from_cache(mock_file, mock_getmtime, mock_exists):
|
|
29
|
+
mock_exists.return_value = True
|
|
30
|
+
mock_getmtime.return_value = time.time() # Very recent cache
|
|
31
|
+
|
|
32
|
+
# Force reload
|
|
33
|
+
mr._load_or_update_cache()
|
|
34
|
+
|
|
35
|
+
assert mr.get_context_window("provider/test-model") == 64000
|
|
36
|
+
assert mr.get_context_window("other_provider/test-model") == 64000
|
|
37
|
+
assert mr._UPDATE_THREAD_STARTED is False # Cache is fresh, no thread spawned
|
|
38
|
+
|
|
39
|
+
@patch("os.path.exists")
|
|
40
|
+
@patch("threading.Thread")
|
|
41
|
+
def test_get_context_window_spawns_thread_if_no_cache(mock_thread, mock_exists):
|
|
42
|
+
mock_exists.return_value = False
|
|
43
|
+
|
|
44
|
+
# Should spawn thread to fetch cache
|
|
45
|
+
mr.get_context_window("some-model")
|
|
46
|
+
|
|
47
|
+
assert mr._UPDATE_THREAD_STARTED is True
|
|
48
|
+
mock_thread.assert_called_once()
|
|
49
|
+
mock_thread.return_value.start.assert_called_once()
|
|
50
|
+
|
|
51
|
+
@patch("os.path.exists")
|
|
52
|
+
@patch("os.path.getmtime")
|
|
53
|
+
@patch("threading.Thread")
|
|
54
|
+
@patch("builtins.open", new_callable=mock_open, read_data='{"old-model": 1000}')
|
|
55
|
+
def test_get_context_window_spawns_thread_if_old_cache(mock_file, mock_thread, mock_getmtime, mock_exists):
|
|
56
|
+
mock_exists.return_value = True
|
|
57
|
+
mock_getmtime.return_value = time.time() - (8 * 24 * 3600) # 8 days old
|
|
58
|
+
|
|
59
|
+
mr.get_context_window("old-model")
|
|
60
|
+
|
|
61
|
+
assert mr._UPDATE_THREAD_STARTED is True
|
|
62
|
+
mock_thread.assert_called_once()
|
|
63
|
+
mock_thread.return_value.start.assert_called_once()
|
|
64
|
+
|
|
65
|
+
@patch("urllib.request.urlopen")
|
|
66
|
+
def test_fetch_and_cache_models(mock_urlopen, tmp_path):
|
|
67
|
+
mock_response = MagicMock()
|
|
68
|
+
mock_response.read.return_value = b'''
|
|
69
|
+
{
|
|
70
|
+
"provider/model-A": {"max_tokens": 1234},
|
|
71
|
+
"provider/model-B": {"max_tokens": 5678, "other": "ignored"},
|
|
72
|
+
"provider/model-C": {"ignored": 10}
|
|
73
|
+
}
|
|
74
|
+
'''
|
|
75
|
+
mock_response.__enter__.return_value = mock_response
|
|
76
|
+
mock_urlopen.return_value = mock_response
|
|
77
|
+
|
|
78
|
+
with patch("nimcode.model_registry._CACHE_FILE", str(tmp_path / "cache.json")):
|
|
79
|
+
mr._fetch_and_cache_models()
|
|
80
|
+
|
|
81
|
+
assert "provider/model-A" in mr._DYNAMIC_CACHE
|
|
82
|
+
assert mr._DYNAMIC_CACHE["provider/model-A"] == 1234
|
|
83
|
+
assert mr._DYNAMIC_CACHE["model-A"] == 1234
|
|
84
|
+
assert "provider/model-C" not in mr._DYNAMIC_CACHE
|
|
85
|
+
|
|
86
|
+
# Verify file is written
|
|
87
|
+
with open(str(tmp_path / "cache.json"), "r") as f:
|
|
88
|
+
data = json.load(f)
|
|
89
|
+
assert data["provider/model-A"] == 1234
|
|
90
|
+
assert data["model-A"] == 1234
|
|
@@ -1 +0,0 @@
|
|
|
1
|
-
__version__ = "0.8.6"
|
|
@@ -1,54 +0,0 @@
|
|
|
1
|
-
"""
|
|
2
|
-
Model registry with context window sizes and metadata.
|
|
3
|
-
Used for auto-compact thresholds and dynamic model listing.
|
|
4
|
-
"""
|
|
5
|
-
import logging
|
|
6
|
-
|
|
7
|
-
logger = logging.getLogger(__name__)
|
|
8
|
-
|
|
9
|
-
# Known NVIDIA NIM models and their context window sizes (in tokens)
|
|
10
|
-
# This acts as a fallback + enrichment layer when API is unavailable
|
|
11
|
-
MODEL_CONTEXT_WINDOWS = {
|
|
12
|
-
# Llama 3.1 family
|
|
13
|
-
"meta/llama-3.1-8b-instruct": 128000,
|
|
14
|
-
"meta/llama-3.3-70b-instruct": 128000,
|
|
15
|
-
"meta/llama-3.1-405b-instruct": 128000,
|
|
16
|
-
# Llama 3.2 family
|
|
17
|
-
"meta/llama-3.2-1b-instruct": 128000,
|
|
18
|
-
"meta/llama-3.2-3b-instruct": 128000,
|
|
19
|
-
"meta/llama-3.2-11b-vision-instruct": 128000,
|
|
20
|
-
"meta/llama-3.2-90b-vision-instruct": 128000,
|
|
21
|
-
# Llama 3.3 family
|
|
22
|
-
"meta/llama-3.3-70b-instruct": 128000,
|
|
23
|
-
# Nemotron
|
|
24
|
-
"nvidia/nemotron-4-340b-instruct": 4096,
|
|
25
|
-
"nvidia/llama-3.1-nemotron-70b-instruct": 128000,
|
|
26
|
-
# Mistral
|
|
27
|
-
"mistralai/mixtral-8x22b-instruct-v0.1": 65536,
|
|
28
|
-
"mistralai/mistral-large-2-instruct": 128000,
|
|
29
|
-
# DeepSeek
|
|
30
|
-
"deepseek-ai/deepseek-r1": 128000,
|
|
31
|
-
# Qwen
|
|
32
|
-
"qwen/qwen2.5-72b-instruct": 128000,
|
|
33
|
-
"qwen/qwq-32b": 128000,
|
|
34
|
-
}
|
|
35
|
-
|
|
36
|
-
# Default context window for unknown models
|
|
37
|
-
DEFAULT_CONTEXT_WINDOW = 8192
|
|
38
|
-
|
|
39
|
-
# Hardcoded fallback model list (used when API is unreachable)
|
|
40
|
-
FALLBACK_MODELS = [
|
|
41
|
-
"meta/llama-3.3-70b-instruct",
|
|
42
|
-
"meta/llama-3.1-8b-instruct",
|
|
43
|
-
"meta/llama-3.1-405b-instruct",
|
|
44
|
-
"meta/llama-3.3-70b-instruct",
|
|
45
|
-
"nvidia/llama-3.1-nemotron-70b-instruct",
|
|
46
|
-
"deepseek-ai/deepseek-r1",
|
|
47
|
-
"mistralai/mixtral-8x22b-instruct-v0.1",
|
|
48
|
-
"qwen/qwen2.5-72b-instruct",
|
|
49
|
-
]
|
|
50
|
-
|
|
51
|
-
|
|
52
|
-
def get_context_window(model_id: str) -> int:
|
|
53
|
-
"""Returns the context window size for a given model ID."""
|
|
54
|
-
return MODEL_CONTEXT_WINDOWS.get(model_id, DEFAULT_CONTEXT_WINDOW)
|
|
@@ -1,38 +0,0 @@
|
|
|
1
|
-
"""Tests for the model_registry module."""
|
|
2
|
-
import pytest
|
|
3
|
-
from nimcode.model_registry import (
|
|
4
|
-
MODEL_CONTEXT_WINDOWS,
|
|
5
|
-
FALLBACK_MODELS,
|
|
6
|
-
DEFAULT_CONTEXT_WINDOW,
|
|
7
|
-
get_context_window,
|
|
8
|
-
)
|
|
9
|
-
|
|
10
|
-
|
|
11
|
-
def test_get_context_window_known_model():
|
|
12
|
-
"""Known models should return their specific context window size."""
|
|
13
|
-
assert get_context_window("meta/llama-3.3-70b-instruct") == 128000
|
|
14
|
-
assert get_context_window("nvidia/nemotron-4-340b-instruct") == 4096
|
|
15
|
-
assert get_context_window("mistralai/mixtral-8x22b-instruct-v0.1") == 65536
|
|
16
|
-
|
|
17
|
-
|
|
18
|
-
def test_get_context_window_unknown_model():
|
|
19
|
-
"""Unknown models should return the default context window."""
|
|
20
|
-
assert get_context_window("some/unknown-model") == DEFAULT_CONTEXT_WINDOW
|
|
21
|
-
assert get_context_window("") == DEFAULT_CONTEXT_WINDOW
|
|
22
|
-
|
|
23
|
-
|
|
24
|
-
def test_fallback_models_not_empty():
|
|
25
|
-
"""Fallback model list should have entries."""
|
|
26
|
-
assert len(FALLBACK_MODELS) > 0
|
|
27
|
-
|
|
28
|
-
|
|
29
|
-
def test_fallback_models_all_in_context_windows():
|
|
30
|
-
"""All fallback models should have a known context window entry."""
|
|
31
|
-
for model in FALLBACK_MODELS:
|
|
32
|
-
assert model in MODEL_CONTEXT_WINDOWS, f"Fallback model {model} missing from context windows"
|
|
33
|
-
|
|
34
|
-
|
|
35
|
-
def test_context_windows_positive():
|
|
36
|
-
"""All context window values should be positive integers."""
|
|
37
|
-
for model, ctx in MODEL_CONTEXT_WINDOWS.items():
|
|
38
|
-
assert isinstance(ctx, int) and ctx > 0, f"Invalid context window for {model}: {ctx}"
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|