Refactor code and add new features

.pylintrc

@@ -278,7 +278,7 @@ valid-metaclass-classmethod-first-arg=mcs
 
 # List of regular expressions of class ancestor names to ignore when counting
 # public methods (see R0903)
-exclude-too-few-public-methods=
+# exclude-too-few-public-methods=
 
 # List of qualified class names to ignore when counting class parents (see
 # R0901)

pyproject.toml

@@ -6,10 +6,19 @@ authors = ["Younis <[email protected]>"]
 readme = "README.md"
 
 [tool.poetry.dependencies]
-python = "^3.10"
+python = ">=3.9.1,<3.12"
 inquirer = "^3.2.4"
 openai = "^1.14.0"
 toml = "^0.10.2"
+esprima = "^4.0.1"
+tree-sitter = "^0.21.0"
+tree-sitter-languages = "^1.10.2"
+langchain = "^0.1.12"
+langchain-community = "^0.0.28"
+langchain-text-splitters = "^0.0.1"
+langchain-openai = "^0.0.8"
+tiktoken = "^0.6.0"
+rich = "^13.7.1"
 
 
 [tool.poetry.group.dev.dependencies]

senior_swe_ai/cache.py

@@ -1,10 +1,39 @@
 """ This module contains functions to create and get the cache directory
 for the application.
 """
+import json
 import os
 import platform
 from pathlib import Path
 
+from altair import Self
+
+
+class VectorCache:
+    """ VectorCache class for storing vector data """
+
+    def __init__(self, filename, vec_ids, commit_hash) -> None:
+        self.filename = filename
+        self.vector_ids = vec_ids
+        self.commit_hash = commit_hash
+
+    def to_json(self) -> str:
+        """Convert the object to json"""
+        return {
+            "filename": self.filename,
+            "commit_hash": self.commit_hash,
+            "vector_ids": self.vector_ids,
+        }
+
+    @classmethod
+    def from_json(cls, data: dict[str, str]) -> Self:
+        """Create a VectorCache object from json"""
+        return cls(
+            filename=data.get("filename"),
+            vec_ids=data.get("vector_ids"),
+            commit_hash=data.get("commit_hash"),
+        )
+
 
 def get_cache_path() -> str:
     """Get the cache directory path for the application."""
@@ -28,3 +57,22 @@ def create_cache_dir() -> None:
     if not os.path.exists(get_cache_path()):
         path = Path(get_cache_path())
         path.mkdir(parents=True, exist_ok=True)
+
+
+def load_vec_cache(filename: str) -> dict[str, VectorCache]:
+    """Load the vector cache from the given file."""
+    with open(get_cache_path() + f'/{filename}', 'r', encoding='utf-8') as f:
+        vec = json.load(f)
+    vec_cache = {}
+    for key, value in vec.items():
+        vec_cache[key] = VectorCache.from_json(value)
+    return vec_cache
+
+
+def save_vec_cache(vector_cache, filename) -> None:
+    """Save the vector cache to the given file."""
+    with open(
+        get_cache_path() + "/" + filename, "w", encoding="utf-8"
+    ) as vector_cache_file:
+        json.dump(vector_cache, default=VectorCache.to_json,
+                  fp=vector_cache_file)

senior_swe_ai/cli.py

@@ -1,20 +1,39 @@
 """ SeniorSWE cli tool utilize AI to help you with your project """
 from argparse import ArgumentParser, Namespace
+import os
 import sys
-from typing import Any
-from senior_swe_ai.git_process import is_git_repo, get_repo_name, get_repo_root
+from typing import List
+from langchain.memory import ConversationSummaryMemory
+from langchain.chains.conversational_retrieval.base import (
+    BaseConversationalRetrievalChain, ConversationalRetrievalChain
+)
+import inquirer
+from langchain_core.documents.base import Document
+from langchain_openai import OpenAIEmbeddings, ChatOpenAI
+from senior_swe_ai.file_handler import parse_code_files
+from senior_swe_ai.git_process import (
+    is_git_repo, get_repo_name, get_repo_root, recursive_load_files
+)
 from senior_swe_ai.conf import config_init, load_conf, append_conf
+from senior_swe_ai.cache import create_cache_dir, get_cache_path, save_vec_cache
+from senior_swe_ai.vec_store import VectorStore
+from senior_swe_ai.consts import FaissModel, faiss_installed
 
 
 def main() -> None:
     """ __main__ """
+    py_version: tuple[int, int] = sys.version_info[:2]
+    if py_version < (3, 9) or py_version > (3, 12):
+        print('This app requires Python ^3.9.x or >3.12.x')
+        sys.exit(1)
+
     parser = ArgumentParser(
         description='SeniorSWE cli tool utilize AI to help you with your project'
     )
 
     parser.add_argument(
         'options', choices=['init', 'chat'],
-        help="'init': initialize the app. 'chat': chat with the AI"
+        help="'init': initialize the app. 'chat': chat with desired codebase."
     )
 
     args: Namespace = parser.parse_args()
@@ -28,18 +47,95 @@ def main() -> None:
         print('The current directory is not a git repository')
         sys.exit(1)
 
-    repo_name: str = get_repo_name()
+    repo_name: str = get_repo_name().upper()
     repo_root: str = get_repo_root()
 
     append_conf({'repo_name': repo_name, 'repo_root': repo_root})
 
     try:
-        conf: dict[Any, Any] = load_conf()
+        conf: dict[str, str] = load_conf()
     except FileNotFoundError:
         config_init()
         append_conf({'repo_name': repo_name, 'repo_root': repo_root})
         conf = load_conf()
 
+    create_cache_dir()
+
+    embed_mdl = OpenAIEmbeddings(
+        model=conf['embed_model'], api_key=conf['api_key'])
+
+    vec_store = VectorStore(embed_mdl, repo_name)
+
+    if not os.path.exists(get_cache_path() + f'/{repo_name}.faiss'):
+        is_faiss_installed: bool = faiss_installed()
+        if not is_faiss_installed:
+            question = [
+                inquirer.List(
+                    'install',
+                    message='FAISS is not installed. Do you want to install it?',
+                    choices=['Yes', 'No'],
+                    default='Yes'
+                )
+            ]
+            answer: dict[str, str] = inquirer.prompt(question)
+            if answer['install'] == 'Yes':
+                question = [
+                    inquirer.List(
+                        "faiss-installation",
+                        message="Please select the appropriate option to install FAISS. \
+                            Use gpu if your system supports CUDA",
+                        choices=[
+                            FaissModel.FAISS_CPU.value,
+                            FaissModel.FAISS_GPU.value,
+                        ],
+                        default=FaissModel.FAISS_CPU.value,
+                    )
+                ]
+                answer: dict[str, str] = inquirer.prompt(question)
+                if answer['faiss-installation'] == 'faiss-cpu':
+                    os.system('pip install faiss-cpu')
+                else:
+                    os.system('pip install faiss-gpu')
+            else:
+                print('FAISS is required for this app to work')
+                sys.exit(1)
+        # all desired files in the git repository tree
+        files: list[str] = recursive_load_files()
+        docs: List[Document] = parse_code_files(files)
+        vec_store.idx_docs(docs)
+        save_vec_cache(vec_store.vec_cache, f'{repo_name}.json')
+
+    vec_store.load_docs()
+    chat_mdl = ChatOpenAI(model=conf['chat_model'], api_key=conf['api_key'], temperature=0.9,
+                          max_tokens=2048)
+    mem = ConversationSummaryMemory(
+        llm=chat_mdl, memory_key='chat_history', return_messages=True
+    )
+    qa: BaseConversationalRetrievalChain = ConversationalRetrievalChain.from_llm(
+        chat_mdl, retriever=vec_store.retrieval, memory=mem)
+
+    try:
+        continue_chat = True
+        while continue_chat:
+            question: str = input(conf['username'] + ': ')
+            answer = qa(question)
+            print(repo_name + ': ' + answer['answer'])
+
+            choice: str = (
+                input(
+                    '[C]ontinue chatting, [R]eset chat history, or [Q]uit? '
+                ).strip().upper()
+            )
+            if choice == 'C':
+                continue
+            if choice == 'R':
+                mem.clear()
+                continue
+            if choice == 'Q':
+                continue_chat = False
+    except KeyboardInterrupt:
+        print('\n✌')
+
 
 if __name__ == '__main__':
     main()

senior_swe_ai/conf.py

@@ -52,7 +52,7 @@ def config_init() -> None:
 
         while api_validate is False:
             api_key: str = getpass.getpass('Enter your OpenAI API key: ')
-            api_validate: bool = validate_api_key(api_key)
+            api_validate: bool = validate_api_key(api_key.strip())
             if api_validate is False:
                 print('Invalid API key. Please try again.')
 
@@ -71,7 +71,7 @@ def config_init() -> None:
 
     conf: dict[str, str] = {
         'api_key': api_key,
-        'username': get_username(),
+        'username': get_username().upper(),
         'embed_model': answers['embed_model'],
         'chat_model': 'gpt-3.5-turbo'
     }
@@ -83,17 +83,19 @@ def validate_api_key(api_key: str) -> bool:
 
     try:
         # Make a simple request to the API
-        client = openai.OpenAI()
-        openai.api_key = api_key
+        client = openai.OpenAI(api_key=api_key)
+        # openai.api_key = api_key
         client.embeddings.create(
             input="A test request to validate the API key",
             model="text-embedding-ada-002"
         )
         return True
-    except openai.AuthenticationError:
+    except openai.AuthenticationError as e:
         # If an AuthenticationError is raised, the API key is invalid
+        print(f"AuthenticationError: {e}")
         return False
-    except openai.OpenAIError:
+    except openai.OpenAIError as e:
+        print(f"OpenAIError: {e}")
         return False
 
 
@@ -107,16 +109,26 @@ def save_conf(conf) -> None:
 def append_conf(conf: dict[Any, Any]) -> None:
     """ Append the configuration to the file """
     conf_file_path: str = get_config_path()
-    with open(conf_file_path, 'a', encoding='utf-8') as conf_file:
-        toml.dump(conf, conf_file)
+    conf_item: dict[Any, Any] = load_conf()
+
+    # Update the existing configuration with the new configuration
+    conf_item.update(conf)
+
+    # Write the updated configuration back to the file
+    with open(conf_file_path, 'w', encoding='utf-8') as conf_file:
+        toml.dump(conf_item, conf_file)
 
 
 def load_conf() -> dict[Any, Any]:
     """ Load the configuration from the file """
     conf_file_path: str = get_config_path()
-    with open(conf_file_path, 'r', encoding='utf-8') as conf_file:
-        conf: dict[Any, Any] = toml.load(conf_file)
-    return conf
+    try:
+        with open(conf_file_path, 'r', encoding='utf-8') as conf_file:
+            conf: dict[Any, Any] = toml.load(conf_file)
+            return conf
+    except FileNotFoundError as e:
+        raise FileNotFoundError(
+            'Configuration file not found, `init` the app first.') from e
 
 
 def get_username() -> str: