import streamlit as st import requests import json import uuid from datetime import datetime from typing import Optional, List, Dict import os # Configure page st.set_page_config( page_title="RAG Document Assistant", page_icon="📚", layout="wide" ) # API Configuration API_BASE_URL = "https://ajoy0071998-multilingual-rag.hf.space" # Change to your server URL # Session State Initialization def init_session_state(): if "token" not in st.session_state: st.session_state.token = None if "user" not in st.session_state: st.session_state.user = None if "documents" not in st.session_state: st.session_state.documents = [] if "conversation_id" not in st.session_state: st.session_state.conversation_id = None if "messages" not in st.session_state: st.session_state.messages = [] if "current_doc_ids" not in st.session_state: st.session_state.current_doc_ids = [] if "selected_docs" not in st.session_state: st.session_state.selected_docs = [] # API Calls def api_request(method: str, endpoint: str, data: Optional[Dict] = None, files: Optional[Dict] = None): """Make API requests with authentication""" headers = {} if st.session_state.token: headers["Authorization"] = f"Bearer {st.session_state.token}" url = f"{API_BASE_URL}{endpoint}" try: if method == "GET": response = requests.get(url, headers=headers) elif method == "POST": if files: response = requests.post(url, headers=headers, files=files, data=data) else: response = requests.post(url, headers=headers, json=data) elif method == "DELETE": response = requests.delete(url, headers=headers) else: return None return response except requests.exceptions.ConnectionError: st.error("❌ Cannot connect to server. Please make sure the server is running.") return None def login(email: str, password: str): """Login user and get token""" response = api_request("POST", "/auth/login", {"email": email, "password": password}) if response and response.status_code == 200: data = response.json() st.session_state.token = data["access_token"] return True return False def register_user(username: str, email: str, password: str): """Register new user""" response = api_request("POST", "/auth/register", { "username": username, "email": email, "password": password }) return response and response.status_code == 200 def get_user_info(): """Get current user info""" response = api_request("GET", "/auth/me") if response and response.status_code == 200: return response.json() return None def get_user_documents(): """Get user's documents""" response = api_request("GET", "/api/v1/documents/") if response and response.status_code == 200: data = response.json() return data.get("documents", []) return [] def upload_document(file): """Upload a document""" files = {"file": file} response = api_request("POST", "/api/v1/rag/ingest", files=files) if response and response.status_code == 200: return response.json() return None def delete_document(document_id: str): """Delete a document""" response = api_request("DELETE", f"/api/v1/documents/{document_id}") return response and response.status_code == 200 def send_query(query: str, doc_ids: List[str], conversation_id: Optional[str] = None): """Send a query to the RAG system""" data = { "query": query, "doc_ids": doc_ids, "conversation_id": conversation_id } response = api_request("POST", "/api/v1/rag/response", data) if response and response.status_code == 200: return response.json() return None # Login/Register Page def show_login_page(): st.title("📚 RAG Document Assistant") st.subheader("Login or Register") tab1, tab2 = st.tabs(["Login", "Register"]) with tab1: with st.form("login_form"): email = st.text_input("Email") password = st.text_input("Password", type="password") submitted = st.form_submit_button("Login") if submitted: if email and password: with st.spinner("Logging in..."): if login(email, password): user_info = get_user_info() if user_info: st.session_state.user = user_info st.rerun() else: st.error("Invalid email or password") else: st.warning("Please fill all fields") with tab2: with st.form("register_form"): username = st.text_input("Username") email = st.text_input("Email") password = st.text_input("Password", type="password") confirm_password = st.text_input("Confirm Password", type="password") submitted = st.form_submit_button("Register") if submitted: if username and email and password and confirm_password: if password != confirm_password: st.error("Passwords do not match") else: with st.spinner("Registering..."): if register_user(username, email, password): st.success("Registration successful! Please login.") else: st.error("Registration failed. User might already exist.") else: st.warning("Please fill all fields") # Main App def main(): init_session_state() # Check if user is logged in if not st.session_state.token: show_login_page() return # Get user info if not available if not st.session_state.user: user_info = get_user_info() if user_info: st.session_state.user = user_info else: # Token might be invalid st.session_state.token = None st.rerun() # Sidebar - Document Management with st.sidebar: st.header(f"👤 {st.session_state.user.get('username', 'User')}") st.caption(f"📧 {st.session_state.user.get('email', '')}") st.divider() # Upload Document st.subheader("📤 Upload Document") uploaded_file = st.file_uploader( "Choose a file", type=["pdf", "doc", "docx"], help="Supported formats: PDF, DOC, DOCX" ) if uploaded_file: if st.button("Upload", type="primary", use_container_width=True): with st.spinner("Uploading document..."): result = upload_document(uploaded_file) if result: st.success("Document uploaded successfully!") st.session_state.documents = get_user_documents() st.rerun() else: st.error("Failed to upload document") st.divider() # Document List st.subheader("📄 Your Documents") # Refresh button if st.button("🔄 Refresh List", use_container_width=True): st.session_state.documents = get_user_documents() st.rerun() # Get documents if not st.session_state.documents: st.session_state.documents = get_user_documents() if st.session_state.documents: for doc in st.session_state.documents: col1, col2 = st.columns([4, 1]) with col1: # Checkbox for document selection doc_id = str(doc["id"]) is_selected = doc_id in st.session_state.selected_docs if st.checkbox( f"📎 {doc['file_name']}", value=is_selected, key=f"doc_{doc_id}" ): if doc_id not in st.session_state.selected_docs: st.session_state.selected_docs.append(doc_id) else: if doc_id in st.session_state.selected_docs: st.session_state.selected_docs.remove(doc_id) with col2: # Delete button if st.button("🗑ī¸", key=f"del_{doc_id}", help="Delete document"): if delete_document(doc_id): st.success(f"Deleted {doc['file_name']}") st.session_state.documents = get_user_documents() if doc_id in st.session_state.selected_docs: st.session_state.selected_docs.remove(doc_id) st.rerun() else: st.error("Failed to delete document") # Show document details with st.expander(f"ℹī¸ Details - {doc['file_name']}"): st.caption(f"📏 Size: {doc['file_size'] / 1024:.1f} KB") st.caption(f"📝 Chunks: {doc.get('chunks', 0)}") st.caption(f"🌐 Language: {doc.get('primary_language', 'unknown')}") st.caption(f"📅 Uploaded: {doc.get('created_at', '')[:10]}") else: st.info("No documents uploaded yet") st.divider() # Logout if st.button("đŸšĒ Logout", use_container_width=True): st.session_state.clear() st.rerun() # Main Content Area # Header with conversation controls col1, col2, col3 = st.columns([2, 2, 1]) with col1: st.title("đŸ’Ŧ RAG Chat") with col2: if st.session_state.conversation_id: st.info(f"Conversation: {st.session_state.conversation_id[:8]}...") with col3: if st.button("🔄 New Conversation", use_container_width=True): st.session_state.messages = [] st.session_state.conversation_id = None st.rerun() # Selected Documents Info if st.session_state.selected_docs: selected_count = len(st.session_state.selected_docs) st.info(f"📌 Querying across {selected_count} selected document(s)") else: st.warning("⚠ī¸ Please select at least one document to query") st.divider() # Display chat messages chat_container = st.container() with chat_container: for message in st.session_state.messages: with st.chat_message(message["role"]): st.markdown(message["content"]) if "timestamp" in message: st.caption(message["timestamp"]) # Chat input at the bottom if st.session_state.selected_docs: user_input = st.chat_input("Ask a question about your documents...") if user_input: # Add user message to chat st.session_state.messages.append({ "role": "user", "content": user_input, "timestamp": datetime.now().strftime("%H:%M") }) # Show user message immediately with st.chat_message("user"): st.markdown(user_input) # Get response with st.chat_message("assistant"): with st.spinner("Thinking..."): response = send_query( query=user_input, doc_ids=st.session_state.selected_docs, conversation_id=st.session_state.conversation_id ) if response: # Update conversation ID for subsequent messages if "conversation_id" in response: st.session_state.conversation_id = response["conversation_id"] # Display response st.markdown(response["answer"]) # Store response in messages st.session_state.messages.append({ "role": "assistant", "content": response["answer"], "timestamp": datetime.now().strftime("%H:%M") }) else: st.error("Failed to get response. Please check if the server is running.") else: st.info("📚 Select documents from the sidebar to start asking questions") if __name__ == "__main__": main()