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()