app.py

"""
Streamlit App for RAG Pipeline using FastAPI
This app uses the REST API endpoints instead of calling functions directly.
"""

import streamlit as st
import requests
from typing import Dict, Any, Optional
import os

# API Configuration
API_BASE_URL = os.getenv("API_BASE_URL", "http://localhost:8000")

# Page configuration
st.set_page_config(
    page_title="RAG Pipeline - PDF Query System (API)",
    page_icon="📚",
    layout="wide"
)

# Initialize session state
if 'session_id' not in st.session_state:
    st.session_state.session_id = None
if 'documents_processed' not in st.session_state:
    st.session_state.documents_processed = False
if 'chat_history' not in st.session_state:
    st.session_state.chat_history = []
if 'system_status' not in st.session_state:
    st.session_state.system_status = None


def check_api_connection():
    """Check if API is accessible."""
    try:
        response = requests.get(f"{API_BASE_URL}/", timeout=5)
        return response.status_code == 200
    except:
        return False


def get_system_status():
    """Get system status from API."""
    try:
        response = requests.get(f"{API_BASE_URL}/status", timeout=5)
        if response.status_code == 200:
            return response.json()
        return None
    except Exception as e:
        st.error(f"Error connecting to API: {str(e)}")
        return None


def upload_documents_api(files, chunk_size, chunk_overlap):
    """Upload and process documents via API."""
    try:
        files_data = []
        for file in files:
            files_data.append(('files', (file.name, file.getvalue(), 'application/pdf')))
        
        data = {
            'chunk_size': chunk_size,
            'chunk_overlap': chunk_overlap
        }
        
        response = requests.post(
            f"{API_BASE_URL}/upload",
            files=files_data,
            data=data,
            timeout=300  # 5 minutes for large files
        )
        
        if response.status_code == 200:
            return response.json()
        else:
            error_detail = response.json().get('detail', 'Unknown error')
            raise Exception(f"API Error: {error_detail}")
    except requests.exceptions.RequestException as e:
        raise Exception(f"Connection error: {str(e)}")


def query_api(query, session_id, top_k, use_memory, metadata_filters=None):
    """Query documents via API."""
    payload = {
        "query": query,
        "session_id": session_id,
        "top_k": top_k,
        "use_memory": use_memory
    }
    
    if metadata_filters:
        payload["metadata_filters"] = metadata_filters
    
    try:
        response = requests.post(
            f"{API_BASE_URL}/query",
            json=payload,
            timeout=60
        )
        
        if response.status_code == 200:
            return response.json()
        else:
            error_detail = response.json().get('detail', 'Unknown error')
            raise Exception(f"API Error: {error_detail}")
    except requests.exceptions.RequestException as e:
        raise Exception(f"Connection error: {str(e)}")


def get_chat_history_api(session_id):
    """Get chat history from API."""
    try:
        response = requests.get(f"{API_BASE_URL}/chat-history/{session_id}", timeout=5)
        if response.status_code == 200:
            return response.json()
        return None
    except:
        return None


def clear_chat_history_api(session_id):
    """Clear chat history via API."""
    try:
        response = requests.delete(f"{API_BASE_URL}/chat-history/{session_id}", timeout=5)
        return response.status_code == 200
    except:
        return False


def get_metadata_options():
    """Get available metadata fields from API status or query."""
    # This would require an API endpoint to get metadata fields
    # For now, we'll use a simplified approach
    status = get_system_status()
    if status and status.get('vector_store_count', 0) > 0:
        # We can't directly access vectorstore from API, so we'll use a workaround
        # Try to get metadata from a sample query or add a new API endpoint
        return {}
    return {}


def main():
    """Main Streamlit app."""
    st.title("📚 RAG Pipeline - PDF Query System (API)")
    st.markdown("Upload PDF files and query them using Retrieval-Augmented Generation (RAG) via REST API")
    
    # Check API connection
    if not check_api_connection():
        st.error(f"❌ Cannot connect to API at {API_BASE_URL}")
        st.info("Please make sure the API server is running:")
        st.code("uvicorn api:app --reload --host 0.0.0.0 --port 8000")
        st.stop()
    
    # Get system status
    if st.session_state.system_status is None:
        st.session_state.system_status = get_system_status()
    
    # Generate session ID if not exists
    if st.session_state.session_id is None:
        import uuid
        st.session_state.session_id = str(uuid.uuid4())
    
    # Sidebar for configuration
    with st.sidebar:
        st.header("⚙️ Configuration")
        
        st.subheader("API Settings")
        api_url = st.text_input(
            "API Base URL",
            value=API_BASE_URL,
            help="Base URL of the FastAPI server"
        )
        if api_url != API_BASE_URL:
            st.session_state.api_base_url = api_url
            st.rerun()
        
        if st.button("🔄 Refresh Status", type="secondary"):
            st.session_state.system_status = get_system_status()
            st.rerun()
        
        st.subheader("Chunking Parameters")
        chunk_size = st.slider(
            "Chunk Size",
            min_value=200,
            max_value=2000,
            value=800,
            step=100,
            help="Size of each text chunk in characters"
        )
        chunk_overlap = st.slider(
            "Chunk Overlap",
            min_value=0,
            max_value=500,
            value=200,
            step=50,
            help="Number of overlapping characters between chunks"
        )
        
        st.subheader("Query Parameters")
        top_k = st.slider(
            "Top K Results",
            min_value=1,
            max_value=10,
            value=5,
            help="Number of document chunks to retrieve"
        )
        
        st.subheader("🔍 Metadata Filters")
        st.caption("Filter documents by metadata for faster, more accurate retrieval")
        use_filters = st.checkbox("Enable Metadata Filtering", value=False)
        
        metadata_filters = {}
        if use_filters:
            st.info("💡 Metadata filtering via API - enter filter values below")
            with st.expander("Advanced: Custom Metadata Filter"):
                filter_key = st.text_input("Metadata Key", placeholder="e.g., source, page")
                filter_value = st.text_input("Metadata Value", placeholder="e.g., document.pdf or 1,2,3")
                if filter_key and filter_value:
                    # Try to parse as list if comma-separated
                    if ',' in filter_value:
                        try:
                            # Try to parse as integers
                            metadata_filters[filter_key] = [int(v.strip()) for v in filter_value.split(',')]
                        except:
                            # Keep as strings
                            metadata_filters[filter_key] = [v.strip() for v in filter_value.split(',')]
                    else:
                        try:
                            # Try integer
                            metadata_filters[filter_key] = int(filter_value)
                        except:
                            # Keep as string
                            metadata_filters[filter_key] = filter_value
        
        st.divider()
        
        st.subheader("System Status")
        status = st.session_state.system_status
        if status:
            if status.get('documents_processed'):
                st.success("✅ Documents Processed")
                st.info(f"🗄️ {status.get('vector_store_count', 0)} documents in vector store")
                if status.get('chunks_available'):
                    st.info(f"📄 {status['chunks_available']} chunks available")
            else:
                st.info("⏳ No documents processed")
            
            if status.get('embedding_model'):
                st.caption(f"Model: {status['embedding_model']}")
        else:
            st.warning("⚠️ Could not fetch system status")
        
        st.divider()
        
        if st.button("🗑️ Clear Chat History", type="secondary"):
            if st.session_state.session_id:
                if clear_chat_history_api(st.session_state.session_id):
                    st.session_state.chat_history = []
                    st.success("Chat history cleared!")
                    st.rerun()
        
        if st.button("🔄 Reset System", type="secondary"):
            try:
                response = requests.post(f"{API_BASE_URL}/reset", timeout=5)
                if response.status_code == 200:
                    st.session_state.documents_processed = False
                    st.session_state.chat_history = []
                    st.session_state.system_status = get_system_status()
                    st.success("System reset!")
                    st.rerun()
            except Exception as e:
                st.error(f"Error resetting system: {str(e)}")
    
    # Main content area
    tab1, tab2 = st.tabs(["📤 Upload & Process", "💬 Chat"])
    
    with tab1:
        st.header("Upload PDF Files")
        st.markdown("Upload one or more PDF files to process and add to the knowledge base via API.")
        
        uploaded_files = st.file_uploader(
            "Choose PDF files",
            type=['pdf'],
            accept_multiple_files=True,
            help="You can upload multiple PDF files at once"
        )
        
        if uploaded_files:
            st.info(f"📎 {len(uploaded_files)} file(s) selected")
            
            # Display file names
            with st.expander("View uploaded files"):
                for file in uploaded_files:
                    st.write(f"- {file.name} ({file.size:,} bytes)")
            
            if st.button("🚀 Process Documents", type="primary"):
                with st.spinner("Uploading and processing documents via API..."):
                    try:
                        result = upload_documents_api(uploaded_files, chunk_size, chunk_overlap)
                        if result.get('success'):
                            st.success("✅ Documents processed successfully!")
                            st.json(result)
                            st.session_state.documents_processed = True
                            st.session_state.system_status = get_system_status()
                        else:
                            st.error("❌ Failed to process documents.")
                    except Exception as e:
                        st.error(f"❌ Error: {str(e)}")
    
    with tab2:
        st.header("💬 Chat with Documents")
        st.markdown("Ask questions about the uploaded PDF documents. The chat remembers previous conversations.")
        
        # Check if documents are processed
        status = st.session_state.system_status
        if not status or not status.get('documents_processed'):
            st.warning("⚠️ Please upload and process documents first in the 'Upload & Process' tab.")
        else:
            # Display chat history
            chat_container = st.container()
            with chat_container:
                if st.session_state.chat_history:
                    for message in st.session_state.chat_history:
                        role = message.get("role", "user")
                        content = message.get("content", "")
                        
                        if role == "user":
                            with st.chat_message("user"):
                                st.write(content)
                        elif role == "assistant":
                            with st.chat_message("assistant"):
                                st.write(content)
                                # Show sources if available
                                if "sources" in message:
                                    with st.expander("📄 Sources"):
                                        for i, source in enumerate(message["sources"], 1):
                                            st.markdown(f"**Source {i}** (Score: {source.get('score', 0):.4f})")
                                            st.caption(f"Preview: {source.get('preview', '')[:200]}...")
                else:
                    st.info("👋 Start a conversation by asking a question below!")
            
            # Chat input
            query = st.chat_input(
                "Ask a question about the documents...",
                key="chat_input"
            )
            
            # Handle query
            if query:
                # Add user message to chat history
                st.session_state.chat_history.append({
                    "role": "user",
                    "content": query
                })
                
                with st.spinner("Thinking..."):
                    try:
                        # Query via API
                        result = query_api(
                            query=query,
                            session_id=st.session_state.session_id,
                            top_k=top_k,
                            use_memory=True,
                            metadata_filters=metadata_filters if use_filters and metadata_filters else None
                        )
                        
                        # Add assistant response to chat history
                        st.session_state.chat_history.append({
                            "role": "assistant",
                            "content": result["answer"],
                            "sources": result.get("sources", [])
                        })
                        
                        # Update session ID if API generated a new one
                        if result.get("session_id"):
                            st.session_state.session_id = result["session_id"]
                        
                        st.rerun()
                        
                    except Exception as e:
                        st.error(f"Error processing query: {str(e)}")
                        # Remove the user message if there was an error
                        if st.session_state.chat_history and st.session_state.chat_history[-1]["role"] == "user":
                            st.session_state.chat_history.pop()
            
            # Example queries
            if not st.session_state.chat_history:
                st.divider()
                st.subheader("💡 Example Queries")
                example_queries = [
                    "What is the main topic of the document?",
                    "Summarize the key points",
                    "What are the main findings?",
                ]
                
                cols = st.columns(len(example_queries))
                for i, example in enumerate(example_queries):
                    with cols[i]:
                        if st.button(f"📝 {example[:30]}...", key=f"example_{i}"):
                            st.session_state.chat_history.append({
                                "role": "user",
                                "content": example
                            })
                            st.rerun()


if __name__ == "__main__":
    main()