Update main.py

main.py

@@ -568,35 +568,97 @@ async def query_system(
         # Get session data
         session = sessions[session_id]
         retriever = session["retriever"]
-        chat_history = session["chat_history"]
+        
+        # Get or initialize chat history
+        chat_history = chat_manager.initialize_chat_history(session_id)
         
         # Create chain
         chain = create_chain(retriever)
         
-        # Get chat history from MongoDB in LangChain format
+        # Extract messages properly for LangChain format
         messages = chat_history.messages
-        langchain_chat_history = [(messages[i].content, messages[i+1].content) 
-                                for i in range(0, len(messages)-1, 2) if i+1 < len(messages)]
         
-        # Query the chain
-        result = chain.invoke({"question": request.query, "chat_history": langchain_chat_history})
+        # Process chat history safely
+        langchain_chat_history = []
         
-        # Update chat history
-        chat_history.add_user_message(request.query)
-        chat_history.add_ai_message(result["answer"])
+        # If messages exist, process them
+        if messages:
+            # Group messages by pairs (user, AI)
+            # This approach is safer than assuming perfect alternating pattern
+            i = 0
+            while i < len(messages) - 1:
+                user_message = messages[i].content
+                ai_message = messages[i+1].content
+                langchain_chat_history.append((user_message, ai_message))
+                i += 2
+                
+        # Add debugging information
+        print(f"Chat history length: {len(langchain_chat_history)}")
+        print(f"Query: {request.query}")
         
-        # Prepare source documents
-        source_docs = [doc.page_content[:100] + "..." for doc in result.get("source_documents", [])]
+        try:
+            # Query the chain
+            result = chain.invoke({
+                "question": request.query,
+                "chat_history": langchain_chat_history
+            })
+            
+            # Extract answer from result
+            answer = result.get("answer", "I couldn't find an answer to your question.")
+            
+            # Update chat history
+            chat_history.add_user_message(request.query)
+            chat_history.add_ai_message(answer)
+            
+            # Prepare source documents with proper error handling
+            source_docs = []
+            if "source_documents" in result and result["source_documents"]:
+                for doc in result["source_documents"]:
+                    try:
+                        # Different LangChain versions might structure documents differently
+                        if hasattr(doc, 'page_content'):
+                            # Regular Document object
+                            content = doc.page_content[:100] + "..." if len(doc.page_content) > 100 else doc.page_content
+                            source_docs.append(content)
+                        elif isinstance(doc, dict) and 'page_content' in doc:
+                            # Dictionary format
+                            content = doc['page_content'][:100] + "..." if len(doc['page_content']) > 100 else doc['page_content']
+                            source_docs.append(content)
+                        elif isinstance(doc, str):
+                            # String format
+                            content = doc[:100] + "..." if len(doc) > 100 else doc
+                            source_docs.append(content)
+                    except Exception as doc_error:
+                        print(f"Error processing source document: {str(doc_error)}")
+            
+            return {
+                "answer": answer,
+                "session_id": session_id,
+                "source_documents": source_docs
+            }
         
-        return {
-            "answer": result["answer"],
-            "session_id": session_id,
-            "source_documents": source_docs
-        }
+        except Exception as chain_error:
+            print(f"Chain invocation error: {str(chain_error)}")
+            # Provide a more graceful fallback
+            fallback_answer = "I apologize, but I encountered an error while processing your question. Please try rephrasing your query or asking about a different topic."
+            
+            # Update chat history even in case of error
+            chat_history.add_user_message(request.query)
+            chat_history.add_ai_message(fallback_answer)
+            
+            return {
+                "answer": fallback_answer,
+                "session_id": session_id,
+                "source_documents": []
+            }
     
     except Exception as e:
+        print(f"Query system error: {str(e)}")
+        import traceback
+        traceback.print_exc()
         raise HTTPException(status_code=500, detail=f"Error querying system: {str(e)}")
 
+        
 @app.get("/sessions", response_model=List[Dict[str, Any]])
 async def get_user_sessions(current_user: User = Depends(get_current_user)):
     """