Clean push without binary history

.gitattributes

@@ -0,0 +1,2 @@
+fonts/dejavu-sans/*.ttf filter=lfs diff=lfs merge=lfs -text
+fonts/dejavu-sans/*.pkl filter=lfs diff=lfs merge=lfs -text

.github/workflows/auto_sync.yml

@@ -0,0 +1,20 @@
+name: Sync to Hugging Face hub
+on:
+  push:
+    branches: [main]
+  workflow_dispatch:
+
+jobs:
+  sync-to-hub:
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v3
+        with:
+          fetch-depth: 0
+          lfs: true
+      - name: Push to hub
+        env:
+          HF_TOKEN: ${{ secrets.HF_TOKEN }}
+        run: git push --force https://mubi0613:$HF_TOKEN@huggingface.co/spaces/mubi0613/SecureVault-AI main
+
+        ----------------------------------------------------------------------------------------------------------------

.gitignore

@@ -0,0 +1,16 @@
+# Local Vault Data 
+notes.json
+vault_config.json
+usage_stats.json
+feedback_log.json
+.env
+
+# Python Environment
+__pycache__/
+*.py[cod]
+*$py.class
+venv/
+.env
+
+# Streamlit logs
+.streamlit/

LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2026 mubi0613
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

README.md

@@ -0,0 +1,69 @@
+# 🛡️ SecureVault AI: Privacy-First RAG Knowledge Base
+
+A production-grade, secure note-taking application that integrates **Retrieval-Augmented Generation (RAG)** with **Zero-Knowledge Encryption**. Built for users who need AI-powered insights without sacrificing data privacy.
+
+## 🚀 Key Engineering Features
+* **Hybrid Security Model:** Uses **PBKDF2** key derivation and **AES-256 (Fernet)** encryption. Data is only decrypted in-memory during an active session.
+* **Local RAG Engine:** Utilizes **FAISS** (Facebook AI Similarity Search) and `all-MiniLM-L6-v2` embeddings for lightning-fast, local semantic search—no private data ever leaves your machine.
+* **AI Observability:** Integrated **Token & Cost Tracker** and a **Feedback Loop** (Correct/Wrong logging) to audit retrieval accuracy and identify gaps in the knowledge base.
+* **Fail-Safe Recovery:** Implements a hashed **Recovery Key** system for account resets while maintaining the mathematical integrity of encrypted data.
+* **Smart Exports:** Unicode-compliant PDF and DOCX generation with specialized font embedding for multi-language support.
+
+## 🛠️ Technical Stack
+* **Frontend:** Streamlit (Custom CSS & Session Management)
+* **AI/NLP:** Sentence-Transformers, TextBlob, FAISS
+* **Cryptography:** Python `cryptography` library (SHA-256 & PBKDF2)
+* **Data:** JSON-based persistent storage
+
+---
+
+## 🔐 Security Architecture
+1. **PIN Hashing:** The Master PIN is never stored; only its **SHA-256** hash is kept for authentication.
+2. **Zero-Knowledge:** The PIN is used to derive the encryption key. If the PIN is lost, "Secret" notes remain mathematically inaccessible.
+3. **Auto-Lock Engine:** A session-based activity monitor automatically wipes the memory and locks the vault after 150 seconds of inactivity.
+
+
+
+---
+
+## 📦 Installation & Setup
+
+### 1. Clone the Repository
+```bash
+git clone [https://github.com/mubi0613/SecureVault-AI.git](https://github.com/mubi0613/SecureVault-AI.git)
+cd SecureVault-AI
+
+2. Install Dependencies
+Ensure you have Python 3.9+ installed, then run:
+```bash
+pip install -r requirements.txt
+
+3. Run the Application
+```bash
+streamlit run main.py
+
+📖 Usage Guide
+Initialize: Set your Master PIN and save your Recovery Key.
+
+Create Notes: Use the sidebar to add notes. Toggle "Mark as Secret" to apply AES-256 encryption.
+
+AI Search: Open the "Ask Your Vault" expander to query your notes using natural language.
+
+Manage: Edit or delete notes, and export them as PDF/DOCX for offline use.
+
+Monitor: Check the "AI Resources" section in the sidebar to view estimated token usage and costs.
+
+📂 Project Structure
+Plaintext
+SecureVault-AI/
+├── main.py            # Streamlit UI & Session Management
+├── vault_logic.py     # Cryptography, RAG, & File IO Logic
+├── requirements.txt   # Project Dependencies
+├── .gitignore         
+├── fonts/             # Custom fonts for PDF Unicode support
+└── README.md          # Documentation
+⚠️ Important Note on Security
+This application is designed with a Zero-Knowledge philosophy. The developers cannot recover your notes if you lose both your Master PIN and your Recovery Key. Data integrity is the sole responsibility of the user.
+
+📜 License
+Distributed under the MIT License. See LICENSE for more information.

main.py

@@ -0,0 +1,269 @@
+import streamlit as st
+from datetime import datetime
+import vault_logic as vl
+import time
+import json
+import os
+
+# --- CONFIGURATION ---
+AUTO_LOCK_SECONDS = 150  # 2 Minutes
+
+st.set_page_config(page_title="AI Vault", layout="wide")
+
+# --- 0. FIRST TIME SETUP ---
+if not vl.is_vault_initialized():
+    st.title("🔒 Initialize Your Secure Vault")
+    
+    # 1. Primary Action: The PIN Box
+    st.info("👋 Welcome! To get started, please set your Master PIN below.")
+    
+    with st.form("setup_form"):
+        new_pin = st.text_input("Create Master PIN", type="password", help="Minimum 4 digits")
+        conf_pin = st.text_input("Confirm Master PIN", type="password")
+        submit = st.form_submit_button("Setup Vault", type="primary")
+    
+    # 2. Professional Spacing
+    st.write("")
+    st.write("") 
+    st.divider()
+    
+    # 3. Security & Recovery Information
+    if 'generated_recovery' not in st.session_state:
+        st.session_state.generated_recovery = vl.generate_recovery_key()
+
+    st.warning("⚠️ **CRITICAL: Save Your Recovery Key**")
+    st.caption("RECOVERY - KEY : This is the ONLY way to reset your vault access if you forget your PIN.")
+    st.code(st.session_state.generated_recovery, language=None)
+    
+    # Important Security Note (Liability Waiver)
+    st.markdown("""
+        <div style="background-color: #ff4b4b11; padding: 20px; border-radius: 10px; border: 1px solid #ff4b4b;">
+            <p style="color: #ff4b4b; font-weight: bold; margin-bottom: 5px;">🛑 Important Security Notice:</p>
+            <p style="font-size: 0.9rem; margin-bottom: 10px;">
+                This vault utilizes <b>Zero-Knowledge AES-Encryption</b>. Your PIN is never stored; it is only hashed. 
+                While the Recovery Key can reset app access, it <b>cannot decrypt</b> existing secret data 
+                locked with a different PIN.
+            </p>
+            <p style="font-size: 0.8rem; opacity: 0.8;">
+                <i>The developer holds no responsibility for data loss resulting from forgotten credentials.</i>
+            </p>
+        </div>
+    """, unsafe_allow_html=True)
+            
+    if submit:        
+        if new_pin == conf_pin and len(new_pin) >= 4:
+            vl.initialize_vault(new_pin, st.session_state.generated_recovery)
+            st.success("Vault Securely Initialized! Refreshing...")
+            time.sleep(1.5)
+            st.rerun()
+        else:
+            st.error("PINs must match and be at least 4 digits.")
+    st.stop()
+
+# --- 1. SESSION STATE ---
+if 'notes' not in st.session_state:
+    st.session_state.notes = vl.load_notes()
+if 'edit_note_id' not in st.session_state:
+    st.session_state.edit_note_id = None
+if 'vault_unlocked' not in st.session_state:
+    st.session_state.vault_unlocked = False
+if 'master_pin' not in st.session_state:
+    st.session_state.master_pin = "" # We store the actual PIN in session only while unlocked
+if 'temp_content' not in st.session_state:
+    st.session_state.temp_content = ""
+if 'form_iteration' not in st.session_state:
+    st.session_state.form_iteration = 0
+if 'last_activity' not in st.session_state:
+    st.session_state.last_activity = time.time()
+if 'show_lock_alert' not in st.session_state:
+    st.session_state.show_lock_alert = False
+
+# --- 2. AUTO-LOCK ENGINE ---
+def update_activity():
+    st.session_state.last_activity = time.time()
+    st.session_state.show_lock_alert = False 
+
+if st.session_state.vault_unlocked:
+    elapsed_time = time.time() - st.session_state.last_activity
+    if elapsed_time > AUTO_LOCK_SECONDS:
+        st.session_state.vault_unlocked = False
+        st.session_state.edit_note_id = None
+        st.session_state.temp_content = "" 
+        st.session_state.show_lock_alert = True 
+        st.rerun() 
+
+# --- 3. CSS & ALERTS ---
+st.markdown("""
+    <style>
+        .block-container {padding-top: 2rem !important;}
+        [data-testid="stSidebarNav"] {display: none;}
+    </style>
+""", unsafe_allow_html=True)
+
+if st.session_state.show_lock_alert:
+    st.error("🚨 **THE VAULT IS AUTO LOCKED**")
+
+# --- 4. SIDEBAR ---
+with st.sidebar:
+    st.title("➕ Add/Edit Note")
+    
+    c_title, c_content, c_secret = "", "", False
+    if st.session_state.edit_note_id:
+        n = next((x for x in st.session_state.notes if x['id'] == st.session_state.edit_note_id), None)
+        if n: 
+            c_title = n['title']
+            c_content = vl.decrypt_data(n['content'], st.session_state.master_pin) if (n.get('secret') and st.session_state.vault_unlocked) else n['content']
+            c_secret = n.get('secret', False)
+
+    default_content = st.session_state.temp_content if st.session_state.temp_content else c_content
+
+    with st.form(key=f"note_form_{st.session_state.form_iteration}", clear_on_submit=False):
+        new_t = st.text_input("Title", value=c_title)
+        new_c = st.text_area("Content", value=default_content, height=150)
+        m_secret = st.checkbox("🤫 Mark as Secret", value=c_secret) if st.session_state.vault_unlocked else False
+        
+        if st.form_submit_button("✨ AI Summarize"):
+            update_activity()
+            if new_c:
+                summary_result = vl.ai_summarize_text(new_c)
+                st.session_state.temp_content = summary_result
+                st.rerun() 
+        
+        if st.form_submit_button("Save", type="primary", use_container_width=True):
+            update_activity()
+            if new_t or new_c:
+                ts = datetime.now().strftime("%Y-%m-%d %H:%M")
+                final_content = vl.encrypt_data(new_c, st.session_state.master_pin) if m_secret else new_c
+                if st.session_state.edit_note_id:
+                    for n in st.session_state.notes:
+                        if n['id'] == st.session_state.edit_note_id:
+                            n.update({"title": new_t, "content": final_content, "timestamp": ts, "secret": m_secret})
+                    st.session_state.edit_note_id = None
+                else:
+                    st.session_state.notes.insert(0, {"id": int(datetime.now().timestamp()), "title": new_t, "content": final_content, "timestamp": ts, "secret": m_secret})
+                
+                vl.save_notes(st.session_state.notes)
+                st.session_state.temp_content = "" 
+                st.session_state.form_iteration += 1 
+                st.rerun()
+
+    st.divider()
+    st.write("### 🔐 Vault Security")
+    if not st.session_state.vault_unlocked:
+        pin_input = st.text_input("Enter Stealth PIN", type="password", key="pin_entry")
+        if pin_input:
+            if vl.verify_pin(pin_input):
+                update_activity()
+                st.session_state.vault_unlocked = True
+                st.session_state.master_pin = pin_input
+                st.rerun()
+            else: st.error("Incorrect PIN")
+
+        # NEW: Recovery Logic
+        with st.expander("Forgot PIN?"):
+            recovery_in = st.text_input("Enter Recovery Key")
+            if st.button("Reset Vault PIN"):
+                # Logic: If recovery key matches, delete config so they can start over
+                # Note: Old encrypted notes stay encrypted (this is the cost of security!)
+                if vl.verify_recovery_key(recovery_in): # You'll add this to vault_logic
+                    os.remove("vault_config.json")
+                    st.warning("Vault Identity Reset. Please refresh to set a new PIN.")
+                    st.rerun()
+
+    else:
+        st.success("🔓 Vault is Open")
+        time_left = int(AUTO_LOCK_SECONDS - (time.time() - st.session_state.last_activity))
+        st.caption(f"Auto-locking in {max(0, time_left)}s")
+        if st.button("🔒 Close Vault", use_container_width=True):
+            st.session_state.vault_unlocked = False
+            st.session_state.master_pin = ""
+            st.rerun()
+
+    st.divider()
+    st.subheader("📊 AI Resources")
+    if os.path.exists("usage_stats.json"):
+        with open("usage_stats.json", "r") as f:
+            stats = json.load(f)
+        col1, col2 = st.columns(2)
+        col1.metric("Tokens", f"{int(stats['total_tokens'])}")
+        col2.metric("Cost", f"${stats['total_cost']:.5f}")
+    else:
+        st.info("No AI usage data yet.")
+
+# --- 5. MAIN PAGE ---
+st.title(vl.get_page_heading(st.session_state.vault_unlocked))
+
+# --- 6. RAG CHAT INTERFACE ---
+if st.session_state.vault_unlocked:
+    with st.expander("💬 Ask Your Vault (AI Search)", expanded=False):
+        user_query = st.text_input("Ask a question about your notes:", placeholder="e.g., What are my goals for 2026?")
+        
+        if user_query:
+            update_activity()
+            # 1. Prepare Decrypted Context
+            decrypted_texts = []
+            for n in st.session_state.notes:
+                content = vl.decrypt_data(n['content'], st.session_state.master_pin) if n.get('secret') else n['content']
+                decrypted_texts.append(content)
+            
+            # 2. RAG Logic
+            index, text_data = vl.create_vector_index([{"content": t} for t in decrypted_texts])
+            context_chunks = vl.query_vault(user_query, index, text_data)
+            
+            # 3. AI Answer
+            context_str = "\n".join(context_chunks)
+            prompt = f"Using ONLY these notes:\n{context_str}\n\nAnswer this question: {user_query}"
+            
+            with st.spinner("Searching Vault..."):
+                answer = vl.get_gemini_response(user_query, context_str)
+                st.info(f"**AI Answer:**\n{answer}")
+
+                # 4. Feedback System ---
+                st.write("---")
+                st.write("Was this answer helpful?")
+                f_col1, f_col2 = st.columns([1, 5])
+                
+                if f_col1.button("✅ Yes", key="fb_yes"):
+                    vl.log_feedback(user_query, answer, context_str, "Correct")
+                    st.success("Thanks for the feedback!")
+                    
+                if f_col2.button("❌ No / Wrong", key="fb_no"):
+                    vl.log_feedback(user_query, answer, context_str, "Wrong")
+                    st.warning("Logged as a failure.")
+
+st.divider()
+search = st.text_input("🔍 Search...", placeholder="Filter notes...")
+
+# --- 7. Display Grid ---
+cols = st.columns(3)
+filtered = vl.get_filtered_notes(st.session_state.notes, st.session_state.vault_unlocked, search)
+
+for idx, note in enumerate(filtered):
+    with cols[idx % 3]: 
+        with st.container(border=True):
+            display_content = vl.decrypt_data(note['content'], st.session_state.master_pin) if (note.get('secret') and st.session_state.vault_unlocked) else note['content']
+            
+            st.subheader(f"🔒 {note['title']}" if note.get('secret') else note['title'])
+            st.write(display_content[:200] + "..." if len(display_content) > 200 else display_content)
+            st.caption(f"🕒 {note['timestamp']}")
+            
+            eb, db = st.columns(2)
+            with eb:
+                if st.button("✏️ Edit", key=f"e_{note['id']}"):
+                    update_activity()
+                    st.session_state.edit_note_id = note['id']
+                    st.rerun()
+            with db:
+                if st.button("🗑️ Delete", key=f"d_{note['id']}"):
+                    st.session_state.notes = [x for x in st.session_state.notes if x['id'] != note['id']]
+                    vl.save_notes(st.session_state.notes)
+                    st.rerun()
+
+            st.divider()
+            p_col, d_col = st.columns(2)
+            with p_col:
+                pdf_bytes = vl.create_pdf(note['title'], display_content)
+                st.download_button("📄 Download PDF", data=pdf_bytes, file_name=f"{note['title']}.pdf", key=f"pdf_{note['id']}", use_container_width=True)
+            with d_col:
+                docx_bytes = vl.create_docx(note['title'], display_content)
+                st.download_button("📝 Download DOCX", data=docx_bytes, file_name=f"{note['title']}.docx", key=f"docx_{note['id']}", use_container_width=True)

requirements.txt

@@ -0,0 +1,11 @@
+streamlit>=1.42.0
+cryptography
+textblob
+fpdf2
+python-docx
+faiss-cpu
+numpy
+sentence-transformers
+google-genai
+python-dotenv
+nltk

simple_notes.json

@@ -0,0 +1 @@
+[]

vault_logic.py

@@ -0,0 +1,290 @@
+import json
+import os
+import base64
+import hashlib
+import time
+from datetime import datetime
+from cryptography.fernet import Fernet
+from cryptography.hazmat.primitives import hashes
+from cryptography.hazmat.primitives.kdf.pbkdf2 import PBKDF2HMAC
+from textblob import TextBlob
+from fpdf import FPDF
+from docx import Document
+import io
+import faiss
+import numpy as np
+from sentence_transformers import SentenceTransformer
+import secrets
+import string
+from google import genai  # Gemini
+from dotenv import load_dotenv
+import nltk
+
+try:
+    nltk.data.find('tokenizers/punkt_tab')
+except LookupError:
+    nltk.download('punkt_tab')
+
+# This line reads the .env file and makes the API key available
+load_dotenv(override=True)  # 'override=True' forces it to refresh the key if you changed .env
+
+NOTES_FILE = "notes.json"
+CONFIG_FILE = "vault_config.json"
+
+# Load a small, fast model for embeddings (runs locally)
+embed_model = SentenceTransformer('all-MiniLM-L6-v2')
+
+# --- 0. PIN & KEY LOGIC ---
+def get_pin_hash(pin: str):
+    """Creates a secure SHA-256 hash of the PIN"""
+    return hashlib.sha256(pin.encode()).hexdigest()
+
+def is_vault_initialized():
+    """Checks if a user has set a PIN yet"""
+    return os.path.exists(CONFIG_FILE)
+
+def generate_recovery_key():
+    """Generates a random 16-character recovery string"""
+    chars = string.ascii_uppercase + string.digits
+    return ''.join(secrets.choice(chars) for _ in range(16))
+
+def initialize_vault(pin: str, recovery_key: str):
+    """Saves the user's initial PIN hash and recovery key hash"""
+    with open(CONFIG_FILE, "w") as f:
+        json.dump({"pin_hash": get_pin_hash(pin),
+                   "recovery_hash": get_pin_hash(recovery_key) # We hash the recovery key too!
+                   }, f)
+
+def verify_recovery_key(input_key: str):
+    """Verifies the recovery key against the stored hash"""
+    if not is_vault_initialized():
+        return False
+    with open(CONFIG_FILE, "r") as f:
+        stored_recovery_hash = json.load(f).get("recovery_hash")
+    return get_pin_hash(input_key) == stored_recovery_hash
+
+def verify_pin(input_pin: str):
+    """Verifies input against stored hash"""
+    if not is_vault_initialized():
+        return False
+    with open(CONFIG_FILE, "r") as f:
+        stored_hash = json.load(f).get("pin_hash")
+    return get_pin_hash(input_pin) == stored_hash
+
+# --- 1. KEY GENERATION (From your PIN) ---
+def generate_key(pin: str):
+    """Derives a functional encryption key from the user's PIN"""
+    password = pin.encode()
+    # We use a static salt for this local project (Professional apps use random salts)
+    salt = b'stable_salt_123' 
+    kdf = PBKDF2HMAC(
+        algorithm=hashes.SHA256(),
+        length=32,
+        salt=salt,
+        iterations=100000,
+    )
+    key = base64.urlsafe_b64encode(kdf.derive(password))
+    return Fernet(key)
+
+# --- 2. ENCRYPTION / DECRYPTION ---
+def encrypt_data(data_string, pin):
+    """Turns readable text into scrambled code"""
+    f = generate_key(pin)
+    return f.encrypt(data_string.encode()).decode()
+
+def decrypt_data(encrypted_string, pin):
+    """Turns scrambled code back into readable text"""
+    try:
+        f = generate_key(pin)
+        return f.decrypt(encrypted_string.encode()).decode()
+    except Exception:
+        return "[Decryption Error: Check PIN]"
+
+# --- 3. UPDATED LOAD/SAVE ---
+def load_notes():
+    if os.path.exists(NOTES_FILE):
+        with open(NOTES_FILE, "r") as f:
+            return json.load(f)
+    return []
+
+def save_notes(notes_list):
+    with open(NOTES_FILE, "w") as f:
+        json.dump(notes_list, f, indent=4)
+
+def get_page_heading(is_unlocked):
+    return "🛡️ Safe Vault" if is_unlocked else "📝 My Notes"
+
+def get_filtered_notes(all_notes, is_unlocked, search_query):
+    visible = all_notes
+    if not is_unlocked:
+        visible = [n for n in all_notes if not n.get('secret', False)]
+    
+    query = search_query.lower()
+    return [n for n in visible if query in n['title'].lower() or query in n['content'].lower()]
+
+# --- 4. AI FEATURES ---
+def ai_summarize_text(text):
+    """Uses NLP to extract key points from long notes."""
+    if len(text) < 50:
+        return text # Don't summarize very short notes
+    
+    blob = TextBlob(text)
+    # Extracting sentences and picking the top 2 for a summary
+    sentences = [str(s) for s in blob.sentences]
+    
+    if len(sentences) > 2:
+        summary = "✨ AI Summary:\n" + "\n".join([f"• {s}" for s in sentences[:2]])
+        
+        # NEW: Log the usage here!
+        track_usage(text, type="input") 
+        track_usage(summary, type="output")
+        return summary
+    return text
+
+# --- 5. EXPORTS ---
+# The Logic Functions : These handle the actual file creation.
+def create_pdf(title, content):
+    try:
+        pdf = FPDF()
+        pdf.add_page()
+
+        # Updated to match your folder: fonts -> dejavu-sans -> DejaVuSans.ttf
+        font_path = os.path.join("fonts", "dejavu-sans", "DejaVuSans.ttf")
+        
+        if os.path.exists(font_path):
+            pdf.add_font('DejaVu', '', font_path, uni=True)
+            pdf.set_font('DejaVu', '', 14)
+        else:
+            # This is where it was falling back before, causing the "?" symbols
+            pdf.set_font("Arial", "B", 16)
+
+        # Write Title
+        pdf.multi_cell(0, 10, title)
+        pdf.ln(5)
+        
+        # Write Content
+        if os.path.exists(font_path):
+            pdf.set_font('DejaVu', '', 12)
+        else:
+            pdf.set_font("Arial", "", 12)
+            
+        pdf.multi_cell(0, 10, content)
+        
+        # Output as bytes
+        return pdf.output(dest='S').encode('latin-1')
+    except Exception as e:
+        return f"PDF Error: {str(e)}".encode('latin-1')
+    
+def create_docx(title, content):
+    doc = Document()
+    doc.add_heading(title, 0)
+    doc.add_paragraph(content)
+    bio = io.BytesIO()
+    doc.save(bio)
+    return bio.getvalue()
+
+# --- 6. AI & RESOURCE TRACKER ---
+def track_usage(text, type="input"):
+    # Rough estimation: 1 token approx 4 characters
+    tokens = len(text) / 4
+    cost = (tokens / 1000) * 0.000125 # Estimate for Gemini 1.5 Flash
+    
+    # Save to a local JSON file to persist data
+    usage_file = "usage_stats.json"
+    stats = {"total_tokens": 0, "total_cost": 0.0}
+    
+    if os.path.exists(usage_file):
+        with open(usage_file, "r") as f:
+            stats = json.load(f)
+            
+    stats["total_tokens"] += tokens
+    stats["total_cost"] += cost
+    
+    with open(usage_file, "w") as f:
+        json.dump(stats, f)
+    
+    return stats
+
+# --- 7. RAG & FEEDBACK ---
+def create_vector_index(notes):
+    """Turns notes into a searchable mathematical index"""
+    if not notes:
+        return None, []
+    
+    # Extract only the text content
+    text_data = [n['content'] for n in notes]
+    
+    # Convert text to vectors (embeddings)
+    embeddings = embed_model.encode(text_data)
+    
+    # Create the FAISS index
+    dimension = embeddings.shape[1]
+    index = faiss.IndexFlatL2(dimension)
+    index.add(np.array(embeddings).astype('float32'))
+    
+    return index, text_data
+
+def query_vault(query, index, text_data, top_k=2):
+    """Finds the most relevant notes for a question"""
+    if index is None:
+        return "No notes found to search."
+        
+    # Convert question to vector
+    query_vector = embed_model.encode([query])
+    
+    # Search the index
+    distances, indices = index.search(np.array(query_vector).astype('float32'), top_k)
+    
+    # Pull the relevant text chunks
+    results = [text_data[i] for i in indices[0] if i != -1]
+    return results
+
+def log_feedback(query, answer, context, status):
+    log_file = "feedback_log.json"
+    entry = {
+        "timestamp": datetime.now().strftime("%Y-%m-%d %H:%M:%S"),
+        "query": query,
+        "answer": answer,
+        "context_used": context,
+        "status": status  # "Correct" or "Wrong"
+    }
+    
+    logs = []
+    if os.path.exists(log_file):
+        with open(log_file, "r") as f:
+            logs = json.load(f)
+            
+    logs.append(entry)
+    with open(log_file, "w") as f:
+        json.dump(logs, f, indent=4)
+
+# --- 8. GEMINI AI ENGINE ---
+def get_gemini_response(user_query, context_str):
+    """Connects to Google Gemini API for free AI logic"""
+    api_key = os.getenv("GEMINI_API_KEY")
+    if not api_key:
+        return "⚠️ Error: GEMINI_API_KEY not found in Hugging Face Secrets."
+    
+    try:
+        # New initialization style
+        client = genai.Client(api_key=api_key)
+        
+        prompt = f"""
+        You are a secure vault assistant. Use the following retrieved notes to answer.
+        Notes Context: {context_str}
+        User Question: {user_query}
+        """
+        
+        # Inside your loop or before you call the model
+        time.sleep(2) # Wait 2 seconds between requests
+        # Simplified model call
+        response = client.models.generate_content(
+            model="gemini-2.5-flash-lite", 
+            contents=prompt
+        )
+        
+        track_usage(prompt, type="input")
+        track_usage(response.text, type="output")
+        return response.text
+    except Exception as e:
+        return f"❌ AI Engine Error: {str(e)}"