Fix print statements in console, Add drop command

Author: pranavvp16

octo/cli.py

@@ -1,5 +1,5 @@
 from octo import Octo
-from octo import DESCRIPTION, VERSION, ASCII_ART
+from octo import ASCII_ART
 from octo import __version__
 
 from rich.console import Console
@@ -59,55 +59,49 @@ def main():
     except:
         if args.command == "start":
             with console.status("[bold green]Starting Mindsdb...", spinner="dots2"):
-                try:
                     octo.start_local()
-                except:
-                    rprint(
-                        "[bold][red]Error: Mindsdb is not installed. Use 'pip install mindsdb' to install."
-                    )
-                    return
-                rprint(f"[bold][green]Started Mindsdb local server")
+                    console.print(f"[bold][green]Started Mindsdb local server")
         else:
-            rprint(
+            console.print(
                 "[bold][red]Error: Mindsdb server is not running. Use 'octo start' to start the server."
             )
             return
 
     if args.command == "stop":
         octo.stop_local()
-        rprint(f"[bold][green]Stopped Mindsdb local server")
+        console.print(f"[bold][green]Stopped Mindsdb local server")
 
     if args.command == "init":
         try:
             owner = args.repo.split("/")[0]
             repo = args.repo.split("/")[1]
         except:
-            rprint(
+            console.print(
                 "[bold][red]Error: Invalid repository name, format should be 'owner/repo'"
             )
             return
         with console.status("[bold green]Initializing repository...", spinner="dots2"):
-            rprint(f"[bold][green] Creating model {owner}/{repo}")
+            console.print(f"[bold][green] Creating model {owner}/{repo}")
             message = octo.init(repo, owner, args.branch, all_files=args.all)
-            rprint(message)
+            console.print(message)
 
     if args.command == "drop":
         try:
             owner = args.repo.split("/")[0]
             repo = args.repo.split("/")[1]
         except:
-            rprint(
+            console.print(
                 "[bold][red]Error: Invalid repository name, format should be 'owner/repo'"
             )
             return
         project = octo._get_project()
         model_names = [i.name for i in project.list_models()]
         if f"{owner}_{repo}" not in model_names:
-            rprint(f"[bold][red]Error: Model {owner}/{repo} does not exist")
+            console.print(f"[bold][red]Error: Model {owner}/{repo} does not exist")
             return
         else:
             project.drop_model(f"{owner}_{repo}")
-            rprint(f"[bold][green]Deleted model {owner}/{repo}")
+            console.print(f"[bold][green]Deleted model {owner}/{repo}")
 
     if args.version:
         text = ASCII_ART.format(__version__)
@@ -123,11 +117,11 @@ def main():
             )
             return
         message = octo.checkout(owner, repo)
-        console.log(message)
+        console.print(message)
 
     if args.command == "status":
         message = octo.status()
-        console.log(message)
+        console.print(message)
 
     if args.command == "tell":
         df = pd.DataFrame({"questions": [args.action]})
@@ -136,7 +130,8 @@ def main():
             answer = pred_df["answer"].iloc[0]
             # Filter string
             answer = answer.replace("\n", "")
-            rprint(answer)
+            console.print(f"\n{answer}")
+            return
 
     if args.command == "drop":
         try:
@@ -148,15 +143,15 @@ def main():
             )
             return
         message = octo.drop(owner, repo)
-        rprint(message)
+        console.print(message)
 
     if args.command == "list":
         model_names = octo.list_models()
         # Format the string
         model_names = [i.replace("_", "/") for i in model_names]
         with console.status("[bold green]Fetching models...", spinner="dots2"):
             for i in model_names:
-                rprint(i)
+                console.print(i)
 
 
 

octo/github_handle/mindsdb_octo.py

@@ -29,7 +29,7 @@ def stop_local(self):
         """
         # Run shell command to stop the local mindsdb server
         command = 'pkill -f "python"'
-        subprocess.run(command, shell=True, check=True)
+        subprocess.run(command, shell=True, check=True, stdout=subprocess.DEVNULL)
 
     def create_model(self, owner, repo, branch, all_files=False):
         project = self._get_project()
@@ -61,13 +61,12 @@ def create_model(self, owner, repo, branch, all_files=False):
             )
         while True:
             if model.get_status() not in ("generating", "training"):
-                print("\nFinished adding repository to knowledge base")
                 break
 
         if model.get_status() == "error":
-            return f"[bold][red]{model.data['error']}"
+            return f"[bold][red]{model.data['error']}", True
 
-        return f"[bold][green]Model created successfully"
+        return f"[bold][green]Model created successfully", False
 
     def predict(self, df, owner, repo):
         project = self._get_project()

octo/github_handle/octo.py

@@ -18,10 +18,11 @@ def init(self, repo, owner, branch, all_files=False):
             self.owner = owner
             self.branch = branch
             self.repo_initialized = True
-            message = self.create_model(
+            message, error = self.create_model(
                 self.owner, self.repo, self.branch, all_files=all_files
             )
-            self._save_state()
+            if not error:
+                self._save_state()
         elif (
             self.repo.lower() == repo
             and self.owner.lower() == owner
@@ -34,10 +35,11 @@ def init(self, repo, owner, branch, all_files=False):
             self.owner = owner
             self.branch = branch
             self.repo_initialized = True
-            message = self.create_model(
+            message, error = self.create_model(
                 self.owner, self.repo, self.branch, all_files=all_files
             )
-            self._save_state()
+            if not error:
+                self._save_state()
         return message
 
     def checkout(self, owner, repo):
@@ -64,12 +66,11 @@ def tell(self, df):
 
     def drop(self, owner, repo):
         project = self._get_project()
-        model_names = [i.name for i in project.list_models()]
-        if f"{owner}_{repo}" not in model_names:
-            return f"[bold][red]Error: Model {owner}/{repo} does not exist"
-        else:
+        try:
             project.drop_model(f"{owner}_{repo}")
             return f"[bold][green]Deleted model {owner}/{repo}"
+        except:
+            return f"[bold][red]Error: Model {owner}/{repo} does not exist"
 
     def _save_state(self):
         state = {

octo/templates/query_dicts.py

@@ -5,11 +5,10 @@
     "branch": "<branch>",
     "reader": "GithubRepositoryReader",
     "filer_type": "include",
-    "filer_extensions": [".md",],
+    "filer_extensions": [".md"],
     "input_column": "questions",
     "openai_api_key": "<openai_key>",
     "github_token": "<github_token>",
-    "model_name": "gpt-4"
 }
 
 # Include all files in the repo
@@ -22,5 +21,4 @@
     "input_column": "questions",
     "openai_api_key": "<openai_key>",
     "github_token": "<github_token>",
-    "model_name": "gpt-4"
 }

requirements.txt

@@ -1,6 +1,5 @@
 rich
 git+https://github.com/pranavvp16/mindsdb.git@llama_index
-mindsdb_sql==0.7.4
 mindsdb_sdk
 openai==1.6.1
 llama_index==0.9.27

setup.py

@@ -11,9 +11,11 @@
     },
     # install_requires=[
     #     'rich',
-    #     'mindsdb_sdk',
-    #     'openai',
-    #     'mindsdb'
-    # ],
-    # dependency_links=["https://github.com/pranavvp16/mindsdb.git@llama_index#egg=mindsdb-llama_index"]
+    #     'mindsdb_llama_index>=23.12.4.2',
+    #     'mindsdb_sdk==2.0.0',
+    #     'openai==1.6.1',
+    #     'llama_index==0.9.27',
+    #     'llama_hub==0.0.67',
+    #],
+    #dependency_links=["https://github.com/pranavvp16/mindsdb.git@llama_index#egg=mindsdb_llama_index-23.12.4.2"]
 )