Fix

rag_processor.py

@@ -1,8 +1,7 @@
 # rag_processor.py
 
 import os
-from dotenv import load_dotenv
-from operator import itemgetter # <--- ADD THIS IMPORT
+from operator import itemgetter
 
 # LLM
 from langchain_groq import ChatGroq
@@ -10,90 +9,58 @@ from langchain_groq import ChatGroq
 # Prompting
 from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
 
-# Chains
-from langchain_core.runnables import RunnableParallel, RunnablePassthrough
-from langchain_core.output_parsers import StrOutputParser
+# Chains and Memory
 from langchain_core.runnables.history import RunnableWithMessageHistory
+from langchain_core.output_parsers import StrOutputParser
+
+def format_docs(docs):
+    """A helper function to format retrieved documents into a single string."""
+    return "\n\n".join(doc.page_content for doc in docs)
 
 def create_rag_chain(retriever, get_session_history_func):
     """
-    Creates an advanced Retrieval-Augmented Generation (RAG) chain with hybrid search,
-    query rewriting, answer refinement, and conversational memory.
-
-    Args:
-        retriever: A configured LangChain retriever object.
-        get_session_history_func: A function to get the chat history for a session.
-
-    Returns:
-        A LangChain runnable object representing the RAG chain with memory.
-        
-    Raises:
-        ValueError: If the GROQ_API_KEY is missing.
+    Creates a simplified and robust Retrieval-Augmented Generation (RAG) chain with memory.
+    This version uses a single, efficient call to the LLM per query.
     """
-    # Load environment variables from .env file
+    # --- 1. Get the API Key from Environment Secrets ---
+    # This correctly reads the secret you set on the Hugging Face Space.
     api_key = os.getenv("GROQ_API_KEY")
     if not api_key:
-        raise ValueError("GROQ_API_KEY not found in environment variables.")
-
-    # --- 1. Initialize the LLM ---
-    # Updated model_name to a standard, high-performance Groq model
-    llm = ChatGroq(model_name="llama-3.1-8b-instant", api_key=api_key, temperature=1)
+        raise ValueError("GROQ_API_KEY secret not found in environment variables.")
 
-    # --- 2. Create Query Rewriting Chain 🧠 ---
-    print("\nSetting up query rewriting chain...")
-    rewrite_template = """You are an expert at rewriting user questions for a vector database.
-You are here to help the user with their document.
-Based on the chat history, reformulate the follow-up question to be a standalone question.
-This new query should be optimized to find the most relevant documents in a knowledge base.
-Do NOT answer the question, only provide the rewritten, optimized question.
+    # --- 2. Initialize the LLM ---
+    llm = ChatGroq(model_name="llama3-8b-8192", api_key=api_key, temperature=0.7)
 
-Chat History:
-{chat_history}
-
-Follow-up Question: {question}
-Standalone Question:"""
-    rewrite_prompt = ChatPromptTemplate.from_messages([
-        ("system", rewrite_template),
-        MessagesPlaceholder(variable_name="chat_history"),
-        ("human", "Based on our conversation, reformulate this question to be a standalone query: {question}")
-    ])
-    query_rewriter = rewrite_prompt | llm | StrOutputParser()
-    
-    # --- 3. Create Main RAG Chain with Memory ---
-    print("\nSetting up main RAG chain...")
-    rag_template = """You are an expert assistant named `Cognichat`.Whenver user ask you about who you are , simply say you are `Cognichat`.
-    You are developed by Ritesh and Alish.
-    Your job is to provide accurate and helpful answers based ONLY on the provided context.
-If the information is not in the context, clearly state that you don't know the answer.
-Provide a clear and concise answer.
+    # --- 3. Define the Conversational RAG Prompt ---
+    # This single prompt handles context, chat history, and the user's question.
+    rag_template = """You are an expert assistant named `CogniChat`, developed by Ritesh and Alish.
+Your job is to provide accurate and helpful answers based ONLY on the provided context.
+If the information to answer the question is not in the context, clearly state that the document does not contain the answer.
+Be concise and clear in your responses. Use formatting like bold text or bullet points if it helps clarity.
 
 Context:
-{context}"""
+{context}
+"""
     rag_prompt = ChatPromptTemplate.from_messages([
         ("system", rag_template),
         MessagesPlaceholder(variable_name="chat_history"),
         ("human", "{question}"),
     ])
 
-    # ============================ FIX IS HERE ============================
-    # Parallel process to fetch context and correctly pass through question and history.
-    # We use itemgetter to select the specific keys from the input dictionary.
-    setup_and_retrieval = RunnableParallel({
-        "context": query_rewriter | retriever,
-        "question": itemgetter("question"),
-        "chat_history": itemgetter("chat_history"),
-    })
-    # =====================================================================
-    
-    # The initial RAG chain
+    # --- 4. Create the Core RAG Chain ---
+    # This is a standard and efficient way to build a RAG chain.
     conversational_rag_chain = (
-        setup_and_retrieval
+        {
+            "context": itemgetter("question") | retriever | format_docs,
+            "question": itemgetter("question"),
+            "chat_history": itemgetter("chat_history"),
+        }
         | rag_prompt
         | llm
         | StrOutputParser()
     )
 
-    # Wrap the chain with memory management
+    # --- 5. Wrap the Chain with Memory Management ---
     chain_with_memory = RunnableWithMessageHistory(
         conversational_rag_chain,
         get_session_history_func,
@@ -101,26 +68,5 @@ Context:
         history_messages_key="chat_history",
     )
 
-    # --- 4. Create Answer Refinement Chain ✨ ---
-    print("\nSetting up answer refinement chain...")
-    refine_template = """You are an expert at editing and refining content.
-Your task is to take a given answer and improve its clarity, structure, and readability.
-Use formatting such as bold text, bullet points, or numbered lists where it enhances the explanation.
-Do not add any new information that wasn't in the original answer.
-
-Original Answer:
-{answer}
-
-Refined Answer:"""
-    refine_prompt = ChatPromptTemplate.from_template(refine_template)
-    refinement_chain = refine_prompt | llm | StrOutputParser()
-
-    # --- 5. Combine Everything into the Final Chain ---
-    # The final chain passes the output of the memory-enabled chain to the refinement chain
-    # Note: We need to adapt the input for the refinement chain
-    final_chain = (
-        lambda input_dict: {"answer": chain_with_memory.invoke(input_dict, config=input_dict.get('config'))}
-    ) | refinement_chain
-    
-    print("\nFinalizing the complete chain with memory...")
-    return final_chain
+    print("\n✅ Simplified RAG chain with memory created successfully.")
+    return chain_with_memory