fixed RAG state message issue

agents/rag_agent.py

@@ -64,6 +64,7 @@ class AgenticRAGState(MessagesState):
     is_sufficient: bool = False
     retry_count: int = 0  # Track number of retries to prevent infinite loops
     max_retries: int = 3  # Maximum number of query rewrites allowed
+    current_query_index: int = 0  # Track which message is the current query
 
 
 class AgenticRAGChat(ChatInterface):
@@ -207,17 +208,38 @@ class AgenticRAGChat(ChatInterface):
         """Evaluate the documents retrieved from the retriever tool."""
         print("Evaluating documents...")
         
-        # Get original user question and retrieved docs
-        user_question = state["messages"][0].content
+        # Check if we've hit max retries
+        if state.get("retry_count", 0) >= state.get("max_retries", 3):
+            print(f"Max retries ({state.get('max_retries', 3)}) reached. Forcing synthesis with available documents.")
+            return {
+                "is_sufficient": True,  # Force synthesis even if not perfect
+                "feedback": "Maximum retries reached. Using available documents."
+            }
+        
+        # Get the CURRENT user question, not the first message in history
+        # Use the current_query_index to get the right message
+        current_query_index = state.get("current_query_index", 0)
+        
+        # Find the current query message
+        user_messages = [msg for msg in state["messages"] if isinstance(msg, HumanMessage)]
+        if current_query_index < len(state["messages"]):
+            user_question = state["messages"][current_query_index].content
+        else:
+            # Fallback: get the last user message
+            user_question = user_messages[-1].content if user_messages else state["messages"][-1].content
+        
+        # Get the retrieved documents (should be the last message)
         retrieved_docs = state["messages"][-1].content
         
+        print(f"Evaluating for query: '{user_question[:50]}...'")  # Debug print
+        
         chain = DOCUMENT_EVALUATOR_PROMPT | self.evaluator_llm
         evaluation = chain.invoke({
             "question": user_question, 
             "retrieved_docs": retrieved_docs
         })
         
-        print(f"Evaluation result: {evaluation}")
+        print(f"Evaluation result: {evaluation} (Retry {state.get('retry_count', 0)}/{state.get('max_retries', 3)})")
         return {
             "is_sufficient": evaluation.is_sufficient,
             "feedback": evaluation.feedback
@@ -227,9 +249,22 @@ class AgenticRAGChat(ChatInterface):
         """Synthesize the final answer from retrieved documents."""
         print("Synthesizing answer...")
         
-        user_question = state["messages"][0].content
+        # Get the CURRENT user question using the index
+        current_query_index = state.get("current_query_index", 0)
+        
+        # Find the current query message
+        user_messages = [msg for msg in state["messages"] if isinstance(msg, HumanMessage)]
+        if current_query_index < len(state["messages"]):
+            user_question = state["messages"][current_query_index].content
+        else:
+            # Fallback: get the last user message
+            user_question = user_messages[-1].content if user_messages else state["messages"][-1].content
+        
+        # Get the retrieved documents
         retrieved_docs = state["messages"][-1].content
         
+        print(f"Synthesizing answer for: '{user_question[:50]}...'")  # Debug print
+        
         chain = DOCUMENT_SYNTHESIZER_PROMPT | self.llm
         answer = chain.invoke({
             "question": user_question, 
@@ -242,10 +277,25 @@ class AgenticRAGChat(ChatInterface):
         """Rewrite the query based on evaluation feedback."""
         print("Rewriting query...")
         
-        user_question = state["messages"][0].content
+        # Increment retry count
+        current_retry = state.get("retry_count", 0)
+        
+        # Get the CURRENT user question using the index
+        current_query_index = state.get("current_query_index", 0)
+        
+        # Find the current query message
+        user_messages = [msg for msg in state["messages"] if isinstance(msg, HumanMessage)]
+        if current_query_index < len(state["messages"]):
+            user_question = state["messages"][current_query_index].content
+        else:
+            # Fallback: get the last user message
+            user_question = user_messages[-1].content if user_messages else state["messages"][-1].content
+        
         retrieved_docs = state["messages"][-1].content
         feedback = state["feedback"]
         
+        print(f"Rewriting query for: '{user_question[:50]}...'")  # Debug print
+        
         chain = QUERY_REWRITER_PROMPT | self.llm
         new_query = chain.invoke({
             "question": user_question, 
@@ -253,8 +303,11 @@ class AgenticRAGChat(ChatInterface):
             "retrieved_docs": retrieved_docs
         })
         
-        print(f"Rewritten query: {new_query.content}")
-        return {"messages": [new_query]}
+        print(f"Rewritten query (Attempt {current_retry + 1}/{state.get('max_retries', 3)}): {new_query.content}")
+        return {
+            "messages": [new_query],
+            "retry_count": current_retry + 1  # Increment retry count
+        }
     
     def _create_graph(self) -> Any:
         """Create the agentic RAG graph."""
@@ -331,30 +384,29 @@ class AgenticRAGChat(ChatInterface):
         # Convert chat history to messages
         history_messages = self._convert_history_to_messages(chat_history)
     
+        # Mark the position where the current query starts
+        # This is important for the evaluator to know which is the actual query
+        history_length = len(history_messages)
+        
         # Add the current message
-        history_messages.append(HumanMessage(content=message))
+        current_query_message = HumanMessage(content=message)
+        history_messages.append(current_query_message)
         
         # Create initial state with full conversation history
+        # Store the index of the current query for reference
         state = AgenticRAGState(
-            messages=history_messages,  # Include full history instead of just current message
+            messages=history_messages,
             feedback="",
             is_sufficient=False,
             retry_count=0,
-            max_retries=3
+            max_retries=3,
+            # Add this to track the current query index
+            current_query_index=history_length  # This is the index of the current query
         )
         
-        # state = AgenticRAGState(
-        #     messages=[HumanMessage(content=message)],
-        #     feedback="",
-        #     is_sufficient=False,
-        #     retry_count=0,
-        #     max_retries=3  # Limit to 3 retries to prevent infinite loops
-        # )
-        
         try:
-            # Run the workflow
             # Run the workflow with increased recursion limit
-            config = {"recursion_limit": 30}  # Increased but reasonable limit
+            config = {"recursion_limit": 30}
             result = self.graph.invoke(state, config=config)
             
             print("\n=== RAG QUERY COMPLETED ===")
@@ -371,10 +423,8 @@ class AgenticRAGChat(ChatInterface):
                 
         except Exception as e:
             print(f"Error in RAG processing: {e}")
-            # Provide a more helpful fallback response
             if "recursion" in str(e).lower():
                 return ("I had difficulty finding the exact information you're looking for in the documents. "
-                       "Based on the available documents, I can see references to various offices and services, "
-                       "but I couldn't find specific details about Mission Support Services. "
-                       "You might want to try asking about a specific aspect or department.")
+                    "Based on the available documents, I can see references to various topics, "
+                    "but I couldn't find specific details. You might want to try asking about a specific aspect.")
             return f"I encountered an error while searching for information: {str(e)}"

app.py

@@ -81,9 +81,9 @@ def create_demo():
         examples=[
             "What is 847 * 293?",
             "What's today's date?",
-            "What's the weather in San Francisco?",
-            "Explain quantum computing in simple terms",
-            "Research the impact of AI on healthcare",
+            # "What's the weather in San Francisco?",
+            # "Explain quantum computing in simple terms",
+            # "Research the impact of AI on healthcare",
         ],
         theme=gr.themes.Soft(),
         analytics_enabled=False,
@@ -95,14 +95,4 @@ def create_demo():
 if __name__ == "__main__":
     # Create and launch the demo
     demo = create_demo()
-    
-    # Check if running in Hugging Face Spaces
-    if os.environ.get("SPACE_ID"):
-        # Hugging Face Spaces configuration
-        demo.launch(
-            server_name="0.0.0.0",
-            server_port=int(os.environ.get("PORT", 7860))
-        )
-    else:
-        # Local development - use simple defaults
-        demo.launch()
+    demo.launch()