Update app.py

app.py

@@ -243,19 +243,26 @@ def retry_last_response(history, use_web_search, model, temperature, num_calls):
     
     return chatbot_interface(last_user_msg, history, use_web_search, model, temperature, num_calls)
 
-def respond(message, history, model, temperature, num_calls, use_web_search, selected_docs):
+def respond(message, history, model, temperature, num_calls, use_web_search, selected_docs, instruction_key):
     logging.info(f"User Query: {message}")
     logging.info(f"Model Used: {model}")
     logging.info(f"Search Type: {'Web Search' if use_web_search else 'PDF Search'}")
-
     logging.info(f"Selected Documents: {selected_docs}")
+    logging.info(f"Instruction Key: {instruction_key}")
 
     try:
+        if instruction_key and instruction_key != "None":
+            # This is a summary generation request
+            instruction = INSTRUCTION_PROMPTS[instruction_key]
+            context_str = get_context_for_summary(selected_docs)
+            message = f"{instruction}\n\nUsing the following context from the PDF documents:\n{context_str}\nGenerate a detailed summary."
+            use_web_search = False  # Ensure we use PDF search for summaries
+
         if use_web_search:
             for main_content, sources in get_response_with_search(message, model, num_calls=num_calls, temperature=temperature):
                 response = f"{main_content}\n\n{sources}"
                 first_line = response.split('\n')[0] if response else ''
-#                logging.info(f"Generated Response (first line): {first_line}")
+                logging.info(f"Generated Response (first line): {first_line}")
                 yield response
         else:
             embed = get_embeddings()
@@ -289,17 +296,25 @@ def respond(message, history, model, temperature, num_calls, use_web_search, sel
                     first_line = partial_response.split('\n')[0] if partial_response else ''
                     logging.info(f"Generated Response (first line): {first_line}")
                     yield partial_response
+
     except Exception as e:
         logging.error(f"Error with {model}: {str(e)}")
         if "microsoft/Phi-3-mini-4k-instruct" in model:
             logging.info("Falling back to Mistral model due to Phi-3 error")
             fallback_model = "mistralai/Mistral-7B-Instruct-v0.3"
-            yield from respond(message, history, fallback_model, temperature, num_calls, use_web_search, selected_docs)
+            yield from respond(message, history, fallback_model, temperature, num_calls, use_web_search, selected_docs, instruction_key)
         else:
             yield f"An error occurred with the {model} model: {str(e)}. Please try again or select a different model."
 
 logging.basicConfig(level=logging.DEBUG)
 
+INSTRUCTION_PROMPTS = {
+    "Asset Managers": "Summarize the key financial metrics, assets under management, and performance highlights for this asset management company.",
+    "Consumer Finance Companies": "Provide a summary of the company's loan portfolio, interest income, credit quality, and key operational metrics.",
+    "Mortgage REITs": "Summarize the REIT's mortgage-backed securities portfolio, net interest income, book value per share, and dividend yield.",
+    # Add more instruction prompts as needed
+}
+
 def get_response_from_cloudflare(prompt, context, query, num_calls=3, temperature=0.2, search_type="pdf"):
     headers = {
         "Authorization": f"Bearer {API_TOKEN}",
@@ -467,6 +482,43 @@ css = """
 }
 """
 
+def get_context_for_summary(selected_docs):
+    embed = get_embeddings()
+    if os.path.exists("faiss_database"):
+        database = FAISS.load_local("faiss_database", embed, allow_dangerous_deserialization=True)
+        retriever = database.as_retriever(search_kwargs={"k": 5})  # Retrieve top 5 most relevant chunks
+        
+        # Create a generic query that covers common financial summary topics
+        generic_query = "financial performance revenue profit assets liabilities cash flow key metrics highlights"
+        
+        relevant_docs = retriever.get_relevant_documents(generic_query)
+        filtered_docs = [doc for doc in relevant_docs if doc.metadata["source"] in selected_docs]
+        
+        if not filtered_docs:
+            return "No relevant information found in the selected documents for summary generation."
+        
+        context_str = "\n".join([doc.page_content for doc in filtered_docs])
+        return context_str
+    else:
+        return "No documents available for summary generation."
+
+def get_context_for_query(query, selected_docs):
+    embed = get_embeddings()
+    if os.path.exists("faiss_database"):
+        database = FAISS.load_local("faiss_database", embed, allow_dangerous_deserialization=True)
+        retriever = database.as_retriever(search_kwargs={"k": 3})  # Retrieve top 3 most relevant chunks
+        
+        relevant_docs = retriever.get_relevant_documents(query)
+        filtered_docs = [doc for doc in relevant_docs if doc.metadata["source"] in selected_docs]
+        
+        if not filtered_docs:
+            return "No relevant information found in the selected documents for the given query."
+        
+        context_str = "\n".join([doc.page_content for doc in filtered_docs])
+        return context_str
+    else:
+        return "No documents available to answer the query."
+
 uploaded_documents = []
 
 def display_documents():
@@ -483,6 +535,8 @@ use_web_search = gr.Checkbox(label="Use Web Search", value=True)
 
 custom_placeholder = "Ask a question (Note: You can toggle between Web Search and PDF Chat in Additional Inputs below)"
 
+instruction_choices = ["None"] + list(INSTRUCTION_PROMPTS.keys())
+
 demo = gr.ChatInterface(
     respond,
     additional_inputs=[
@@ -490,10 +544,11 @@ demo = gr.ChatInterface(
         gr.Slider(minimum=0.1, maximum=1.0, value=0.2, step=0.1, label="Temperature"),
         gr.Slider(minimum=1, maximum=5, value=1, step=1, label="Number of API Calls"),
         use_web_search,
-        document_selector
+        document_selector,
+        gr.Dropdown(choices=instruction_choices, label="Select Entity Type for Summary", value="None")
     ],
     title="AI-powered Web Search and PDF Chat Assistant",
-    description="Chat with your PDFs or use web search to answer questions. Toggle between Web Search and PDF Chat in Additional Inputs below.",
+    description="Chat with your PDFs, use web search to answer questions, or generate summaries. Select an Entity Type for Summary to generate a specific summary.",
     theme=gr.themes.Soft(
         primary_hue="orange",
         secondary_hue="amber",
@@ -522,11 +577,11 @@ demo = gr.ChatInterface(
     analytics_enabled=False,
     textbox=gr.Textbox(placeholder=custom_placeholder, container=False, scale=7),
     chatbot = gr.Chatbot(  
-    show_copy_button=True,
-    likeable=True,
-    layout="bubble",
-    height=400,
-)
+        show_copy_button=True,
+        likeable=True,
+        layout="bubble",
+        height=400,
+    )
 )
 
 # Add file upload functionality