ArthyP/technical-rag-assistant
0
1#!/usr/bin/env python32"""3Technical Documentation RAG System - Streamlit Interface4 5A professional web interface for the RAG system with answer generation,6optimized for technical documentation Q&A.7"""8 9import os10# Set environment variables before importing streamlit11os.environ['HOME'] = '/app'12os.environ['STREAMLIT_CONFIG_DIR'] = '/app/.streamlit'13os.environ['STREAMLIT_BROWSER_GATHER_USAGE_STATS'] = 'false'14 15import streamlit as st16import sys17from pathlib import Path18import time19import traceback20from typing import List, Dict, Any21import json22 23# Add project root to path24project_root = Path(__file__).parent25sys.path.insert(0, str(project_root))26 27# Import directly since we're in the project directory28sys.path.insert(0, str(Path(__file__).parent))29from src.rag_with_generation import RAGWithGeneration30 31 32# Page configuration33st.set_page_config(34 page_title="Technical Documentation RAG Assistant",35 page_icon="🔍",36 layout="wide",37 initial_sidebar_state="expanded"38)39 40# Custom CSS for professional appearance41st.markdown("""42<style>43 .main-header {44 font-size: 2.5rem;45 font-weight: bold;46 color: #1f77b4;47 text-align: center;48 margin-bottom: 2rem;49 padding: 1rem;50 background: linear-gradient(90deg, #f0f8ff, #e6f3ff);51 border-radius: 10px;52 border-left: 5px solid #1f77b4;53 }54 55 .system-stats {56 background-color: #f8f9fa;57 padding: 1rem;58 border-radius: 8px;59 border-left: 4px solid #28a745;60 margin: 1rem 0;61 }62 63 .error-box {64 background-color: #f8d7da;65 color: #721c24;66 padding: 1rem;67 border-radius: 8px;68 border-left: 4px solid #dc3545;69 margin: 1rem 0;70 }71 72 .metrics-container {73 display: flex;74 justify-content: space-around;75 margin: 1rem 0;76 }77 78 .metric-box {79 background: white;80 padding: 1rem;81 border-radius: 8px;82 box-shadow: 0 2px 4px rgba(0,0,0,0.1);83 text-align: center;84 min-width: 120px;85 }86 87 .citation-box {88 background-color: #e8f4fd;89 padding: 0.8rem;90 border-radius: 6px;91 border-left: 3px solid #2196F3;92 margin: 0.5rem 0;93 font-size: 0.9rem;94 }95 96 .sample-query {97 background-color: #f0f8ff;98 padding: 0.8rem;99 border-radius: 6px;100 margin: 0.5rem 0;101 cursor: pointer;102 border-left: 3px solid #4CAF50;103 }104 105 .sample-query:hover {106 background-color: #e6f3ff;107 }108</style>109""", unsafe_allow_html=True)110 111 112def initialize_rag_system(api_token=None, model_name=None):113 """Initialize the RAG system with HuggingFace API."""114 try:115 # Check for token in environment first116 import os117 token = api_token or os.getenv("HUGGINGFACE_API_TOKEN")118 119 # Use selected model or default based on Pro vs Free tier120 if not model_name:121 if token:122 model_name = "mistralai/Mistral-7B-Instruct-v0.2" # Pro: Best for technical Q&A123 else:124 model_name = "gpt2-medium" # Free tier: Best available125 126 rag_system = RAGWithGeneration(127 model_name=model_name,128 api_token=token, # Use provided token or env variable129 temperature=0.3,130 max_tokens=512131 )132 return rag_system, None133 except Exception as e:134 error_msg = f"RAG system initialization failed: {str(e)}"135 return None, error_msg136 137 138def display_header():139 """Display application header with branding."""140 st.markdown("""141 <div class="main-header">142 🔍 Technical Documentation RAG Assistant143 </div>144 """, unsafe_allow_html=True)145 146 st.markdown("""147 <div style="text-align: center; margin-bottom: 2rem; color: #666;">148 Advanced hybrid retrieval with local LLM answer generation<br>149 <strong>Built for Swiss ML Engineering Excellence</strong>150 </div>151 """, unsafe_allow_html=True)152 153 154def display_system_status(rag_system):155 """Display system status and metrics in sidebar."""156 with st.sidebar:157 st.markdown("### 🏥 System Status")158 159 # Basic system info160 chunk_count = len(rag_system.chunks) if rag_system.chunks else 0161 source_count = len(set(chunk.get('source', '') for chunk in rag_system.chunks)) if rag_system.chunks else 0162 163 col1, col2 = st.columns(2)164 with col1:165 st.metric("📄 Documents", source_count)166 with col2:167 st.metric("🧩 Chunks", chunk_count)168 169 # Model info170 st.markdown("### 🤖 Model Status")171 if rag_system.answer_generator.api_token:172 st.success("✅ HuggingFace API (Authenticated)")173 else:174 st.info("ℹ️ HuggingFace API (Free Tier)")175 st.info("🔍 Hybrid Search Active")176 177 # API Configuration178 st.markdown("### ⚙️ API Configuration")179 with st.expander("HuggingFace Configuration"):180 st.markdown("""181 **Using HF Token:**182 1. Set as Space Secret: `HUGGINGFACE_API_TOKEN`183 2. Or paste token below184 185 **Benefits of token:**186 - Higher rate limits187 - Better models (Llama 2, Falcon)188 - Faster response times189 """)190 191 token_input = st.text_input(192 "HF Token", 193 type="password", 194 help="Your HuggingFace API token",195 key="hf_token_input"196 )197 198 # Model selection - Pro tier models from your guide199 if rag_system.answer_generator.api_token:200 model_options = [201 "mistralai/Mistral-7B-Instruct-v0.2", # Best for technical Q&A202 "codellama/CodeLlama-7b-Instruct-hf", # Perfect for code docs203 "meta-llama/Llama-2-7b-chat-hf", # Well-rounded204 "codellama/CodeLlama-13b-Instruct-hf", # Higher quality (slower)205 "meta-llama/Llama-2-13b-chat-hf", # Better reasoning206 "microsoft/DialoGPT-large", # Conversational fallback207 "tiiuae/falcon-7b-instruct", # Efficient option208 "gpt2-medium" # Emergency fallback209 ]210 else:211 model_options = [212 "gpt2-medium", # Best bet for free tier 213 "gpt2", # Always available214 "distilgpt2" # Fastest option215 ]216 217 current_model = rag_system.answer_generator.model_name218 selected_model = st.selectbox(219 "Model",220 model_options,221 index=model_options.index(current_model) if current_model in model_options else 0,222 key="model_select"223 )224 225 col1, col2 = st.columns(2)226 with col1:227 if st.button("Update Configuration"):228 if token_input or selected_model != current_model:229 # Reinitialize with new settings230 st.session_state['api_token'] = token_input if token_input else st.session_state.get('api_token')231 st.session_state['selected_model'] = selected_model232 st.session_state['rag_system'] = None233 st.rerun()234 235 with col2:236 if st.button("Test Pro Models"):237 # Test all Pro models from your guide238 pro_models = [239 "mistralai/Mistral-7B-Instruct-v0.2",240 "codellama/CodeLlama-7b-Instruct-hf", 241 "meta-llama/Llama-2-7b-chat-hf",242 "codellama/CodeLlama-13b-Instruct-hf",243 "meta-llama/Llama-2-13b-chat-hf",244 "microsoft/DialoGPT-large",245 "tiiuae/falcon-7b-instruct"246 ]247 248 test_token = token_input if token_input else st.session_state.get('api_token')249 250 with st.spinner("Testing Pro models..."):251 results = {}252 for model in pro_models:253 try:254 import requests255 import os256 token = test_token or os.getenv("HUGGINGFACE_API_TOKEN")257 258 headers = {"Content-Type": "application/json"}259 if token:260 headers["Authorization"] = f"Bearer {token}"261 262 response = requests.post(263 f"https://api-inference.huggingface.co/models/{model}",264 headers=headers,265 json={"inputs": "What is RISC-V?", "parameters": {"max_new_tokens": 50}},266 timeout=15267 )268 269 if response.status_code == 200:270 results[model] = "✅ Available"271 elif response.status_code == 404:272 results[model] = "❌ Not found"273 elif response.status_code == 503:274 results[model] = "⏳ Loading"275 else:276 results[model] = f"❌ Error {response.status_code}"277 278 except Exception as e:279 results[model] = f"❌ Failed: {str(e)[:30]}"280 281 # Display results282 st.subheader("🧪 Pro Model Test Results:")283 for model, status in results.items():284 model_short = model.split('/')[-1]285 st.write(f"**{model_short}**: {status}")286 287 if chunk_count > 0:288 st.markdown("### 📊 Index Statistics")289 st.markdown(f"""290 - **Indexed Documents**: {source_count}291 - **Total Chunks**: {chunk_count}292 - **Search Method**: Hybrid (Semantic + BM25)293 - **Embeddings**: 384-dim MiniLM-L6294 """)295 296 297def handle_query_interface(rag_system):298 """Handle the main query interface."""299 if not rag_system.chunks:300 st.warning("⚠️ No documents indexed yet. Please upload documents in the 'Manage Documents' tab.")301 return302 303 # Query input304 query = st.text_input(305 "Enter your question:",306 placeholder="e.g., What is RISC-V and what are its main features?",307 key="main_query"308 )309 310 # Advanced options311 with st.expander("🔧 Advanced Options"):312 col1, col2, col3 = st.columns(3)313 314 with col1:315 top_k = st.slider("Results to retrieve", 3, 10, 5)316 with col2:317 dense_weight = st.slider("Semantic weight", 0.5, 1.0, 0.7, 0.1)318 with col3:319 use_fallback = st.checkbox("Use fallback model", False)320 321 if st.button("🔍 Search & Generate Answer", type="primary"):322 if not query.strip():323 st.error("Please enter a question.")324 return325 326 try:327 # Execute query with timing328 start_time = time.time()329 330 with st.spinner("🔍 Searching documents and generating answer..."):331 result = rag_system.query_with_answer(332 question=query,333 top_k=top_k,334 use_hybrid=True,335 dense_weight=dense_weight,336 use_fallback_llm=use_fallback337 )338 339 total_time = time.time() - start_time340 341 # Display results342 display_query_results(result, total_time)343 344 except Exception as e:345 st.error(f"❌ Query failed: {str(e)}")346 st.markdown(f"**Error details:** {traceback.format_exc()}")347 348 349def display_query_results(result: Dict, total_time: float):350 """Display query results with metrics and citations."""351 352 # Performance metrics353 col1, col2, col3, col4 = st.columns(4)354 355 with col1:356 st.metric("⏱️ Total Time", f"{total_time:.2f}s")357 with col2:358 st.metric("🎯 Confidence", f"{result['confidence']:.1%}")359 with col3:360 st.metric("📄 Sources", len(result['sources']))361 with col4:362 retrieval_time = result['retrieval_stats']['retrieval_time']363 st.metric("🔍 Retrieval", f"{retrieval_time:.2f}s")364 365 # Answer366 st.markdown("### 💬 Generated Answer")367 st.markdown(f"""368 <div style="background-color: #f8f9fa; padding: 1.5rem; border-radius: 10px; border-left: 5px solid #28a745; color: #333333;">369 {result['answer']}370 </div>371 """, unsafe_allow_html=True)372 373 # Citations374 if result['citations']:375 st.markdown("### 📚 Sources & Citations")376 377 for i, citation in enumerate(result['citations'], 1):378 st.markdown(f"""379 <div class="citation-box">380 <strong>[{i}]</strong> {citation['source']} (Page {citation['page']})<br>381 <small><em>Relevance: {citation['relevance']:.1%}</em></small><br>382 <small>"{citation['snippet']}"</small>383 </div>384 """, unsafe_allow_html=True)385 386 # Technical details387 with st.expander("🔬 Technical Details"):388 st.json({389 "retrieval_method": result['retrieval_stats']['method'],390 "chunks_retrieved": result['retrieval_stats']['chunks_retrieved'],391 "dense_weight": result['retrieval_stats'].get('dense_weight', 'N/A'),392 "model_used": result['generation_stats']['model'],393 "generation_time": f"{result['generation_stats']['generation_time']:.3f}s"394 })395 396 397def handle_document_upload(rag_system):398 """Handle document upload and indexing."""399 st.subheader("📤 Upload Documents")400 401 uploaded_files = st.file_uploader(402 "Upload PDF documents",403 type=['pdf'],404 accept_multiple_files=True,405 help="Upload technical documentation, manuals, or research papers"406 )407 408 if uploaded_files:409 for uploaded_file in uploaded_files:410 if st.button(f"Index {uploaded_file.name}", key=f"index_{uploaded_file.name}"):411 try:412 # Save uploaded file temporarily in app directory413 import tempfile414 import os415 416 # Create temp directory in app folder417 temp_dir = Path("/app/temp_uploads")418 temp_dir.mkdir(exist_ok=True)419 420 temp_path = temp_dir / uploaded_file.name421 with open(temp_path, "wb") as f:422 f.write(uploaded_file.getvalue())423 424 # Index the document425 st.write(f"🔄 Starting to process {uploaded_file.name}...")426 st.write(f"📁 File saved to: {temp_path}")427 st.write(f"📏 File size: {temp_path.stat().st_size} bytes")428 429 # Capture print output for debugging430 import io431 import sys432 433 captured_output = io.StringIO()434 sys.stdout = captured_output435 436 try:437 with st.spinner(f"Processing {uploaded_file.name}..."):438 chunk_count = rag_system.index_document(temp_path)439 440 finally:441 # Restore stdout442 sys.stdout = sys.__stdout__443 444 # Show captured output445 output = captured_output.getvalue()446 if output:447 st.text_area("Processing Log:", output, height=150)448 449 st.success(f"✅ {uploaded_file.name} indexed! {chunk_count} chunks added.")450 451 # Clean up temp file452 try:453 temp_path.unlink()454 except:455 pass456 457 st.rerun()458 459 except Exception as e:460 st.error(f"❌ Failed to index {uploaded_file.name}: {str(e)}")461 import traceback462 st.error(f"Details: {traceback.format_exc()}")463 464 465def display_sample_queries():466 """Display sample queries for demonstration."""467 st.subheader("💡 Sample Questions")468 st.markdown("Click on any question to try it:")469 470 sample_queries = [471 "What is RISC-V and what are its main features?",472 "How does RISC-V compare to ARM and x86 architectures?",473 "What are the different RISC-V instruction formats?",474 "Explain RISC-V base integer instructions",475 "What are the benefits of using RISC-V in embedded systems?",476 "How does RISC-V handle memory management?",477 "What are RISC-V privileged instructions?",478 "Describe RISC-V calling conventions"479 ]480 481 for query in sample_queries:482 if st.button(query, key=f"sample_{hash(query)}"):483 st.session_state['sample_query'] = query484 st.rerun()485 486 487def main():488 """Main Streamlit application."""489 490 # Initialize session state491 if 'rag_system' not in st.session_state:492 st.session_state['rag_system'] = None493 st.session_state['init_error'] = None494 if 'api_token' not in st.session_state:495 st.session_state['api_token'] = None496 497 # Display header498 display_header()499 500 # Initialize RAG system501 if st.session_state['rag_system'] is None:502 with st.spinner("Initializing RAG system..."):503 selected_model = st.session_state.get('selected_model')504 rag_system, error = initialize_rag_system(505 st.session_state.get('api_token'),506 selected_model507 )508 st.session_state['rag_system'] = rag_system509 st.session_state['init_error'] = error510 511 rag_system = st.session_state['rag_system']512 init_error = st.session_state['init_error']513 514 # Check for initialization errors515 if init_error:516 st.markdown(f"""517 <div class="error-box">518 ❌ <strong>Failed to initialize RAG system:</strong><br>519 {init_error}<br><br>520 <strong>System uses HuggingFace Inference API</strong><br>521 If you see network errors, please check your internet connection.522 </div>523 """, unsafe_allow_html=True)524 return525 526 if rag_system is None:527 st.error("Failed to initialize RAG system. Please check the logs.")528 return529 530 # Display system status in sidebar531 display_system_status(rag_system)532 533 # Main interface534 tab1, tab2, tab3 = st.tabs(["🤔 Ask Questions", "📄 Manage Documents", "💡 Examples"])535 536 with tab1:537 # Handle sample query selection538 if 'sample_query' in st.session_state:539 st.text_input(540 "Enter your question:",541 value=st.session_state['sample_query'],542 key="main_query"543 )544 del st.session_state['sample_query']545 546 handle_query_interface(rag_system)547 548 with tab2:549 handle_document_upload(rag_system)550 551 # Option to load test document552 st.subheader("📖 Test Document")553 test_pdf_path = Path("data/test/riscv-base-instructions.pdf")554 555 if test_pdf_path.exists():556 if st.button("Load RISC-V Test Document"):557 try:558 with st.spinner("Loading test document..."):559 st.write(f"🔄 Processing test document: {test_pdf_path}")560 st.write(f"📏 File size: {test_pdf_path.stat().st_size} bytes")561 562 chunk_count = rag_system.index_document(test_pdf_path)563 st.success(f"✅ Test document loaded! {chunk_count} chunks indexed.")564 st.rerun()565 except Exception as e:566 st.error(f"Failed to load test document: {e}")567 import traceback568 st.error(f"Details: {traceback.format_exc()}")569 else:570 st.info("Test document not found at data/test/riscv-base-instructions.pdf")571 572 with tab3:573 display_sample_queries()574 575 # Footer576 st.markdown("---")577 st.markdown("""578 <div style="text-align: center; color: #666; font-size: 0.9rem;">579 Technical Documentation RAG Assistant | Powered by HuggingFace API & RISC-V Documentation<br>580 Built for ML Engineer Portfolio | Swiss Tech Market Focus581 </div>582 """, unsafe_allow_html=True)583 584 585if __name__ == "__main__":586 main()