Initial Draft

app.py

@@ -13,6 +13,13 @@ from langchain.prompts.prompt import PromptTemplate
 from langchain_community.llms import LlamaCpp
 from langchain.chains import RetrievalQA
 
+from dotenv import load_dotenv
+import google.generativeai as genai
+
+# Loading Google Gemini
+load_dotenv()
+genai.configure(api_key=os.getenv("GOOGLE_API_KEY"))
+
 # Upload pdf file into 'pdf-data' folder if it does not exist
 def fn_upload_pdf(mv_pdf_input_file, mv_processing_message):
     """Upload pdf file into 'pdf-data' folder if it does not exist"""
@@ -229,6 +236,57 @@ def fn_generate_QnA_response(mv_selected_model, mv_user_question, lv_vector_stor
 
     return lv_llm_response
 
+# Function return API based QA Response using Vector Store
+def fn_generate_API_QnA_response(mv_selected_model, mv_user_question, lv_vector_store, mv_processing_message):
+    """Returns QA Response using Vector Store"""
+
+    lv_template   = """Instruction:
+                    You are an AI assistant for answering questions about the provided context.
+                    You are given the following extracted parts of a long document and a question. Provide a detailed answer.
+                    If you don't know the answer, just say "Hmm, I'm not sure." Don't try to make up an answer.
+                    =======
+                    {context}
+                    =======
+                    Question: {question}
+                    Output:\n"""
+    lv_qa_prompt = PromptTemplate(
+                                template=lv_template,
+                                input_variables=["question", "context"]
+                              )
+
+    lv_vector_search_result = lv_vector_store.similarity_search(mv_user_question, k=2)
+    # print("Vector Search Result - ")
+    # print(lv_vector_search_result)
+
+    # -- Creating formatted document result
+    lv_document_context = ""
+    lv_count = 0
+    for lv_result in lv_vector_search_result:
+        # print("Concatenating Result of page - " + str(lv_count) + " with content of document page no - "+str(lv_result.metadata["page"]))
+        lv_document_context += lv_result.page_content 
+        lv_count += 1
+
+    # print("Formatted Document Search Result - ")
+    # print(lv_document_context)
+
+    lv_qa_formatted_prompt = lv_qa_prompt.format(  
+                                            question=mv_user_question,
+                                            context=lv_document_context
+                                         )
+    
+    if mv_selected_model == 'Google Gemini-pro':
+        lv_model = genai.GenerativeModel('gemini-pro')
+
+        print("Step4: Generating LLM response")
+        fn_display_user_messages("Step4: Generating LLM response","Info", mv_processing_message)
+    
+        lv_llm_response = lv_model.generate_content(lv_qa_prompt)
+    
+        print("Step5: LLM response generated")
+        fn_display_user_messages("Step5: LLM response generated","Info", mv_processing_message)
+    
+        return lv_llm_response
+    
 # Main Function
 def main():
     
@@ -248,7 +306,8 @@ def main():
                                         [
                                             'microsoft/phi-2',
                                             'google/gemma-2b',
-                                            'mistralai/Mistral-7B-Instruct-v0.2'
+                                            'mistralai/Mistral-7B-Instruct-v0.2',
+                                            'Google Gemini-pro'
                                         ]
                                       )
     
@@ -268,7 +327,10 @@ def main():
     st.text("")
 
     # -- Downloading Model Files
-    fn_download_llm_models(mv_selected_model, mv_processing_message)
+    if mv_selected_model != "Google Gemini-pro":
+        fn_download_llm_models(mv_selected_model, mv_processing_message)
+    else:
+        print("Call Google API")
 
     # -- Processing PDF
     if (mv_pdf_input_file is not None):
@@ -290,8 +352,11 @@ def main():
            st.session_state.messages.append({"role": "user", "content": mv_user_question})
 
            # -- Generating LLM response
-           lv_response = fn_generate_QnA_response(mv_selected_model, mv_user_question, lv_vector_store, mv_processing_message)
-
+           if mv_selected_model != "Google Gemini-pro":
+               lv_response = fn_generate_QnA_response(mv_selected_model, mv_user_question, lv_vector_store, mv_processing_message)
+           else:
+               lv_response = fn_generate_API_QnA_response(mv_selected_model, mv_user_question, lv_vector_store, mv_processing_message)
+           
            # -- Adding assistant response to chat history
            st.session_state.messages.append({"role": "assistant", "content": lv_response})