Simplify

src/config.py

@@ -2,8 +2,8 @@ import os
 from pathlib import Path
 
 OPENAI_API_KEY = os.getenv("OPENAI_API_KEY", "")
-OPENAI_MODEL = "gpt-4o-mini"  # Cheaper model variant
-USE_CACHE = True  # Enable response caching
+OPENAI_MODEL = "gpt-4o-mini" 
+USE_CACHE = True  
 
 CHROMA_DB_PATH = "./chroma_db"
 DOCSTORE_PATH = "./docstore"
@@ -12,10 +12,10 @@ PROCESSED_FILES_LOG = "./processed_files.txt"
 EMBEDDING_MODEL = "sentence-transformers/all-mpnet-base-v2"
 EMBEDDING_DIM = 768
 
-MAX_CHUNK_SIZE = 500  # Smaller chunks = fewer tokens
-CHUNK_OVERLAP = 50    # Less overlap = fewer chunks
-TEMPERATURE = 0.3     # Lower = faster, cheaper
-MAX_TOKENS = 500      # Limit response size (vs 1500)
+MAX_CHUNK_SIZE = 500  
+CHUNK_OVERLAP = 50    
+TEMPERATURE = 0.3     
+MAX_TOKENS = 500     
 
 LANGUAGE = "russian"
 
@@ -26,6 +26,6 @@ UPLOAD_FOLDER = "./uploaded_pdfs"
 Path(UPLOAD_FOLDER).mkdir(exist_ok=True)
 MAX_PDF_SIZE_MB = 50
 
-BATCH_SEARCH_RESULTS = 3
-CACHE_RESPONSES = True
-SUMMARIZE_FIRST = True 
+BATCH_SEARCH_RESULTS = 3 
+CACHE_RESPONSES = True    
+SUMMARIZE_FIRST = True    

src/pdf_parser.py

@@ -9,13 +9,16 @@ from PIL import Image
 import pytesseract
 from config import DOCSTORE_PATH, PROCESSED_FILES_LOG
 
+
 class PDFParser:
     def __init__(self, debug: bool = True):
         self.docstore_path = Path(DOCSTORE_PATH)
         self.docstore_path.mkdir(exist_ok=True)
         self.processed_files = self._load_processed_files()
         self.debug = debug
+        
         self._configure_tesseract()
+        
         if self.debug:
             print("PDFParser initialized")
 
@@ -31,13 +34,13 @@ class PDFParser:
             print(f"[PDF Parser] {label}")
             if isinstance(data, dict):
                 for key, val in data.items():
-                    print(f" {key}: {val}")
+                    print(f"  {key}: {val}")
             elif isinstance(data, (list, tuple)):
-                print(f" Count: {len(data)}")
+                print(f"  Count: {len(data)}")
                 for i, item in enumerate(data[:3]):
-                    print(f" [{i}]: {str(item)[:100]}")
+                    print(f"  [{i}]: {str(item)[:100]}")
             else:
-                print(f" {data}")
+                print(f"  {data}")
 
     def _load_processed_files(self) -> Dict[str, str]:
         if os.path.exists(PROCESSED_FILES_LOG):
@@ -66,12 +69,14 @@ class PDFParser:
                 reader = PyPDF2.PdfReader(file)
                 page_count = len(reader.pages)
                 self._debug_print("PDF Text Extraction", f"Total pages: {page_count}")
+                
                 for page_num, page in enumerate(reader.pages):
                     page_text = page.extract_text()
                     text += page_text + "\n"
                     self._debug_print(f"Page {page_num+1} Text Length", len(page_text))
         except Exception as e:
             self._debug_print("ERROR extracting text", str(e))
+        
         self._debug_print("Total Text Extracted", len(text))
         return text
 
@@ -79,24 +84,33 @@ class PDFParser:
         images_data = []
         try:
             self._debug_print("Image Extraction Started", f"File: {pdf_path}")
+            
             images = convert_from_path(pdf_path, dpi=150)
-            self._debug_print("PDF to Images Conversion", f"Total images: {len(images)}")
+            self._debug_print("PDF to Images", f"Total images: {len(images)}")
+            
             for idx, image in enumerate(images):
                 self._debug_print(f"Processing Image {idx}", f"Size: {image.size}")
+                
                 image_path = self.docstore_path / f"{doc_id}_image_{idx}.png"
                 image.save(image_path)
                 self._debug_print(f"Image {idx} Saved", str(image_path))
+                
                 self._debug_print(f"Image {idx} OCR")
+                
                 try:
                     ocr_text = pytesseract.image_to_string(image, lang='rus')
+                    
                     ocr_text = ocr_text.strip()
+                    
                     if not ocr_text or len(ocr_text) < 5:
-                        self._debug_print(f"Image {idx} OCR Result", f"EMPTY ({len(ocr_text)} chars)")
+                        self._debug_print(f"Image {idx} OCR Result", f"EMPTY or very short ({len(ocr_text)} chars)")
                     else:
                         self._debug_print(f"Image {idx} OCR Result", f"Success - {len(ocr_text)} chars: {ocr_text[:150]}")
+                    
                 except Exception as ocr_error:
                     self._debug_print(f"Image {idx} OCR ERROR", str(ocr_error))
-                    ocr_text = f"Image {idx}: OCR failed - {str(ocr_error)}"
+                    ocr_text = f"[Image {idx}: OCR failed - {str(ocr_error)}]"
+                
                 images_data.append({
                     'page': idx,
                     'path': str(image_path),
@@ -105,6 +119,7 @@ class PDFParser:
                 })
         except Exception as e:
             self._debug_print("ERROR extracting images", str(e))
+        
         self._debug_print("Image Extraction Complete", f"Total: {len(images_data)}")
         return images_data
 
@@ -113,7 +128,9 @@ class PDFParser:
         try:
             text = self._extract_text_from_pdf(pdf_path)
             lines = text.split('\n')
+            
             self._debug_print("Table Detection", f"Scanning {len(lines)} lines")
+            
             current_table = []
             for line in lines:
                 if '|' in line or '\t' in line:
@@ -125,37 +142,48 @@ class PDFParser:
                             'description': f"Table {len(tables_data) + 1}"
                         })
                     current_table = []
+            
             if current_table and len(current_table) > 1:
                 tables_data.append({
                     'content': '\n'.join(current_table),
                     'description': f"Table {len(tables_data) + 1}"
                 })
+            
             self._debug_print("Tables Found", len(tables_data))
         except Exception as e:
             self._debug_print("ERROR extracting tables", str(e))
+        
         return tables_data
 
     def parse_pdf(self, pdf_path: str) -> Tuple[str, List[Dict], List[Dict]]:
         file_hash = self._get_file_hash(pdf_path)
         doc_id = Path(pdf_path).stem
+        
         self._debug_print("PDF Parsing Started", f"File: {doc_id}")
+        
         if doc_id in self.processed_files:
             if self.processed_files[doc_id] == file_hash:
                 self._debug_print("Status", f"File {doc_id} already processed")
                 return self._load_extracted_data(doc_id)
+        
         print(f"Processing PDF: {doc_id}")
+        
         text = self._extract_text_from_pdf(pdf_path)
         images = self._extract_images_from_pdf(pdf_path, doc_id)
         tables = self._extract_tables_from_pdf(pdf_path, doc_id)
+        
         self._debug_print("Extraction Summary", {
             'text_length': len(text),
             'images_count': len(images),
             'tables_count': len(tables),
             'images_with_ocr': sum(1 for img in images if img.get('ocr_text', '').strip())
         })
+        
         self._save_extracted_data(doc_id, text, images, tables)
+        
         self.processed_files[doc_id] = file_hash
         self._save_processed_files()
+        
         return text, images, tables
 
     def _save_extracted_data(self, doc_id: str, text: str, images: List[Dict], tables: List[Dict]):
@@ -167,6 +195,7 @@ class PDFParser:
         data_path = self.docstore_path / f"{doc_id}_data.json"
         with open(data_path, 'w', encoding='utf-8') as f:
             json.dump(data, f, ensure_ascii=False, indent=2)
+        
         self._debug_print("Data Saved", str(data_path))
 
     def _load_extracted_data(self, doc_id: str) -> Tuple[str, List[Dict], List[Dict]]:
@@ -176,4 +205,15 @@ class PDFParser:
                 data = json.load(f)
             return data['text'], data['images'], data['tables']
         except:
-            return "", [], []
+            return "", [], []
+
+    def get_all_documents(self) -> Dict:
+        all_docs = {}
+        for json_file in self.docstore_path.glob("*_data.json"):
+            doc_id = json_file.stem.replace("_data", "")
+            try:
+                with open(json_file, 'r', encoding='utf-8') as f:
+                    all_docs[doc_id] = json.load(f)
+            except:
+                pass
+        return all_docs

src/rag_system.py

@@ -1,25 +1,32 @@
 from typing import List, Dict
 from langchain_openai import ChatOpenAI
-from langchain_core.messages import HumanMessage
+from langchain_core.messages import HumanMessage, SystemMessage
 import base64
 import os
 from pathlib import Path
 from config import (
-    OPENAI_API_KEY, OPENAI_MODEL, TEMPERATURE, MAX_TOKENS,
+    OPENAI_API_KEY, OPENAI_MODEL, TEMPERATURE, MAX_TOKENS, 
     LANGUAGE, CHROMA_DB_PATH
 )
 
+
 class VisualMultimodalRAG:
+    
     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",
+            model_name="gpt-4o-mini", 
             api_key=api_key,
             temperature=TEMPERATURE,
             max_tokens=MAX_TOKENS,
         )
+        
         self.language = LANGUAGE
+        self.visual_summaries_log = []
+        
         if self.debug:
             print("VisualMultimodalRAG initialized")
 
@@ -27,10 +34,10 @@ class VisualMultimodalRAG:
         if self.debug:
             print(f"DEBUG [{label}]:")
             if isinstance(data, (list, dict)):
-                print(f" Type: {type(data).__name__}")
-                print(f" Content: {str(data)[:300]}...")
+                print(f"  Type: {type(data).__name__}")
+                print(f"  Content: {str(data)[:300]}...")
             else:
-                print(f" {data}")
+                print(f"  {data}")
 
     def _image_to_base64(self, image_path: str) -> str:
         try:
@@ -43,11 +50,13 @@ class VisualMultimodalRAG:
 
     def analyze_image_visually(self, image_path: str, image_idx: int) -> str:
         if not os.path.exists(image_path):
-            return f"Image {image_idx}: File not found - {image_path}"
+            return f"[Image {image_idx}: File not found - {image_path}]"
+        
         try:
             image_base64 = self._image_to_base64(image_path)
             if not image_base64:
-                return f"Image {image_idx}: Could not convert to base64"
+                return f"[Image {image_idx}: Could not convert to base64]"
+            
             file_ext = Path(image_path).suffix.lower()
             media_type_map = {
                 '.jpg': 'image/jpeg',
@@ -57,7 +66,9 @@ class VisualMultimodalRAG:
                 '.webp': 'image/webp'
             }
             media_type = media_type_map.get(file_ext, 'image/png')
-            print(f"Analyzing image {image_idx} visually (as {media_type})...")
+            
+            print(f"Analyzing image {image_idx}...")
+            
             message = HumanMessage(
                 content=[
                     {
@@ -80,25 +91,34 @@ Analysis:"""
                     }
                 ],
             )
+            
             response = self.llm.invoke([message])
             analysis = response.content.strip()
+            
             if self.debug:
                 self._debug_print(f"Image {image_idx} Visual Analysis", analysis)
+            
             print(f"Image {image_idx} analyzed successfully")
             return analysis
+            
         except Exception as e:
-            error_msg = f"Image {image_idx}: Vision analysis failed - {str(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]:
+
         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
+            
             visual_analysis = self.analyze_image_visually(image_path, idx)
+            
             visual_analyses.append({
                 'type': 'image_visual',
                 'image_index': idx,
@@ -106,24 +126,32 @@ Analysis:"""
                 'visual_analysis': visual_analysis,
                 'ocr_text': image.get('ocr_text', '')
             })
+        
         return visual_analyses
 
     def summarize_text_chunks(self, text: str, chunk_size: int = 1500) -> List[Dict]:
         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:
                 continue
+            
             try:
                 prompt = f"""Summarize this text chunk in {self.language}.
 Be brief and meaningful. Extract key points, facts, and main ideas.
+
 Text Chunk:
 {chunk}
+
 Summary:"""
+                
                 message = HumanMessage(content=prompt)
                 response = self.llm.invoke([message])
                 summary = response.content.strip()
+                
                 chunks.append({
                     'type': 'text_chunk',
                     'chunk_index': len(chunks),
@@ -131,27 +159,37 @@ Summary:"""
                     'summary': summary,
                     'chunk_length': len(chunk)
                 })
+                
                 if self.debug:
                     self._debug_print(f"Text Chunk {len(chunks)-1} Summary", summary)
+                
             except Exception as e:
                 print(f"Error summarizing text chunk: {e}")
+        
         return chunks
 
     def summarize_tables(self, tables: List[Dict]) -> List[Dict]:
         summaries = []
+        
         for idx, table in enumerate(tables):
             table_content = table.get('content', '')
+            
             if not table_content or len(table_content.strip()) < 10:
                 continue
+            
             try:
                 prompt = f"""Analyze and summarize this table/structured data in {self.language}.
 Extract key insights, row/column meanings, and important figures. Be brief and meaningful.
+
 Table Content:
 {table_content}
+
 Summary:"""
+                
                 message = HumanMessage(content=prompt)
                 response = self.llm.invoke([message])
                 summary = response.content.strip()
+                
                 summaries.append({
                     'type': 'table',
                     'table_index': idx,
@@ -159,21 +197,27 @@ Summary:"""
                     'summary': summary,
                     'table_length': len(table_content)
                 })
+                
                 if self.debug:
                     self._debug_print(f"Table {idx} Summary", summary)
+                
             except Exception as e:
                 print(f"Error summarizing table {idx}: {e}")
+        
         return summaries
 
     def process_and_store_document(
-        self,
-        text: str,
+        self, 
+        text: str, 
         images: List[Dict],
         tables: List[Dict],
         vector_store,
         doc_id: str
     ) -> Dict:
-        print("PROCESSING WITH VISUAL IMAGE ANALYSIS: " + doc_id)
+
+        print(f"PROCESSING WITH VISUAL IMAGE ANALYSIS: {doc_id}")
+
+        
         results = {
             'doc_id': doc_id,
             'image_visual_analyses': [],
@@ -181,42 +225,54 @@ Summary:"""
             'table_summaries': [],
             'total_stored': 0
         }
-        print("VISUAL IMAGE ANALYSIS (" + str(len(images)) + " total)")
+        
+
+        print(f"VISUAL IMAGE ANALYSIS ({len(images)} total)")
+
+        
         image_analyses = self.analyze_images_visually(images)
         results['image_visual_analyses'] = image_analyses
+        
         image_docs = {
             'text': ' | '.join([
-                f"Image {a['image_index']}: {a['visual_analysis']}"
+                f"Image {a['image_index']}: {a['visual_analysis']}" 
                 for a in image_analyses
             ]),
             'images': [],
             'tables': []
         }
+        
         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]}...")
+            print(f"   Image {analysis['image_index']} (visual analysis)")
+            print(f"     Path: {analysis['image_path']}")
+            print(f"     Analysis: {analysis['visual_analysis'][:100]}...")
+        
         if image_analyses:
             try:
                 vector_store.add_documents(
-                    image_docs,
+                    image_docs, 
                     f"{doc_id}_images_visual"
                 )
                 results['total_stored'] += len(image_analyses)
-                print(f"Stored {len(image_analyses)} image visual analyses")
+                print(f" Stored {len(image_analyses)} image visual analyses")
             except Exception as e:
-                print(f"Error storing image analyses: {e}")
-        print("TEXT CHUNK SUMMARIZATION")
+                print(f" Error storing image analyses: {e}")
+        
+        print(f" TEXT CHUNK SUMMARIZATION")
+        
         text_summaries = self.summarize_text_chunks(text)
         results['text_summaries'] = text_summaries
+        
         text_docs = {
-            'text': ' | '.join([f"Chunk {s['chunk_index']}: {s['summary']}"
-            for s in text_summaries]),
+            'text': ' | '.join([f"Chunk {s['chunk_index']}: {s['summary']}" 
+                               for s in text_summaries]),
             'images': [],
             'tables': []
         }
+        
         for summary in text_summaries:
-            print(f" Chunk {summary['chunk_index']}: {summary['summary'][:50]}...")
+            print(f"  Chunk {summary['chunk_index']}: {summary['summary'][:50]}...")
+        
         if text_summaries:
             try:
                 vector_store.add_documents(
@@ -224,20 +280,25 @@ Summary:"""
                     f"{doc_id}_text_chunks"
                 )
                 results['total_stored'] += len(text_summaries)
-                print(f"Stored {len(text_summaries)} text chunk summaries")
+                print(f" Stored {len(text_summaries)} text chunk summaries")
             except Exception as e:
-                print(f"Error storing text summaries: {e}")
-        print("TABLE SUMMARIZATION (" + str(len(tables)) + " total)")
+                print(f" Error storing text summaries: {e}")
+        
+        print(f" TABLE SUMMARIZATION ({len(tables)} total)")
+        
         table_summaries = self.summarize_tables(tables)
         results['table_summaries'] = table_summaries
+        
         table_docs = {
-            'text': ' | '.join([f"Table {s['table_index']}: {s['summary']}"
-            for s in table_summaries]),
+            'text': ' | '.join([f"Table {s['table_index']}: {s['summary']}" 
+                               for s in table_summaries]),
             'images': [],
             'tables': []
         }
+        
         for summary in table_summaries:
-            print(f" Table {summary['table_index']}: {summary['summary'][:50]}...")
+            print(f"   Table {summary['table_index']}: {summary['summary'][:50]}...")
+        
         if table_summaries:
             try:
                 vector_store.add_documents(
@@ -245,17 +306,21 @@ Summary:"""
                     f"{doc_id}_tables"
                 )
                 results['total_stored'] += len(table_summaries)
-                print(f"Stored {len(table_summaries)} table summaries")
+                print(f" Stored {len(table_summaries)} table summaries")
             except Exception as e:
-                print(f"Error storing table summaries: {e}")
-        print("STORAGE SUMMARY")
-        print(f" Images analyzed and stored: {len(image_analyses)}")
-        print(f" Text chunks summarized and stored: {len(text_summaries)}")
-        print(f" Tables summarized and stored: {len(table_summaries)}")
-        print(f" Total items stored in vector: {results['total_stored']}")
+                print(f" Error storing table summaries: {e}")
+        
+        print(f" STORAGE SUMMARY")
+        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 in vector: {results['total_stored']}")
+        
+        self.visual_summaries_log.append(results)
         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):
@@ -264,41 +329,56 @@ Summary:"""
             start = end - overlap
         return chunks
 
+    def get_visual_summaries_log(self) -> List[Dict]:
+        """Get all visual analysis logs"""
+        return self.visual_summaries_log
+
+
 class AnsweringRAG:
+    
     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",
+            model_name="gpt-4o-mini",  
             api_key=api_key,
             temperature=TEMPERATURE,
             max_tokens=MAX_TOKENS,
         )
+        
         self.language = LANGUAGE
+        self.answer_log = []
+        
         if self.debug:
-            print("AnsweringRAG initialized")
+            print("AnsweringRAG initialized ")
 
     def _debug_print(self, label: str, data: any):
         if self.debug:
-            print(f"DEBUG [{label}]:")
+            print(f" DEBUG [{label}]:")
             if isinstance(data, (list, dict)):
-                print(f" Type: {type(data).__name__}")
-                print(f" Content: {str(data)[:300]}...")
+                print(f"  Type: {type(data).__name__}")
+                print(f"  Content: {str(data)[:300]}...")
             else:
-                print(f" {data}")
+                print(f"  {data}")
 
     def analyze_and_answer(
-        self,
-        question: str,
+        self, 
+        question: str, 
         search_results: List[Dict]
     ) -> Dict:
-        print("ANALYZING QUESTION & GENERATING ANSWER")
+
+        
+        print(f"ANALYZING QUESTION & GENERATING ANSWER")
+        
         print(f"Question: {question}")
         print(f"Search Results Found: {len(search_results)}")
+        
         if not search_results:
-            print("No search results found!")
-            answer = f"""No relevant information in the document to answer question: "{question}"
+            print(f"No search results found!")
+            answer = f"""No relevant information in the document to answer your question: "{question}"
 """
+            
             result = {
                 'question': question,
                 'answer': answer,
@@ -306,23 +386,33 @@ class AnsweringRAG:
                 'confidence': 'low',
                 'search_results': []
             }
+            self.answer_log.append(result)
             return result
+        
         context_parts = []
         for idx, result in enumerate(search_results, 1):
             content = result.get('content', '')
+            metadata = result.get('metadata', {})
             content_type = result.get('type', 'unknown')
             distance = result.get('distance', 0)
             relevance = 1 - distance if distance else 0
+            
             context_parts.append(f"""
 [Source {idx} - {content_type.upper()} (relevance: {relevance:.1%})]
 {content}""")
+        
         full_context = "\n".join(context_parts)
+        
         self._debug_print("Context Prepared", f"{len(context_parts)} sources, {len(full_context)} chars")
+        
         analysis_prompt = f"""You are a helpful assistant analyzing document content to answer user questions.
+
 USER QUESTION:
 "{question}"
+
 RELEVANT CONTENT FROM DOCUMENT:
 {full_context}
+
 INSTRUCTIONS:
 1. Analyze the provided content carefully
 2. Extract information relevant to the question
@@ -331,29 +421,37 @@ INSTRUCTIONS:
 5. Be specific and cite the content when relevant
 6. Structure your answer clearly with key points
 ANSWER:"""
-        print("Analyzing search results...")
-        print(f" Context size: {len(full_context)} characters")
-        print(f" Sources: {len(search_results)}")
+        
+        print(f"Analyzing search results...")
+        print(f"   Context size: {len(full_context)} characters")
+        print(f"   Sources: {len(search_results)}")
+        
         try:
             message = HumanMessage(content=analysis_prompt)
             response = self.llm.invoke([message])
             answer = response.content.strip()
+            
             confidence = self._estimate_confidence(len(search_results), answer)
-            print("Answer generated successfully")
-            print(f" Confidence: {confidence}")
-            print(f" Answer length: {len(answer)} characters")
+            
+            print(f" Answer generated successfully")
+            print(f"   Confidence: {confidence}")
+            print(f"   Answer length: {len(answer)} characters")
+            
             result = {
                 'question': question,
                 'answer': answer,
                 'sources_used': len(search_results),
                 'confidence': confidence,
-                'search_results': search_results,
-                'formatted_sources': self._format_sources(search_results)
+                'search_results': search_results
             }
+            
+            self.answer_log.append(result)
             return result
+            
         except Exception as e:
-            print(f"Error generating answer: {e}")
-            answer = "Error while analyzing the search results"
+            print(f" Error generating answer: {e}")
+            answer = f"I encountered an error while analyzing the search results. Please try again."
+            
             result = {
                 'question': question,
                 'answer': answer,
@@ -362,24 +460,55 @@ ANSWER:"""
                 'error': str(e),
                 'search_results': search_results
             }
+            
+            self.answer_log.append(result)
             return result
 
     def _estimate_confidence(self, sources_count: int, answer: str) -> str:
         answer_length = len(answer)
+        
         if sources_count >= 3 and answer_length > 500:
             return "high"
+        
         elif sources_count >= 2 and answer_length > 200:
             return "medium"
+        
         else:
             return "low"
 
-    def _format_sources(self, search_results: List[Dict]) -> List[Dict]:
+    def get_answer_with_sources(
+        self,
+        question: str,
+        search_results: List[Dict]
+    ) -> Dict:
+        
+        result = self.analyze_and_answer(question, search_results)
+        
         formatted_sources = []
-        for idx, source in enumerate(search_results, 1):
+        for idx, source in enumerate(result['search_results'], 1):
             formatted_sources.append({
                 'index': idx,
                 'type': source.get('type', 'unknown'),
                 'content': source.get('content', ''),
                 'relevance': 1 - source.get('distance', 0) if source.get('distance') else 0
             })
-        return formatted_sources
+        
+        result['formatted_sources'] = formatted_sources
+        return result
+
+    def get_answer_log(self) -> List[Dict]:
+        return self.answer_log
+
+    def print_answer_with_sources(self, result: Dict, max_source_length: int = 300):
+        
+        print(f"ANSWER TO: {result['question']}")
+        
+        print(f"ANSWER (Confidence: {result['confidence'].upper()}):")
+        print(result['answer'])
+        
+        if result.get('formatted_sources'):
+            print(f"SOURCES USED ({len(result['formatted_sources'])} total):")
+            for source in result['formatted_sources']:
+                print(f"\n[Source {source['index']} - {source['type'].upper()} ({source['relevance']:.0%} relevant)]")
+                print(f"{source['content'][:max_source_length]}...")
+        

src/vector_store.py

@@ -9,9 +9,9 @@ from config import CHROMA_DB_PATH, EMBEDDING_MODEL, EMBEDDING_DIM
 
 class CLIPEmbedder:
     def __init__(self, model_name: str = EMBEDDING_MODEL):
-        print(f"Loading embedding model: {model_name}")
+        print(f" Loading embedding model: {model_name}")
         self.model = SentenceTransformer(model_name)
-        print(f"Model loaded successfully")
+        print(f" Model loaded successfully")
     
     def embed(self, text: str) -> List[float]:
         try:
@@ -35,15 +35,15 @@ class VectorStore:
         self.persist_directory = CHROMA_DB_PATH
         self.embedder = CLIPEmbedder()
         
-        print(f"Initializing ChromaDB at: {self.persist_directory}")
+        print(f" Initializing ChromaDB at: {self.persist_directory}")
         
         try:
             self.client = chromadb.PersistentClient(
                 path=self.persist_directory
             )
-            print(f"ChromaDB initialized")
+            print(f" ChromaDB initialized")
         except Exception as e:
-            print(f"Error initializing ChromaDB: {e}")
+            print(f" Error initializing ChromaDB: {e}")
             self.client = chromadb.PersistentClient(
                 path=self.persist_directory
             )
@@ -54,7 +54,7 @@ class VectorStore:
                 metadata={"hnsw:space": "cosine"}
             )
             count = self.collection.count()
-            print(f"Collection loaded: {count} items in store")
+            print(f" Collection loaded: {count} items in store")
         except Exception as e:
             print(f"Error with collection: {e}")
             self.collection = self.client.get_or_create_collection(
@@ -66,7 +66,7 @@ class VectorStore:
         metadatas = []
         ids = []
         
-        print(f"Adding documents for: {doc_id}")
+        print(f" Adding documents for: {doc_id}")
         
         if 'text' in documents and documents['text']:
             chunks = self._chunk_text(documents['text'], chunk_size=1000, overlap=200)
@@ -78,7 +78,7 @@ class VectorStore:
                     'chunk_idx': str(idx)
                 })
                 ids.append(f"{doc_id}_text_{idx}")
-            print(f"Text: {len(chunks)} chunks")
+            print(f"   Text: {len(chunks)} chunks")
         
         if 'images' in documents:
             image_count = 0
@@ -94,7 +94,7 @@ class VectorStore:
                     ids.append(f"{doc_id}_image_{idx}")
                     image_count += 1
             if image_count > 0:
-                print(f"Images: {image_count} with OCR text")
+                print(f"   Images: {image_count} with OCR text")
         
         if 'tables' in documents:
             table_count = 0
@@ -109,10 +109,10 @@ class VectorStore:
                     ids.append(f"{doc_id}_table_{idx}")
                     table_count += 1
             if table_count > 0:
-                print(f"Tables: {table_count}")
+                print(f"   Tables: {table_count}")
         
         if texts:
-            print(f"Generating {len(texts)} embeddings...")
+            print(f"   Generating {len(texts)} embeddings...")
             embeddings = self.embedder.embed_batch(texts)
             
             try:
@@ -122,11 +122,13 @@ class VectorStore:
                     embeddings=embeddings,
                     metadatas=metadatas
                 )
-                print(f"Successfully added {len(texts)} items to vector store")
+                print(f" Successfully added {len(texts)} items to vector store")
+                print(f" Data persisted automatically to: {self.persist_directory}")
             except Exception as e:
-                print(f"Error adding to collection: {e}")
+                print(f" Error adding to collection: {e}")
 
     def search(self, query: str, n_results: int = 5) -> List[Dict]:
+        """Search vector store for similar documents"""
         try:
             query_embedding = self.embedder.embed(query)
             
@@ -180,13 +182,13 @@ class VectorStore:
             results = self.collection.get(where={'doc_id': doc_id})
             if results['ids']:
                 self.collection.delete(ids=results['ids'])
-                print(f"Deleted {len(results['ids'])} documents for {doc_id}")
-                print(f"Changes persisted automatically")
+                print(f" Deleted {len(results['ids'])} documents for {doc_id}")
         except Exception as e:
             print(f"Error deleting documents: {e}")
 
     def persist(self):
-        print("Vector store is using auto-persist")
+
+        print(" Vector store is using auto-persist")
 
     def clear_all(self):
         try:
@@ -195,6 +197,6 @@ class VectorStore:
                 name="multimodal_rag",
                 metadata={"hnsw:space": "cosine"}
             )
-            print("Collection cleared and reset")
+            print(" Collection cleared and reset")
         except Exception as e:
             print(f"Error clearing collection: {e}")