bharathmunakala/Role_Base_Access_Control
0
1__import__('pysqlite3')
2import sys
3sys.modules['sqlite3'] = sys.modules.pop('pysqlite3')
4
5import sqlite3
6
7import os
8import streamlit as st
9import chromadb
10from typing import Dict, Optional, Any
11from pathlib import Path
12from dotenv import load_dotenv
13from llama_index.core import VectorStoreIndex, StorageContext, Settings
14from llama_index.vector_stores.chroma import ChromaVectorStore
15from llama_index.llms.groq import Groq
16from llama_index.embeddings.cohere import CohereEmbedding
17
18
19from arize.otel import register
20from openinference.instrumentation.llama_index import LlamaIndexInstrumentor
21
22# Setup OTel via Arize's convenience function
23tracer_provider = register(
24 space_id=os.getenv("ARIZE_SPACE_ID"),
25 api_key=os.getenv("ARIZE_API_KEY"),
26 project_name="rbacrag" # Choose a project name
27)
28
29# Instrument LlamaIndex
30LlamaIndexInstrumentor().instrument(tracer_provider=tracer_provider)
31
32# Import database module
33from database import db, initialize_users
34
35# Load environment variables
36load_dotenv()
37
38# Initialize default users
39initialize_users()
40
41# Role-based access control for documents
42ROLE_ACCESS = {
43 "hr": ["hr", "general"],
44 "engineering": ["engineering", "general"],
45 "finance": ["finance", "general"],
46 "marketing": ["marketing", "general"]
47}
48
49# Initialize session state
50def initialize_session_state():
51 """Initialize or reset the session state"""
52 if "authenticated" not in st.session_state:
53 st.session_state.authenticated = False
54 if "username" not in st.session_state:
55 st.session_state.username = None
56 if "role" not in st.session_state:
57 st.session_state.role = None
58 if "messages" not in st.session_state:
59 st.session_state.messages = []
60 if "vector_index" not in st.session_state:
61 st.session_state.vector_index = None
62 if "query_engine" not in st.session_state:
63 st.session_state.query_engine = None
64
65# Set page config
66st.set_page_config(
67 page_title="Departmental RAG System",
68 page_icon="🔒",
69 layout="centered",
70 initial_sidebar_state="collapsed"
71)
72
73# Initialize session state
74initialize_session_state()
75
76# Authentication functions
77def login(username: str, password: str) -> bool:
78 """
79 Authenticate user and set session state
80
81 Args:
82 username: The username to authenticate
83 password: The password to verify
84
85 Returns:
86 bool: True if authentication was successful, False otherwise
87 """
88 try:
89 user = db.verify_user(username, password)
90 if user:
91 st.session_state.authenticated = True
92 st.session_state.username = user["username"]
93 st.session_state.role = user["role"]
94 st.session_state.messages = [
95 {"role": "assistant", "content": f"Welcome, {user['username']}! How can I assist you today?"}
96 ]
97 st.rerun() # Rerun to update the UI
98 return True
99 return False
100 except Exception as e:
101 st.error(f"An error occurred during login: {str(e)}")
102 return False
103
104def logout():
105 """
106 Log out the current user and clear session state
107 """
108 username = st.session_state.get('username', 'Unknown')
109 st.session_state.clear()
110 initialize_session_state()
111 st.success(f"Successfully logged out {username}")
112 st.rerun() # Rerun to update the UI
113
114@st.cache_resource
115def load_vector_index(role: str):
116 """Load the ChromaDB index for the user's role"""
117 try:
118 # Initialize Cohere embeddings
119 cohere_api_key = os.getenv("COHERE_API_KEY")
120 if not cohere_api_key:
121 raise ValueError("COHERE_API_KEY not found in environment variables")
122
123 embed_model = CohereEmbedding(
124 cohere_api_key=cohere_api_key,
125 model_name="embed-english-v3.0",
126 input_type="search_document"
127 )
128 Settings.embed_model = embed_model
129
130 # Initialize Chroma client
131 persist_dir = f"./chroma_db/{role}"
132 chroma_client = chromadb.PersistentClient(path=persist_dir)
133
134 # Get the collection
135 chroma_collection = chroma_client.get_collection("documents")
136
137 # Create vector store
138 vector_store = ChromaVectorStore(chroma_collection=chroma_collection)
139
140 # Create storage context
141 storage_context = StorageContext.from_defaults(vector_store=vector_store)
142
143 # Load the index
144 index = VectorStoreIndex.from_vector_store(
145 vector_store=vector_store,
146 storage_context=storage_context,
147 embed_model=embed_model
148 )
149
150 return index
151 except Exception as e:
152 st.error(f"Error loading vector index: {str(e)}")
153 st.stop()
154
155def chat_interface():
156 """Main chat interface"""
157 # Add styled heading
158 st.markdown(f"<h2 style='color: #1407fa;'>💬 {st.session_state.role.capitalize()} Department Chat</h3>", unsafe_allow_html=True)
159
160 # Display chat messages
161 for message in st.session_state.messages:
162 with st.chat_message(message["role"]):
163 st.markdown(message["content"])
164
165 # Load the appropriate index for the user's role
166 index = load_vector_index(st.session_state.role)
167
168 # Initialize Groq LLM
169 try:
170 llm = Groq(
171 model="llama3-8b-8192",
172 api_key=os.getenv("GROQ_API_KEY"),
173 temperature=0.5,
174 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."
175 )
176
177 # Create query engine with the LLM
178 query_engine = index.as_query_engine(
179 llm=llm,
180 similarity_top_k=3,
181 response_mode="compact"
182 )
183 except Exception as e:
184 st.error(f"Error initializing LLM: {str(e)}")
185 st.warning("Falling back to default LLM settings. Some features may be limited.")
186 query_engine = index.as_query_engine(
187 similarity_top_k=3,
188 response_mode="compact"
189 )
190
191 # Chat input
192 if prompt := st.chat_input(f"Ask about {st.session_state.role} documents..."):
193 # Add user message to chat history
194 st.session_state.messages.append({"role": "user", "content": prompt})
195
196 # Display user message
197 with st.chat_message("user"):
198 st.markdown(prompt)
199
200 # Get and display assistant response
201 with st.chat_message("assistant"):
202 message_placeholder = st.empty()
203 full_response = ""
204
205 try:
206 # Get response from query engine
207 response = query_engine.query(prompt)
208 full_response = str(response)
209 message_placeholder.markdown(full_response)
210 except Exception as e:
211 error_msg = f"Error generating response: {str(e)}"
212 message_placeholder.error(error_msg)
213 full_response = error_msg
214
215 # Add assistant response to chat history
216 st.session_state.messages.append({"role": "assistant", "content": full_response})
217
218def show_login_form():
219 """Display the beautiful login form"""
220 st.markdown(
221 """
222 <style>
223 .main {
224 background-color: #1a1a2e;
225 color: white;
226 }
227 .stTextInput > div > div > input {
228 background-color: #2a2a3e;
229 color: white;
230 border: 1px solid #4a4a6a;
231 border-radius: 8px;
232 }
233 .stTextInput > div > div > input::placeholder {
234 color: #a0a0b0 !important;
235 opacity: 1 !important;
236 }
237 .stButton > button {
238 background-color: #e94560;
239 color: white;
240 border: none;
241 border-radius: 8px;
242 padding: 10px 20px;
243 font-size: 16px;
244 width: 100%;
245 }
246 .stButton > button:hover {
247 background-color: #d83450;
248 }
249 h1, h2, h3, h4, h5, h6 {
250 color: white;
251 }
252 .st-emotion-cache-1r6slb0 {
253 border: 1px solid #4a4a6a;
254 border-radius: 12px;
255 padding: 2rem;
256 background-color: #232339;
257 }
258 </style>
259 """,
260 unsafe_allow_html=True
261 )
262 st.markdown('<div style="text-align: center; margin-top: -80px; margin-bottom: 30px;"><h1 style="font-size: 3rem;">🔒</h1></div>', unsafe_allow_html=True)
263 st.markdown('<h1 style="text-align: center; margin-bottom: 20px;">Department Portal</h1>', unsafe_allow_html=True)
264 st.markdown('<p style="text-align: center; color: #a0a0b0; margin-bottom: 30px;">Sign in to access your department\'s knowledge base</p>', unsafe_allow_html=True)
265
266 with st.container():
267 with st.form("login_form", border=True):
268 username = st.text_input("Username", placeholder="Enter your username")
269 password = st.text_input("Password", type="password", placeholder="Enter your password")
270 login_button = st.form_submit_button("Sign In")
271
272 if login_button:
273 if not username or not password:
274 st.error("Please enter both username and password")
275 elif login(username, password):
276 st.success(f"Welcome, {username}! Redirecting...")
277 else:
278 st.error("Invalid username or password")
279
280 with st.expander("Need demo credentials?"):
281 st.markdown("""
282 - **Engineering:** `Tony` / `password123`
283 - **Marketing:** `Bruce` / `securepass`
284 - **Finance:** `Sam` / `financepass`
285 - **HR:** `Natasha` / `hrpass123`
286 """)
287
288 st.markdown('<p style="text-align: center; margin-top: 2rem; color: #a0a0b0;">2025 Department RAG System</p>', unsafe_allow_html=True)
289
290
291def main():
292 """
293 Main application entry point
294 Handles routing between login and main application
295 """
296 # Sidebar for logout and user info
297 if st.session_state.authenticated:
298 st.set_page_config(layout="wide", initial_sidebar_state="expanded")
299 with st.sidebar:
300 st.markdown(f"### Welcome, {st.session_state.username}")
301 st.markdown(f"**Role:** {st.session_state.role.capitalize()}")
302
303 if st.button("Logout", key="logout_btn"):
304 logout()
305 return
306
307 st.markdown("---")
308 st.markdown("### About")
309 st.markdown("""
310 This is a secure departmental RAG system that provides
311 role-based access to information across different departments.
312 """)
313
314 # Main content area
315 if not st.session_state.authenticated:
316 show_login_form()
317 else:
318 chat_interface()
319
320if __name__ == "__main__":
321 main()
322