Merge pull request #14 from kubetoolsca/pre-release

Pre release

krs/krs.py

@@ -4,7 +4,7 @@
 from krs.main import KrsMain
 from krs.utils.constants import KRSSTATE_PICKLE_FILEPATH, KRS_DATA_DIRECTORY
 
-app = typer.Typer()
+app = typer.Typer(help="krs: A command line interface to scan your Kubernetes Cluster, detect errors, provide resolutions using LLMs and recommend latest tools for your cluster")
 krs = KrsMain()
 
 def check_initialized():
@@ -16,11 +16,11 @@ def check_initialized():
     os.mkdir(KRS_DATA_DIRECTORY)
 
 @app.command()
-def init():
+def init(kubeconfig: str = typer.Option('~/.kube/config', help="Custom path for kubeconfig file if not default")):
     """
     Initializes the services and loads the scanner.
     """
-    krs.initialize()
+    krs.initialize(kubeconfig)
     typer.echo("Services initialized and scanner loaded.")
 
 @app.command()
@@ -71,13 +71,14 @@ def recommend():
     krs.generate_recommendations()
 
 @app.command()
-def health(change_model: bool = typer.Option(False, help="Option to reinitialize/change the LLM, if set to True")):
+def health(change_model: bool = typer.Option(False, help="Option to reinitialize/change the LLM, if set to True"),
+           device: str = typer.Option('cpu', help='Option to run Huggingface models on GPU by entering the option as "gpu"')):
     """
-    Starts an interactive terminal to chat with user.
+    Starts an interactive terminal using an LLM of your choice to detect and fix issues with your cluster
     """
     check_initialized()
     typer.echo("\nStarting interactive terminal...\n")
-    krs.health_check(change_model)
+    krs.health_check(change_model, device)
 
 @app.command()
 def export():

krs/main.py

@@ -2,7 +2,6 @@
 from krs.utils.cluster_scanner import KubetoolsScanner
 from krs.utils.llm_client import KrsGPTClient
 from krs.utils.functional import extract_log_entries, CustomJSONEncoder
-from termcolor import colored
 import os, pickle, time, json
 from tabulate import tabulate
 from krs.utils.constants import (KRSSTATE_PICKLE_FILEPATH, LLMSTATE_PICKLE_FILEPATH, POD_INFO_FILEPATH, KRS_DATA_DIRECTORY)
@@ -173,7 +172,7 @@ def print_recommendations(self):
         print(tabulate(recommendations, headers=["Category", "Recommendation", "Tool Name", "CNCF Status"], tablefmt="grid"))
 
 
-    def health_check(self, change_model=False):
+    def health_check(self, change_model=False, device='cpu'):
 
         if os.path.exists(LLMSTATE_PICKLE_FILEPATH) and not change_model:
             continue_previous_chat = input("\nDo you want to continue fixing the previously selected pod ? (y/n): >> ")
@@ -184,13 +183,13 @@ def health_check(self, change_model=False):
                     break
 
             if continue_previous_chat=='y':
-                krsllmclient = KrsGPTClient()
+                krsllmclient = KrsGPTClient(device=device)
                 self.continue_chat = True
             else:
-                krsllmclient = KrsGPTClient(reset_history=True)
+                krsllmclient = KrsGPTClient(reset_history=True, device=device)
 
         else:
-            krsllmclient = KrsGPTClient(reinitialize=True)
+            krsllmclient = KrsGPTClient(reinitialize=True, device=device)
             self.continue_chat = False
 
         if not self.continue_chat:
@@ -244,7 +243,7 @@ def get_logs_from_pod(self, namespace_index, pod_index):
         try:
             namespace_index -= 1
             pod_index -= 1
-            namespace = list(self.pod_info.keys())[namespace_index]
+            namespace = list(self.list_namespaces())[namespace_index]
             return list(self.pod_info[namespace][pod_index]['info']['Logs'].values())[0]
         except KeyError as e:
             print("\nKindly enter a value from the available namespaces and pods")

krs/requirements.txt

@@ -1,3 +1,5 @@
 typer==0.12.3
 requests==2.32.2
-kubernetes==29.0.0
+kubernetes==29.0.0
+tabulate==0.9.0
+

krs/utils/llm_client.py

@@ -5,7 +5,7 @@
 
 class KrsGPTClient:
 
-    def __init__(self, reinitialize=False, reset_history=False):
+    def __init__(self, reinitialize=False, reset_history=False, device='cpu'):
 
         self.reinitialize = reinitialize
         self.client = None
@@ -16,6 +16,7 @@ def __init__(self, reinitialize=False, reset_history=False):
         self.continue_chat = False
         self.history = []
         self.max_tokens = MAX_OUTPUT_TOKENS
+        self.device = device
 
 
         if not self.reinitialize:
@@ -137,7 +138,7 @@ def init_huggingface_client(self, reinitialize=False):
         try:
             self.tokenizer = AutoTokenizer.from_pretrained(self.model)
             self.model_hf = AutoModelForCausalLM.from_pretrained(self.model)
-            self.pipeline = pipeline('text-generation', model=self.model_hf, tokenizer=self.tokenizer)
+            self.pipeline = pipeline('text-generation', model=self.model_hf, tokenizer=self.tokenizer, device=0 if self.device == 'gpu' else -1)
 
         except OSError as e:
             print("\nError loading model: ", e)