Update app.py

app.py

@@ -90,55 +90,74 @@ def process_documents(selected_files):
 # ✅ Query document
 
 
+import os
+import time
+import logging
+from gtts import gTTS
+from langchain.chains import RetrievalQA
+from langchain_google_genai import ChatGoogleGenerativeAI
+
+# ✅ Ensure temp_file_map exists
+temp_file_map = {}
+
 def query_document(question):
-    
     if vector_store is None:
         return "❌ No documents processed.", None
 
     # ✅ Fetch stored documents  
-    stored_docs = vector_store.get()["documents"]  
+    stored_docs = vector_store.get()["documents"]
 
     # ✅ Calculate total word count safely  
-    total_words = sum(len(doc.split()) if isinstance(doc, str) else len(doc.page_content.split()) for doc in stored_docs)  
-
-    # ✅ Dynamically adjust k based on document size  
-    if total_words < 500:  
-        k_value = 3  
-    elif total_words < 2000:  
-        k_value = 5  
-    else:  
-        k_value = 10  
-
-    retriever = vector_store.as_retriever(search_type="similarity", search_kwargs={"k": k_value})  
-
-    # ✅ Improved prompt for detailed response  
-    detailed_prompt = f"""  
-    Provide a **detailed and structured answer** to the following question.  
-    - Use relevant **examples, key points, and explanations**.  
-    - If applicable, provide **step-by-step analysis** or comparisons.  
-    - Ensure **clarity and completeness**.  
+    total_words = sum(len(doc.split()) if isinstance(doc, str) else len(doc.page_content.split()) for doc in stored_docs)
+
+    # ✅ Categorize file size  
+    if total_words < 500:
+        file_size_category = "small"
+        k_value = 3
+    elif total_words < 2000:
+        file_size_category = "medium"
+        k_value = 5
+    else:
+        file_size_category = "large"
+        k_value = 10
+
+    retriever = vector_store.as_retriever(search_type="similarity", search_kwargs={"k": k_value})
+
+    # ✅ Adjust response detail based on file size  
+    if file_size_category == "small":
+        prompt_prefix = "Provide a **concise** response focusing on key points."
+    elif file_size_category == "medium":
+        prompt_prefix = "Provide a **detailed response** with examples and key insights."
+    else:
+        prompt_prefix = "Provide a **comprehensive and structured response**, including step-by-step analysis and explanations."
+
+    # ✅ Final prompt  
+    detailed_prompt = f"""{prompt_prefix}  
+    - Ensure clarity and completeness.  
+    - Highlight the most relevant information.  
     **Question:** {question}  
-    """  
+    """
 
-    # ✅ Dynamically select model based on document size  
-    if total_words < 1000:  
-        model_name = "gemini-2.0-pro-exp-02-05"  # More detailed responses for small files  
-    else:  
-        model_name = "gemini-2.0-flash"  # Faster processing for large documents  
+    # ✅ Dynamically select model based on file size  
+    if file_size_category in ["small", "medium"]:
+        model_name = "gemini-2.0-pro-exp-02-05"
+    else:
+        model_name = "gemini-2.0-flash"
 
-    logging.info(f"🧠 Using Model: {model_name} for processing")  
+    logging.info(f"🧠 Using Model: {model_name} for {file_size_category} file.")
 
-    model = ChatGoogleGenerativeAI(model=model_name, google_api_key=GOOGLE_API_KEY)  
-    qa_chain = RetrievalQA.from_chain_type(llm=model, retriever=retriever)  
-    response = qa_chain.invoke({"query": detailed_prompt})["result"]  
+    model = ChatGoogleGenerativeAI(model=model_name, google_api_key=GOOGLE_API_KEY)
+    qa_chain = RetrievalQA.from_chain_type(llm=model, retriever=retriever)
+    response = qa_chain.invoke({"query": detailed_prompt})["result"]
 
     # ✅ Convert response to speech  
-    tts = gTTS(text=response, lang="en")  
-    temp_audio_path = os.path.join(temp_dir, "response.mp3")  
-    tts.save(temp_audio_path)  
-    temp_file_map["response.mp3"] = time.time()  
+    tts = gTTS(text=response, lang="en")
+    temp_audio_path = os.path.join(temp_dir, "response.mp3")
+    tts.save(temp_audio_path)
+    temp_file_map["response.mp3"] = time.time()
+
+    return response, temp_audio_path
 
-    return response, temp_audio_path 
 
 
 # ✅ Gradio UI