app.py

import os
import streamlit as st
import google.genai as genai
from PyPDF2 import PdfReader
from dotenv import load_dotenv

from spire.doc import *
from spire.doc.common import *


# Load environment variables
load_dotenv()

# Global constants
MODEL_NAME = os.getenv("GEMINI_MODEL_NAME", "models/gemini-2.5-flash-preview-05-20")
print(f"Using model: {MODEL_NAME}")


# Configure Streamlit page
st.set_page_config(
    page_title="Document Analysis Chatbot",
    page_icon="📄",
    layout="wide",
)

# Initialize session state
if "messages" not in st.session_state:
    st.session_state.messages = []
if "document_content" not in st.session_state:
    st.session_state.document_content = ""
if "document_summary" not in st.session_state:
    st.session_state.document_summary = ""
if "data_protection_rules" not in st.session_state:
    st.session_state.data_protection_rules = ""
if "genai_client" not in st.session_state:
    st.session_state.genai_client = None
if "common_prefix" not in st.session_state:
    st.session_state.common_prefix = ""


def setup_gemini():
    """Setup Gemini API configuration with 2.5 Flash for implicit caching"""
    # Try to load API key from environment first
    api_key = os.getenv("GEMINI_API_KEY")

    if not api_key or api_key == "your_gemini_api_key_here":
        # Fallback to Streamlit input if not found in .env
        api_key = st.sidebar.text_input(
            "Enter your Gemini API Key:",
            type="password",
            help="API key not found in .env file. Please enter manually or update your .env file.",
        )
    else:
        st.sidebar.success("✅ API Key loaded from .env file")

    if api_key and api_key != "your_gemini_api_key_here":
        try:
            # Initialize the genai client with Gemini 2.5 Flash
            client = genai.Client(api_key=api_key)
            st.session_state.genai_client = client
            return client
        except Exception as e:
            st.sidebar.error(f"Error initializing Gemini client: {str(e)}")
            return None
    return None


def count_tokens(client, text: str) -> int:
    """Count tokens in text using Gemini API"""
    try:
        response = client.models.count_tokens(model=MODEL_NAME, contents=[text])
        return response.total_tokens
    except Exception as e:
        # Fallback estimation: roughly 4 characters per token
        return len(text) // 4


def create_common_prefix(document_content: str, data_protection_rules: str) -> str:
    """Create a common prefix for implicit caching optimization"""
    prefix = f"""BỐI CẢNH TÀI LIỆU VÀ QUY ĐỊNH:

TÀI LIỆU CHÍNH CẦN PHÂN TÍCH:
{document_content}

---

NGHỊ ĐỊNH BẢO VỆ DỮ LIỆU CÁ NHÂN - NGHỊ ĐỊNH 13/2023/NĐ-CP:
{data_protection_rules}

---

Bạn là một trợ lý AI chuyên nghiệp trong việc phân tích tài liệu và kiểm tra tuân thủ pháp luật. Hãy sử dụng toàn bộ thông tin trên để trả lời các câu hỏi một cách chính xác và chi tiết.

"""
    return prefix


def generate_with_implicit_cache(client, prefix: str, prompt: str, model_name: str = MODEL_NAME):
    """Generate content optimized for implicit caching by placing large content at the beginning"""
    try:
        # Combine prefix (large common content) with specific prompt
        full_prompt = f"{prefix}\n{prompt}"

        response = client.models.generate_content(model=model_name, contents=[full_prompt])

        # Display cache hit information if available
        if hasattr(response, "usage_metadata") and response.usage_metadata:
            total_tokens = getattr(response.usage_metadata, "total_token_count", 0) or 0
            prompt_tokens = getattr(response.usage_metadata, "prompt_token_count", 0) or 0
            candidates_tokens = getattr(response.usage_metadata, "candidates_token_count", 0) or 0
            cached_tokens = getattr(response.usage_metadata, "cached_content_token_count", 0) or 0

            if cached_tokens > 0:
                st.info(f"💰 Implicit Cache Hit: {cached_tokens:,} tokens cached | Total: {total_tokens:,} tokens")
            elif total_tokens > 0:
                # Show token usage even without cache hits
                st.info(
                    f"📊 Token Usage: {prompt_tokens:,} prompt + {candidates_tokens:,} response = {total_tokens:,} total"
                )

        return response.text

    except Exception as e:
        st.error(f"Error generating content: {str(e)}")
        return f"Lỗi khi tạo nội dung: {str(e)}. Vui lòng thử lại."


def load_data_protection_rules():
    """Load data protection rules from local file"""
    with open("documents/13_2023_ND-CP_465185.txt", "r", encoding="utf-8") as f:
        text = f.read()
    return text


def extract_text_from_file(uploaded_file):
    """Extract text from uploaded file (PDF, MARKDOWN, DOC/DOCX, or TXT)"""
    try:
        # Reset file pointer to beginning
        uploaded_file.seek(0)

        file_extension = uploaded_file.name.split(".")[-1].lower()

        if file_extension == "pdf":
            pdf_reader = PdfReader(uploaded_file)
            text = ""
            for page in pdf_reader.pages:
                text += page.extract_text()

        elif file_extension in ("doc", "docx"):
            # Save uploaded file temporarily for Spire.Doc
            import tempfile

            with tempfile.NamedTemporaryFile(delete=False, suffix=f".{file_extension}") as tmp_file:
                # Write the uploaded file content to temp file
                tmp_file.write(uploaded_file.read())
                tmp_file.flush()

                # Now load from the temporary file path
                doc = Document()
                doc.LoadFromFile(tmp_file.name)
                text = (
                    doc.GetText()
                    .replace("Evaluation Warning: The document was created with Spire.Doc for Python.", "")
                    .strip()
                )

                # Clean up temp file
                os.unlink(tmp_file.name)

        elif file_extension in ("md", "txt"):
            text = str(uploaded_file.read(), "utf-8")

        else:
            st.error(f"Unsupported file format: {file_extension}")
            return None

        return text

    except Exception as e:
        st.error(f"Error extracting text from file: {str(e)}")
        return None


def summarize_document(client, prefix: str) -> str:
    """Generate summary of the document using implicit caching"""
    prompt = """
    YÊU CẦU TÓM TẮT TÀI LIỆU CHÍNH CẦN PHÂN TÍCH (KHÔNG BAO GỒM NGHỊ ĐỊNH BẢO VỆ DỮ LIỆU):

    Hãy tóm tắt nội dung tài liệu một cách chi tiết và có cấu trúc:

    1. **Tóm tắt tổng quát** (2-3 câu)
    2. **Các phần chính trong tài liệu** với tiêu đề và nội dung chính
    3. **Những điểm quan trọng cần lưu ý** 
    4. **Cấu trúc và tổ chức nội dung**
    5. **Các khái niệm và thuật ngữ chính**
    6. **Kết luận và khuyến nghị** (nếu có)

    Vui lòng phân tích toàn bộ tài liệu để đưa ra tóm tắt đầy đủ, chính xác và có cấu trúc rõ ràng.
    """

    return generate_with_implicit_cache(client, prefix, prompt)


def check_data_protection_compliance(client, prefix: str) -> str:
    """Check compliance using implicit caching with common prefix"""
    prompt = """
    YÊU CẦU KIỂM TRA TUÂN THỦ:

    Hãy thực hiện phân tích tuân thủ bảo vệ dữ liệu cá nhân một cách chi tiết và có hệ thống:

    **PHÂN TÍCH TUÂN THỦ NGHỊ ĐỊNH 13/2023/NĐ-CP**

    Thực hiện so sánh giữa nội dung tài liệu với các quy định bảo vệ dữ liệu cá nhân:

    1. **Đánh giá tổng quan**
       - Có vi phạm các quy định hay không?
       - Mức độ nghiêm trọng (Cao/Trung bình/Thấp)

    2. **Phân tích chi tiết các vi phạm (nếu có)**
       - Điều khoản cụ thể bị vi phạm
       - Nội dung vi phạm trong tài liệu
       - Trích dẫn quy định tương ứng

    3. **Các vấn đề cần lưu ý**
       - Thu thập dữ liệu không đúng quy định
       - Thiếu thông báo hoặc đồng ý
       - Vấn đề bảo mật và lưu trữ
       - Quyền của chủ thể dữ liệu

    4. **Khuyến nghị khắc phục**
       - Các biện pháp cụ thể cần thực hiện
       - Thứ tự ưu tiên xử lý
       - Biện pháp phòng ngừa

    5. **Kết luận và đánh giá rủi ro**

    Phân tích toàn bộ nội dung tài liệu một cách kỹ lưỡng, không bỏ sót bất kỳ phần nào.
    """

    return generate_with_implicit_cache(client, prefix, prompt)


def verify_information(client, prefix: str, user_claim: str) -> str:
    """Verify information using implicit caching"""
    prompt = f"""
    YÊU CẦU XÁC MINH THÔNG TIN:

    **Thông tin cần kiểm tra:** "{user_claim}"

    **Phân tích và đánh giá:**

    1. **Kết luận:** ĐÚNG / SAI / KHÔNG XÁC ĐỊNH / MỘT PHẦN ĐÚNG

    2. **Phân tích chi tiết:**
       - Bằng chứng hỗ trợ từ tài liệu
       - Trích dẫn chính xác các phần liên quan
       - Vị trí xuất hiện trong tài liệu

    3. **Đánh giá độ tin cậy:**
       - Mức độ tin cậy: Cao / Trung bình / Thấp
       - Lý do đánh giá

    4. **Thông tin bổ sung:**
       - Ngữ cảnh liên quan
       - Thông tin mâu thuẫn (nếu có)
       - Giải thích chi tiết

    5. **Kết luận cuối cùng:**
       - Tóm tắt kết quả xác minh
       - Khuyến nghị về việc sử dụng thông tin này

    Vui lòng tìm kiếm toàn bộ tài liệu để đưa ra kết luận chính xác và đáng tin cậy nhất.
    """

    return generate_with_implicit_cache(client, prefix, prompt)


def chat_with_document(client, prefix: str, question: str) -> str:
    """Answer questions using implicit caching"""
    prompt = f"""
    YÊU CẦU TRẢ LỜI CÂU HỎI:

    **Câu hỏi:** {question}

    **Cấu trúc trả lời:**

    1. **Trả lời trực tiếp:**
       - Câu trả lời ngắn gọn, rõ ràng

    2. **Thông tin chi tiết:**
       - Giải thích đầy đủ và bối cảnh
       - Trích dẫn các phần liên quan trong tài liệu

    3. **Phân tích bổ sung:**
       - Các góc nhìn khác nhau (nếu có)
       - Thông tin liên quan hữu ích

    4. **Kết luận và khuyến nghị:**
       - Tóm tắt câu trả lời
       - Đề xuất hành động (nếu phù hợp)

    **Lưu ý:** Sử dụng toàn bộ thông tin có sẵn trong tài liệu để đưa ra câu trả lời hoàn chỉnh, chính xác và hữu ích nhất.
    """

    return generate_with_implicit_cache(client, prefix, prompt)


# Main app
def main():
    st.title("📄 Document Analysis Chatbot")
    st.markdown("*Powered by Gemini 2.5 Flash with Automatic Cost Optimization*")
    # st.markdown("---")

    # Load data protection rules automatically on startup
    if not st.session_state.data_protection_rules:
        with st.spinner("Đang tải quy định bảo vệ dữ liệu..."):
            st.session_state.data_protection_rules = load_data_protection_rules()

    # Setup Gemini client
    client = setup_gemini()

    if not client:
        st.warning("⚠️ Vui lòng cấu hình Gemini API Key để bắt đầu!")
        st.sidebar.markdown(
            """
        ### 🔧 Cấu hình API Key:
        
        **Cách 1 (Khuyến nghị):** Tạo file `.env`
        ```
        GEMINI_API_KEY=your_api_key_here
        ```
        
        **Cách 2:** Nhập trực tiếp ở sidebar
        
        ### 📝 Lấy API Key:
        1. Truy cập [Google AI Studio](https://aistudio.google.com/app/apikey)
        2. Tạo API Key mới
        3. Sao chép và cấu hình
        """
        )
        return

    # File upload section
    st.sidebar.header("📁 Upload Tài liệu")
    uploaded_file = st.sidebar.file_uploader(
        "Chọn tài liệu cần phân tích",
        type=["pdf", "doc", "docx", "md", "txt"],
        help="Hỗ trợ định dạng PDF, DOC/DOCX, MARKDOWN, TXT",
    )

    # Show data protection status
    st.sidebar.header("⚖️ Quy định bảo vệ dữ liệu")
    if "File quy định gốc không được tìm thấy" in st.session_state.data_protection_rules:
        st.sidebar.warning("⚠️ Sử dụng quy định mặc định")
        st.sidebar.info("Đặt file '13_2023_ND-CP.doc' trong thư mục gốc để tải quy định đầy đủ")
    else:
        st.sidebar.success("✅ Đã tải Nghị định 13/2023/NĐ-CP")

    # Process uploaded files and create common prefix for implicit caching
    if uploaded_file:
        with st.spinner("Đang xử lý tài liệu..."):
            content = extract_text_from_file(uploaded_file)
            if content:
                st.session_state.document_content = content
                st.sidebar.success(f"✅ Đã tải tài liệu: {uploaded_file.name}")

                # Create common prefix for implicit caching optimization
                st.session_state.common_prefix = create_common_prefix(
                    st.session_state.document_content, st.session_state.data_protection_rules
                )

                # Show token count for optimization insight
                total_tokens = count_tokens(client, st.session_state.common_prefix)
                if total_tokens >= 1024:
                    st.sidebar.success(f"🚀 Optimized for implicit caching ({total_tokens:,} tokens)")
                else:
                    st.sidebar.info(f"📊 Document loaded ({total_tokens:,} tokens)")

    # Main interface
    if st.session_state.document_content:
        # Create tabs for different features
        tab1, tab2, tab3, tab4 = st.tabs(["💬 Chat", "📋 Tóm tắt", "✅ Xác minh", "⚖️ Kiểm tra tuân thủ"])

        with tab1:
            st.header("💬 Đặt câu hỏi về tài liệu")

            # Create a container for chat messages
            prompt = st.chat_input("Đặt câu hỏi về tài liệu...")

            with st.container(height=350):
                # Display chat history
                for message in st.session_state.messages:
                    with st.chat_message(message["role"]):
                        st.markdown(message["content"])

                # Chat input
                if prompt:
                    # Add user message
                    st.session_state.messages.append({"role": "user", "content": prompt})
                    with st.chat_message("user"):
                        st.markdown(prompt)

                    # Get AI response using implicit caching
                    with st.chat_message("assistant"):
                        with st.spinner("Đang suy nghĩ..."):
                            response = chat_with_document(client, st.session_state.common_prefix, prompt)
                            st.markdown(response)
                            st.session_state.messages.append({"role": "assistant", "content": response})

                # Force rerun to show new messages and auto-scroll

        with tab2:
            # # Auto-generate summary
            # if not st.session_state.document_summary:
            #     with st.spinner("Tạo tóm tắt với implicit caching..."):
            #         st.session_state.document_summary = summarize_document(client, st.session_state.common_prefix)
            # else:
            st.header("📋 Tóm tắt tài liệu")
            if (
                st.session_state.document_summary
                and "Error processing (Assignment 2 Final Project LLMs02): " not in st.session_state.document_summary
            ):
                st.markdown(st.session_state.document_summary)
            else:
                if st.button("Tạo tóm tắt tài liệu"):
                    with st.spinner("Đang tóm tắt tài liệu với implicit caching..."):
                        summary = summarize_document(client, st.session_state.document_content)
                        st.session_state.document_summary = summary
                        st.markdown(summary)

        with tab3:
            st.header("✅ Xác minh thông tin")
            claim = st.text_area("Nhập thông tin cần xác minh:", height=100)

            if st.button("Kiểm tra thông tin"):
                if claim:
                    with st.spinner("Đang xác minh với implicit caching..."):
                        verification = verify_information(client, st.session_state.common_prefix, claim)
                        st.markdown("### Kết quả xác minh:")
                        st.markdown(verification)
                else:
                    st.warning("Vui lòng nhập thông tin cần xác minh!")

        with tab4:
            st.header("⚖️ Kiểm tra tuân thủ bảo vệ dữ liệu")
            if st.button("Kiểm tra vi phạm"):
                with st.spinner("Đang phân tích tuân thủ với implicit caching..."):
                    compliance = check_data_protection_compliance(client, st.session_state.common_prefix)
                    st.markdown("### Kết quả kiểm tra:")
                    st.markdown(compliance)

    else:
        st.info("👆 Vui lòng upload tài liệu trong sidebar để bắt đầu!")

        # Show implicit caching benefits
        st.markdown(
            """
        ### 🚀 Tính năng Implicit Caching tự động với Gemini 2.5 Flash:
        
        #### 💰 **Tiết kiệm chi phí tự động**
        - **Không cần cấu hình**: Implicit caching được bật mặc định
        - **Tiết kiệm thông minh**: Tự động phát hiện và tái sử dụng nội dung chung
        - **Chi phí giảm**: Chỉ tính phí cho token mới, không tính lại token đã cache
        
        #### ⚡ **Tối ưu hóa hiệu suất**
        - **Nội dung lớn ở đầu**: Đặt tài liệu và quy định ở vị trí đầu tiên
        - **Prefix chung**: Sử dụng cùng một prefix cho tất cả yêu cầu
        - **Batching thông minh**: Gửi các yêu cầu tương tự trong thời gian ngắn
        
        #### 📊 **Theo dõi hiệu quả**
        - **Token tracking**: Hiển thị số token được cache hit
        - **Cost monitoring**: Theo dõi tiết kiệm chi phí thực tế
        - **Performance metrics**: Đo lường tốc độ và hiệu quả
        
        #### 🎯 **Tối ưu cho phân tích tài liệu**
        - **Full document processing**: Xử lý toàn bộ tài liệu không giới hạn
        - **Consistent context**: Duy trì bối cảnh nhất quán qua các truy vấn
        - **Smart prefixing**: Tự động tối ưu thứ tự nội dung
        """
        )


if __name__ == "__main__":
    main()