Feat: Implement Conversational Memory (Contextual Rewriting)

src/app.py

@@ -103,7 +103,16 @@ if prompt := st.chat_input("Ask about any satellite (e.g., 'What is Gaofen 1?').
                     if health["warning"]:
                         st.warning(f"⚠️ Low Memory Warning: Only {health['available_mb']:.0f}MB available. Query might be slow.")
                     
-                    response, docs = engine.query(prompt)
+                    # Construct Chat History
+                    # We need pairs of (User, AI) from session_state.messages
+                    # Excluding the current new prompt which is already appended but not part of 'history' yet for this context
+                    chat_history = []
+                    msgs = st.session_state.messages[:-1] 
+                    for i in range(0, len(msgs) - 1, 2):
+                        if msgs[i]["role"] == "user" and msgs[i+1]["role"] == "assistant":
+                            chat_history.append((msgs[i]["content"], msgs[i+1]["content"]))
+                    
+                    response, docs = engine.query(prompt, chat_history=chat_history)
 
                     st.markdown(response)
 

src/rag_engine.py

@@ -79,18 +79,50 @@ class SatelliteRAG:
             api_key=settings.GROQ_API_KEY
         )
 
+    def _rewrite_query(self, question: str, chat_history: List[Tuple[str, str]]) -> str:
+        """Rewrite question based on history to be standalone."""
+        if not chat_history:
+            return question
+            
+        logger.info("Rewriting question with conversational context...")
+        
+        template = """Given the following conversation and a follow up question, rephrase the follow up question to be a standalone question.
+        
+        Chat History:
+        {history}
+        
+        Follow Up Input: {question}
+        Standalone Question:"""
+        
+        try:
+            prompt = ChatPromptTemplate.from_template(template)
+            chain = prompt | self.llm | StrOutputParser()
+            
+            # Format history as a string
+            history_str = "\n".join([f"User: {q}\nAssistant: {a}" for q, a in chat_history])
+            
+            standalone_question = chain.invoke({"history": history_str, "question": question})
+            logger.info(f"Rephrased '{question}' -> '{standalone_question}'")
+            return standalone_question
+        except Exception as e:
+            logger.error(f"Failed to rewrite question: {e}")
+            return question
+
     @retry(
         stop=stop_after_attempt(3),
         wait=wait_exponential(multiplier=1, min=2, max=10),
         reraise=True
     )
-    def query(self, question: str) -> Tuple[str, List[Document]]:
+    def query(self, question: str, chat_history: List[Tuple[str, str]] = []) -> Tuple[str, List[Document]]:
         """
         Query the RAG system.
         Retries up to 3 times on failure (e.g. API Rate Limits).
         """
+        # 0. Contextual Rewriting
+        standalone_question = self._rewrite_query(question, chat_history)
+
         # Retrieval
-        logger.info(f"Starting query process for: {question}")
+        logger.info(f"Starting query process for: {standalone_question}")
         try:
             # Force GC to clear any previous large objects
             import gc
@@ -101,7 +133,7 @@ class SatelliteRAG:
             retriever = self.vector_store.as_retriever(search_kwargs={"k": 4})
             
             logger.info("Step 2: Invoking retriever (Embedding inference)...")
-            docs = retriever.invoke(question)
+            docs = retriever.invoke(standalone_question)
             logger.info(f"Step 3: Retrieval successful. Found {len(docs)} chunks.")
             
             context_text = "\n\n".join([d.page_content for d in docs])
@@ -127,6 +159,7 @@ class SatelliteRAG:
             prompt = ChatPromptTemplate.from_template(template)
             chain = prompt | self.llm | StrOutputParser()
             
+            # Use original question for answer generation to keep tone, but context is from standalone
             response = chain.invoke({"context": context_text, "question": question})
             logger.info("Step 5: LLM generation successful.")
             return response, docs