aniket47/document-intelligence-chatbot
0
1import streamlit as st2import os3from typing import List, Dict4import time5 6# Import custom components7from components.document_processor import DocumentProcessor8from components.vector_store import VectorStore9from components.query_router import QueryRouter, QueryType10from components.web_search import WebSearcher11from components.huggingface_client import HuggingFaceClient12 13# Page configuration14st.set_page_config(15 page_title="Universal Document Intelligence Chatbot",16 layout="wide",17 initial_sidebar_state="expanded"18)19 20@st.cache_resource21def get_hf_client():22 """Get or create HuggingFace client with caching"""23 try:24 print("Initializing cached HuggingFace client...")25 client = HuggingFaceClient()26 # Force model loading27 success = client._load_model()28 print(f"Model loading success: {success}")29 print(f"Model is_loaded: {client.is_loaded}")30 return client, success31 except Exception as e:32 print(f"Failed to initialize HuggingFace client: {str(e)}")33 return None, False34 35class DocumentChatbot:36 """37 Main chatbot application class38 """39 40 def __init__(self, serper_api_key: str = None):41 self.doc_processor = DocumentProcessor()42 self.vector_store = VectorStore()43 self.query_router = QueryRouter()44 self.web_searcher = None45 46 # Get cached HuggingFace client47 self.hf_client, self.model_loaded = get_hf_client()48 49 # Initialize web searcher if API key is available50 self.init_web_search(serper_api_key)51 52 def init_web_search(self, api_key: str = None):53 """Initialize or reinitialize web search with provided API key"""54 try:55 self.web_searcher = WebSearcher(api_key=api_key)56 return True57 except ValueError as e:58 self.web_searcher = None59 return False60 61 # Load existing index if available62 self.vector_store.load_index()63 64 def is_ai_model_available(self):65 """Check if AI model is available"""66 return self.hf_client is not None and self.hf_client.is_loaded67 68 def process_uploaded_files(self, uploaded_files):69 """Process uploaded PDF files"""70 if not uploaded_files:71 return72 73 with st.spinner("Processing uploaded documents..."):74 all_chunks = []75 76 for uploaded_file in uploaded_files:77 try:78 # Process the PDF79 chunks = self.doc_processor.process_document(uploaded_file)80 all_chunks.extend(chunks)81 82 st.success(f"Processed {uploaded_file.name}: {len(chunks)} chunks")83 84 except Exception as e:85 st.error(f"Error processing {uploaded_file.name}: {str(e)}")86 87 if all_chunks:88 # Add to vector store89 self.vector_store.add_documents(all_chunks)90 self.vector_store.save_index()91 92 st.success(f"Successfully processed {len(all_chunks)} document chunks!")93 94 # Update session state95 st.session_state.documents_loaded = True96 st.session_state.vector_stats = self.vector_store.get_stats()97 98 def search_documents(self, query: str, k: int = 5) -> List[Dict]:99 """Search documents using vector similarity"""100 if self.vector_store.index is None or len(self.vector_store.documents) == 0:101 print(f"No documents available - index: {self.vector_store.index is not None}, docs: {len(self.vector_store.documents) if hasattr(self.vector_store, 'documents') else 'N/A'}")102 return []103 104 results = self.vector_store.search(query, k=k)105 print(f"Document search for '{query}': found {len(results)} results")106 if results:107 scores = [r.get('score', 0) for r in results]108 print(f"Score range: {min(scores):.3f} - {max(scores):.3f}")109 return results110 111 def get_web_search_results(self, query: str) -> List[Dict]:112 """Get web search results"""113 if not self.web_searcher:114 return []115 116 try:117 return self.web_searcher.search_and_format(query, num_results=3)118 except Exception as e:119 st.error(f"Web search error: {str(e)}")120 return []121 122 def generate_response(self, query: str) -> Dict:123 """Generate response using smart routing and HuggingFace for LLM responses"""124 response = {125 'query': query,126 'sources': [],127 'answer': '',128 'routing_info': '',129 'search_strategy': 'unknown'130 }131 132 # Search documents first, but respect query routing133 doc_results = self.search_documents(query)134 135 # NEW: Use semantic-based routing instead of keyword-based136 routing_analysis = self.query_router.analyze_query_semantic(query, self.vector_store, similarity_threshold=0.15)137 138 print(f"DEBUG: Semantic routing result: {routing_analysis}")139 140 # SMART ROUTING: Use semantic similarity to determine strategy141 if routing_analysis['suggested_route'] == QueryType.WEB_SEARCH:142 # Query is not relevant to documents - use web search143 response['search_strategy'] = 'web_search'144 response['routing_info'] = f"Strategy: web_search (reason: {routing_analysis['reasoning'][0] if routing_analysis['reasoning'] else 'semantic analysis'})"145 print(f"DEBUG: Using web search for query: '{query}' (similarity: {routing_analysis.get('similarity_score', 0):.3f})")146 web_results = self.get_web_search_results(query)147 print(f"DEBUG: Web search returned {len(web_results) if web_results else 0} results")148 149 if web_results:150 # Create context from web results151 context = "Web search results:\n"152 for i, result in enumerate(web_results[:3], 1):153 context += f"{i}. {result['title']}: {result['snippet']}\n"154 response['sources'].append({155 'type': 'web',156 'title': result['title'],157 'snippet': result['snippet'],158 'link': result.get('link', ''),159 'source': result.get('source', '')160 })161 162 print(f"DEBUG: Web context created, length: {len(context)}")163 164 # Generate response using HuggingFace165 if self.is_ai_model_available():166 system_prompt = "You are a helpful AI assistant that answers questions based on web search results. Be accurate and cite sources when appropriate."167 ai_response = self.hf_client.generate_response(query, context, system_prompt)168 169 if len(ai_response.strip()) < 50 or "not sure" in ai_response.lower():170 response['answer'] = f"**🌐 Web Search Results:**\n{context}\n\n**🤖 AI Analysis:**\n{ai_response}"171 else:172 response['answer'] = f"**🤖 AI Analysis:**\n{ai_response}\n\n**🌐 Web Search Results:**\n{context}"173 response['ai_model_used'] = True174 else:175 response['answer'] = f"**🌐 Web Search Results:**\n{context}"176 response['ai_model_used'] = False177 178 print(f"DEBUG: Returning web search response")179 return response180 else:181 print("DEBUG: No web results, falling back to document search")182 183 # If semantic routing suggests documents, use them184 elif routing_analysis['suggested_route'] == QueryType.DOCUMENT_ONLY and doc_results and len(doc_results) > 0:185 best_score = max([r.get('score', 0) for r in doc_results])186 187 print(f"DEBUG: Using documents based on semantic routing: {len(doc_results)} results, best score: {best_score:.3f}")188 189 response['search_strategy'] = 'document_search'190 response['routing_info'] = f"Strategy: document_search (semantic similarity: {routing_analysis.get('similarity_score', 0):.3f}, found {len(doc_results)} matches)"191 192 # Create context from document results193 context = "Relevant information from your documents:\n"194 for i, result in enumerate(doc_results[:3], 1):195 doc = result['document']196 score = result['score']197 context += f"{i}. From {doc['metadata']['filename']} (relevance: {score:.2f}):\n{doc['text']}\n\n"198 199 response['sources'].append({200 'type': 'document',201 'filename': doc['metadata']['filename'],202 'text': doc['text'],203 'score': score,204 'chunk_id': doc['metadata'].get('chunk_index', 0)205 })206 207 # Generate response using HuggingFace208 if self.is_ai_model_available():209 system_prompt = "You are a helpful AI assistant that answers questions based on provided document context. Be accurate and cite the source documents when appropriate."210 print(f"DEBUG: Generating AI response for query: '{query[:50]}...'")211 print(f"DEBUG: Context length: {len(context)}")212 ai_response = self.hf_client.generate_response(query, context, system_prompt)213 print(f"DEBUG: AI response received: '{ai_response[:100]}...'")214 print(f"DEBUG: AI response length: {len(ai_response.strip())}")215 216 # Always combine AI response with document context for better user experience217 if ai_response and len(ai_response.strip()) > 5:218 response['answer'] = f"**🤖 AI Summary:**\n{ai_response}\n\n**📄 Source Documents:**\n{context}"219 response['ai_model_used'] = True220 else:221 # Fallback if AI response is empty222 response['answer'] = f"**📄 Source Documents:**\n{context}"223 response['ai_model_used'] = False224 else:225 print("DEBUG: AI model not available, using fallback")226 # Fallback response if HuggingFace is not available227 response['answer'] = f"**📄 Source Documents:**\n{context}"228 response['ai_model_used'] = False229 230 return response231 232 # Fallback: Use web search if no relevant documents found233 print("DEBUG: Using web search fallback")234 response['search_strategy'] = 'web_search'235 response['routing_info'] = f"Strategy: web_search (no relevant documents found or documents not relevant enough)"236 web_results = self.get_web_search_results(query)237 238 if web_results:239 # Create context from web results240 context = "Web search results:\n"241 for i, result in enumerate(web_results[:3], 1):242 context += f"{i}. {result['title']}: {result['snippet']}\n"243 response['sources'].append({244 'type': 'web',245 'title': result['title'],246 'snippet': result['snippet'],247 'link': result.get('link', ''),248 'source': result.get('source', '')249 })250 251 # Generate response using HuggingFace252 if self.is_ai_model_available():253 system_prompt = "You are a helpful AI assistant. Answer the user's question based on the provided web search results. Be informative and cite your sources."254 ai_response = self.hf_client.generate_response(query, context, system_prompt)255 256 if len(ai_response.strip()) < 50 or "not sure" in ai_response.lower():257 response['answer'] = f"**🌐 Web Search Results:**\n{context}\n\n**🤖 AI Analysis:**\n{ai_response}"258 else:259 response['answer'] = f"**🤖 AI Analysis:**\n{ai_response}\n\n**🌐 Web Search Results:**\n{context}"260 response['ai_model_used'] = True261 else:262 response['answer'] = f"**🌐 Web Search Results:**\n{context}"263 response['ai_model_used'] = False264 else:265 response['answer'] = "I couldn't find relevant information in your documents or through web search. Please try rephrasing your question or upload more relevant documents."266 267 return response268 269def main():270 """Main application function"""271 272 # Initialize session state273 if 'chatbot' not in st.session_state:274 # Try to get API key from environment variable first275 env_api_key = os.getenv("SERPER_API_KEY")276 st.session_state.chatbot = DocumentChatbot(serper_api_key=env_api_key)277 278 if 'chat_history' not in st.session_state:279 st.session_state.chat_history = []280 281 if 'documents_loaded' not in st.session_state:282 st.session_state.documents_loaded = False283 284 # Header285 st.title("Universal Document Intelligence Chatbot")286 st.markdown("*Upload documents and ask questions - get answers from your files or the web*")287 288 # Sidebar for document management289 with st.sidebar:290 st.header("Document Management")291 292 # File upload293 uploaded_files = st.file_uploader(294 "Upload PDF documents",295 type=['pdf'],296 accept_multiple_files=True,297 help="Upload PDF files to create a knowledge base"298 )299 300 # Process uploaded files301 if uploaded_files:302 if st.button("Process Documents", type="primary"):303 st.session_state.chatbot.process_uploaded_files(uploaded_files)304 305 # Display statistics306 if st.session_state.documents_loaded:307 st.subheader("Knowledge Base Stats")308 stats = st.session_state.chatbot.vector_store.get_stats()309 st.metric("Documents", stats['total_documents'])310 st.metric("Vector Dimension", stats['dimension'])311 st.info(f"Model: {stats['model_name']}")312 313 # Clear documents314 if st.session_state.documents_loaded:315 if st.button("Clear All Documents", type="secondary"):316 st.session_state.chatbot.vector_store.clear_index()317 st.session_state.documents_loaded = False318 st.session_state.chat_history = []319 st.success("Documents cleared!")320 st.rerun()321 322 # AI Model status323 st.subheader("AI Model Status")324 if st.session_state.chatbot.hf_client and st.session_state.chatbot.hf_client.is_available():325 st.success("✅ AI model loaded")326 else:327 st.warning("⚠️ AI model loading...")328 st.info("Models are being downloaded. This may take a few minutes on first run.")329 330 # Web Search Configuration331 st.subheader("🌐 Web Search")332 333 # Check if web search is already enabled334 web_search_enabled = st.session_state.chatbot.web_searcher is not None335 336 if web_search_enabled:337 st.success("✅ Web search enabled")338 if st.button("🔄 Change API Key"):339 st.session_state.show_api_input = True340 st.rerun()341 else:342 st.warning("⚠️ Web search disabled")343 344 # Show API key input field345 if not web_search_enabled or st.session_state.get('show_api_input', False):346 st.markdown("---")347 st.markdown("**Enter your Serper API Key:**")348 st.caption("Get a free API key at [serper.dev](https://serper.dev/) (2,500 searches/month free)")349 350 api_key = st.text_input(351 "Serper API Key",352 type="password",353 placeholder="Enter your API key here",354 help="Your API key is not stored and only used during this session",355 key="serper_api_key_input"356 )357 358 if api_key:359 if st.button("Enable Web Search", type="primary"):360 success = st.session_state.chatbot.init_web_search(api_key)361 if success:362 st.success("✅ Web search enabled!")363 st.session_state.show_api_input = False364 st.rerun()365 else:366 st.error("❌ Invalid API key. Please check and try again.")367 368 if not api_key:369 st.info("💡 Web search is optional. The chatbot works with documents only.")370 371 st.markdown("---")372 373 # Main chat interface374 st.header("Chat Interface")375 376 # Display chat history377 for i, chat in enumerate(st.session_state.chat_history):378 with st.chat_message("user"):379 st.write(chat['query'])380 381 with st.chat_message("assistant"):382 st.write(chat['answer'])383 384 # Show routing info385 if chat.get('routing_info'):386 with st.expander("Search Strategy"):387 st.info(chat['routing_info'])388 389 # Show sources390 if chat.get('sources'):391 with st.expander(f"Sources ({len(chat['sources'])} found)"):392 for j, source in enumerate(chat['sources'], 1):393 if source['type'] == 'document':394 st.markdown(f"**{j}. Document Source:**")395 st.markdown(f"- **File:** {source['filename']}")396 st.markdown(f"- **Relevance:** {source['score']:.2f}")397 st.markdown(f"- **Text:** {source['text'][:200]}...")398 elif source['type'] == 'web':399 st.markdown(f"**{j}. Web Source:**")400 st.markdown(f"- **Title:** {source['title']}")401 st.markdown(f"- **Source:** {source.get('source', 'Unknown')}")402 if source.get('link'):403 st.markdown(f"- **Link:** {source['link']}")404 405 # Query input406 query = st.chat_input("Ask a question about your documents or anything else...")407 408 if query:409 # Add user message to chat410 with st.chat_message("user"):411 st.write(query)412 413 # Generate response414 with st.chat_message("assistant"):415 with st.spinner("Thinking..."):416 response = st.session_state.chatbot.generate_response(query)417 418 st.write(response['answer'])419 420 # Show routing info421 if response.get('routing_info'):422 with st.expander("Search Strategy"):423 st.info(response['routing_info'])424 st.caption(f"Strategy used: {response['search_strategy']}")425 426 # Show sources427 if response.get('sources'):428 with st.expander(f"Sources ({len(response['sources'])} found)"):429 for j, source in enumerate(response['sources'], 1):430 if source['type'] == 'document':431 st.markdown(f"**{j}. Document Source:**")432 st.markdown(f"- **File:** {source['filename']}")433 st.markdown(f"- **Relevance:** {source['score']:.2f}")434 st.markdown(f"- **Text:** {source['text'][:200]}...")435 elif source['type'] == 'web':436 st.markdown(f"**{j}. Web Source:**")437 st.markdown(f"- **Title:** {source['title']}")438 st.markdown(f"- **Source:** {source.get('source', 'Unknown')}")439 if source.get('link'):440 st.markdown(f"- **Link:** {source['link']}")441 442 # Add to chat history443 st.session_state.chat_history.append({444 'query': query,445 'answer': response['answer'],446 'routing_info': response.get('routing_info'),447 'sources': response.get('sources', []),448 'search_strategy': response.get('search_strategy')449 })450 451 # Instructions452 if not st.session_state.chat_history:453 st.markdown("""454 ### Getting Started:455 456 1. **Upload PDFs** - Use the sidebar to add your documents457 2. **Click Process** - This creates a searchable knowledge base458 3. **Start Chatting** - Ask questions in the box below459 460 ### What you can ask:461 462 **About your documents:**463 - "What does the report say about..."464 - "Summarize the main points"465 - "Find information about X"466 467 **General questions:**468 - "What's the latest news on..."469 - "How does X work?"470 - "Compare A and B"471 472 The chatbot automatically decides whether to search your documents or the web.473 """)474 475if __name__ == "__main__":476 main()