add frontend in gradio

lib/utils.py

@@ -35,12 +35,13 @@ def select_model():
     for i, model in enumerate(models):
         print(f"{i + 1}. {model}")
 
-    while True:
-        try:
-            choice = int(input("Select a model by number: ")) - 1
-            if 0 <= choice < len(models):
-                return models[choice]
-            else:
-                print("Invalid choice. Please select a number from the list.")
-        except ValueError:
-            print("Invalid input. Please enter a number.")
+    return models[0] #MH - temp, just use openAI
+    # while True:
+    #     try:
+    #         choice = int(input("Select a model by number: ")) - 1
+    #         if 0 <= choice < len(models):
+    #             return models[choice]
+    #         else:
+    #             print("Invalid choice. Please select a number from the list.")
+    #     except ValueError:
+    #         print("Invalid input. Please enter a number.")

main.py

@@ -12,20 +12,44 @@ from lib.chain import create_retriever, create_qa_chain
 from lib.utils import read_prompt, load_LLM, select_model
 from lib.models import MODELS_MAP
 
+import time
+import gradio as gr
+
+def slow_echo(message, history):
+    for i in range(len(message)):
+        time.sleep(0.05)
+        yield message[: i + 1]
+
 # set_debug(True)
 
-def main():
+def build():
+    with gr.Blocks() as demo:
+        repo_url = gr.Textbox(label="Repo URL", placeholder="Enter the repository URL here...")
+        submit_btn = gr.Button("Submit Repo URL")
+
+        user_input = gr.Textbox(label="User Input", placeholder="Enter your question here...")
+        chat_output = gr.Textbox(label="Chat Output", placeholder="The answer will appear here...")
+        # add a status textbox
+
+        def update_repo_url(new_url):
+            updated_url = main(new_url)
+            return updated_url
+
+        def generate_answer(user_input):
+            answer = qa_chain.invoke(user_input)
+            print(f"Answer: {answer}")
+            return answer['output']
+
+        submit_btn.click(update_repo_url, inputs=repo_url, outputs=repo_url)
+        user_input_submit_btn = gr.Button("Submit Question")
+        user_input_submit_btn.click(generate_answer, inputs=user_input, outputs=chat_output)
+
+    demo.launch()
+
+def main(repo_url):
     # Prompt user to select the model
     model_name = select_model()
     model_info = MODELS_MAP[model_name]
-
-    # Parse the command line arguments
-    parser = argparse.ArgumentParser(description="GitHub Repo QA CLI Application")
-    parser.add_argument("repo_url", type=str, help="URL of the GitHub repository")
-    args = parser.parse_args()
-
-    # Extract the repository name from the URL
-    repo_url = args.repo_url
     repo_name = repo_url.split("/")[-1].replace(".git", "")
 
     # Compute the path to the data folder relative to the script's directory
@@ -53,16 +77,10 @@ def main():
     print(f"Creating retrieval QA chain using {model_name}...")
     llm = load_LLM(model_name)
     retriever = create_retriever(model_name, db_dir, document_chunks)
+    global qa_chain
     qa_chain = create_qa_chain(llm, retriever, prompts_text)
-
-    print("You can start asking questions. Type 'exit' to quit.")
-    while True:
-        question = input("Question: ")
-        if question.lower() == "exit":
-            break
-        answer = qa_chain.invoke(question)
-        print(f"Answer: {answer['output']}")
-
+    print(f"Ready to chat!")
+    return repo_url
 
 if __name__ == "__main__":
-    main()
+    build()

prompt_templates/README.md

@@ -1,4 +1,6 @@
 `python -m venv .venv`
 `source .venv/bin/activate`
 `pip3 install -r requirements.txt`
-`python3 main.py https://github.com/streamlit/streamlit`
+`python3 main.py https://github.com/streamlit/streamlit`
+
+where the URL is the repo you want to ask questions about

requirements.txt

@@ -17,4 +17,5 @@ tree_sitter
 tree_sitter_languages
 pysqlite3-binary
 git
-gradio
+gradio
+time