Merge pull request #1 from berylliumsec/update_readme

linting and readme

.github/workflows/publish.yml

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

README.md

@@ -21,13 +21,13 @@ Welcome to Neutron.
 
 **First i would like to thank the All-Mighty God who is the source of all knowledge, without Him, this would not be possible.**
 
-**Disclaimer: AI can make mistakes, consider cross-checking suggestions.**
+## **Disclaimer: AI can make mistakes, consider cross-checking suggestions.**
 
 ## [Click Here to Watch Neutron in Action](https://youtu.be/v5X8TNPsMbM)
 
 ## 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.
+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).
 
 ## Compatibility
 

src/neutron/client.py

@@ -1,5 +1,6 @@
 import argparse
 import os
+
 import requests
 
 
@@ -9,7 +10,9 @@ def send_request(question: str, server_url: str = "http://localhost:8000"):
 
     payload = {"question": question}
     # Retrieve the token from the NEUTRON_TOKEN environment variable
-    token = os.environ.get("NEUTRON_TOKEN", "default_token")  # Use a default value if NEUTRON_TOKEN is not set
+    token = os.environ.get(
+        "NEUTRON_TOKEN", "default_token"
+    )  # Use a default value if NEUTRON_TOKEN is not set
     headers = {"Authorization": token}  # Use the token from the environment variable
     try:
         response = requests.post(endpoint_url, json=payload, headers=headers)
@@ -18,15 +21,24 @@ def send_request(question: str, server_url: str = "http://localhost:8000"):
             print("Response:", response.json().get("response", "No response received"))
         else:
             # If an error happens, FastAPI will still return JSON formatted error messages
-            print("Error:", response.status_code, response.json().get("detail", "Unknown error"))
+            print(
+                "Error:",
+                response.status_code,
+                response.json().get("detail", "Unknown error"),
+            )
     except Exception as e:
         print(f"An error occurred while sending the request: {e}")
 
 
 def main():
     parser = argparse.ArgumentParser(description="Send a question to the AI server.")
-    parser.add_argument('question', type=str, help='The question to ask the AI server.')
-    parser.add_argument('--server_url', type=str, default="http://localhost:8000", help='The URL of the AI server, defaults to http://localhost:8000')
+    parser.add_argument("question", type=str, help="The question to ask the AI server.")
+    parser.add_argument(
+        "--server_url",
+        type=str,
+        default="http://localhost:8000",
+        help="The URL of the AI server, defaults to http://localhost:8000",
+    )
 
     args = parser.parse_args()
 

src/neutron/interactive_model.py

@@ -12,6 +12,7 @@
 from langchain_core.runnables import RunnablePassthrough
 from transformers import (AutoModelForCausalLM, AutoTokenizer,
                           BitsAndBytesConfig, pipeline)
+
 from neutron import utilities
 
 transformers.logging.set_verbosity_error()
@@ -44,7 +45,8 @@ def __init__(self):
             print(f"total memory available {total_memory_gb}")
 
         if total_memory_gb >= 24:
-            self.model = AutoModelForCausalLM.from_pretrained(utilities.return_path("neutron_model"),
+            self.model = AutoModelForCausalLM.from_pretrained(
+                utilities.return_path("neutron_model"),
                 quantization_config=bnb_config,
                 device_map={"": 0},
             )

src/neutron/server.py

@@ -6,10 +6,8 @@
 from fastapi.middleware.cors import CORSMiddleware
 from pydantic import BaseModel
 
-
 from neutron.interactive_model import InteractiveModel
 
-
 # Set up argument parsing
 parser = argparse.ArgumentParser(description="Run the FastAPI server.")
 parser.add_argument(

src/neutron/utilities.py

@@ -53,9 +53,9 @@ def get_input_with_default(message, default_text=None):
 
 def ensure_model_folder_exists(model_directory, auto_update=True):
     if model_directory == "neutron_model":
-                s3_url = model_s3_url
+        s3_url = model_s3_url
     else:
-        s3_url= chroma_s3_url
+        s3_url = chroma_s3_url
     try:
         metadata_file = os.path.join(model_directory, "metadata.json")
 
@@ -64,16 +64,15 @@ def ensure_model_folder_exists(model_directory, auto_update=True):
         except Exception as e:
             logging.error(f"Error getting local metadata: {e}")
 
-
         try:
             s3_etag = get_s3_file_etag(s3_url)
         except Exception as e:
             logging.error(f"Error getting S3 file etag: {e}")
 
-
         if s3_etag is None:
-            logging.warning("No S3 etag found. Possibly no internet connection or issue with S3.")
-
+            logging.warning(
+                "No S3 etag found. Possibly no internet connection or issue with S3."
+            )
 
         # Check if the model directory exists and has the same etag (metadata)
         if folder_exists_and_not_empty(model_directory) and local_etag == s3_etag:
@@ -95,7 +94,9 @@ def ensure_model_folder_exists(model_directory, auto_update=True):
                 logging.info("User chose not to update the model directory.")
                 return  # Exit if user chooses not to update
         else:
-            logging.info("Auto-update is enabled. Downloading new version if necessary...")
+            logging.info(
+                "Auto-update is enabled. Downloading new version if necessary..."
+            )
 
         # Proceed with the removal of the existing model directory and the download of the new version
         if os.path.exists(model_directory):
@@ -106,9 +107,10 @@ def ensure_model_folder_exists(model_directory, auto_update=True):
                 logging.error(f"Error removing existing model directory: {e}")
                 return
 
-        logging.info(f"{model_directory} not found or is outdated. Downloading and unzipping...")
+        logging.info(
+            f"{model_directory} not found or is outdated. Downloading and unzipping..."
+        )
         try:
-
             download_and_unzip(s3_url, f"{model_directory}.zip")
             # Save new metadata
             save_local_metadata(metadata_file, s3_etag)
@@ -117,6 +119,7 @@ def ensure_model_folder_exists(model_directory, auto_update=True):
     except Exception as e:
         logging.error(f"Unexpected error: {e}")
 
+
 def resource_path(relative_path):
     """Get absolute path to resource, works for dev and for PyInstaller"""
     if getattr(sys, "frozen", False):
@@ -217,9 +220,7 @@ def is_internet_available(host="8.8.8.8", port=53, timeout=3):
 
 def get_s3_file_etag(s3_url):
     if not is_internet_available():
-        logging.error(
-            "No internet connection available. Skipping version check."
-        )
+        logging.error("No internet connection available. Skipping version check.")
         return None
     response = requests.head(s3_url)
     return response.headers.get("ETag")
@@ -276,6 +277,4 @@ def check_new_pypi_version(package_name="neutron-ai"):
                 "yellow",
             )
     except Exception as e:
-        logging.error(
-            f"An error occurred while checking for the latest version: {e}"
-        )
+        logging.error(f"An error occurred while checking for the latest version: {e}")