Merge pull request #123 from SaintsSec/vincent/cleanup

Code cleanup and improved token and config management

chips/SSG/navi_chip_creator.py

@@ -1,10 +1,12 @@
-import navi_internal
-import sys
 import os
 import re
 import subprocess  # nosec
+import sys
 import webbrowser
+
 from colorama import Fore
+
+import navi_internal
 from navi_shell import restart_navi
 
 command: str = "chip-create"

chips/SSG/navi_chip_installer.py

@@ -1,13 +1,14 @@
 import os
-import sys
-import requests
-import zipfile
 import shutil
 import subprocess  # nosec
+import sys
 import uuid
-import navi_internal
+import zipfile
 
+import requests
 from colorama import Fore
+
+import navi_internal
 from navi import get_parameters
 from navi_shell import restart_navi
 
@@ -328,9 +329,9 @@ def update_chip(chip_name: str) -> None:
 
 def help_text() -> None:
     navi.print_message("Chip Manager\n"
-                  "chips [install | uninstall | search | update] [app/query]\n\n"
-                  "List currently installed chips\n"
-                  "chips list")
+                       "chips [install | uninstall | search | update] [app/query]\n\n"
+                       "List currently installed chips\n"
+                       "chips list")
 
 
 def run(arguments=None) -> None:

chips/SSG/navi_memories.py

@@ -1,6 +1,7 @@
-import navi_internal
 import os
 
+import navi_internal
+
 command: str = "memory"
 use: str = "Manage chat memory sessions."
 aliases: list = ['session', 'memory']
@@ -22,13 +23,15 @@
 
 help_params: tuple = ('-help', '-h')
 
+
 def print_params() -> None:
     """Print available parameters and descriptions."""
     print(f"{'Parameter':<20} | {'Description'}")
     print("-" * 50)
     for param, description in params.items():
         print(f"{param:<20} | {description}")
 
+
 def list_sessions(active_session) -> None:
     sessions = os.listdir(memory_dir)
     print("Available Sessions:")
@@ -39,12 +42,14 @@ def list_sessions(active_session) -> None:
         else:
             print(f"  {session_name}")
 
+
 def does_session_exist(session_name) -> bool:
     if not os.path.exists(os.path.join(memory_dir, f"{session_name}.json")):
         print(f"Session '{session_name}' does not exist.")
         return False
     return True
 
+
 def run(arguments=None) -> None:
     navi_instance = navi_internal.navi_instance
     active_session = navi_instance.get_active_session()
@@ -74,4 +79,4 @@ def run(arguments=None) -> None:
         case x if x[0] in help_params:
             print_params()
         case _:
-            navi_instance.print_message("Invalid arguments. Use '-help' for more information.")
+            navi_instance.print_message("Invalid arguments. Use '-help' for more information.")

chips/SSG/navi_settings.py

@@ -34,7 +34,7 @@ def modify_config(key, value):
                 file.write(line)
                 continue
 
-            current_key, current_value = line.strip().split("=", 1)
+            current_key, _ = line.strip().split("=", 1)
             if current_key == key:
                 file.write(f"{key}={value}\n")
                 modified = True
@@ -69,14 +69,28 @@ def read_config(path_to_config):
 
 
 def settings_init():
+    import getpass
+
+    default_config = read_config(default_config_path)
+
+    # Check if the config file exists
     if os.path.exists(config_path):
-        return read_config(config_path)
+        # Read the existing config
+        user_config = read_config(config_path)
+        # Add any missing keys from the default config
+        for key, default_value in default_config.items():
+            if key not in user_config:
+                print(f"Adding missing key: {key} with default value: {default_value}")
+                user_config[key] = default_value
+                modify_config(key, default_value)
+        return user_config
     else:
-        import getpass
-        default_config = read_config(default_config_path)
+        # If the config file doesn't exist, create it using defaults
+        print("Config file not found. Creating a new one with default settings.")
         create_config(default_config)
         modify_config("username", getpass.getuser())
-        return read_config(config_path)
+
+    return read_config(config_path)
 
 
 def run(arguments=None):

chips/SSG/navi_spec.py

@@ -1,10 +1,12 @@
 #!/bin/python3
 import platform
+import socket
+from datetime import datetime
+
 import psutil
 import requests
-import socket
+
 import navi_internal
-from datetime import datetime
 
 # Navi Command System Variables
 command = "navi_specs"
@@ -78,7 +80,8 @@ def run(arguments=None):
 
     # Main functions
     response_message = navi_instance.llm_chat(
-        "Give me a really simple quip about getting my systems specs. Dont include commands or references to operating systems.", True)
+        "Give me a really simple quip about getting my systems specs. Dont include commands or references to operating systems.",
+        True)
     clean_text = str(response_message).replace("(", "").replace(")", "").replace(", 200", "").replace("\"", "").replace(
         "\\n", "")
     output = clean_text + "\n"

chips/SSG/navi_system.py

@@ -1,7 +1,8 @@
 #!/bin/python3
+import shlex
 import subprocess  # nosec
+
 import navi_internal
-import shlex
 from navi import get_command_path
 
 command = "navi_sys"
@@ -13,7 +14,8 @@ def run(arguments=None):
     navi_command = str(arguments).replace("TERMINAL OUTPUT", "", 1).strip()
     base_command = navi_command.split()[0]
     if get_command_path(base_command) is not None:
-        navi_instance.print_message(f"\nDo I have your permission to use your **shell** to execute the following: \n\n{navi_command}\n")
+        navi_instance.print_message(
+            f"\nDo I have your permission to use your **shell** to execute the following: \n\n{navi_command}\n")
         user_input = input("Do you want me to continue (y/n): ").strip().lower()
         if user_input == 'y':
             result = subprocess.run(

chips/SSG/navi_token_management.py

@@ -0,0 +1,112 @@
+import navi_internal
+
+token_limit_max: int = 4096
+navi_settings: dict = {}
+
+command: str = "token-config"
+use: str = "Adjust token limits for RAG and chat memory."
+aliases: list = ['set_tokens', 'token_limits']
+params: dict = {
+    '-rag': 'Set the token limit for RAG (e.g., -rag 1024)',
+    '-chat': 'Set the token limit for chat memory (e.g., -chat 3072)',
+    '-show': 'Display the current token limit partition.',
+    '-help': 'Display help information.',
+}
+
+help_params: tuple = ('-help', '-h')
+
+
+def print_params() -> None:
+    print(f"{'Parameter':<10} | {'Description'}")
+    print("-" * 40)
+
+    for param, description in params.items():
+        print(f"{param:<10} | {description}")
+
+
+def set_token_limit(rag: int = None, chat: int = None) -> str:
+    global navi_settings
+    token_limit_rag = int(navi_settings["token_limit_rag"])
+    token_limit_chat = int(navi_settings["token_limit_chat"])
+
+    # Calculate new limits if provided
+    new_rag = token_limit_rag if rag is None else rag
+    new_chat = token_limit_chat if chat is None else chat
+
+    # Validate the total allocation
+    if new_rag + new_chat > token_limit_max:
+        return f"Error: Total allocation exceeds the maximum token limit of {token_limit_max}."
+
+    from navi_shell import modify_navi_settings, get_navi_settings
+    modify_navi_settings("token_limit_rag", new_rag)
+    modify_navi_settings("token_limit_chat", new_chat)
+
+    # Refresh the navi_settings dictionary
+    navi_settings = get_navi_settings()
+
+    return f"Token limits updated. RAG: {new_rag}, Chat: {new_chat}"
+
+
+def show_token_limits() -> str:
+    rag = int(navi_settings["token_limit_rag"])
+    chat = int(navi_settings["token_limit_chat"])
+    return f"Current Token Limits:\n- RAG: {rag}\n- Chat: {chat}\n- Total: {rag + chat} (Max: {token_limit_max})"
+
+
+def run(arguments=None) -> None:
+    navi_instance = navi_internal.navi_instance
+    global token_limit_max
+    token_limit_max = navi_instance.get_max_token_limit()
+    arg_array = arguments.text.split()
+
+    arg_array.pop(0)
+
+    from navi_shell import get_navi_settings
+    global navi_settings
+    navi_settings = get_navi_settings()
+
+    # Parse parameters
+    if arg_array:
+        rag_limit = None
+        chat_limit = None
+        show_only = False
+
+        for i in range(len(arg_array)):
+            arg = arg_array[i]
+
+            match arg:
+                case '-rag':
+                    try:
+                        # Fetch the next value and skip it in the loop
+                        rag_limit = int(arg_array[i + 1])
+                        arg_array[i + 1] = None  # Mark as processed
+                    except (IndexError, ValueError):
+                        navi_instance.print_message("Error: Invalid value for -rag.")
+                        return
+                case '-chat':
+                    try:
+                        # Fetch the next value and skip it in the loop
+                        chat_limit = int(arg_array[i + 1])
+                        arg_array[i + 1] = None  # Mark as processed
+                    except (IndexError, ValueError):
+                        navi_instance.print_message("Error: Invalid value for -chat.")
+                        return
+                case '-show':
+                    show_only = True
+                case x if x in help_params:
+                    print_params()
+                    return
+                case _:
+                    # Skip any previously processed value
+                    if arg is not None:
+                        navi_instance.print_message(f"Invalid parameter: {arg}")
+                        return
+
+        if show_only:
+            navi_instance.print_message(show_token_limits())
+            return
+
+        result = set_token_limit(rag=rag_limit, chat=chat_limit)
+        navi_instance.print_message(result)
+    else:
+        print_params()

chips/__init__.py

@@ -1,10 +1,9 @@
 """Init commands."""
-from os.path import dirname, basename, isfile, join, exists
-import os
-import sys
-import glob
 import importlib
 import logging
+import os
+import sys
+from os.path import dirname, join, exists
 from pathlib import Path
 
 package_dir = Path(__file__).parent

default_config

@@ -4,4 +4,6 @@ use_local_model=True
 dont_check_for_updates=False
 update_branch=main
 session=DEFAULT_SESSION
-overwrite_session=True
+overwrite_session=True
+token_limit_rag=2048
+token_limit_chat=2048

install/local_model.py

@@ -1,15 +1,12 @@
 import os
+import platform
 import shutil
-import subprocess
+import subprocess  # nosec
 import tempfile
 from getpass import getpass
 from typing import Tuple
 
 
-import subprocess
-import platform
-import os
-
 def install_ollama() -> bool:
     try:
         if platform.system() in ["Linux"]:
@@ -46,8 +43,8 @@ def install_ollama() -> bool:
             # Download the installer
             subprocess.run(
                 ["powershell", "-Command",
-                 "Invoke-WebRequest -Uri https://ollama.com/download/OllamaSetup.exe -OutFile ollama_installer.exe"],
-                shell=True,
+                 "Invoke-WebRequest", "-Uri", "https://ollama.com/download/OllamaSetup.exe",
+                 "-OutFile", "ollama_installer.exe"],
                 check=True,
                 capture_output=True,
                 text=True,
@@ -56,8 +53,7 @@ def install_ollama() -> bool:
             # Run the installer
             subprocess.run(
                 ["powershell", "-Command",
-                 "Start-Process -FilePath ./ollama_installer.exe"],
-                shell=True,
+                 "Start-Process", "-FilePath", "./ollama_installer.exe", "-Wait"],
                 check=True,
                 capture_output=True,
                 text=True,
@@ -89,7 +85,6 @@ def install_ollama() -> bool:
                 print(f"Failed to remove installer: {e}")
 
 
-
 def ollama_installed() -> bool:
     try:
         subprocess.run(

navi.py

@@ -1,5 +1,5 @@
-import re
 import platform
+import re
 import subprocess  # nosec
 
 

navi_banner.py

@@ -4,7 +4,6 @@
 
 versionNum = get_navi_version()
 
-
 # The cover art:
 art = rf"""{breakline}
 |               _   __            _                 |