Update src/enhanced_rag_system.py

src/enhanced_rag_system.py

@@ -6,22 +6,15 @@ from typing import List, Tuple, Dict, Optional
 from langchain.schema import Document
 import re
 import json
-import warnings
-warnings.filterwarnings('ignore')
 
-# Import vector store components with better error handling
+# Import vector store components
 try:
     import chromadb
     from chromadb.config import Settings
     from sentence_transformers import SentenceTransformer
     VECTOR_STORE_AVAILABLE = True
-    print("✅ ChromaDB and SentenceTransformers imported successfully")
-except ImportError as e:
-    VECTOR_STORE_AVAILABLE = False
-    print(f"⚠️ Vector store import error: {e}")
-except Exception as e:
+except ImportError:
     VECTOR_STORE_AVAILABLE = False
-    print(f"⚠️ Vector store initialization error: {e}")
 
 # Import LLM components
 try:
@@ -29,12 +22,8 @@ try:
     LLM_AVAILABLE = bool(os.getenv("OPENAI_API_KEY"))
     if LLM_AVAILABLE:
         openai.api_key = os.getenv("OPENAI_API_KEY")
-        print("✅ OpenAI API key found and configured")
-    else:
-        print("⚠️ OpenAI API key not found in environment")
 except ImportError:
     LLM_AVAILABLE = False
-    print("⚠️ OpenAI library not available")
 
 # Import our custom modules
 from document_processor import DocumentProcessor
@@ -100,31 +89,17 @@ class EnhancedRAGSystem:
             print("⚠️ Using fallback mode with template responses")
     
     def _initialize_vector_store(self):
-        """Initialize ChromaDB vector store with better error handling"""
+        """Initialize ChromaDB vector store"""
         if not VECTOR_STORE_AVAILABLE:
-            print("⚠️ ChromaDB/SentenceTransformers not available, using in-memory search")
+            print("⚠️ ChromaDB not available, using in-memory search")
             return
         
         try:
-            print("🔧 Initializing ChromaDB...")
-            
-            # Try different ChromaDB configurations for HuggingFace compatibility
-            try:
-                # First try: PersistentClient (newer API)
-                self.chroma_client = chromadb.PersistentClient(path="./chroma_db")
-                print("✅ Using ChromaDB PersistentClient")
-            except Exception as e1:
-                try:
-                    # Second try: Client with settings (older API)
-                    self.chroma_client = chromadb.Client(Settings(
-                        chroma_db_impl="duckdb+parquet",
-                        persist_directory="./chroma_db"
-                    ))
-                    print("✅ Using ChromaDB Client with Settings")
-                except Exception as e2:
-                    # Third try: Simple client
-                    self.chroma_client = chromadb.Client()
-                    print("✅ Using ChromaDB in-memory client")
+            # Initialize ChromaDB client
+            self.chroma_client = chromadb.Client(Settings(
+                chroma_db_impl="duckdb+parquet",
+                persist_directory="./chroma_db"
+            ))
             
             # Get or create collection
             collection_name = "finsolve_documents"
@@ -138,25 +113,15 @@ class EnhancedRAGSystem:
                 )
                 print(f"✅ Created new ChromaDB collection: {collection_name}")
             
-            # Initialize embedding model with smaller model for HuggingFace
-            try:
-                self.embedding_model = SentenceTransformer("all-MiniLM-L6-v2")
-                print("✅ Loaded sentence transformer model: all-MiniLM-L6-v2")
-            except Exception as e:
-                # Fallback to even smaller model
-                try:
-                    self.embedding_model = SentenceTransformer("paraphrase-MiniLM-L3-v2")
-                    print("✅ Loaded fallback sentence transformer model: paraphrase-MiniLM-L3-v2")
-                except Exception as e2:
-                    print(f"❌ Failed to load embedding model: {e2}")
-                    raise e2
+            # Initialize embedding model
+            self.embedding_model = SentenceTransformer("sentence-transformers/all-MiniLM-L6-v2")
+            print("✅ Loaded sentence transformer model")
             
             self.vector_store_initialized = True
             
         except Exception as e:
             print(f"⚠️ ChromaDB initialization failed: {str(e)}")
             print("⚠️ Falling back to in-memory search")
-            self.vector_store_initialized = False
     
     def _initialize_llm(self):
         """Initialize OpenAI LLM"""
@@ -165,7 +130,7 @@ class EnhancedRAGSystem:
             return
         
         try:
-            # Test OpenAI connection with updated API
+            # Test OpenAI connection
             response = openai.ChatCompletion.create(
                 model=self.llm_model,
                 messages=[{"role": "user", "content": "Hello"}],
@@ -182,15 +147,10 @@ class EnhancedRAGSystem:
     
     def _load_documents_to_vector_store(self):
         """Load documents into ChromaDB vector store"""
-        if not self.vector_store_initialized or not self.embedding_model:
+        if not self.vector_store_initialized or self.collection.count() > 0:
             return
         
         try:
-            # Check if documents already loaded
-            if self.collection.count() > 0:
-                print(f"✅ ChromaDB already contains {self.collection.count()} documents")
-                return
-            
             print("📄 Loading documents into vector store...")
             
             texts = []
@@ -205,7 +165,7 @@ class EnhancedRAGSystem:
                     "title": doc.metadata.get("title", "Document"),
                     "department": doc.metadata.get("department", "General"),
                     "type": doc.metadata.get("type", "Document"),
-                    "chunk_id": str(doc.metadata.get("chunk_id", 0)),
+                    "chunk_id": doc.metadata.get("chunk_id", 0),
                     "source": doc.metadata.get("source", "unknown")
                 }
                 
@@ -213,23 +173,16 @@ class EnhancedRAGSystem:
                 metadatas.append(metadata)
                 ids.append(doc_id)
             
-            # Generate embeddings in batches to avoid memory issues
-            batch_size = 10
-            for i in range(0, len(texts), batch_size):
-                batch_texts = texts[i:i+batch_size]
-                batch_metadatas = metadatas[i:i+batch_size] 
-                batch_ids = ids[i:i+batch_size]
-                
-                # Generate embeddings
-                embeddings = self.embedding_model.encode(batch_texts).tolist()
-                
-                # Add to ChromaDB
-                self.collection.add(
-                    embeddings=embeddings,
-                    documents=batch_texts,
-                    metadatas=batch_metadatas,
-                    ids=batch_ids
-                )
+            # Generate embeddings
+            embeddings = self.embedding_model.encode(texts).tolist()
+            
+            # Add to ChromaDB
+            self.collection.add(
+                embeddings=embeddings,
+                documents=texts,
+                metadatas=metadatas,
+                ids=ids
+            )
             
             print(f"✅ Loaded {len(self.documents)} documents into ChromaDB")
             
@@ -242,10 +195,10 @@ class EnhancedRAGSystem:
         print("🤖 FINSOLVE RAG SYSTEM STATUS")
         print("="*50)
         print(f"✅ Python: Core system initialized")
-        print(f"{'✅' if self.vector_store_initialized else '⚠️'} ChromaDB Vector Store: {'Ready' if self.vector_store_initialized else 'Fallback mode'}")
-        print(f"{'✅' if self.llm_initialized else '⚠️'} OpenAI LLM: {'OpenAI GPT' if self.llm_initialized else 'Template mode'}")
+        print(f"{'✅' if self.vector_store_initialized else '⚠️'} ChromaDB Vector Store: {'Available' if self.vector_store_initialized else 'Fallback mode'}")
+        print(f"{'✅' if self.llm_initialized else '⚠️'} OpenAI LLM: {'Available' if self.llm_initialized else 'Template mode'}")
         print(f"✅ Streamlit: UI active")
-        print(f"🔄 FastAPI: Simulated API")
+        print(f"✅ FastAPI: Simulated endpoints")
         print(f"✅ Authentication: JWT-style RBAC")
         print(f"✅ NLP: Intent classification + {'LLM' if self.llm_initialized else 'Templates'}")
         print(f"✅ RAG: Vector retrieval + context augmentation")
@@ -647,6 +600,34 @@ Please contact your administrator if you need access to additional information."
             
             return None
             
+        except Exception as e:
+            print(f"❌ Error creating visualization: {str(e)}")
+            return None
+    
+    def _create_data_table(self, content: str, query_intent: str) -> Optional[str]:
+        """Create data tables from content"""
+        try:
+            if query_intent == "finance":
+                data = {
+                    'Metric': ['Q4 Revenue', 'Annual Revenue', 'Net Income', 'Gross Margin', 'ROI'],
+                    'Value': ['$2.6B', '$9.4B', '$325M', '64%', '15%'],
+                    'YoY Growth': ['+35%', '+28%', '+18%', '+6%', '+3%']
+                }
+                df = pd.DataFrame(data)
+                return df.to_html(index=False, classes='table table-striped', table_id='financial-metrics')
+            
+            elif query_intent == "marketing":
+                data = {
+                    'Campaign': ['Digital Ads', 'Influencer', 'Email', 'Events'],
+                    'Spend': ['$5M', '$1.5M', '$0.2M', '$2M'],
+                    'ROI': ['3.5x', '4.2x', '2.0x', '5.0x'],
+                    'Leads': ['180K', '60K', '25K', '300']
+                }
+                df = pd.DataFrame(data)
+                return df.to_html(index=False, classes='table table-striped', table_id='marketing-metrics')
+            
+            return None
+            
         except Exception as e:
             print(f"❌ Error creating table: {str(e)}")
             return None
@@ -732,20 +713,11 @@ Please contact your administrator if you need access to additional information."
                 "streamlit": "✅ Active", 
                 "vector_store": "✅ ChromaDB" if self.vector_store_initialized else "⚠️ Fallback",
                 "llm": f"✅ {self.llm_model}" if self.llm_initialized else "⚠️ Templates",
-                "fastapi": "✅ Real FastAPI" if self._check_fastapi_running() else "🔄 Simulated",
+                "fastapi": "✅ Simulated",
                 "authentication": "✅ JWT-style RBAC"
             }
         }
     
-    def _check_fastapi_running(self) -> bool:
-        """Check if FastAPI server is running"""
-        try:
-            import requests
-            response = requests.get("http://localhost:8000/health", timeout=2)
-            return response.status_code == 200
-        except:
-            return False
-    
     def get_available_documents_for_role(self, role: str) -> List[Dict]:
         """Get list of documents available for a specific role"""
         accessible_docs = self.auth_system.get_accessible_documents(role)
@@ -759,32 +731,4 @@ Please contact your administrator if you need access to additional information."
                     **doc_info[doc_name]
                 })
         
-        return available❌ Error creating visualization: {str(e)}")
-            return None
-    
-    def _create_data_table(self, content: str, query_intent: str) -> Optional[str]:
-        """Create data tables from content"""
-        try:
-            if query_intent == "finance":
-                data = {
-                    'Metric': ['Q4 Revenue', 'Annual Revenue', 'Net Income', 'Gross Margin', 'ROI'],
-                    'Value': ['$2.6B', '$9.4B', '$325M', '64%', '15%'],
-                    'YoY Growth': ['+35%', '+28%', '+18%', '+6%', '+3%']
-                }
-                df = pd.DataFrame(data)
-                return df.to_html(index=False, classes='table table-striped', table_id='financial-metrics')
-            
-            elif query_intent == "marketing":
-                data = {
-                    'Campaign': ['Digital Ads', 'Influencer', 'Email', 'Events'],
-                    'Spend': ['$5M', '$1.5M', '$0.2M', '$2M'],
-                    'ROI': ['3.5x', '4.2x', '2.0x', '5.0x'],
-                    'Leads': ['180K', '60K', '25K', '300']
-                }
-                df = pd.DataFrame(data)
-                return df.to_html(index=False, classes='table table-striped', table_id='marketing-metrics')
-            
-            return None
-            
-        except Exception as e:
-            print(f"
+        return available