Changes for Citation

agents/rag_agent.py

@@ -80,58 +80,68 @@ def load_rag_agent(vectorstore_path: str = "vectorstore/"):
     # Grounded prompt
     prompt_template = """You are a helpful HDFC Bank policy assistant.
 
-Use ONLY the context below to answer the customer's question.
+    Use ONLY the context below to answer the customer's question.
 
-IMPORTANT:
-- Always include the sources at the end of your answer.
-- Also include a short explanation titled "Why this answer?"
-- Explain briefly how the answer was derived from context
-- The sources are provided in the context.
-- Format sources exactly like this:
+    IMPORTANT:
+    - Always include the sources at the end of your answer.
+    - Also include a short explanation titled "Why this answer?"
+    - Explain briefly how the answer was derived from context
+    - The sources are provided in the context.
+    - Format sources exactly like this:
 
-Sources:
-- file1.pdf
-- file2.pdf
+    Sources:
+    - file1.pdf
+    - file2.pdf
 
-If the answer is not in the context, say:
-"I don't have enough information in the policy documents to answer this. Please contact HDFC Bank directly."
+    If the answer is not in the context, say:
+    "I don't have enough information in the policy documents to answer this. Please contact HDFC Bank directly."
 
-Context:
-{context}
+    Context:
+    {context}
 
-Customer Question: {question}
+    Customer Question: {question}
 
-Answer:"""
+    Answer:"""
 
     prompt = PromptTemplate(
         template=prompt_template,
         input_variables=["context", "question"]
     )
 
-    def format_docs(docs):
-        formatted = []
+    def extract_sources(docs):
         sources = []
 
         for doc in docs:
             source = doc.metadata.get("source", "Unknown")
             filename = os.path.basename(source)
-
             sources.append(filename)
-            formatted.append(doc.page_content)
 
-        context = "\n\n".join(formatted)
+        return list(set(sources))
 
-        unique_sources = list(set(sources))
-        source_text = "\n\nSources:\n" + "\n".join(f"- {s}" for s in unique_sources)
+    def run_rag(question):
+        # 1. Retrieve documents
+        docs = retriever.get_relevant_documents(question)
 
-        return context + source_text
+        # 2. Create context for LLM
+        context = "\n\n".join(doc.page_content for doc in docs)
 
-    # LCEL chain
-    rag_chain = (
-        {"context": retriever | format_docs, "question": RunnablePassthrough()}
-        | prompt
-        | llm
-        | StrOutputParser()
-    )
+        # 3. Extract sources separately
+        sources = []
+        for doc in docs:
+            source = doc.metadata.get("source", "Unknown")
+            filename = os.path.basename(source)
+            sources.append(filename)
+
+        sources = list(set(sources))
+
+        # 4. Call LLM
+        response = llm.invoke(
+            prompt.format(context=context, question=question)
+        )
+
+        return {
+            "answer": response.content,
+            "sources": sources
+        }
 
-    return rag_chain
+    return run_rag

app.py

@@ -77,27 +77,25 @@ if query := st.chat_input("Ask your banking question here..."):
                 "agent_used": "",
                 "response": ""
             })
-
-        response = format_currency(result["response"])
+        
+        response = result["answer"]
+        sources = result["sources"]
         agent_used = result["agent_used"].upper()
 
-        # Show which agent handled it
+        # Show agent
         if agent_used == "RAG":
-            st.caption(" Answered by: Policy Agent (RAG)")
+            st.caption("Answered by: Policy Agent (RAG)")
         else:
-            st.caption(" Answered by: Data Agent (SQL)")
-
-        if "Sources:" in response:
-            answer, sources = response.split("Sources:")
+            st.caption("Answered by: Data Agent (SQL)")
 
         # Show answer
-            st.markdown(answer)
+        st.markdown(response)
 
-        # Show sources nicely
+        # Show sources
+        if sources:
             st.markdown("### Sources")
-            for s in sources.strip().split("\n"):
-                if s.strip():
-                   st.markdown(f"- {s.replace('-', '').strip()}")
+            for s in sources:
+                st.markdown(f"- {s}")
         else:
             st.markdown(response)