Update app.py

app.py

@@ -1,21 +1,42 @@
 import os
 import gradio as gr
-from langchain_community.document_loaders import PyPDFLoader, TextLoader, Docx2txtLoader
-from langchain_text_splitters import RecursiveCharacterTextSplitter
-from langchain_huggingface import HuggingFaceEmbeddings
-from langchain_community.vectorstores import FAISS
-from langchain_groq import ChatGroq
+
+# Classic & Community Imports
 from langchain_classic.chains import ConversationalRetrievalChain
 from langchain_classic.memory import ConversationBufferMemory
-
-# --- NEW IMPORTS FOR HYBRID SEARCH ---
+from langchain_groq import ChatGroq
+from langchain_community.vectorstores import FAISS
+from langchain_huggingface import HuggingFaceEmbeddings
+from langchain_community.document_loaders import PyPDFLoader, TextLoader, Docx2txtLoader
+from langchain_text_splitters import RecursiveCharacterTextSplitter
 from langchain_community.retrievers import BM25Retriever
 from langchain.retrievers import EnsembleRetriever
+from langchain.prompts import PromptTemplate
+
+# --- 1. SETUP API & SYSTEM PROMPT ---
+# Hugging Face uses os.getenv for secrets
+api_key = os.getenv("GROQ_API")
+
+STRICT_PROMPT_TEMPLATE = """You are a strict document-based assistant. 
+Use the following pieces of context to answer the user's question. 
+
+RESTRICTIONS:
+1. ONLY use the information provided in the context below.
+2. If the answer is not contained within the context, specifically say: "I'm sorry, but the provided documents do not contain information to answer this question."
+3. Do NOT use your own outside knowledge.
+
+Context:
+{context}
+
+Question: {question}
+Helpful Answer:"""
 
-# 1. SETUP API
-api_key = os.environ.get("GROQ_API")
+STRICT_PROMPT = PromptTemplate(
+    template=STRICT_PROMPT_TEMPLATE, 
+    input_variables=["context", "question"]
+)
 
-# 2. FILE LOADING LOGIC
+# --- 2. LOADING LOGIC ---
 def load_any(path: str):
     p = path.lower()
     if p.endswith(".pdf"): return PyPDFLoader(path).load()
@@ -23,41 +44,39 @@ def load_any(path: str):
     if p.endswith(".docx"): return Docx2txtLoader(path).load()
     return []
 
-# 3. HYBRID PROCESSING FUNCTION
-def process_files(files):
+# --- 3. HYBRID PROCESSING ---
+def process_files(files, response_length):
     if not files or not api_key:
-        return None, "⚠️ Missing files or API key."
+        return None, "⚠️ Missing files or GROQ_API key in Secrets."
 
     try:
-        # Load all documents
         docs = []
         for file_obj in files:
             docs.extend(load_any(file_obj.name))
         
-        if not docs:
-            return None, "⚠️ No readable text found."
-
-        # Split into chunks
-        splitter = RecursiveCharacterTextSplitter(chunk_size=700, chunk_overlap=100)
+        splitter = RecursiveCharacterTextSplitter(chunk_size=800, chunk_overlap=100)
         chunks = splitter.split_documents(docs)
 
-        # A. Semantic Search (FAISS)
+        # Hybrid Retrievers
         embeddings = HuggingFaceEmbeddings(model_name="sentence-transformers/all-MiniLM-L6-v2")
         faiss_db = FAISS.from_documents(chunks, embeddings)
         faiss_retriever = faiss_db.as_retriever(search_kwargs={"k": 3})
-
-        # B. Keyword Search (BM25) - THIS IS THE MULTI-RETRIEVER ADDITION
+        
         bm25_retriever = BM25Retriever.from_documents(chunks)
         bm25_retriever.k = 3
 
-        # C. Ensemble (Hybrid Search)
         ensemble_retriever = EnsembleRetriever(
             retrievers=[faiss_retriever, bm25_retriever],
-            weights=[0.6, 0.4] # 60% Semantic, 40% Keyword
+            weights=[0.5, 0.5]
         )
 
-        # D. Classic Chain Setup
-        llm = ChatGroq(groq_api_key=api_key, model="llama-3.3-70b-versatile", temperature=0)
+        llm = ChatGroq(
+            groq_api_key=api_key, 
+            model="llama-3.3-70b-versatile", 
+            temperature=0,
+            max_tokens=int(response_length)
+        )
+        
         memory = ConversationBufferMemory(
             memory_key="chat_history", 
             return_messages=True, 
@@ -66,18 +85,19 @@ def process_files(files):
         
         chain = ConversationalRetrievalChain.from_llm(
             llm=llm,
-            retriever=ensemble_retriever, # Use Hybrid Retriever
+            retriever=ensemble_retriever,
+            combine_docs_chain_kwargs={"prompt": STRICT_PROMPT},
             memory=memory,
             return_source_documents=True,
             output_key="answer"
         )
         
-        return chain, f"✅ Hybrid Multi-RAG Ready! ({len(chunks)} chunks)"
+        return chain, f"✅ Knowledge base built! Max answer length: {response_length} tokens."
 
     except Exception as e:
         return None, f"❌ Error: {str(e)}"
 
-# 4. CHAT FUNCTION
+# --- 4. CHAT FUNCTION ---
 def chat_function(message, history, chain):
     if not chain:
         return "⚠️ Build the chatbot first."
@@ -85,27 +105,27 @@ def chat_function(message, history, chain):
     res = chain.invoke({"question": message})
     answer = res["answer"]
     
-    # Format Sources
     sources = list(set([os.path.basename(d.metadata.get("source", "unknown")) for d in res.get("source_documents", [])]))
-    source_text = "\n\n---\n**Sources:** " + ", ".join(sources)
+    source_display = "\n\n----- \n**Sources used:** " + ", ".join(sources)
     
-    return answer + source_text
+    return answer + source_display
 
-# 5. UI
-with gr.Blocks(title="Hybrid RAG") as demo:
-    gr.Markdown("# 🚀 Hybrid Multi-RAG Chatbot")
+# --- 5. UI BUILDING ---
+with gr.Blocks(theme=gr.themes.Soft()) as demo:
+    gr.Markdown("# 🛡️ Strict Hybrid Multi-RAG")
     chain_state = gr.State(None)
     
     with gr.Row():
         with gr.Column(scale=1):
-            file_input = gr.File(file_count="multiple", label="Upload Docs")
-            build_btn = gr.Button("Build Hybrid RAG", variant="primary")
+            file_input = gr.File(file_count="multiple", label="1. Upload Documents")
+            len_slider = gr.Slider(minimum=100, maximum=4000, value=1000, step=100, label="2. Response Length")
+            build_btn = gr.Button("3. Build Restricted Chatbot", variant="primary")
             status = gr.Textbox(label="Status", interactive=False)
         
         with gr.Column(scale=2):
             gr.ChatInterface(fn=chat_function, additional_inputs=[chain_state])
 
-    build_btn.click(process_files, inputs=[file_input], outputs=[chain_state, status])
+    build_btn.click(process_files, inputs=[file_input, len_slider], outputs=[chain_state, status])
 
 if __name__ == "__main__":
     demo.launch()