Update query_index.py

query_index.py

@@ -1,8 +1,8 @@
-
 import os
 import logging
 from typing import Dict, List, Optional
 from dotenv import load_dotenv
+from llama_index.llms.openai import OpenAI
 
 from llama_index.core import (
     StorageContext,
@@ -43,7 +43,10 @@ class MultimodalRAGSystem:
         self.config = MultimodalRAGConfig()
         self.index = None
         self.query_engine = None
-        
+        self.rewrite_llm = OpenAI(
+            model="gpt-4o-mini",  
+            temperature=0.0
+        )
         self._initialize_system()
         
     def _initialize_system(self):
@@ -87,48 +90,47 @@ class MultimodalRAGSystem:
         
         logger.info(f"System Ready! Model: {self.config.LLM_MODEL}")
 
-    def _rewrite_question(self, question: str, chat_history: List[Dict[str, str]]) -> str:
+    def ask(self, query_str: str, chat_history: Optional[List[Dict[str, str]]] = None) -> Dict:
         """
-        Rewrite a follow-up question into a standalone question
-        using recent conversation history.
+        Query the RAG system with optional chat history for context.
+        
+        Args:
+            query_str: The user's question
+            chat_history: List of dicts with 'role' and 'content' keys
+            
+        Returns:
+            Dict with 'answer', 'images', and 'texts' keys
         """
-        if not chat_history:
-            return question
-    
-        history_text = "\n".join(
-            f"{turn['role'].capitalize()}: {turn['content']}"
-            for turn in chat_history[-4:]  # last 2 turns
-        )
+        if not self.query_engine:
+            raise RuntimeError("Query engine not initialized")
     
-        prompt = f"""
-    You are rewriting a follow-up question into a standalone question.
+        logger.info(f"Original question: {query_str}")
     
-    Conversation:
-    {history_text}
+        # Rewrite follow-up into standalone question if history exists
+        standalone_question = query_str
     
-    Follow-up question:
-    {question}
+        if chat_history and len(chat_history) > 0:
+            # Convert chat history to context string
+            history_text = "\n".join(
+                f"{turn['role'].capitalize()}: {turn['content']}"
+                for turn in chat_history[-4:]  # last 2 turns (4 messages)
+            )
+            
+            rewrite_prompt = (
+                "Given the previous conversation and the follow-up question, "
+                "rewrite the follow-up question as a standalone question that includes all necessary context.\n\n"
+                f"Conversation history:\n{history_text}\n\n"
+                f"Follow-up question:\n{query_str}\n\n"
+                "Rewrite this as a standalone question that can be understood without the conversation history. "
+                "Only output the rewritten question, nothing else.\n\n"
+                "Standalone question:"
+            )
     
-    Standalone question:
-    """
+            standalone_question = self.rewrite_llm.complete(
+                rewrite_prompt
+            ).text.strip()
     
-        rewritten = self.query_engine._llm.complete(prompt)
-        return rewritten.text.strip()
-
-    def ask(self, question: str, chat_history: Optional[List[Dict[str, str]]] = None) -> Dict:
-        """
-        Ask a question and return answer + source images
-        """
-        if not self.query_engine:
-             raise RuntimeError("Query engine not initialized")
-
-        logger.info(f"Original question: {question}")
-
-        if chat_history:
-            standalone_question = self._rewrite_question(question, chat_history)
             logger.info(f"Rewritten question: {standalone_question}")
-        else:
-            standalone_question = question
         
         response = self.query_engine.query(standalone_question)
 
@@ -196,12 +198,24 @@ class MultimodalRAGSystem:
 def main():
     try:
         rag = MultimodalRAGSystem()
+        chat_hist = []
         while True:
             q = input("Query (q to quit): ")
             if q.lower() == 'q': break
-            print(rag.ask(q))
+            
+            result = rag.ask(q, chat_history=chat_hist)
+            print(f"\nAnswer: {result['answer']}\n")
+            
+            # Update history
+            chat_hist.append({"role": "user", "content": q})
+            chat_hist.append({"role": "assistant", "content": result['answer']})
+            
+            # Keep history reasonable
+            if len(chat_hist) > 6:
+                chat_hist = chat_hist[-6:]
+                
     except Exception as e:
         print(f"Error: {e}")
 
 if __name__ == "__main__":
-    main()
+    main()