Update app.py

app.py

@@ -1,6 +1,5 @@
-
+# app.py - Fixed version with proper error handling
 import sqlite3
-
 import os
 import streamlit as st
 import chromadb
@@ -12,28 +11,41 @@ from llama_index.vector_stores.chroma import ChromaVectorStore
 from llama_index.llms.groq import Groq
 from llama_index.embeddings.cohere import CohereEmbedding
 
+# Load environment variables first
+load_dotenv()
 
-from arize.otel import register
-from openinference.instrumentation.llama_index import LlamaIndexInstrumentor
+# Disable ChromaDB telemetry to remove the warning
+os.environ["ANONYMIZED_TELEMETRY"] = "False"
 
-# Setup OTel via Arize's convenience function
-tracer_provider = register(
-    space_id=os.getenv("ARIZE_SPACE_ID"),
-    api_key=os.getenv("ARIZE_API_KEY"),
-    project_name="rbacrag" # Choose a project name
-)
-
-# Instrument LlamaIndex
-LlamaIndexInstrumentor().instrument(tracer_provider=tracer_provider)
+# Setup OTel via Arize's convenience function with error handling
+try:
+    from arize.otel import register
+    from openinference.instrumentation.llama_index import LlamaIndexInstrumentor
+    
+    if os.getenv("ARIZE_SPACE_ID") and os.getenv("ARIZE_API_KEY"):
+        tracer_provider = register(
+            space_id=os.getenv("ARIZE_SPACE_ID"),
+            api_key=os.getenv("ARIZE_API_KEY"),
+            project_name="rbacrag"
+        )
+        LlamaIndexInstrumentor().instrument(tracer_provider=tracer_provider)
+    else:
+        print("Arize credentials not found, skipping instrumentation")
+except Exception as e:
+    print(f"Warning: Arize instrumentation failed: {e}")
 
 # Import database module
 from database import db, initialize_users
 
-# Load environment variables
-load_dotenv()
-
-# Initialize default users
-initialize_users()
+# Initialize default users with better error handling
+try:
+    success_count, error_count = initialize_users()
+    if error_count > 0:
+        print(f"Database initialization completed with {error_count} errors (likely users already exist)")
+    else:
+        print(f"Database initialization successful: {success_count} users ready")
+except Exception as e:
+    print(f"Error during user initialization: {e}")
 
 # Role-based access control for documents
 ROLE_ACCESS = {
@@ -43,7 +55,6 @@ ROLE_ACCESS = {
     "marketing": ["marketing", "general"]
 }
 
-# Initialize session state
 def initialize_session_state():
     """Initialize or reset the session state"""
     if "authenticated" not in st.session_state:
@@ -70,7 +81,6 @@ st.set_page_config(
 # Initialize session state
 initialize_session_state()
 
-# Authentication functions
 def login(username: str, password: str) -> bool:
     """
     Authenticate user and set session state
@@ -91,7 +101,7 @@ def login(username: str, password: str) -> bool:
             st.session_state.messages = [
                 {"role": "assistant", "content": f"Welcome, {user['username']}! How can I assist you today?"}
             ]
-            st.rerun()  # Rerun to update the UI
+            st.rerun()
             return True
         return False
     except Exception as e:
@@ -99,23 +109,23 @@ def login(username: str, password: str) -> bool:
         return False
 
 def logout():
-    """
-    Log out the current user and clear session state
-    """
+    """Log out the current user and clear session state"""
     username = st.session_state.get('username', 'Unknown')
     st.session_state.clear()
     initialize_session_state()
     st.success(f"Successfully logged out {username}")
-    st.rerun()  # Rerun to update the UI
+    st.rerun()
 
 @st.cache_resource
 def load_vector_index(role: str):
-    """Load the ChromaDB index for the user's role"""
+    """Load the ChromaDB index for the user's role with enhanced error handling"""
     try:
         # Initialize Cohere embeddings
         cohere_api_key = os.getenv("COHERE_API_KEY")
         if not cohere_api_key:
-            raise ValueError("COHERE_API_KEY not found in environment variables")
+            st.error("❌ COHERE_API_KEY not found in environment variables")
+            st.info("Please set your Cohere API key in the .env file")
+            st.stop()
             
         embed_model = CohereEmbedding(
             cohere_api_key=cohere_api_key,
@@ -124,12 +134,59 @@ def load_vector_index(role: str):
         )
         Settings.embed_model = embed_model
         
-        # Initialize Chroma client
+        # Docker-compatible ChromaDB initialization
         persist_dir = f"./chroma_db/{role}"
-        chroma_client = chromadb.PersistentClient(path=persist_dir)
         
-        # Get the collection
-        chroma_collection = chroma_client.get_collection("documents")
+        # Ensure directory exists
+        Path(persist_dir).mkdir(parents=True, exist_ok=True)
+        
+        # Initialize Chroma client with telemetry disabled
+        try:
+            chroma_client = chromadb.PersistentClient(
+                path=persist_dir,
+                settings=chromadb.Settings(
+                    anonymized_telemetry=False,
+                    allow_reset=True
+                )
+            )
+        except Exception as e:
+            st.warning(f"Failed to connect to persistent ChromaDB: {e}")
+            st.info("Attempting to create new collection...")
+            
+            # Try to reset and recreate
+            try:
+                chroma_client = chromadb.PersistentClient(path=persist_dir)
+                chroma_client.reset()
+                chroma_client = chromadb.PersistentClient(
+                    path=persist_dir,
+                    settings=chromadb.Settings(
+                        anonymized_telemetry=False,
+                        allow_reset=True
+                    )
+                )
+            except:
+                # Fallback to in-memory client
+                st.warning("⚠️ Using in-memory ChromaDB (data will not persist)")
+                chroma_client = chromadb.Client(
+                    settings=chromadb.Settings(anonymized_telemetry=False)
+                )
+        
+        # Try to get existing collection, create if it doesn't exist
+        collection_name = "documents"
+        try:
+            chroma_collection = chroma_client.get_collection(collection_name)
+            st.success(f"✅ Connected to existing collection for {role} role")
+        except Exception:
+            st.warning(f"⚠️ Collection '{collection_name}' not found for role '{role}'. Creating empty collection.")
+            try:
+                chroma_collection = chroma_client.create_collection(
+                    name=collection_name,
+                    metadata={"hnsw:space": "cosine"}
+                )
+                st.info("📝 Created new empty collection. You may need to add documents first.")
+            except Exception as create_error:
+                st.error(f"❌ Failed to create collection: {create_error}")
+                st.stop()
         
         # Create vector store
         vector_store = ChromaVectorStore(chroma_collection=chroma_collection)
@@ -137,22 +194,36 @@ def load_vector_index(role: str):
         # Create storage context
         storage_context = StorageContext.from_defaults(vector_store=vector_store)
         
-        # Load the index
-        index = VectorStoreIndex.from_vector_store(
-            vector_store=vector_store,
-            storage_context=storage_context,
-            embed_model=embed_model
-        )
+        # Check if collection has documents
+        if chroma_collection.count() == 0:
+            st.warning(f"📭 No documents found in {role} collection.")
+            st.info("The system will work, but responses will be limited without documents.")
+            # Create empty index for now
+            index = VectorStoreIndex([], storage_context=storage_context, embed_model=embed_model)
+        else:
+            st.info(f"📚 Found {chroma_collection.count()} documents in {role} collection")
+            # Load the index
+            index = VectorStoreIndex.from_vector_store(
+                vector_store=vector_store,
+                storage_context=storage_context,
+                embed_model=embed_model
+            )
         
         return index
+        
     except Exception as e:
-        st.error(f"Error loading vector index: {str(e)}")
+        st.error(f"❌ Error loading vector index: {str(e)}")
+        st.info("**Possible solutions:**")
+        st.info("1. Check that ChromaDB collections exist for this role")
+        st.info("2. Verify database files are properly mounted in Docker")
+        st.info("3. Check permissions on the database directory")
+        st.info("4. Ensure COHERE_API_KEY is set correctly")
         st.stop()
 
 def chat_interface():
     """Main chat interface"""
     # Add styled heading
-    st.markdown(f"<h2 style='color: #1407fa;'>💬 {st.session_state.role.capitalize()} Department Chat</h3>", unsafe_allow_html=True)
+    st.markdown(f"<h2 style='color: #1407fa;'>💬 {st.session_state.role.capitalize()} Department Chat</h2>", unsafe_allow_html=True)
     
     # Display chat messages
     for message in st.session_state.messages:
@@ -164,9 +235,15 @@ def chat_interface():
     
     # Initialize Groq LLM
     try:
+        groq_api_key = os.getenv("GROQ_API_KEY")
+        if not groq_api_key:
+            st.error("❌ GROQ_API_KEY not found in environment variables")
+            st.info("Please set your Groq API key in the .env file")
+            st.stop()
+            
         llm = Groq(
             model="llama3-8b-8192", 
-            api_key=os.getenv("GROQ_API_KEY"),
+            api_key=groq_api_key,
             temperature=0.5,
             system_prompt=f"You are a helpful assistant specialized in {st.session_state.role} department documents. Answer the user queries with the help of the provided context with high accuracy and precision."
         )
@@ -178,8 +255,8 @@ def chat_interface():
             response_mode="compact"
         )
     except Exception as e:
-        st.error(f"Error initializing LLM: {str(e)}")
-        st.warning("Falling back to default LLM settings. Some features may be limited.")
+        st.error(f"❌ Error initializing LLM: {str(e)}")
+        st.warning("⚠️ Falling back to default LLM settings. Some features may be limited.")
         query_engine = index.as_query_engine(
             similarity_top_k=3,
             response_mode="compact"
@@ -205,7 +282,7 @@ def chat_interface():
                 full_response = str(response)
                 message_placeholder.markdown(full_response)
             except Exception as e:
-                error_msg = f"Error generating response: {str(e)}"
+                error_msg = f"❌ Error generating response: {str(e)}"
                 message_placeholder.error(error_msg)
                 full_response = error_msg
         
@@ -284,7 +361,6 @@ def show_login_form():
 
     st.markdown('<p style="text-align: center; margin-top: 2rem; color: #a0a0b0;">2025 Department RAG System</p>', unsafe_allow_html=True)
 
-
 def main():
     """
     Main application entry point
@@ -292,7 +368,6 @@ def main():
     """
     # Sidebar for logout and user info
     if st.session_state.authenticated:
-        st.set_page_config(layout="wide", initial_sidebar_state="expanded")
         with st.sidebar:
             st.markdown(f"### Welcome, {st.session_state.username}")
             st.markdown(f"**Role:** {st.session_state.role.capitalize()}")
@@ -307,6 +382,16 @@ def main():
             This is a secure departmental RAG system that provides
             role-based access to information across different departments.
             """)
+            
+            # Show database status
+            try:
+                users = db.list_users()
+                st.markdown("---")
+                st.markdown("### System Status")
+                st.markdown(f"✅ Database: {len(users)} users")
+                st.markdown("✅ Authentication: Active")
+            except:
+                st.markdown("⚠️ Database: Connection issues")
     
     # Main content area
     if not st.session_state.authenticated: