Merge pull request #3 from berylliumsec/nebula_pro_awareness

Adding more memory information

.github/workflows/publish.yml

@@ -120,7 +120,7 @@ jobs:
           prerelease: false
           body: |
             Release Notes:
-              - Updating Readme and Linting
+              - Updating Readme and adding memory usage display
       - name: Upload Release Assets
         if: github.event.pull_request.merged == true
         run: |

README.md

@@ -23,12 +23,26 @@ Welcome to Neutron.
 
 ## **Disclaimer: AI can make mistakes, consider cross-checking suggestions.**
 
-## [Click Here to Watch Neutron in Action](https://youtu.be/v5X8TNPsMbM)
+## 🌐 Introducing Nebula Pro: A New Era in Ethical Hacking 🌐
+
+🚀 We're thrilled to unveil a sneak peek of Nebula Pro, our latest innovation designed to empower ethical hackers with advanced, AI-driven capabilities. After months of dedicated development, we have launched the preview version. Some of the exciting features are:
+
+- AI Powered Autonomous Mode
+- AI Powered Suggestions
+- AI Powered Note Taking
+
+**Neutron will become a part of Nebula Pro's free tier**
+# 📺 [Click Here to Get Access To Nebula Pro Now](https://www.berylliumsec.com/nebula-pro-waitlist) 🚀
+
+
 
 ## Why Neutron?
 
 The purpose of Neutron is straightforward: to provide security professionals  with access to a free AI assistant that can be invoked directly from their command line interface. It was built as part of the free tier of [Nebula Pro](https://www.berylliumsec.com/nebula-pro-waitlist).
 
+## [Click Here to Watch Neutron in Action](https://youtu.be/v5X8TNPsMbM)
+
+
 ## Compatibility
 
 Neutron has been extensively tested and optimized for Linux platforms. As of now, its functionality on Windows or macOS is not guaranteed, and it may not operate as expected.
@@ -37,9 +51,9 @@ Neutron has been extensively tested and optimized for Linux platforms. As of now
 
 - Storage: A minimum of 50GB is recommended.
 
-- RAM: A minimum of 32GB RAM memory is recommended.
+- RAM: A minimum of 16GB RAM memory is recommended.
 
-- Graphics Processing Unit (GPU): While not mandatory, having at least 24GB of GPU memory is recommended for optimal performance.
+- Graphics Processing Unit (GPU) (NOT MANDATORY, Neutron can run on CPU): While not mandatory, having at least 24GB of GPU memory is recommended for optimal performance.
 
 
 **PYPI based distribution requirement(s)**

src/neutron/interactive_model.py

@@ -42,7 +42,9 @@ def __init__(self):
             total_memory_gb = torch.cuda.get_device_properties(0).total_memory / (
                 1024**3
             )  # Convert bytes to GB
-            print(f"total memory available {total_memory_gb}")
+            print(f"total GPU memory available {total_memory_gb}")
+            if total_memory_gb < 24:
+                print("There isnt enough GPU memory, will use CPU")
 
         if total_memory_gb >= 24:
             self.model = AutoModelForCausalLM.from_pretrained(

src/neutron/server.py

@@ -2,6 +2,7 @@
 import os
 from typing import Any, Dict
 
+import psutil
 from fastapi import FastAPI, HTTPException, Request, status
 from fastapi.middleware.cors import CORSMiddleware
 from pydantic import BaseModel
@@ -32,6 +33,17 @@
 )
 
 model = InteractiveModel()  # Initializes the model with its default configuration
+# Get current process ID
+pid = os.getpid()
+# Get the process info using psutil
+process = psutil.Process(pid)
+# Get memory usage (in bytes)
+memory_use = process.memory_info().rss
+memory_use_gb = memory_use / 1024 / 1024 / 1024
+print(f"Memory used by Neutron: {memory_use_gb} GB")
+print(
+    "Embrace the future of AI Powered Ethical Hacking with Nebula Pro ->> https://www.berylliumsec.com/nebula-pro-waitlist "
+)
 
 
 def check_auth(token: str) -> bool:
@@ -58,6 +70,9 @@ def ask(request: Request, query: Query) -> Dict[str, Any]:
 
     try:
         response = model.invoke(query.question)
+        memory_use = process.memory_info().rss
+        memory_use_gb = memory_use / 1024 / 1024 / 1024
+        print(f"Memory used by Neutron: {memory_use_gb} GB")
         return {"response": response}
     except KeyError as e:
         raise HTTPException(status_code=status.HTTP_400_BAD_REQUEST, detail={e})

src/neutron/utilities.py

@@ -241,27 +241,23 @@ def get_latest_pypi_version(package_name):
         if response.status_code == 200:
             return response.json()["info"]["version"]
     except requests.exceptions.RequestException as e:
-        logging.error(f"Failed to get latest version information: {e}", "red")
+        logging.error(f"Failed to get latest version information: {e}")
     return None
 
 
 def check_new_pypi_version(package_name="neutron-ai"):
     """Check if a newer version of the package is available on PyPI."""
     if not is_internet_available():
-        logging.error(
-            "No internet connection available. Skipping version check.", "red"
-        )
+        logging.error("No internet connection available. Skipping version check.")
         return
 
     try:
         installed_version = version(package_name)
     except Exception as e:
-        logging.error(
-            f"Error retrieving installed version of {package_name}: {e}", "red"
-        )
+        logging.error(f"Error retrieving installed version of {package_name}: {e}")
         return
 
-    logging.info(f"Installed version: {installed_version}", "green")
+    print(f"Installed version: {installed_version}")
 
     try:
         latest_version = get_latest_pypi_version(package_name)
@@ -272,9 +268,8 @@ def check_new_pypi_version(package_name="neutron-ai"):
             return
 
         if latest_version > installed_version:
-            logging.info(
-                f"A newer version ({latest_version}) of {package_name} is available on PyPI. Please consider updating to access the latest features!",
-                "yellow",
+            print(
+                f"A newer version ({latest_version}) of {package_name} is available on PyPI. Please consider updating to access the latest features!"
             )
     except Exception as e:
         logging.error(f"An error occurred while checking for the latest version: {e}")