Update src/rag_system.py

src/rag_system.py

@@ -1,42 +1,46 @@
 """
-Enhanced RAG System - Individual Summarization + Vector Store Persistence
-Summarizes each image, text chunk, and table separately, then stores results
+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
-import hashlib
+import base64
+import os
+from pathlib import Path
 from config import (
     OPENAI_API_KEY, OPENAI_MODEL, TEMPERATURE, MAX_TOKENS, 
-    LANGUAGE, CACHE_RESPONSES, BATCH_SEARCH_RESULTS
+    LANGUAGE, CHROMA_DB_PATH
 )
 
 
 class MultimodalRAG:
     """
     RAG system that:
-    1. Summarizes each component individually
-    2. Stores summaries in vector store
-    3. Enables fine-grained semantic search
+    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=OPENAI_MODEL,
+            model_name="gpt-4o",  # CRITICAL: gpt-4o has vision
             api_key=api_key,
             temperature=TEMPERATURE,
             max_tokens=MAX_TOKENS,
         )
         
-        self.conversation_history = []
         self.language = LANGUAGE
-        self.summaries_log = []
+        self.visual_summaries_log = []
         
         if self.debug:
-            print("✅ EnhancedMultimodalRAG initialized")
+            print("✅ VisualMultimodalRAG initialized with gpt-4o (vision model)")
 
     def _debug_print(self, label: str, data: any):
         """Print debug information"""
@@ -48,54 +52,134 @@ class MultimodalRAG:
             else:
                 print(f"  {data}")
 
-    def summarize_image(self, image_ocr_text: str, image_idx: int) -> str:
+    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')
+            return image_data
+        except Exception as e:
+            print(f"Error converting image to base64: {e}")
+            return None
+
+    def analyze_image_visually(self, image_path: str, image_idx: int) -> str:
         """
-        Summarize a single image's OCR text
-        Returns concise summary focused on image content
+        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 image_ocr_text or len(image_ocr_text.strip()) < 5:
-            return f"[Image {image_idx}: No readable text or empty content]"
+        if not os.path.exists(image_path):
+            return f"[Image {image_idx}: File not found - {image_path}]"
         
         try:
-            prompt = f"""Summarize this text extracted from an image in {self.language}.
-Keep it concise but informative. Focus on key information, data, and visual elements.
+            # 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',
+                '.jpeg': 'image/jpeg',
+                '.png': 'image/png',
+                '.gif': 'image/gif',
+                '.webp': 'image/webp'
+            }
+            media_type = media_type_map.get(file_ext, 'image/png')
+            
+            print(f"🔍 Analyzing image {image_idx} visually (as {media_type})...")
+            
+            # Create message with image
+            message = HumanMessage(
+                content=[
+                    {
+                        "type": "image_url",
+                        "image_url": {
+                            "url": f"data:{media_type};base64,{image_base64}",
+                        },
+                    },
+                    {
+                        "type": "text",
+                        "text": f"""Analyze this image in detail in {self.language}.
 
-Image OCR Text:
-{image_ocr_text}
+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
 
-Summary (2-3 sentences maximum):"""
+Be specific and detailed. Focus on visual information that cannot be extracted from text alone.
+
+Analysis:"""
+                    }
+                ],
+            )
             
-            message = HumanMessage(content=prompt)
+            # Call gpt-4o with vision
             response = self.llm.invoke([message])
-            summary = response.content.strip()
+            analysis = response.content.strip()
             
             if self.debug:
-                self._debug_print(f"Image {image_idx} Summary", summary)
+                self._debug_print(f"Image {image_idx} Visual Analysis", analysis)
+            
+            print(f"✅ Image {image_idx} analyzed successfully")
+            return analysis
             
-            return summary
         except Exception as e:
-            error_msg = f"[Image {image_idx}: Summarization failed - {str(e)}]"
-            print(f"Error summarizing image {image_idx}: {e}")
+            error_msg = f"[Image {image_idx}: Vision analysis failed - {str(e)}]"
+            print(f"❌ Error analyzing image {image_idx}: {e}")
             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):
+            image_path = image.get('path', '')
+            
+            if not image_path:
+                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({
+                'type': 'image_visual',
+                'image_index': idx,
+                'image_path': image_path,
+                'visual_analysis': visual_analysis,
+                'ocr_text': image.get('ocr_text', '')  # Keep OCR as backup
+            })
+        
+        return visual_analyses
+
     def summarize_text_chunks(self, text: str, chunk_size: int = 1500) -> List[Dict]:
         """
         Chunk text and summarize each chunk individually
-        Returns list of {chunk_text, summary, type, index}
         """
         chunks = []
-        
-        # Split text into chunks
         text_chunks = self._chunk_text(text, chunk_size=chunk_size, overlap=300)
         
         self._debug_print("Text Chunking", f"Created {len(text_chunks)} chunks")
         
         for idx, chunk in enumerate(text_chunks):
-            if len(chunk.strip()) < 50:  # Skip very small chunks
+            if len(chunk.strip()) < 50:
                 continue
             
             try:
-                # Summarize chunk
                 prompt = f"""Summarize this text chunk in {self.language}.
 Keep it concise. Extract key points, facts, and main ideas.
 
@@ -111,7 +195,7 @@ Summary (2-3 sentences maximum):"""
                 chunks.append({
                     'type': 'text_chunk',
                     'chunk_index': len(chunks),
-                    'original_text': chunk[:500],  # Store first 500 chars
+                    'original_text': chunk[:500],
                     'summary': summary,
                     'chunk_length': len(chunk)
                 })
@@ -127,7 +211,6 @@ Summary (2-3 sentences maximum):"""
     def summarize_tables(self, tables: List[Dict]) -> List[Dict]:
         """
         Summarize each table individually
-        Returns list of {table_content, summary, type, index}
         """
         summaries = []
         
@@ -138,7 +221,6 @@ Summary (2-3 sentences maximum):"""
                 continue
             
             try:
-                # Summarize table
                 prompt = f"""Analyze and summarize this table/structured data in {self.language}.
 Extract key insights, row/column meanings, and important figures.
 
@@ -167,27 +249,6 @@ Summary (2-3 sentences maximum):"""
         
         return summaries
 
-    def summarize_images(self, images: List[Dict]) -> List[Dict]:
-        """
-        Summarize each image individually
-        Returns list of {image_index, ocr_text, summary, type}
-        """
-        summaries = []
-        
-        for idx, image in enumerate(images):
-            ocr_text = image.get('ocr_text', '')
-            summary = self.summarize_image(ocr_text, idx)
-            
-            summaries.append({
-                'type': 'image',
-                'image_index': idx,
-                'original_ocr': ocr_text[:500],
-                'summary': summary,
-                'ocr_length': len(ocr_text)
-            })
-        
-        return summaries
-
     def process_and_store_document(
         self, 
         text: str, 
@@ -197,58 +258,61 @@ Summary (2-3 sentences maximum):"""
         doc_id: str
     ) -> Dict:
         """
-        Main function: Summarize all components and store in vector store
-        Returns summary statistics
+        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 AND STORING: {doc_id}")
+        print(f"PROCESSING WITH VISUAL IMAGE ANALYSIS: {doc_id}")
         print(f"{'='*70}")
         
         results = {
             'doc_id': doc_id,
-            'image_summaries': [],
+            'image_visual_analyses': [],
             'text_summaries': [],
             'table_summaries': [],
             'total_stored': 0
         }
         
-        # 1. Summarize and store images
-        print(f"\n🖼️ PROCESSING IMAGES ({len(images)} total)")
+        # 1. Analyze images VISUALLY using gpt-4o
+        print(f"\n🖼️ VISUAL IMAGE ANALYSIS (gpt-4o vision) ({len(images)} total)")
         print(f"{'─'*70}")
         
-        image_summaries = self.summarize_images(images)
-        results['image_summaries'] = image_summaries
+        image_analyses = self.analyze_images_visually(images)
+        results['image_visual_analyses'] = image_analyses
         
-        # Store each image summary in vector store
+        # Store each image analysis in vector store
         image_docs = {
-            'text': ' | '.join([f"Image {s['image_index']}: {s['summary']}" 
-                               for s in image_summaries]),
+            'text': ' | '.join([
+                f"Image {a['image_index']}: {a['visual_analysis']}" 
+                for a in image_analyses
+            ]),
             'images': [],
             'tables': []
         }
         
-        for summary in image_summaries:
-            print(f"  ✅ Image {summary['image_index']}: {summary['summary'][:50]}...")
+        for analysis in image_analyses:
+            print(f"  ✅ Image {analysis['image_index']} (visual analysis)")
+            print(f"     Path: {analysis['image_path']}")
+            print(f"     Analysis: {analysis['visual_analysis'][:100]}...")
         
-        if image_summaries:
+        if image_analyses:
             try:
                 vector_store.add_documents(
                     image_docs, 
-                    f"{doc_id}_images"
+                    f"{doc_id}_images_visual"
                 )
-                results['total_stored'] += len(image_summaries)
-                print(f"✅ Stored {len(image_summaries)} image summaries")
+                results['total_stored'] += len(image_analyses)
+                print(f"✅ Stored {len(image_analyses)} image visual analyses")
             except Exception as e:
-                print(f"❌ Error storing image summaries: {e}")
+                print(f"❌ Error storing image analyses: {e}")
         
         # 2. Summarize and store text chunks
-        print(f"\n📝 PROCESSING TEXT CHUNKS")
+        print(f"\n📝 TEXT CHUNK SUMMARIZATION")
         print(f"{'─'*70}")
         
         text_summaries = self.summarize_text_chunks(text)
         results['text_summaries'] = text_summaries
         
-        # Store each text chunk summary in vector store
         text_docs = {
             'text': ' | '.join([f"Chunk {s['chunk_index']}: {s['summary']}" 
                                for s in text_summaries]),
@@ -271,13 +335,12 @@ Summary (2-3 sentences maximum):"""
                 print(f"❌ Error storing text summaries: {e}")
         
         # 3. Summarize and store tables
-        print(f"\n📋 PROCESSING TABLES ({len(tables)} total)")
+        print(f"\n📋 TABLE SUMMARIZATION ({len(tables)} total)")
         print(f"{'─'*70}")
         
         table_summaries = self.summarize_tables(tables)
         results['table_summaries'] = table_summaries
         
-        # Store each table summary in vector store
         table_docs = {
             'text': ' | '.join([f"Table {s['table_index']}: {s['summary']}" 
                                for s in table_summaries]),
@@ -303,13 +366,13 @@ Summary (2-3 sentences maximum):"""
         print(f"\n{'='*70}")
         print(f"📊 STORAGE SUMMARY")
         print(f"{'='*70}")
-        print(f"  Images summarized & stored: {len(image_summaries)}")
+        print(f"  Images analyzed visually & stored: {len(image_analyses)}")
         print(f"  Text chunks summarized & stored: {len(text_summaries)}")
         print(f"  Tables summarized & stored: {len(table_summaries)}")
-        print(f"  Total items stored: {results['total_stored']}")
+        print(f"  Total items stored in vector: {results['total_stored']}")
         print(f"{'='*70}")
         
-        self.summaries_log.append(results)
+        self.visual_summaries_log.append(results)
         return results
 
     def _chunk_text(self, text: str, chunk_size: int = 1500, overlap: int = 300) -> List[str]:
@@ -322,12 +385,6 @@ Summary (2-3 sentences maximum):"""
             start = end - overlap
         return chunks
 
-    def get_summaries_log(self) -> List[Dict]:
-        """Get all processing logs"""
-        return self.summaries_log
-
-    def clear_history(self):
-        """Clear conversation history"""
-        self.conversation_history = []
-        if self.debug:
-            print("✅ Conversation history cleared")
+    def get_visual_summaries_log(self) -> List[Dict]:
+        """Get all visual analysis logs"""
+        return self.visual_summaries_log