app.py

"""
Intelligent Document Analyzer - AI-Powered Document Intelligence Application

This application demonstrates document intelligence capabilities by:
- Uploading PDF/text documents
- Generating AI-powered summaries and key insights
- Extracting risk flags and important entities
- Enabling interactive Q&A on uploaded documents

Built with multi-agent architecture for planning, review, and improvement cycles.
"""

import streamlit as st
import os
from pathlib import Path
from typing import Optional, List, Dict, Any
import json
from datetime import datetime
from io import BytesIO

import pdfplumber
try:
    from PyPDF2 import PdfReader
except ImportError:
    PdfReader = None

st.set_page_config(
    page_title="Document Analyzer",
    page_icon="🔍",
    layout="wide",
    initial_sidebar_state="expanded"
)


def load_external_js():
    """Load external JavaScript and CSS styling."""
    # Get the directory where this script is located
    script_dir = Path(__file__).parent
    
    # Load external JavaScript using absolute path
    js_path = script_dir / "static" / "app.js"
    with open(js_path, "r") as f:
        st.markdown(f"<script>{f.read()}</script>", unsafe_allow_html=True)


class DocumentProcessor:
    """Handles document extraction and preprocessing."""
    
    @staticmethod
    def extract_text_from_pdf(file_content=None, file_path: str = None) -> str:
        """Extract text from PDF files."""
        text = ""
        
        # Pre-read UploadedFile content once to avoid consuming the file pointer
        content_bytes = None
        if file_content and not isinstance(file_content, bytes):
            content_bytes = file_content.read()
        
        if pdfplumber is not None:
            try:
                if file_content:
                    # Handle both bytes and UploadedFile objects
                    if isinstance(file_content, bytes):
                        with pdfplumber.open(BytesIO(file_content)) as pdf:
                            for page in pdf.pages:
                                page_text = page.extract_text()
                                if page_text:
                                    text += page_text + "\n"
                    else:
                        # Use pre-read content
                        with pdfplumber.open(BytesIO(content_bytes)) as pdf:
                            for page in pdf.pages:
                                page_text = page.extract_text()
                                if page_text:
                                    text += page_text + "\n"
                else:
                    with pdfplumber.open(file_path) as pdf:
                        for page in pdf.pages:
                            page_text = page.extract_text()
                            if page_text:
                                text += page_text + "\n"
            except Exception as e:
                st.warning(f"pdfplumber failed: {e}, trying alternative...")
        
        # Fallback to PyPDF2
        if not text and PdfReader is not None:
            try:
                if file_content:
                    # Handle both bytes and UploadedFile objects
                    if isinstance(file_content, bytes):
                        reader = PdfReader(BytesIO(file_content))
                    else:
                        # Use pre-read content instead of reading again
                        reader = PdfReader(BytesIO(content_bytes))
                    for page in reader.pages:
                        text += page.extract_text() + "\n"
                else:
                    reader = PdfReader(file_path)
                    for page in reader.pages:
                        text += page.extract_text() + "\n"
            except Exception as e:
                st.error(f"PDF extraction failed: {e}")
        
        return text.strip()
    
    @staticmethod
    def extract_text_from_txt(file_content=None, file_path: str = None) -> str:
        """Extract text from plain text files."""
        if file_content:
            return file_content.decode('utf-8')
        elif file_path:
            with open(file_path, 'r', encoding='utf-8') as f:
                return f.read()
        return ""
    
    @staticmethod
    def preprocess_text(text: str) -> str:
        """Clean and preprocess extracted text."""
        # Remove excessive whitespace
        import re
        text = re.sub(r'\s+', ' ', text)
        text = re.sub(r'\n\s*\n', '\n\n', text)
        return text.strip()


class MultiAgentOrchestrator:
    """
    Multi-agent system for document analysis with planning, review, and improvement cycles.
    
    Agents:
    - Planner Agent: Determines analysis strategy and breaks down tasks
    - Analyzer Agent: Performs deep content analysis and extraction
    - Reviewer Agent: Validates findings and checks for completeness
    - Improver Agent: Refines outputs based on reviewer feedback
    """
    
    def __init__(self, api_key: str = None, model: str = "claude-haiku-4-5-20251001", store_prompts: bool = True):
        # Store API key with fallback to environment variable
        self.api_key = api_key or os.getenv("ANTHROPIC_API_KEY")
        self.model = model
        self.conversation_history: List[Dict] = []
        self.store_prompts = store_prompts
        self.last_prompt_sent = None
        self.last_api_response = None
        
    def _call_llm(self, system_prompt: str, user_prompt: str) -> str:
        """Call LLM with given prompts."""
        # Store prompt for display if key is provided
        if self.store_prompts and self.api_key:
            self.last_prompt_sent = {
                "system": system_prompt,
                "user": user_prompt
            }
        
        try:
            from anthropic import Anthropic
            client = Anthropic(api_key=self.api_key)
            
            response = client.messages.create(
                model=self.model,
                max_tokens=2000,
                temperature=0.3,
                system=system_prompt,
                messages=[{"role": "user", "content": user_prompt}]
            )
            self.last_api_response = response.content[0].text
            return self.last_api_response
        except Exception as e:
            # Fallback to mock analysis for demo purposes
            return self._mock_analysis(system_prompt, user_prompt)
    
    def _call_llm_stream(self, system_prompt: str, user_prompt: str):
        """Call LLM with streaming response."""
        try:
            from anthropic import Anthropic
            client = Anthropic(api_key=self.api_key)
            
            # Store prompt for display if key is provided
            if self.store_prompts and self.api_key:
                self.last_prompt_sent = {
                    "system": system_prompt,
                    "user": user_prompt
                }
            
            with client.messages.stream(
                model=self.model,
                max_tokens=2000,
                temperature=0.3,
                system=system_prompt,
                messages=[{"role": "user", "content": user_prompt}]
            ) as stream:
                for text in stream.text_stream:
                    yield text
                
                # Store the complete response
                self.last_api_response = stream.get_final_message().content[0].text
                
        except Exception as e:
            # Fallback to mock streaming analysis for demo purposes
            yield from self._mock_analysis_stream(system_prompt, user_prompt)
    
    def _mock_analysis(self, system_prompt: str, user_prompt: str) -> str:
        """Mock analysis when no API key is available."""
        # Check if this is a Q&A question (contains question words or ends with ?)
        is_question = any(
            word in user_prompt.lower() 
            for word in ["what", "how", "why", "when", "where", "who", "which", "can you", "could you", "is there", "are there"]
        ) or user_prompt.strip().endswith("?")
        
        # Check if system prompt indicates Q&A mode
        is_qa_mode = "q&a" in system_prompt.lower() or "answer questions" in system_prompt.lower()
        
        if is_question or is_qa_mode:
            return f"""Based on the document content, here's what I found regarding your question:

**Key Findings:**

The document contains relevant information that addresses your inquiry. Based on my analysis of the provided text:

1. **Primary Information**: The document discusses operational procedures and strategic considerations with detailed explanations of processes and methodologies.

2. **Important Details**: Several key points are highlighted throughout the document, including timelines, responsibilities, and expected outcomes.

3. **Actionable Items**: The content includes specific recommendations and next steps that should be considered.

**Summary Answer:**
The information you're looking for appears to be covered in the main body of the document. For more specific details about this topic, I would recommend reviewing the sections on operational procedures and strategic planning.

*Note: This is a mock response since no Anthropic API key was provided. With an API key configured, I would provide a more precise answer based on actual AI analysis.*"""
        elif "summary" in system_prompt.lower() or "summarize" in user_prompt.lower():
            return """## Executive Summary

This document appears to be a professional business/technical document containing important information about operations, policies, or analysis. The content demonstrates structured communication with clear sections and actionable insights.

### Key Points Identified:
1. Primary focus on operational efficiency and strategic planning
2. Multiple stakeholders mentioned with distinct roles
3. Risk considerations are addressed throughout
4. Recommendations include specific action items

## Document Characteristics
- **Structure**: Well-organized with clear headings
- **Tone**: Professional and analytical
- **Complexity**: Medium to high technical depth
- **Actionability**: Contains concrete recommendations"""
        elif "risk" in system_prompt.lower():
            return """## Risk Analysis

### Identified Risk Factors:

**🟡 Medium Risk Items:**
- Operational dependencies on external systems
- Potential compliance gaps in documented processes
- Resource allocation constraints

**🟢 Low Risk Items:**
- Standard business continuity measures in place
- Documentation appears current and maintained

### Recommendations:
1. Review operational dependencies quarterly
2. Update compliance documentation as needed
3. Consider resource buffer for critical operations"""
        elif "insight" in system_prompt.lower():
            return """## Key Insights Extracted

### Strategic Insights:
1. **Efficiency Focus**: Document emphasizes process optimization and waste reduction
2. **Stakeholder Alignment**: Multiple parties need coordinated action
3. **Risk-Aware Planning**: Decisions consider potential downsides

### Tactical Insights:
1. Clear timelines and milestones established
2. Resource requirements are quantified
3. Success metrics are defined

### Actionable Takeaways:
- Prioritize high-impact, low-effort initiatives first
- Establish regular review cadence for progress tracking
- Document lessons learned for future reference"""
        else:
            return """## Analysis Results

The document has been analyzed using multi-agent AI systems. Key findings include structured information suitable for decision-making purposes. The content demonstrates professional communication standards and contains actionable recommendations."""
    
    def _mock_analysis_stream(self, system_prompt: str, user_prompt: str):
        """Mock streaming analysis when no API key is available."""
        result = self._mock_analysis(system_prompt, user_prompt)
        # Simulate streaming by yielding character by character
        for char in result:
            yield char


class PlannerAgent(MultiAgentOrchestrator):
    """Plans the analysis strategy for a given document."""
    
    def create_analysis_plan(self, document_text: str) -> Dict[str, Any]:
        """Create a structured plan for analyzing the document."""
        system_prompt = """You are a Document Analysis Planner. Your role is to:
1. Assess the document type and structure
2. Identify key sections and their importance
3. Determine what analysis approaches would be most valuable
4. Create a step-by-step analysis plan

Output should be in JSON format with keys: document_type, main_sections, priority_areas, analysis_approach."""
        
        user_prompt = f"Analyze this document and create an analysis plan:\n\n{document_text[:5000]}"
        
        response = self._call_llm(system_prompt, user_prompt)
        try:
            # Try to parse as JSON
            import re
            json_match = re.search(r'\{.*\}', response, re.DOTALL)
            if json_match:
                return json.loads(json_match.group())
        except:
            pass
        
        return {
            "document_type": "general",
            "main_sections": ["introduction", "body", "conclusion"],
            "priority_areas": ["key_findings", "recommendations"],
            "analysis_approach": "comprehensive"
        }


class AnalyzerAgent(MultiAgentOrchestrator):
    """Performs deep content analysis on documents."""
    
    def generate_summary(self, document_text: str) -> str:
        """Generate a comprehensive summary of the document."""
        system_prompt = """You are a Document Analysis Expert. Create a detailed executive summary that captures:
- Main purpose and objectives
- Key findings and insights
- Important data points or metrics
- Conclusions and recommendations

Format your response with clear headings and bullet points for readability."""
        
        user_prompt = f"Summarize this document:\n\n{document_text[:8000]}"
        return self._call_llm(system_prompt, user_prompt)
    
    def extract_risk_flags(self, document_text: str) -> List[str]:
        """Extract potential risk factors or concerns from the document."""
        system_prompt = """You are a Risk Analyst. Identify any risk factors, concerns, or areas requiring attention in this document. Categorize by severity (HIGH/MEDIUM/LOW) and provide brief explanations."""
        
        user_prompt = f"Analyze for risks:\n\n{document_text[:8000]}"
        return self._call_llm(system_prompt, user_prompt)
    
    def extract_key_insights(self, document_text: str) -> List[str]:
        """Extract key insights and actionable takeaways."""
        system_prompt = """You are an Insights Extractor. Identify the most valuable insights from this document that would help a decision-maker. Focus on:
- Strategic implications
- Actionable recommendations
- Important patterns or trends
- Critical success factors"""
        
        user_prompt = f"Extract key insights:\n\n{document_text[:8000]}"
        return self._call_llm(system_prompt, user_prompt)


class ReviewerAgent(MultiAgentOrchestrator):
    """Reviews and validates analysis outputs."""
    
    def review_analysis(self, summary: str, risks: str, insights: str) -> Dict[str, Any]:
        """Review the complete analysis for quality and completeness."""
        system_prompt = """You are a Quality Reviewer. Evaluate the document analysis for:
1. Completeness - Are all important aspects covered?
2. Accuracy - Do findings align with typical document patterns?
3. Clarity - Is the output clear and actionable?

Provide feedback on what could be improved."""
        
        user_prompt = f"Review this analysis:\n\nSummary:\n{summary}\n\nRisks:\n{risks}\n\nInsights:\n{insights}"
        review = self._call_llm(system_prompt, user_prompt)
        
        return {
            "quality_score": 85,  # Mock score
            "completeness": "Good coverage of key areas",
            "feedback": review
        }


class ImproverAgent(MultiAgentOrchestrator):
    """Improves analysis based on reviewer feedback."""
    
    def improve_analysis(self, original_summary: str, review_feedback: Dict) -> str:
        """Refine the summary based on reviewer feedback."""
        system_prompt = """You are an Analysis Improver. Enhance the document summary based on reviewer feedback. Make it more comprehensive, clear, and actionable."""
        
        user_prompt = f"Original Summary:\n{original_summary}\n\nReview Feedback:\n{review_feedback.get('feedback', '')}"
        return self._call_llm(system_prompt, user_prompt)


def initialize_session_state():
    """Initialize Streamlit session state variables."""
    if "document_text" not in st.session_state:
        st.session_state.document_text = ""
    if "analysis_results" not in st.session_state:
        st.session_state.analysis_results = None
    if "chat_history" not in st.session_state:
        st.session_state.chat_history = []
    if "api_key" not in st.session_state:
        st.session_state.api_key = ""
    if "anthropic_prompts" not in st.session_state:
        st.session_state.anthropic_prompts = []


def run_full_analysis(document_text: str) -> Dict[str, Any]:
    """Run the complete multi-agent analysis pipeline."""
    
    # Clear previous prompts and store agent instances for prompt retrieval
    st.session_state.anthropic_prompts = []
    agents_list = []
    
    # Initialize agents with prompt storage enabled
    planner = PlannerAgent(api_key=st.session_state.api_key or None, store_prompts=True)
    analyzer = AnalyzerAgent(api_key=st.session_state.api_key or None, store_prompts=True)
    reviewer = ReviewerAgent(api_key=st.session_state.api_key or None, store_prompts=True)
    improver = ImproverAgent(api_key=st.session_state.api_key or None, store_prompts=True)
    agents_list = [planner, analyzer, reviewer, improver]
    
    # Step 1: Planning
    with st.spinner("📋 Planner Agent: Creating analysis strategy..."):
        analysis_plan = planner.create_analysis_plan(document_text)
    
    # Step 2: Analysis
    with st.spinner("🔍 Analyzer Agent: Generating summary and insights..."):
        summary = analyzer.generate_summary(document_text)
        
    with st.spinner("⚠️ Analyzer Agent: Identifying risk factors..."):
        risks = analyzer.extract_risk_flags(document_text)
        
    with st.spinner("💡 Analyzer Agent: Extracting key insights..."):
        insights = analyzer.extract_key_insights(document_text)
    
    # Step 3: Review
    with st.spinner("👁️ Reviewer Agent: Validating analysis quality..."):
        review = reviewer.review_analysis(summary, risks, insights)
    
    # Step 4: Improvement
    with st.spinner("✨ Improver Agent: Refining outputs..."):
        improved_summary = improver.improve_analysis(summary, review)
    
    # Collect all prompts from agents
    prompt_entries = []
    agent_names = ["Planner", "Analyzer (Summary)", "Analyzer (Risks)", "Analyzer (Insights)", "Reviewer", "Improver"]
    
    for i, agent in enumerate(agents_list):
        if hasattr(agent, 'last_prompt_sent') and agent.last_prompt_sent:
            prompt_entries.append({
                "agent": agent_names[i] if i < len(agent_names) else f"Agent {i+1}",
                "system_prompt": agent.last_prompt_sent.get("system", ""),
                "user_prompt": agent.last_prompt_sent.get("user", "")
            })
    
    st.session_state.anthropic_prompts = prompt_entries
    
    return {
        "plan": analysis_plan,
        "summary": improved_summary,
        "risks": risks,
        "insights": insights,
        "review": review,
        "timestamp": datetime.now().strftime("%Y-%m-%d %H:%M:%S"),
        "prompts": prompt_entries
    }


def main():
    """Main application entry point."""
    load_external_js()
    initialize_session_state()
    
    # Sidebar configuration
    with st.sidebar:
        st.header("⚙️ Configuration")
        
        api_key = st.text_input(
            "Anthropic Claude API Key (Optional)",
            type="password",
            help="Provide an Anthropic API key for enhanced analysis. Without it, demo mode will be used."
        )
        if api_key:
            st.session_state.api_key = api_key
        
        model = st.selectbox(
            "Model Selection",
            ["claude-sonnet-4-6", "claude-opus-4-7", "claude-haiku-4-5-20251001"],
            index=0
        )
        
        st.divider()
        
        st.header("📁 Document Info")
        if st.session_state.document_text:
            char_count = len(st.session_state.document_text)
            word_count = len(st.session_state.document_text.split())
            st.metric("Characters", f"{char_count:,}")
            st.metric("Words", f"{word_count:,}")
        
        st.divider()
        
        if st.button("🗑️ Clear Analysis", type="secondary"):
            st.session_state.document_text = ""
            st.session_state.analysis_results = None
            st.session_state.chat_history = []
            st.rerun()
    
    # Main content area
    st.markdown('<p class="main-header">🔍 Intelligent Document Analyzer</p>', unsafe_allow_html=True)
    st.markdown('<p class="sub-header">Upload documents for AI-powered analysis, summaries, and Q&A</p>', unsafe_allow_html=True)
    
    # Display "Using Claude LLM" badge if API key is provided
    if st.session_state.api_key:
        st.success("🤖 **Using Claude LLM** - Anthropic API Key configured", icon="✅")
    
    # File upload section
    uploaded_file = st.file_uploader(
        "Upload a document (PDF or TXT)",
        type=["pdf", "txt"],
        help="Supported formats: PDF, Plain Text"
    )
    
    if uploaded_file is not None:
        # Process the file
        file_type = uploaded_file.name.split(".")[-1].lower()
        
        if file_type == "pdf":
            text = DocumentProcessor.extract_text_from_pdf(file_content=uploaded_file)
        else:
            text = uploaded_file.read().decode("utf-8")
        
        # Store in session state
        st.session_state.document_text = DocumentProcessor.preprocess_text(text)
        st.success(f"✅ Document loaded! {len(st.session_state.document_text.split())} words extracted.")
    
    # Display document preview if available
    if st.session_state.document_text:
        with st.expander("📄 View Document Preview"):
            preview_text = st.session_state.document_text[:5000] + "..." if len(st.session_state.document_text) > 5000 else st.session_state.document_text
            st.text_area("Document Content", value=preview_text, height=200, disabled=True)
        
        # Analysis buttons
        col1, col2 = st.columns([1, 1])
        with col1:
            if st.button("🚀 Run Full Analysis", type="primary", use_container_width=True):
                results = run_full_analysis(st.session_state.document_text)
                st.session_state.analysis_results = results
                st.rerun()
        
        with col2:
            if st.button("⚡ Quick Summary", use_container_width=True):
                analyzer = AnalyzerAgent(api_key=st.session_state.api_key or None)
                summary = analyzer.generate_summary(st.session_state.document_text)
                st.session_state.analysis_results = {"summary": summary, "timestamp": datetime.now().strftime("%Y-%m-%d %H:%M:%S")}
                st.rerun()
        
        # Display results if available
        if st.session_state.analysis_results:
            st.divider()
            
            # Display prompts sent to Anthropic API (if key was provided)
            if st.session_state.api_key and st.session_state.anthropic_prompts:
                st.divider()
                with st.expander(f"📝 View Prompts Sent to Anthropic API ({len(st.session_state.anthropic_prompts)} calls)", expanded=False):
                    for i, prompt_entry in enumerate(st.session_state.anthropic_prompts):
                        st.markdown(f"**{i+1}. {prompt_entry['agent']}**")
                        
                        st.markdown("**🤖 System Prompt:**")
                        st.code(prompt_entry["system_prompt"], language="markdown")
                        
                        st.markdown("**👤 User Prompt:**")
                        # Truncate very long prompts in display
                        user_text = prompt_entry["user_prompt"]
                        if len(user_text) > 1000:
                            st.code(user_text[:997] + "...", language="markdown")
                        else:
                            st.code(user_text, language="markdown")
                        
                        st.divider()
            
            # Summary section
            st.markdown("### 📋 Executive Summary")
            st.markdown(st.session_state.analysis_results.get("summary", "No summary available."))
            
            # Multi-column display for risks and insights
            col1, col2 = st.columns(2)
            
            with col1:
                st.markdown("### ⚠️ Risk Analysis")
                risks = st.session_state.analysis_results.get("risks", "")
                if risks:
                    st.markdown(risks)
                else:
                    st.info("No risk analysis available.")
            
            with col2:
                st.markdown("### 💡 Key Insights")
                insights = st.session_state.analysis_results.get("insights", "")
                if insights:
                    st.markdown(insights)
                else:
                    st.info("No insights extracted yet.")
        
        # Q&A Section with Streaming Support
        st.divider()
        st.markdown("### 💬 Document Q&A")
        
        if st.session_state.api_key:
            st.info("🔄 **Streaming enabled**: Answers will appear in real-time as they are generated by Claude.")
        else:
            st.warning("⚠️ **Demo Mode**: No API key configured. Mock responses will be used.")
        
        # Chat input
        user_question = st.text_input(
            "Ask a question about this document:",
            placeholder="e.g., What are the main recommendations?",
            key="qa_input"
        )
        
        if user_question and st.button("🔎 Ask"):
            # Add to chat history (user message)
            st.session_state.chat_history.append({"role": "user", "content": user_question})
            
            # Prepare the document context and question
            doc_context = st.session_state.document_text[:10000]  # Limit context size
            system_prompt = f"""You are a Document Q&A Assistant. Answer questions based on this document content:

{doc_context}

If the answer is not in the document, state that clearly."""
            
            # Store this prompt if API key is provided
            if st.session_state.api_key:
                st.session_state.anthropic_prompts.append({
                    "agent": "Q&A Assistant",
                    "system_prompt": system_prompt,
                    "user_prompt": user_question
                })
            
            # Use streaming if API key is provided
            if st.session_state.api_key:
                analyzer = AnalyzerAgent(api_key=st.session_state.api_key)
                
                with st.spinner("🤖 Thinking..."):
                    placeholder = st.empty()
                    full_response = ""
                    api_error = False
                    
                    try:
                        for chunk in analyzer._call_llm_stream(system_prompt, user_question):
                            full_response += chunk
                            placeholder.markdown(full_response + "▌")
                        
                        # Check if we got a real response or mock fallback
                        if "Note: This is a mock response since no Anthropic API key was provided" in full_response:
                            api_error = True
                            st.error(f"⚠️ **API Error**: The mock response was returned. API key value: {st.session_state.api_key}. Please check your API key and try again.")
                        
                        placeholder.markdown(full_response)
                        response = full_response
                    except Exception as e:
                        api_error = True
                        st.error(f"⚠️ **API Error**: {str(e)}")
                        placeholder.markdown("Sorry, there was an error connecting to the Anthropic API. Please check your API key and try again.")
                        response = ""
            else:
                # Streaming mock fallback when no API key is provided
                with st.spinner("🤖 Thinking..."):
                    placeholder = st.empty()
                    analyzer = AnalyzerAgent(api_key=None)
                    full_response = ""
                    
                    for chunk in analyzer._call_llm_stream(system_prompt, user_question):
                        full_response += chunk
                        placeholder.markdown(full_response + "▌")
                    
                    placeholder.markdown(full_response)
                    response = full_response
            
            # Add assistant response to chat history
            st.session_state.chat_history.append({"role": "assistant", "content": response})
        
        # Display chat history (last 5 messages)
        if st.session_state.chat_history:
            for msg in st.session_state.chat_history[-5:]:
                if msg["role"] == "user":
                    with st.chat_message("user"):
                        st.write(msg["content"])
                else:
                    with st.chat_message("assistant"):
                        st.write(msg["content"])


if __name__ == "__main__":
    main()