add action arguments

fynnfluegge · Sep 22, 2023 · 3439c12 · 3439c12
1 parent 88d3ec3
commit 3439c12
Show file tree

Hide file tree

Showing 3 changed files with 47 additions and 20 deletions.
diff --git a/codeqai/app.py b/codeqai/app.py
@@ -1,3 +1,4 @@
+import argparse
 import os
 
 from langchain.chains import ConversationalRetrievalChain
@@ -7,12 +8,22 @@
 from yaspin import yaspin
 
 from codeqai import codeparser, repo
-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.vector_store import VectorStore
 
 
 def run():
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "action",
+        choices=["search", "chat", "configure"],
+        help="Action to perform. 'search' will semantically search the codebase. 'chat' will chat with the codebase.",
+    )
+    args = parser.parse_args()
+
+    if args.action == "configure":
+        create_config()
+
     # load config
     config = None
     try:
@@ -47,29 +58,42 @@ def run():
     qa = ConversationalRetrievalChain.from_llm(
         llm, retriever=vector_store.retriever, memory=memory
     )
+
     while True:
-        question = input("🤖 Ask me anything about the codebase: ")
+        choice = None
+        if args.action == "search":
+            search_pattern = input("🤖 Enter a search pattern: ")
+            spinner = yaspin(text="🤖 Processing...", color="green")
+            spinner.start()
+            similarity_result = vector_store.similarity_search(search_pattern)
+            spinner.stop()
+            for doc in similarity_result:
+                # print(doc.metadata["file_name"])
+                # print(doc.metadata["method_name"])
+                # print(doc.page_content)
+                print(doc)
 
-        spinner = yaspin(text=f"🤖 Processing question...")
-        spinner.start()
+            choice = input("(C)ontinue search or (E)xit [C]?").strip().lower()
 
-        result = qa(question)
+        elif args.action == "chat":
+            question = input("🤖 Ask me anything about the codebase: ")
+            spinner = yaspin(text="🤖 Processing...", color="green")
+            spinner.start()
+            result = qa(question)
+            spinner.stop()
+            print(result["answer"])
 
-        spinner.stop()
-        print(result["answer"])
+            choice = (
+                input("(C)ontinue chat, (R)eset chat or (E)xit [C]?").strip().lower()
+            )
 
-        choice = (
-            input("Do you want to (C)ontinue chat, (R)eset chat or (E)xit [C]?")
-            .strip()
-            .lower()
-        )
+            if choice == "r":
+                print("Resetting chat...")
+                memory.clear()
 
         if choice == "" or choice == "c":
             continue
         elif choice == "e":
             break
-        elif choice == "r":
-            print("Resetting chat...")
-            memory.clear()
         else:
             print("Invalid choice. Please enter 'C', 'E', or 'R'.")
diff --git a/codeqai/codeparser.py b/codeqai/codeparser.py
@@ -13,7 +13,7 @@ def parse_code_files(code_files: list[str]) -> list[Document]:
     source_code_documents, docstring_documents = [], []
     source_code_splitter = None
     docstring_splitter = RecursiveCharacterTextSplitter(
-        chunk_size=512, chunk_overlap=128
+        chunk_size=1024, chunk_overlap=128
     )
     for code_file in code_files:
         with open(code_file, "r") as file:
@@ -31,7 +31,7 @@ def parse_code_files(code_files: list[str]) -> list[Document]:
             if langchain_language:
                 source_code_splitter = RecursiveCharacterTextSplitter.from_language(
                     language=langchain_language,
-                    chunk_size=512,
+                    chunk_size=1024,
                     chunk_overlap=128,
                 )
 

diff --git a/codeqai/vector_store.py b/codeqai/vector_store.py
@@ -25,10 +25,13 @@ def __init__(
             self.db.save_local(index_name=self.name, folder_path=get_cache_path())
         self.retriever = self.db.as_retriever(search_type="mmr", search_kwargs={"k": 8})
 
+    def similarity_search(self, query: str):
+        return self.db.similarity_search(query, k=4)
+
     def install_faiss(self):
         try:
-            from faiss import (FAISS_VERSION_MAJOR,  # noqa: F401
-                               FAISS_VERSION_MINOR)
+            from faiss import FAISS_VERSION_MAJOR  # noqa: F401
+            from faiss import FAISS_VERSION_MINOR
         except:  # noqa: E722
             question = [
                 inquirer.Confirm(