app.py

import streamlit as st
import os
from typing import List, Dict
import time

# Import custom components
from components.document_processor import DocumentProcessor
from components.vector_store import VectorStore
from components.query_router import QueryRouter, QueryType
from components.web_search import WebSearcher
from components.huggingface_client import HuggingFaceClient

# Page configuration
st.set_page_config(
    page_title="Universal Document Intelligence Chatbot",
    layout="wide",
    initial_sidebar_state="expanded"
)

@st.cache_resource
def get_hf_client():
    """Get or create HuggingFace client with caching"""
    try:
        print("Initializing cached HuggingFace client...")
        client = HuggingFaceClient()
        # Force model loading
        success = client._load_model()
        print(f"Model loading success: {success}")
        print(f"Model is_loaded: {client.is_loaded}")
        return client, success
    except Exception as e:
        print(f"Failed to initialize HuggingFace client: {str(e)}")
        return None, False

class DocumentChatbot:
    """
    Main chatbot application class
    """
    
    def __init__(self, serper_api_key: str = None):
        self.doc_processor = DocumentProcessor()
        self.vector_store = VectorStore()
        self.query_router = QueryRouter()
        self.web_searcher = None
        
        # Get cached HuggingFace client
        self.hf_client, self.model_loaded = get_hf_client()
        
        # Initialize web searcher if API key is available
        self.init_web_search(serper_api_key)
    
    def init_web_search(self, api_key: str = None):
        """Initialize or reinitialize web search with provided API key"""
        try:
            self.web_searcher = WebSearcher(api_key=api_key)
            return True
        except ValueError as e:
            self.web_searcher = None
            return False
        
        # Load existing index if available
        self.vector_store.load_index()
    
    def is_ai_model_available(self):
        """Check if AI model is available"""
        return self.hf_client is not None and self.hf_client.is_loaded
    
    def process_uploaded_files(self, uploaded_files):
        """Process uploaded PDF files"""
        if not uploaded_files:
            return
        
        with st.spinner("Processing uploaded documents..."):
            all_chunks = []
            
            for uploaded_file in uploaded_files:
                try:
                    # Process the PDF
                    chunks = self.doc_processor.process_document(uploaded_file)
                    all_chunks.extend(chunks)
                    
                    st.success(f"Processed {uploaded_file.name}: {len(chunks)} chunks")
                    
                except Exception as e:
                    st.error(f"Error processing {uploaded_file.name}: {str(e)}")
            
            if all_chunks:
                # Add to vector store
                self.vector_store.add_documents(all_chunks)
                self.vector_store.save_index()
                
                st.success(f"Successfully processed {len(all_chunks)} document chunks!")
                
                # Update session state
                st.session_state.documents_loaded = True
                st.session_state.vector_stats = self.vector_store.get_stats()
    
    def search_documents(self, query: str, k: int = 5) -> List[Dict]:
        """Search documents using vector similarity"""
        if self.vector_store.index is None or len(self.vector_store.documents) == 0:
            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'}")
            return []
        
        results = self.vector_store.search(query, k=k)
        print(f"Document search for '{query}': found {len(results)} results")
        if results:
            scores = [r.get('score', 0) for r in results]
            print(f"Score range: {min(scores):.3f} - {max(scores):.3f}")
        return results
    
    def get_web_search_results(self, query: str) -> List[Dict]:
        """Get web search results"""
        if not self.web_searcher:
            return []
        
        try:
            return self.web_searcher.search_and_format(query, num_results=3)
        except Exception as e:
            st.error(f"Web search error: {str(e)}")
            return []
    
    def generate_response(self, query: str) -> Dict:
        """Generate response using smart routing and HuggingFace for LLM responses"""
        response = {
            'query': query,
            'sources': [],
            'answer': '',
            'routing_info': '',
            'search_strategy': 'unknown'
        }
        
        # Search documents first, but respect query routing
        doc_results = self.search_documents(query)
        
        # NEW: Use semantic-based routing instead of keyword-based
        routing_analysis = self.query_router.analyze_query_semantic(query, self.vector_store, similarity_threshold=0.15)
        
        print(f"DEBUG: Semantic routing result: {routing_analysis}")
        
        # SMART ROUTING: Use semantic similarity to determine strategy
        if routing_analysis['suggested_route'] == QueryType.WEB_SEARCH:
            # Query is not relevant to documents - use web search
            response['search_strategy'] = 'web_search'
            response['routing_info'] = f"Strategy: web_search (reason: {routing_analysis['reasoning'][0] if routing_analysis['reasoning'] else 'semantic analysis'})"
            print(f"DEBUG: Using web search for query: '{query}' (similarity: {routing_analysis.get('similarity_score', 0):.3f})")
            web_results = self.get_web_search_results(query)
            print(f"DEBUG: Web search returned {len(web_results) if web_results else 0} results")
            
            if web_results:
                # Create context from web results
                context = "Web search results:\n"
                for i, result in enumerate(web_results[:3], 1):
                    context += f"{i}. {result['title']}: {result['snippet']}\n"
                    response['sources'].append({
                        'type': 'web',
                        'title': result['title'],
                        'snippet': result['snippet'],
                        'link': result.get('link', ''),
                        'source': result.get('source', '')
                    })
                
                print(f"DEBUG: Web context created, length: {len(context)}")
                
                # Generate response using HuggingFace
                if self.is_ai_model_available():
                    system_prompt = "You are a helpful AI assistant that answers questions based on web search results. Be accurate and cite sources when appropriate."
                    ai_response = self.hf_client.generate_response(query, context, system_prompt)
                    
                    if len(ai_response.strip()) < 50 or "not sure" in ai_response.lower():
                        response['answer'] = f"**🌐 Web Search Results:**\n{context}\n\n**🤖 AI Analysis:**\n{ai_response}"
                    else:
                        response['answer'] = f"**🤖 AI Analysis:**\n{ai_response}\n\n**🌐 Web Search Results:**\n{context}"
                    response['ai_model_used'] = True
                else:
                    response['answer'] = f"**🌐 Web Search Results:**\n{context}"
                    response['ai_model_used'] = False
                
                print(f"DEBUG: Returning web search response")
                return response
            else:
                print("DEBUG: No web results, falling back to document search")
        
        # If semantic routing suggests documents, use them
        elif routing_analysis['suggested_route'] == QueryType.DOCUMENT_ONLY and doc_results and len(doc_results) > 0:
            best_score = max([r.get('score', 0) for r in doc_results])
            
            print(f"DEBUG: Using documents based on semantic routing: {len(doc_results)} results, best score: {best_score:.3f}")
            
            response['search_strategy'] = 'document_search'
            response['routing_info'] = f"Strategy: document_search (semantic similarity: {routing_analysis.get('similarity_score', 0):.3f}, found {len(doc_results)} matches)"
            
            # Create context from document results
            context = "Relevant information from your documents:\n"
            for i, result in enumerate(doc_results[:3], 1):
                doc = result['document']
                score = result['score']
                context += f"{i}. From {doc['metadata']['filename']} (relevance: {score:.2f}):\n{doc['text']}\n\n"
                
                response['sources'].append({
                    'type': 'document',
                    'filename': doc['metadata']['filename'],
                    'text': doc['text'],
                    'score': score,
                    'chunk_id': doc['metadata'].get('chunk_index', 0)
                })
                
            # Generate response using HuggingFace
            if self.is_ai_model_available():
                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."
                print(f"DEBUG: Generating AI response for query: '{query[:50]}...'")
                print(f"DEBUG: Context length: {len(context)}")
                ai_response = self.hf_client.generate_response(query, context, system_prompt)
                print(f"DEBUG: AI response received: '{ai_response[:100]}...'")
                print(f"DEBUG: AI response length: {len(ai_response.strip())}")
                
                # Always combine AI response with document context for better user experience
                if ai_response and len(ai_response.strip()) > 5:
                    response['answer'] = f"**🤖 AI Summary:**\n{ai_response}\n\n**📄 Source Documents:**\n{context}"
                    response['ai_model_used'] = True
                else:
                    # Fallback if AI response is empty
                    response['answer'] = f"**📄 Source Documents:**\n{context}"
                    response['ai_model_used'] = False
            else:
                print("DEBUG: AI model not available, using fallback")
                # Fallback response if HuggingFace is not available
                response['answer'] = f"**📄 Source Documents:**\n{context}"
                response['ai_model_used'] = False
            
            return response
        
        # Fallback: Use web search if no relevant documents found
        print("DEBUG: Using web search fallback")
        response['search_strategy'] = 'web_search'
        response['routing_info'] = f"Strategy: web_search (no relevant documents found or documents not relevant enough)"
        web_results = self.get_web_search_results(query)
        
        if web_results:
            # Create context from web results
            context = "Web search results:\n"
            for i, result in enumerate(web_results[:3], 1):
                context += f"{i}. {result['title']}: {result['snippet']}\n"
                response['sources'].append({
                    'type': 'web',
                    'title': result['title'],
                    'snippet': result['snippet'],
                    'link': result.get('link', ''),
                    'source': result.get('source', '')
                })
            
            # Generate response using HuggingFace
            if self.is_ai_model_available():
                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."
                ai_response = self.hf_client.generate_response(query, context, system_prompt)
                
                if len(ai_response.strip()) < 50 or "not sure" in ai_response.lower():
                    response['answer'] = f"**🌐 Web Search Results:**\n{context}\n\n**🤖 AI Analysis:**\n{ai_response}"
                else:
                    response['answer'] = f"**🤖 AI Analysis:**\n{ai_response}\n\n**🌐 Web Search Results:**\n{context}"
                response['ai_model_used'] = True
            else:
                response['answer'] = f"**🌐 Web Search Results:**\n{context}"
                response['ai_model_used'] = False
        else:
            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."
        
        return response

def main():
    """Main application function"""
    
    # Initialize session state
    if 'chatbot' not in st.session_state:
        # Try to get API key from environment variable first
        env_api_key = os.getenv("SERPER_API_KEY")
        st.session_state.chatbot = DocumentChatbot(serper_api_key=env_api_key)
    
    if 'chat_history' not in st.session_state:
        st.session_state.chat_history = []
    
    if 'documents_loaded' not in st.session_state:
        st.session_state.documents_loaded = False
    
    # Header
    st.title("Universal Document Intelligence Chatbot")
    st.markdown("*Upload documents and ask questions - get answers from your files or the web*")
    
    # Sidebar for document management
    with st.sidebar:
        st.header("Document Management")
        
        # File upload
        uploaded_files = st.file_uploader(
            "Upload PDF documents",
            type=['pdf'],
            accept_multiple_files=True,
            help="Upload PDF files to create a knowledge base"
        )
        
        # Process uploaded files
        if uploaded_files:
            if st.button("Process Documents", type="primary"):
                st.session_state.chatbot.process_uploaded_files(uploaded_files)
        
        # Display statistics
        if st.session_state.documents_loaded:
            st.subheader("Knowledge Base Stats")
            stats = st.session_state.chatbot.vector_store.get_stats()
            st.metric("Documents", stats['total_documents'])
            st.metric("Vector Dimension", stats['dimension'])
            st.info(f"Model: {stats['model_name']}")
        
        # Clear documents
        if st.session_state.documents_loaded:
            if st.button("Clear All Documents", type="secondary"):
                st.session_state.chatbot.vector_store.clear_index()
                st.session_state.documents_loaded = False
                st.session_state.chat_history = []
                st.success("Documents cleared!")
                st.rerun()
        
        # AI Model status
        st.subheader("AI Model Status")
        if st.session_state.chatbot.hf_client and st.session_state.chatbot.hf_client.is_available():
            st.success("✅ AI model loaded")
        else:
            st.warning("⚠️ AI model loading...")
            st.info("Models are being downloaded. This may take a few minutes on first run.")
        
        # Web Search Configuration
        st.subheader("🌐 Web Search")
        
        # Check if web search is already enabled
        web_search_enabled = st.session_state.chatbot.web_searcher is not None
        
        if web_search_enabled:
            st.success("✅ Web search enabled")
            if st.button("🔄 Change API Key"):
                st.session_state.show_api_input = True
                st.rerun()
        else:
            st.warning("⚠️ Web search disabled")
            
        # Show API key input field
        if not web_search_enabled or st.session_state.get('show_api_input', False):
            st.markdown("---")
            st.markdown("**Enter your Serper API Key:**")
            st.caption("Get a free API key at [serper.dev](https://serper.dev/) (2,500 searches/month free)")
            
            api_key = st.text_input(
                "Serper API Key",
                type="password",
                placeholder="Enter your API key here",
                help="Your API key is not stored and only used during this session",
                key="serper_api_key_input"
            )
            
            if api_key:
                if st.button("Enable Web Search", type="primary"):
                    success = st.session_state.chatbot.init_web_search(api_key)
                    if success:
                        st.success("✅ Web search enabled!")
                        st.session_state.show_api_input = False
                        st.rerun()
                    else:
                        st.error("❌ Invalid API key. Please check and try again.")
            
            if not api_key:
                st.info("💡 Web search is optional. The chatbot works with documents only.")
        
        st.markdown("---")
    
    # Main chat interface
    st.header("Chat Interface")
    
    # Display chat history
    for i, chat in enumerate(st.session_state.chat_history):
        with st.chat_message("user"):
            st.write(chat['query'])
        
        with st.chat_message("assistant"):
            st.write(chat['answer'])
            
            # Show routing info
            if chat.get('routing_info'):
                with st.expander("Search Strategy"):
                    st.info(chat['routing_info'])
            
            # Show sources
            if chat.get('sources'):
                with st.expander(f"Sources ({len(chat['sources'])} found)"):
                    for j, source in enumerate(chat['sources'], 1):
                        if source['type'] == 'document':
                            st.markdown(f"**{j}. Document Source:**")
                            st.markdown(f"- **File:** {source['filename']}")
                            st.markdown(f"- **Relevance:** {source['score']:.2f}")
                            st.markdown(f"- **Text:** {source['text'][:200]}...")
                        elif source['type'] == 'web':
                            st.markdown(f"**{j}. Web Source:**")
                            st.markdown(f"- **Title:** {source['title']}")
                            st.markdown(f"- **Source:** {source.get('source', 'Unknown')}")
                            if source.get('link'):
                                st.markdown(f"- **Link:** {source['link']}")
    
    # Query input
    query = st.chat_input("Ask a question about your documents or anything else...")
    
    if query:
        # Add user message to chat
        with st.chat_message("user"):
            st.write(query)
        
        # Generate response
        with st.chat_message("assistant"):
            with st.spinner("Thinking..."):
                response = st.session_state.chatbot.generate_response(query)
            
            st.write(response['answer'])
            
            # Show routing info
            if response.get('routing_info'):
                with st.expander("Search Strategy"):
                    st.info(response['routing_info'])
                    st.caption(f"Strategy used: {response['search_strategy']}")
            
            # Show sources
            if response.get('sources'):
                with st.expander(f"Sources ({len(response['sources'])} found)"):
                    for j, source in enumerate(response['sources'], 1):
                        if source['type'] == 'document':
                            st.markdown(f"**{j}. Document Source:**")
                            st.markdown(f"- **File:** {source['filename']}")
                            st.markdown(f"- **Relevance:** {source['score']:.2f}")
                            st.markdown(f"- **Text:** {source['text'][:200]}...")
                        elif source['type'] == 'web':
                            st.markdown(f"**{j}. Web Source:**")
                            st.markdown(f"- **Title:** {source['title']}")
                            st.markdown(f"- **Source:** {source.get('source', 'Unknown')}")
                            if source.get('link'):
                                st.markdown(f"- **Link:** {source['link']}")
        
        # Add to chat history
        st.session_state.chat_history.append({
            'query': query,
            'answer': response['answer'],
            'routing_info': response.get('routing_info'),
            'sources': response.get('sources', []),
            'search_strategy': response.get('search_strategy')
        })
    
    # Instructions
    if not st.session_state.chat_history:
        st.markdown("""
        ### Getting Started:
        
        1. **Upload PDFs** - Use the sidebar to add your documents
        2. **Click Process** - This creates a searchable knowledge base
        3. **Start Chatting** - Ask questions in the box below
        
        ### What you can ask:
        
        **About your documents:**
        - "What does the report say about..."
        - "Summarize the main points"
        - "Find information about X"
        
        **General questions:**
        - "What's the latest news on..."
        - "How does X work?"
        - "Compare A and B"
        
        The chatbot automatically decides whether to search your documents or the web.
        """)

if __name__ == "__main__":
    main()