app.py

import streamlit as st
import streamlit_authenticator as stauth
import tempfile
import os
from datetime import datetime, date
from dotenv import load_dotenv
from openai import OpenAI
import pypdfium2 as pdfium
import yaml
from yaml.loader import SafeLoader
from docling.datamodel.pipeline_options import PdfPipelineOptions
from docling.document_converter import DocumentConverter, PdfFormatOption
from docling.datamodel.base_models import InputFormat

load_dotenv()

st.set_page_config(
    page_title="PDF RAG Quality Analyzer",
    page_icon="📄",
    layout="wide"
)

QUALITY_PROMPT = """You are an expert at evaluating document quality for RAG (Retrieval-Augmented Generation) pipelines.

Analyze the following markdown content extracted from a PDF and evaluate its suitability for RAG.

Consider these factors:
1. **Text Extraction Quality**: Is the text readable? Are there OCR errors, garbled characters, or missing content?
2. **Structure Preservation**: Are headings, paragraphs, lists, and tables properly formatted?
3. **Content Coherence**: Does the text flow logically? Are sentences complete?
4. **Information Density**: Does the document contain meaningful, retrievable information?
5. **Noise Level**: Is there excessive boilerplate, headers/footers, or irrelevant content?

Provide your assessment in the following format:

## Quality Score: [1-10]/10

## Verdict: [EXCELLENT / GOOD / ACCEPTABLE / POOR / UNUSABLE]

## Summary
[2-3 sentence summary of overall quality]

## Detailed Analysis
- **Text Extraction**: [score/10] - [brief explanation]
- **Structure**: [score/10] - [brief explanation]
- **Coherence**: [score/10] - [brief explanation]
- **Information Density**: [score/10] - [brief explanation]
- **Noise Level**: [score/10] - [brief explanation]

## Recommendations
[Specific suggestions to improve the document quality for RAG, if applicable]

---
MARKDOWN CONTENT TO ANALYZE:
"""


def load_config():
    with open("config.yaml") as file:
        return yaml.load(file, Loader=SafeLoader)


def save_config(config):
    with open("config.yaml", "w") as file:
        yaml.dump(config, file, default_flow_style=False)


def get_user_request_count(username: str) -> int:
    """Get the number of requests made by user today."""
    if "user_requests" not in st.session_state:
        st.session_state["user_requests"] = {}
    
    today = str(date.today())
    user_key = f"{username}_{today}"
    return st.session_state["user_requests"].get(user_key, 0)


def increment_user_request(username: str):
    """Increment the request count for user today."""
    if "user_requests" not in st.session_state:
        st.session_state["user_requests"] = {}
    
    today = str(date.today())
    user_key = f"{username}_{today}"
    st.session_state["user_requests"][user_key] = st.session_state["user_requests"].get(user_key, 0) + 1


def check_rate_limit(username: str, max_requests: int) -> tuple[bool, int]:
    """Check if user is within rate limit. Returns (allowed, remaining)."""
    current = get_user_request_count(username)
    remaining = max_requests - current
    return remaining > 0, max(0, remaining)


def get_pdf_page_count(pdf_path: str) -> int:
    pdf = pdfium.PdfDocument(pdf_path)
    return len(pdf)


def convert_pdf_to_markdown(pdf_path: str, page_range: tuple[int, int] | None = None) -> str:
    pipeline_options = PdfPipelineOptions()
    pipeline_options.do_ocr = True
    pipeline_options.do_table_structure = True
    pipeline_options.table_structure_options.do_cell_matching = True
    pipeline_options.do_code_enrichment = True
    pipeline_options.generate_picture_images = True

    converter = DocumentConverter(
        format_options={
            InputFormat.PDF: PdfFormatOption(pipeline_options=pipeline_options)
        }
    )

    if page_range:
        result = converter.convert(pdf_path, page_range=page_range)
    else:
        result = converter.convert(pdf_path)
    return result.document.export_to_markdown()


def analyze_quality(markdown_content: str) -> str:
    api_key = os.getenv("OPENAI_API_KEY")
    if not api_key:
        raise ValueError("OpenAI API key not configured on server")
    
    client = OpenAI(api_key=api_key)
    
    truncated_content = markdown_content[:15000] if len(markdown_content) > 15000 else markdown_content
    
    response = client.chat.completions.create(
        model="gpt-4o-mini",
        messages=[
            {"role": "system", "content": "You are an expert document quality analyst."},
            {"role": "user", "content": QUALITY_PROMPT + truncated_content}
        ],
        temperature=0.3,
        max_tokens=1500
    )
    
    return response.choices[0].message.content


def main():
    config = load_config()
    
    authenticator = stauth.Authenticate(
        config["credentials"],
        config["cookie"]["name"],
        config["cookie"]["key"],
        config["cookie"]["expiry_days"],
    )
    
    max_requests = config.get("rate_limit", {}).get("max_requests_per_day", 20)
    
    try:
        authenticator.login(location="main", key="login-main")
    except Exception as e:
        st.error(f"Login error: {e}")
    
    if st.session_state.get("authentication_status"):
        username = st.session_state.get("username")
        name = st.session_state.get("name")
        
        with st.sidebar:
            st.write(f"👤 **{name}**")
            authenticator.logout("Logout", "sidebar", key="logout-main")
            
            st.markdown("---")
            
            allowed, remaining = check_rate_limit(username, max_requests)
            st.metric("Requests remaining today", remaining)
            
            st.markdown("---")
            st.markdown("""
            ### How it works
            1. Upload a PDF document
            2. Docling converts it to Markdown
            3. GPT-4o-mini analyzes the quality
            4. Get a detailed RAG suitability report
            """)
        
        st.title("📄 PDF RAG Quality Analyzer")
        st.markdown("*Check if your PDF documents are suitable for RAG pipelines*")
        
        uploaded_file = st.file_uploader(
            "Upload your PDF document",
            type=["pdf"],
            help="Upload a PDF file to analyze its quality for RAG"
        )
        
        if uploaded_file is not None:
            with tempfile.NamedTemporaryFile(delete=False, suffix=".pdf") as tmp_file:
                tmp_file.write(uploaded_file.getvalue())
                tmp_path = tmp_file.name
            
            try:
                total_pages = get_pdf_page_count(tmp_path)
            except Exception:
                total_pages = None
            
            col1, col2 = st.columns([1, 1])
            
            with col1:
                st.subheader("📥 Uploaded File")
                page_info = f"\n\n**Pages:** {total_pages}" if total_pages else ""
                st.info(f"**Filename:** {uploaded_file.name}\n\n**Size:** {uploaded_file.size / 1024:.1f} KB{page_info}")
            
            with col2:
                st.subheader("📄 Page Selection")
                use_page_range = st.checkbox("Analyze specific pages only", value=False)
                
                page_range = None
                if use_page_range and total_pages:
                    col_start, col_end = st.columns(2)
                    with col_start:
                        start_page = st.number_input(
                            "Start page",
                            min_value=1,
                            max_value=total_pages,
                            value=1,
                            step=1
                        )
                    with col_end:
                        end_page = st.number_input(
                            "End page",
                            min_value=1,
                            max_value=total_pages,
                            value=min(total_pages, start_page),
                            step=1
                        )
                    if end_page >= start_page:
                        page_range = (start_page, end_page)
                        st.caption(f"Will analyze pages {start_page} to {end_page}")
                    else:
                        st.warning("End page must be >= start page")
            
            if st.button("🔍 Analyze Document", type="primary", use_container_width=True):
                allowed, remaining = check_rate_limit(username, max_requests)
                
                if not allowed:
                    st.error(f"⚠️ Daily limit reached ({max_requests} requests/day). Try again tomorrow.")
                    os.unlink(tmp_path)
                    return
                
                try:
                    range_text = f" (pages {page_range[0]}-{page_range[1]})" if page_range else ""
                    with st.spinner(f"Converting PDF to Markdown with Docling{range_text}..."):
                        markdown_content = convert_pdf_to_markdown(tmp_path, page_range=page_range)
                    
                    st.session_state["markdown_content"] = markdown_content
                    
                    with st.spinner("Analyzing quality with AI..."):
                        analysis = analyze_quality(markdown_content)
                    
                    st.session_state["analysis"] = analysis
                    increment_user_request(username)
                    
                except Exception as e:
                    st.error(f"Error processing document: {str(e)}")
                finally:
                    os.unlink(tmp_path)
            
            if "analysis" in st.session_state:
                st.markdown("---")
                
                st.subheader("📊 Quality Analysis")
                st.markdown(st.session_state["analysis"])
                
                st.markdown("---")
                
                with st.expander("📝 View Extracted Markdown", expanded=False):
                    st.markdown(
                        f'<div style="max-height: 500px; overflow-y: auto; background: #1e1e1e; padding: 1rem; border-radius: 0.5rem;"><pre style="white-space: pre-wrap; margin: 0;">{st.session_state["markdown_content"]}</pre></div>',
                        unsafe_allow_html=True
                    )
                
                st.download_button(
                    label="⬇️ Download Markdown",
                    data=st.session_state["markdown_content"],
                    file_name=f"{uploaded_file.name.replace('.pdf', '')}.md",
                    mime="text/markdown"
                )
    
    elif st.session_state.get("authentication_status") is False:
        st.error("❌ Username/password is incorrect")
    
    elif st.session_state.get("authentication_status") is None:
        st.title("📄 PDF RAG Quality Analyzer")
        st.info("👆 Please log in to continue")


if __name__ == "__main__":
    main()