From e8dc67df87216254817027f3b9932e4a7a1279a0 Mon Sep 17 00:00:00 2001
From: fynnfluegge <fynnfluegge@gmx.de>
Date: Tue, 26 Sep 2023 11:18:07 +0200
Subject: [PATCH] fix: faiss installation with yaspin spinner

---
 codeqai/app.py          |  9 ++-------
 codeqai/embeddings.py   |  4 ++++
 codeqai/llm.py          |  3 +--
 codeqai/vector_store.py | 13 +++++++++++--
 pyproject.toml          |  2 +-
 5 files changed, 19 insertions(+), 12 deletions(-)

diff --git a/codeqai/app.py b/codeqai/app.py
index faf0362..6daf3fd 100644
--- a/codeqai/app.py
+++ b/codeqai/app.py
@@ -8,7 +8,8 @@
 from yaspin import yaspin
 
 from codeqai import codeparser, repo, utils
-from codeqai.config import create_cache_dir, create_config, get_cache_path, load_config
+from codeqai.config import (create_cache_dir, create_config, get_cache_path,
+                            load_config)
 from codeqai.constants import EmbeddingsModel, LllmHost
 from codeqai.embeddings import Embeddings
 from codeqai.llm import LLM
@@ -55,19 +56,13 @@ def run():
         files = repo.load_files()
         documents = codeparser.parse_code_files(files)
         spinner.stop()
-        spinner = yaspin(text="💾 Indexing vector store...", color="green")
-        spinner.start()
         vector_store = VectorStore(
             repo_name,
             embeddings=embeddings_model.embeddings,
             documents=documents,
         )
-        spinner.stop()
     else:
-        spinner = yaspin(text="💾 Loading vector store...", color="green")
-        spinner.start()
         vector_store = VectorStore(repo_name, embeddings=embeddings_model.embeddings)
-        spinner.stop()
 
     llm = LLM(
         llm_host=LllmHost[config["llm-host"].upper().replace("-", "_")],
diff --git a/codeqai/embeddings.py b/codeqai/embeddings.py
index 2d99cbe..2d20f1a 100644
--- a/codeqai/embeddings.py
+++ b/codeqai/embeddings.py
@@ -68,6 +68,8 @@ def _install_sentence_transformers(self):
                 )
             except subprocess.CalledProcessError as e:
                 print(f"Error during sentence_transformers installation: {e}")
+        else:
+            exit("sentence_transformers is required for local embeddings.")
 
     def _install_instructor_embedding(self):
         question = [
@@ -96,3 +98,5 @@ def _install_instructor_embedding(self):
                 )
             except subprocess.CalledProcessError as e:
                 print(f"Error during sentence_transformers installation: {e}")
+        else:
+            exit("InstructorEmbedding is required for local embeddings.")
diff --git a/codeqai/llm.py b/codeqai/llm.py
index 311a01d..15906a1 100644
--- a/codeqai/llm.py
+++ b/codeqai/llm.py
@@ -136,5 +136,4 @@ def supports_metal():
                         )
 
             else:
-                print("", "Installation cancelled. Exiting.", "")
-                return None
+                exit("llama-cpp-python is required for local LLM.")
diff --git a/codeqai/vector_store.py b/codeqai/vector_store.py
index 8f2cfee..9e444cc 100644
--- a/codeqai/vector_store.py
+++ b/codeqai/vector_store.py
@@ -2,6 +2,7 @@
 from langchain.embeddings.base import Embeddings
 from langchain.schema import Document
 from langchain.vectorstores import FAISS
+from yaspin import yaspin
 
 from codeqai import utils
 from codeqai.config import get_cache_path
@@ -15,14 +16,20 @@ def __init__(
         self.embeddings = embeddings
         self.install_faiss()
         if not documents:
+            spinner = yaspin(text="💾 Loading vector store...", color="green")
+            spinner.start()
             self.db = FAISS.load_local(
                 index_name=self.name,
                 folder_path=get_cache_path(),
                 embeddings=self.embeddings,
             )
+            spinner.stop()
         else:
+            spinner = yaspin(text="💾 Indexing vector store...", color="green")
+            spinner.start()
             self.db = FAISS.from_documents(documents, embeddings)
             self.db.save_local(index_name=self.name, folder_path=get_cache_path())
+            spinner.stop()
         self.retriever = self.db.as_retriever(search_type="mmr", search_kwargs={"k": 8})
 
     def similarity_search(self, query: str):
@@ -30,12 +37,12 @@ def similarity_search(self, query: str):
 
     def install_faiss(self):
         try:
-            import faiss
+            import faiss  # noqa: F401
         except ImportError:
             question = [
                 inquirer.Confirm(
                     "confirm",
-                    message=f"{utils.get_bold_text('FAISS')} not found in this python environment. Do you want to install it now?",
+                    message=f"{utils.get_bold_text('faiss')} package not found in this python environment. Do you want to install it now?",
                     default=True,
                 ),
             ]
@@ -72,3 +79,5 @@ def install_faiss(self):
                         )
                     except subprocess.CalledProcessError as e:
                         print(f"Error during faiss installation: {e}")
+            else:
+                exit("faiss package is required for codeqai to work.")
diff --git a/pyproject.toml b/pyproject.toml
index f154565..dcc48d4 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "codeqai"
-version = "0.0.1"
+version = "0.0.2"
 description = ""
 authors = ["fynnfluegge <fynnfluegge@gmx.de>"]
 readme = "README.md"