Update src/rag_system.py

src/rag_system.py

@@ -1,8 +1,4 @@
-"""
-Enhanced RAG System - Visual Image Analysis
-Sends base64 images directly to GPT-4o for visual analysis (not just OCR)
-Then stores results in vector store
-"""
+
 from typing import List, Dict
 from langchain_openai import ChatOpenAI
 from langchain_core.messages import HumanMessage, SystemMessage
@@ -16,21 +12,13 @@ from config import (
 
 
 class VisualMultimodalRAG:
-    """
-    RAG system that:
-    1. Sends images as base64 to GPT-4o for visual analysis
-    2. Gets detailed visual descriptions and insights
-    3. Stores visual analysis in vector store
-    4. Enables image-based semantic search
-    """
-    
+
     def __init__(self, api_key: str = None, debug: bool = True):
         api_key = api_key or OPENAI_API_KEY
         self.debug = debug
         
-        # Use gpt-4o for vision capabilities
         self.llm = ChatOpenAI(
-            model_name="gpt-4o-mini",  # CRITICAL: gpt-4o has vision
+            model_name="gpt-4o-mini",
             api_key=api_key,
             temperature=TEMPERATURE,
             max_tokens=MAX_TOKENS,
@@ -40,10 +28,9 @@ class VisualMultimodalRAG:
         self.visual_summaries_log = []
         
         if self.debug:
-            print("✅ VisualMultimodalRAG initialized with gpt-4o (vision model)")
+            print("✅ VisualMultimodalRAG initialized")
 
     def _debug_print(self, label: str, data: any):
-        """Print debug information"""
         if self.debug:
             print(f"\n🔍 DEBUG [{label}]:")
             if isinstance(data, (list, dict)):
@@ -53,7 +40,6 @@ class VisualMultimodalRAG:
                 print(f"  {data}")
 
     def _image_to_base64(self, image_path: str) -> str:
-        """Convert image file to base64 string"""
         try:
             with open(image_path, 'rb') as image_file:
                 image_data = base64.b64encode(image_file.read()).decode('utf-8')
@@ -63,28 +49,15 @@ class VisualMultimodalRAG:
             return None
 
     def analyze_image_visually(self, image_path: str, image_idx: int) -> str:
-        """
-        Send actual image (base64) to gpt-4o for visual analysis
-        Returns detailed visual analysis/description
-        
-        gpt-4o can see:
-        - Charts, graphs, diagrams
-        - Tables and structured data
-        - Photos and drawings
-        - Handwritten text
-        - Screenshots
-        - Any visual content
-        """
+
         if not os.path.exists(image_path):
             return f"[Image {image_idx}: File not found - {image_path}]"
         
         try:
-            # Convert image to base64
             image_base64 = self._image_to_base64(image_path)
             if not image_base64:
                 return f"[Image {image_idx}: Could not convert to base64]"
             
-            # Determine image type
             file_ext = Path(image_path).suffix.lower()
             media_type_map = {
                 '.jpg': 'image/jpeg',
@@ -108,23 +81,21 @@ class VisualMultimodalRAG:
                     },
                     {
                         "type": "text",
-                        "text": f"""Analyze this image in detail in {self.language}.
-
-Provide a comprehensive visual analysis including:
-1. **What you see** - Main objects, elements, structure
-2. **Data/Content** - Any numbers, text, charts, graphs
-3. **Purpose** - What this image is showing or representing
-4. **Key insights** - Important patterns, trends, or information
-5. **Connections** - How this relates to document content
-
-Be specific and detailed. Focus on visual information that cannot be extracted from text alone.
+                        "text": f"""You are assistant for analyzing and aggregating information. Analyze this image.
+                        
+Provide a visual analysis that includes:
+1. Main objects and element
+2. Data/Content - Any numbers, text, charts, graphs
+3. What this image is showing or representing
+4. Important patterns, trends, or information
+5. How image relates to document content
+Be brief and meaningful. Focus on visual information that cannot be extracted from text. Response on {self.language}.
 
 Analysis:"""
                     }
                 ],
             )
             
-            # Call gpt-4o with vision
             response = self.llm.invoke([message])
             analysis = response.content.strip()
             
@@ -140,10 +111,7 @@ Analysis:"""
             return error_msg
 
     def analyze_images_visually(self, images: List[Dict]) -> List[Dict]:
-        """
-        Analyze each image visually using gpt-4o vision
-        Returns list of {image_index, visual_analysis, type}
-        """
+
         visual_analyses = []
         
         for idx, image in enumerate(images):
@@ -153,7 +121,6 @@ Analysis:"""
                 print(f"⚠️  Image {idx}: No path provided")
                 continue
             
-            # Analyze image visually (not just OCR)
             visual_analysis = self.analyze_image_visually(image_path, idx)
             
             visual_analyses.append({
@@ -167,9 +134,7 @@ Analysis:"""
         return visual_analyses
 
     def summarize_text_chunks(self, text: str, chunk_size: int = 1500) -> List[Dict]:
-        """
-        Chunk text and summarize each chunk individually
-        """
+
         chunks = []
         text_chunks = self._chunk_text(text, chunk_size=chunk_size, overlap=300)
         
@@ -181,12 +146,12 @@ Analysis:"""
             
             try:
                 prompt = f"""Summarize this text chunk in {self.language}.
-Keep it concise. Extract key points, facts, and main ideas.
+Be brief and meaningful. Extract key points, facts, and main ideas.
 
 Text Chunk:
 {chunk}
 
-Summary (2-3 sentences maximum):"""
+Summary:"""
                 
                 message = HumanMessage(content=prompt)
                 response = self.llm.invoke([message])
@@ -209,9 +174,6 @@ Summary (2-3 sentences maximum):"""
         return chunks
 
     def summarize_tables(self, tables: List[Dict]) -> List[Dict]:
-        """
-        Summarize each table individually
-        """
         summaries = []
         
         for idx, table in enumerate(tables):
@@ -222,12 +184,12 @@ Summary (2-3 sentences maximum):"""
             
             try:
                 prompt = f"""Analyze and summarize this table/structured data in {self.language}.
-Extract key insights, row/column meanings, and important figures.
+Extract key insights, row/column meanings, and important figures. Be brief and meaningful.
 
 Table Content:
 {table_content}
 
-Summary (2-3 sentences maximum):"""
+Summary:"""
                 
                 message = HumanMessage(content=prompt)
                 response = self.llm.invoke([message])
@@ -257,10 +219,7 @@ Summary (2-3 sentences maximum):"""
         vector_store,
         doc_id: str
     ) -> Dict:
-        """
-        Main function: Analyze all components visually and store in vector store
-        Images are analyzed using gpt-4o vision (not just OCR)
-        """
+
         print(f"\n{'='*70}")
         print(f"PROCESSING WITH VISUAL IMAGE ANALYSIS: {doc_id}")
         print(f"{'='*70}")
@@ -273,14 +232,12 @@ Summary (2-3 sentences maximum):"""
             'total_stored': 0
         }
         
-        # 1. Analyze images VISUALLY using gpt-4o
         print(f"\n🖼️ VISUAL IMAGE ANALYSIS (gpt-4o vision) ({len(images)} total)")
         print(f"{'─'*70}")
         
         image_analyses = self.analyze_images_visually(images)
         results['image_visual_analyses'] = image_analyses
         
-        # Store each image analysis in vector store
         image_docs = {
             'text': ' | '.join([
                 f"Image {a['image_index']}: {a['visual_analysis']}" 
@@ -306,7 +263,6 @@ Summary (2-3 sentences maximum):"""
             except Exception as e:
                 print(f"❌ Error storing image analyses: {e}")
         
-        # 2. Summarize and store text chunks
         print(f"\n📝 TEXT CHUNK SUMMARIZATION")
         print(f"{'─'*70}")
         
@@ -334,7 +290,6 @@ Summary (2-3 sentences maximum):"""
             except Exception as e:
                 print(f"❌ Error storing text summaries: {e}")
         
-        # 3. Summarize and store tables
         print(f"\n📋 TABLE SUMMARIZATION ({len(tables)} total)")
         print(f"{'─'*70}")
         
@@ -362,7 +317,6 @@ Summary (2-3 sentences maximum):"""
             except Exception as e:
                 print(f"❌ Error storing table summaries: {e}")
         
-        # 4. Summary statistics
         print(f"\n{'='*70}")
         print(f"📊 STORAGE SUMMARY")
         print(f"{'='*70}")
@@ -376,7 +330,6 @@ Summary (2-3 sentences maximum):"""
         return results
 
     def _chunk_text(self, text: str, chunk_size: int = 1500, overlap: int = 300) -> List[str]:
-        """Split text into overlapping chunks"""
         chunks = []
         start = 0
         while start < len(text):
@@ -386,24 +339,18 @@ Summary (2-3 sentences maximum):"""
         return chunks
 
     def get_visual_summaries_log(self) -> List[Dict]:
-        """Get all visual analysis logs"""
         return self.visual_summaries_log
 
 
 class AnsweringRAG:
-    """
-    RAG system that:
-    1. Searches vector store for relevant content
-    2. ANALYZES search results
-    3. Generates intelligent answers based on context
-    """
+
     
     def __init__(self, api_key: str = None, debug: bool = True):
         api_key = api_key or OPENAI_API_KEY
         self.debug = debug
         
         self.llm = ChatOpenAI(
-            model_name="gpt-4o-mini",  # Use gpt-4o for better understanding
+            model_name="gpt-4o-mini", 
             api_key=api_key,
             temperature=TEMPERATURE,
             max_tokens=MAX_TOKENS,
@@ -413,10 +360,9 @@ class AnsweringRAG:
         self.answer_log = []
         
         if self.debug:
-            print("✅ AnsweringRAG initialized with answer generation")
+            print("✅ AnsweringRAG initialized")
 
     def _debug_print(self, label: str, data: any):
-        """Print debug information"""
         if self.debug:
             print(f"\n🔍 DEBUG [{label}]:")
             if isinstance(data, (list, dict)):
@@ -430,18 +376,7 @@ class AnsweringRAG:
         question: str, 
         search_results: List[Dict]
     ) -> Dict:
-        """
-        Analyze search results and generate intelligent answer
-        
-        Returns:
-        {
-            'question': user question,
-            'answer': detailed answer,
-            'sources_used': number of sources,
-            'confidence': low/medium/high,
-            'search_results': original search results
-        }
-        """
+
         
         print(f"\n{'='*70}")
         print(f"ANALYZING QUESTION & GENERATING ANSWER")
@@ -450,15 +385,10 @@ class AnsweringRAG:
         print(f"\n❓ Question: {question}")
         print(f"📊 Search Results Found: {len(search_results)}")
         
-        # Check if we have search results
         if not search_results:
             print(f"⚠️  No search results found!")
-            answer = f"""I could not find relevant information in the document to answer your question: "{question}"
-
-Try:
-- Using different keywords
-- Breaking the question into smaller parts
-- Asking about other topics in the document"""
+            answer = f"""No relevant information in the document to answer question: "{question}"
+"""
             
             result = {
                 'question': question,
@@ -470,7 +400,6 @@ Try:
             self.answer_log.append(result)
             return result
         
-        # Build context from search results
         context_parts = []
         for idx, result in enumerate(search_results, 1):
             content = result.get('content', '')
@@ -487,11 +416,10 @@ Try:
         
         self._debug_print("Context Prepared", f"{len(context_parts)} sources, {len(full_context)} chars")
         
-        # Build prompt to analyze results and answer question
         analysis_prompt = f"""You are a helpful assistant analyzing document content to answer user questions.
-
+        
 USER QUESTION:
-"{question}"
+{question}
 
 RELEVANT CONTENT FROM DOCUMENT:
 {full_context}
@@ -503,7 +431,6 @@ INSTRUCTIONS:
 4. If the content doesn't fully answer the question, explain what information is available
 5. Be specific and cite the content when relevant
 6. Structure your answer clearly with key points
-
 ANSWER:"""
         
         print(f"\n🔍 Analyzing search results...")
@@ -511,12 +438,10 @@ ANSWER:"""
         print(f"   Sources: {len(search_results)}")
         
         try:
-            # Call LLM to analyze and answer
             message = HumanMessage(content=analysis_prompt)
             response = self.llm.invoke([message])
             answer = response.content.strip()
             
-            # Determine confidence level
             confidence = self._estimate_confidence(len(search_results), answer)
             
             print(f"✅ Answer generated successfully")
@@ -551,18 +476,14 @@ ANSWER:"""
             return result
 
     def _estimate_confidence(self, sources_count: int, answer: str) -> str:
-        """Estimate confidence level of answer"""
         answer_length = len(answer)
         
-        # High confidence: multiple sources, substantial answer
         if sources_count >= 3 and answer_length > 500:
             return "high"
         
-        # Medium confidence: some sources, decent answer
         elif sources_count >= 2 and answer_length > 200:
             return "medium"
         
-        # Low confidence: few sources or short answer
         else:
             return "low"
 
@@ -571,14 +492,9 @@ ANSWER:"""
         question: str,
         search_results: List[Dict]
     ) -> Dict:
-        """
-        Get answer AND properly formatted sources
-        Returns both answer and formatted source citations
-        """
-        
+
         result = self.analyze_and_answer(question, search_results)
         
-        # Format sources for display
         formatted_sources = []
         for idx, source in enumerate(result['search_results'], 1):
             formatted_sources.append({
@@ -592,11 +508,9 @@ ANSWER:"""
         return result
 
     def get_answer_log(self) -> List[Dict]:
-        """Get all answer generation logs"""
         return self.answer_log
 
     def print_answer_with_sources(self, result: Dict, max_source_length: int = 300):
-        """Pretty print answer with sources"""
         
         print(f"\n{'='*70}")
         print(f"ANSWER TO: {result['question']}")