Update src/app.py

refactored to add the document loader, add document summarization, and remove email builder and prompt writer

src/app.py

@@ -1,46 +1,133 @@
 import streamlit as st
 import requests
 import os
-import unicodedata
-import resources 
+import re
+import io
+import zipfile
 import tracker
-import rag_engine 
+import rag_engine
+import doc_loader 
 from openai import OpenAI
 from datetime import datetime
 
 # --- CONFIGURATION ---
 st.set_page_config(page_title="Navy AI Toolkit", page_icon="⚓", layout="wide")
 
-# 1. SETUP CREDENTIALS
-API_URL_ROOT = os.getenv("API_URL")  # For Ollama models
-OPENAI_KEY = os.getenv("OPENAI_API_KEY")  # For GPT-4o
+API_URL_ROOT = os.getenv("API_URL") 
+OPENAI_KEY = os.getenv("OPENAI_API_KEY") 
 
 # --- INITIALIZATION ---
 if "roles" not in st.session_state:
     st.session_state.roles = []
 
-# --- LOGIN / REGISTER LOGIC ---
+# --- FLATTENER LOGIC (Integrated) ---
+class OutlineProcessor:
+    """Parses text outlines for the Flattener tool."""
+    def __init__(self, file_content):
+        self.raw_lines = file_content.split('\n')
+    
+    def _is_list_item(self, line):
+        pattern = r"^\s*(\d+\.|[a-zA-Z]\.|-|\*)\s+"
+        return bool(re.match(pattern, line))
+
+    def _merge_multiline_items(self):
+        merged_lines = []
+        for line in self.raw_lines:
+            stripped = line.strip()
+            if not stripped: continue
+            if not merged_lines:
+                merged_lines.append(line)
+                continue
+            if not self._is_list_item(line):
+                merged_lines[-1] = merged_lines[-1].rstrip() + " " + stripped
+            else:
+                merged_lines.append(line)
+        return merged_lines
+
+    def parse(self):
+        clean_lines = self._merge_multiline_items()
+        stack = []
+        results = []
+        for line in clean_lines:
+            stripped = line.strip()
+            indent = len(line) - len(line.lstrip())
+            while stack and stack[-1]['indent'] >= indent:
+                stack.pop()
+            stack.append({'indent': indent, 'text': stripped})
+            if len(stack) > 1:
+                context_str = " > ".join([item['text'] for item in stack[:-1]])
+            else:
+                context_str = "ROOT"
+            results.append({"context": context_str, "target": stripped})
+        return results
+
+# --- HELPER FUNCTIONS ---
+def query_model_universal(messages, max_tokens, model_choice, user_key=None):
+    """Unified router for both Chat and Tools."""
+    # 1. OpenAI Path
+    if "GPT-4o" in model_choice:
+        key = user_key if user_key else OPENAI_KEY
+        if not key: return "[Error: No OpenAI API Key]", None
+        
+        client = OpenAI(api_key=key)
+        try:
+            resp = client.chat.completions.create(
+                model="gpt-4o", max_tokens=max_tokens, messages=messages, temperature=0.3
+            )
+            usage = {"input": resp.usage.prompt_tokens, "output": resp.usage.completion_tokens}
+            return resp.choices[0].message.content, usage
+        except Exception as e:
+            return f"[OpenAI Error: {e}]", None
+
+    # 2. Local Path
+    else:
+        model_map = {
+            "Granite 4 (IBM)": "granite4:latest",
+            "Llama 3.2 (Meta)": "llama3.2:latest",
+            "Gemma 3 (Google)": "gemma3:latest"
+        }
+        tech_name = model_map.get(model_choice)
+        if not tech_name: return "[Error: Model Map Failed]", None
+        
+        url = f"{API_URL_ROOT}/generate"
+        
+        # Flatten history for Ollama
+        hist = ""
+        sys_msg = "You are a helpful assistant."
+        for m in messages:
+            if m['role']=='system': sys_msg = m['content']
+            elif m['role']=='user': hist += f"User: {m['content']}\n"
+            elif m['role']=='assistant': hist += f"Assistant: {m['content']}\n"
+        hist += "Assistant: "
+        
+        try:
+            r = requests.post(url, json={"text": hist, "persona": sys_msg, "max_tokens": max_tokens, "model": tech_name}, timeout=300)
+            if r.status_code == 200:
+                d = r.json()
+                return d.get("response", ""), d.get("usage", {"input":0,"output":0})
+            return f"[Local Error {r.status_code}]", None
+        except Exception as e:
+            return f"[Conn Error: {e}]", None
+
+def update_sidebar_metrics():
+    # Helper to safely update metrics if placeholder exists
+    if metric_placeholder:
+        stats = tracker.get_daily_stats()
+        u_stats = stats["users"].get(st.session_state.username, {"input":0, "output":0})
+        metric_placeholder.metric("My Tokens Today", u_stats["input"] + u_stats["output"])
+
+# --- LOGIN ---
 if "authentication_status" not in st.session_state or st.session_state["authentication_status"] is None:
-    # If not logged in, show tabs
     login_tab, register_tab = st.tabs(["🔑 Login", "📝 Register"])
-    
     with login_tab:
-        is_logged_in = tracker.check_login()
-        if is_logged_in:
-            # Check if a different user was previously logged in
+        if tracker.check_login():
+            # Session Isolation Logic
             if "last_user" in st.session_state and st.session_state.last_user != st.session_state.username:
-                # WIPE EVERYTHING
                 st.session_state.messages = []
-                st.session_state.email_draft = ""
                 st.session_state.user_openai_key = None
-                
-            # Update the tracker
             st.session_state.last_user = st.session_state.username
-            
-            # Download DB and Refresh
             tracker.download_user_db(st.session_state.username)
-            st.rerun() # Refresh to show the app
-    
+            st.rerun()
     with register_tab:
         st.header("Create Account")
         with st.form("reg_form"):
@@ -56,16 +143,11 @@ if "authentication_status" not in st.session_state or st.session_state["authenti
                     st.success(msg)
                 else:
                     st.error(msg)
+                    
+    if not st.session_state.get("authentication_status"): st.stop()
 
-    # Stop execution if not logged in
-    if not st.session_state.get("authentication_status"):
-        st.stop()
-
-# --- GLOBAL PLACEHOLDERS ---
+# --- SIDEBAR ---
 metric_placeholder = None
-admin_metric_placeholder = None
-
-# --- SIDEBAR (CONSOLIDATED) ---
 with st.sidebar:
     st.header("👤 User Profile")
     st.write(f"Welcome, **{st.session_state.name}**")
@@ -77,8 +159,6 @@ with st.sidebar:
     if "admin" in st.session_state.roles:
         st.divider()
         st.header("🛡️ Admin Tools")
-        admin_metric_placeholder = st.empty()
-        
         log_path = tracker.get_log_path()
         if log_path.exists():
             with open(log_path, "r") as f:
@@ -89,502 +169,217 @@ with st.sidebar:
                 file_name=f"usage_log_{datetime.now().strftime('%Y-%m-%d')}.json",
                 mime="application/json"
             )
-        else:
-            st.warning("No logs found yet.")
-    
-    # Logout
-    if "authenticator" in st.session_state:
-        st.session_state.authenticator.logout(location='sidebar')
     
     st.divider()
     
-    # --- MODEL SELECTOR ---
-    st.header("🧠 Model Selector")
-
+    # Model Selector
+    st.header("🧠 Intelligence")
     model_map = {
-        "Granite 4 (IBM)": "granite4:latest",
+        "Granite 4 (IBM)": "granite4:latest", 
         "Llama 3.2 (Meta)": "llama3.2:latest",
         "Gemma 3 (Google)": "gemma3:latest"
     }
+    opts = list(model_map.keys())
+    model_captions = ["Slower, free, private" for _ in opts]
     
-    model_options = list(model_map.keys())
-    model_captions = ["Slower for now, but free and private" for _ in model_options]
-
-    # 2. CHECK FOR GPT-4o ACCESS (Admin OR User Key)
-    # We moved the input UP so the user can unlock the option immediately
-    
-    # Check if user is admin
+    # Vision Key Input (User or Admin)
     is_admin = "admin" in st.session_state.roles
-    
-    # Input for Non-Admins
-    user_api_key = None
+    user_key = None
     if not is_admin:
-        user_api_key = st.text_input(
+        user_key = st.text_input(
             "🔓 Unlock GPT-4o (Enter API Key)", 
             type="password", 
-            help="Enter your OpenAI API key to access GPT-4o. Press Enter to apply.",
-            key=f"user_key_{st.session_state.username}"
+            key=f"key_{st.session_state.username}",
+            help="Required for Vision Mode and GPT-4o."
         )
-        if user_api_key:
-            st.session_state.user_openai_key = user_api_key
+        if user_key: 
+            st.session_state.user_openai_key = user_key
             st.caption("✅ Key Active")
-        else:
+        else: 
             st.session_state.user_openai_key = None
     else:
+        # Admin defaults to system key, but we ensure state is clean
         st.session_state.user_openai_key = None
-
-    # 3. DYNAMICALLY ADD GPT-4o TO THE LIST
-    # If Admin OR if they just entered a key, show the option
+    
+    # Unlock GPT-4o option
     if is_admin or st.session_state.get("user_openai_key"):
-        model_options.append("GPT-4o (Omni)")
+        opts.append("GPT-4o (Omni)")
         model_captions.append("Fast, smart, sends data to OpenAI")
-
-    # 4. RENDER THE SELECTOR
-    model_choice = st.radio(
-        "Choose your Intelligence:",
-        model_options,
-        captions=model_captions,
-        key="model_selector_radio" 
-    )
+        
+    model_choice = st.radio("Select Model:", opts, captions=model_captions, key="model_selector_radio")
     st.info(f"Connected to: **{model_choice}**")
     
     st.divider()
-    st.header("⚙️ Controls")
-    max_len = st.slider("Max Response Length (Tokens)", 100, 2000, 500)
-
-# --- HELPER FUNCTIONS ---
-def update_sidebar_metrics():
-    """Refreshes the global placeholders defined in the sidebar."""
-    if metric_placeholder is None: 
-        return 
-    
-    stats = tracker.get_daily_stats()
-    user_stats = stats["users"].get(st.session_state.username, {"input":0, "output":0})
-    
-    metric_placeholder.metric("My Tokens Today", user_stats["input"] + user_stats["output"])
-    
-    if "admin" in st.session_state.roles and admin_metric_placeholder is not None:
-        admin_metric_placeholder.metric("Team Total Today", stats["total_tokens"])
+    if st.session_state.authenticator:
+        st.session_state.authenticator.logout(location='sidebar')
 
-# Call metrics once on load
 update_sidebar_metrics()
 
-def query_local_model(messages, max_tokens, model_name): 
-    if not API_URL_ROOT:
-        return "Error: API_URL not set.", None
-    
-    url = API_URL_ROOT + "/generate"
-    
-    # --- FLATTEN MESSAGE HISTORY ---
-    formatted_history = ""
-    system_persona = "You are a helpful assistant." # Default
-    
-    for msg in messages:
-        if msg['role'] == 'system':
-            system_persona = msg['content']
-        elif msg['role'] == 'user':
-            formatted_history += f"User: {msg['content']}\n"
-        elif msg['role'] == 'assistant':
-            formatted_history += f"Assistant: {msg['content']}\n"
-            
-    # Append the "Assistant:" prompt at the end to cue the model
-    formatted_history += "Assistant: "
-
-    payload = {
-        "text": formatted_history, 
-        "persona": system_persona, 
-        "max_tokens": max_tokens,
-        "model": model_name 
-    } 
-    
-    try:
-        response = requests.post(url, json=payload, timeout=300)
-        
-        if response.status_code == 200:
-            response_data = response.json()
-            ans = response_data.get("response", "")
-            usage = response_data.get("usage", {"input":0, "output":0})
-            return ans, usage
-            
-        return f"Error {response.status_code}: {response.text}", None
-        
-    except Exception as e:
-        return f"Connection Error: {e}", None
-
-def query_openai_model(messages, max_tokens):
-    # 1. Check for User Key first
-    api_key_to_use = st.session_state.get("user_openai_key")
-    
-    # 2. Fallback to System Key
-    if not api_key_to_use:
-        api_key_to_use = OPENAI_KEY
-
-    # 3. Final Safety Check
-    if not api_key_to_use:
-        return "Error: No API Key available. Please enter one in the sidebar.", None
-    
-    client = OpenAI(api_key=api_key_to_use)
-    
-    try:
-        response = client.chat.completions.create(
-            model="gpt-4o",
-            max_tokens=max_tokens,
-            messages=messages,
-            temperature=0.3
-        )
-        usage_obj = response.usage
-        usage_dict = {"input": usage_obj.prompt_tokens, "output": usage_obj.completion_tokens}
-        return response.choices[0].message.content, usage_dict
-        
-    except Exception as e:
-        return f"OpenAI Error: {e}", None
+# --- MAIN APP ---
+st.title("⚓ Navy AI Toolkit")
+tab1, tab2 = st.tabs(["💬 Chat Playground", "📂 Knowledge & Tools"])
 
-def clean_text(text):
-    if not text: return ""
-    text = unicodedata.normalize('NFKC', text)
-    replacements = {'“': '"', '”': '"', '‘': "'", '’': "'", '–': '-', '—': '-', '…': '...', '\u00a0': ' '}
-    for old, new in replacements.items():
-        text = text.replace(old, new)
-    return text.strip()
-
-def ask_ai(user_prompt, system_persona, max_tokens):
-    # 1. Standardize Input: Convert the strings into the Message List format
-    messages_payload = [
-        {"role": "system", "content": system_persona},
-        {"role": "user", "content": user_prompt}
-    ]
-
-    # 2. Routing Logic
-    if "GPT-4o" in model_choice:
-        return query_openai_model(messages_payload, max_tokens)
-    else:
-        technical_name = model_map[model_choice]
-        return query_local_model(messages_payload, max_tokens, technical_name)
-
-# --- MAIN UI ---
-st.title("AI Toolkit")
-tab1, tab2, tab3, tab4 = st.tabs(["📧 Email Builder", "💬 Chat Playground", "🛠️ Prompt Architect", "📚 Knowledge Base"])
-
-# --- TAB 1: EMAIL BUILDER ---
+# === TAB 1: CHAT ===
 with tab1:
-    st.header("Structured Email Generator")
-    if "email_draft" not in st.session_state:
-        st.session_state.email_draft = ""
-
-    st.subheader("1. Define the Voice")
-    style_mode = st.radio("How should the AI write?", ["Use a Preset Persona", "Mimic My Style"], horizontal=True)
+    st.header("Discussion & Analysis")
+    if "messages" not in st.session_state: st.session_state.messages = []
     
-    selected_persona_instruction = ""
-    if style_mode == "Use a Preset Persona":
-        persona_name = st.selectbox("Select a Persona", list(resources.TONE_LIBRARY.keys()))
-        selected_persona_instruction = resources.TONE_LIBRARY[persona_name]
-        st.info(f"**System Instruction:** {selected_persona_instruction}")
-    else:
-        st.info("Upload 1-3 text files of your previous emails.")
-        uploaded_style_files = st.file_uploader("Upload Samples (.txt)", type=["txt"], accept_multiple_files=True)
-        if uploaded_style_files:
-            style_context = ""
-            for uploaded_file in uploaded_style_files:
-                string_data = uploaded_file.read().decode("utf-8")
-                style_context += f"---\n{string_data}\n---\n"
-            selected_persona_instruction = f"Analyze these examples and mimic the style:\n{style_context}"
-
-    st.divider()
-    st.subheader("2. Details")
-    c1, c2 = st.columns(2)
-    with c1: recipient = st.text_input("Recipient")
-    with c2: topic = st.text_input("Topic")
+    c1, c2 = st.columns([3, 1])
+    with c1: st.caption(f"Active Model: **{st.session_state.get('model_selector_radio', 'Granite')}**")
+    with c2: use_rag = st.toggle("Enable Knowledge Base", value=False)
     
-    st.caption("Content Source")
-    input_method = st.toggle("Upload notes file?")
-    raw_notes = ""
-    if input_method:
-        notes_file = st.file_uploader("Upload Notes (.txt)", type=["txt"])
-        if notes_file: raw_notes = notes_file.read().decode("utf-8")
-    else:
-        raw_notes = st.text_area("Paste notes:", height=150)
-
-    # Context Bar
-    est_tokens = len(raw_notes) / 4 
-    st.progress(min(est_tokens / 128000, 1.0), text=f"Context: {int(est_tokens)} tokens")
-
-    if st.button("Draft Email", type="primary"):
-        if not raw_notes:
-            st.warning("Please provide notes.")
-        else:
-            clean_notes = clean_text(raw_notes)
-            with st.spinner(f"Drafting with {model_choice}..."):
-                prompt = f"TASK: Write email.\nTO: {recipient}\nTOPIC: {topic}\nSTYLE: {selected_persona_instruction}\nDATA: {clean_notes}"
-                
-                reply, usage = ask_ai(prompt, "You are an expert ghostwriter.", max_len)
-                st.session_state.email_draft = reply
-                
-                if usage:
-                    if "GPT-4o" in model_choice:
-                         m_name = "GPT-4o"
-                    else:
-                         m_name = model_choice.split(" ")[0]
-                    tracker.log_usage(m_name, usage["input"], usage["output"])
-                    update_sidebar_metrics()
-
-    if st.session_state.email_draft:
-        st.subheader("Draft Result")
-        st.text_area("Copy your email:", value=st.session_state.email_draft, height=300)
-
-# --- TAB 2: CHAT PLAYGROUND ---
-with tab2:
-    st.header("Choose Your Model and Start a Discussion")
-    
-    # --- INITIALIZE CHAT MEMORY (MUST BE DONE FIRST) ---
-    if "messages" not in st.session_state:
-        st.session_state.messages = []
-        
-    # --- CONTROLS AND METRICS ---
-    c1, c2, c3 = st.columns([2, 1, 1])
-    with c1:
-        # FIX: Access the correct key from the sidebar widget
-        # We default to the global variable 'model_choice' if state is missing
-        selected_model_name = st.session_state.get('model_selector_radio', model_choice)
-        st.caption(f"Active Model: **{selected_model_name}**")
+    for msg in st.session_state.messages:
+        with st.chat_message(msg["role"]): st.markdown(msg["content"])
         
-    with c2:
-        use_rag = st.toggle("🔌 Enable Knowledge Base", value=False)
-    
-    with c3:
-        # --- NEW FEATURE: DOWNLOAD CHAT ---
-        chat_log = ""
-        for msg in st.session_state.messages:
-            role = "USER" if msg['role'] == 'user' else "ASSISTANT"
-            chat_log += f"[{role}]: {msg['content']}\n\n"
-            
-        if chat_log:
-            st.download_button(
-                label="💾 Save Chat",
-                data=chat_log,
-                file_name="mission_log.txt",
-                mime="text/plain",
-                help="Download the current conversation history."
-            )
-    
-    st.divider()
-    
-    # --- DISPLAY CONVERSATION HISTORY ---
-    for message in st.session_state.messages:
-        with st.chat_message(message["role"]):
-            st.markdown(message["content"])
-
-    # --- CHAT INPUT HANDLING ---
-    if prompt := st.chat_input("Ask a question..."):
-        # 1. Display User Message and save to history
+    if prompt := st.chat_input("Input command..."):
         st.session_state.messages.append({"role": "user", "content": prompt})
-        with st.chat_message("user"):
-            st.markdown(prompt)
-
-        # 2. Default Configuration (Standard AI Mode)
-        system_persona = "You are a helpful AI assistant. Answer the user's question to the best of your ability."
-        final_user_content = prompt
-        retrieved_docs = [] 
+        with st.chat_message("user"): st.markdown(prompt)
+        
+        # RAG Search
+        context_txt = ""
+        sys_p = "You are a helpful assistant."
         
-        # 3. Handle RAG Logic (Only if enabled)
         if use_rag:
-            with st.spinner("🧠 Searching Knowledge Base..."):
-                retrieved_docs = rag_engine.search_knowledge_base(
-                    prompt, 
-                    st.session_state.username
-                )
-                
-                if retrieved_docs:
-                    # RAG SUCCESS: Switch to Strict Navy Persona
-                    system_persona = (
-                        "You are a Navy Document Analyst. Your task is to answer the user's question "
-                        "using ONLY the Context provided below. "
-                        "If the answer is not present in the Context, return ONLY this exact phrase: "
-                        "'I cannot find that information in the provided documents.'"
+            with st.spinner("Searching DB..."):
+                docs = rag_engine.search_knowledge_base(prompt, st.session_state.username)
+                if docs:
+                    sys_p = (
+                        "You are a Navy Document Analyst. Answer using ONLY the Context below. "
+                        "If the answer is not in the context, say 'I cannot find that information'."
                     )
+                    for d in docs: context_txt += f"\n---\n{d.page_content}"
                     
-                    # Format Context
-                    context_text = ""
-                    for doc in retrieved_docs:
-                        score = doc.metadata.get('relevance_score', 'N/A')
-                        src = os.path.basename(doc.metadata.get('source', 'Unknown'))
-                        context_text += f"---\nSOURCE: {src} (Rel: {score})\nTEXT: {doc.page_content}\n"
-                    
-                    # Augment User Prompt
-                    final_user_content = (
-                        f"User Question: {prompt}\n\n"
-                        f"Relevant Context:\n{context_text}\n\n"
-                        "Answer the question using the context provided."
-                    )
+        final_prompt = f"{prompt}\n\nCONTEXT:\n{context_txt}" if context_txt else prompt
         
-        # 4. Construct Payload (Now using the CORRECT persona)
-        messages_payload = [{"role": "system", "content": system_persona}]
-        
-        # --- MEMORY LOGIC: SLIDING WINDOW ---
-        history_depth = 8 
-        recent_history = st.session_state.messages[-(history_depth+1):-1] 
-        messages_payload.extend(recent_history)
-        
-        # Add the final (potentially augmented) user message to payload
-        messages_payload.append({"role": "user", "content": final_user_content})
-        
-        # 5. Generate Response
+        # Generation
         with st.chat_message("assistant"):
-            with st.spinner(f"Thinking with {selected_model_name}..."):
-                # Determine model ID
-                model_id = "" 
-                ollama_map = {
-                    "Granite 4 (IBM)": "granite4:latest",
-                    "Llama 3.2 (Meta)": "llama3.2:latest",
-                    "Gemma 3 (Google)": "gemma3:latest"
-                }
-                for key, val in ollama_map.items():
-                    if key in selected_model_name:
-                        model_id = val
-                        break
+            with st.spinner("Thinking..."):
+                # Memory Window
+                hist = [{"role":"system", "content":sys_p}] + st.session_state.messages[-6:-1] + [{"role":"user", "content":final_prompt}]
                 
-                # ROUTING CHECK
-                if not model_id and "gpt" in selected_model_name.lower():
-                    # If it's the GPT model choice
-                    response, usage = query_openai_model(messages_payload, max_len)
-                elif model_id:
-                    # If it's the local Ollama model
-                    response, usage = query_local_model(messages_payload, max_len, model_id)
-                else:
-                    response, usage = "Error: Could not determine model to use.", None
+                resp, usage = query_model_universal(hist, 2000, model_choice, st.session_state.get("user_openai_key"))
+                st.markdown(resp)
                 
-                st.markdown(response)
-            
-        # 6. Save Assistant Response
-        st.session_state.messages.append({"role": "assistant", "content": response})
+                if usage:
+                    m_name = "GPT-4o" if "GPT-4o" in model_choice else model_choice.split()[0]
+                    tracker.log_usage(m_name, usage["input"], usage["output"])
+                    update_sidebar_metrics()
+                    
+        st.session_state.messages.append({"role": "assistant", "content": resp})
         
-        # 7. Metrics & Context Display
-        if usage:
-            if "GPT-4o" in selected_model_name:
-                 m_name = "GPT-4o"
-            else:
-                 m_name = selected_model_name.split(" ")[0]
-            tracker.log_usage(m_name, usage["input"], usage["output"])
-            update_sidebar_metrics()
-
-        if use_rag and retrieved_docs:
+        if use_rag and context_txt:
             with st.expander("📚 View Context Used"):
-                for i, doc in enumerate(retrieved_docs):
-                    score = doc.metadata.get('relevance_score', 'N/A')
-                    src = os.path.basename(doc.metadata.get('source', 'Unknown'))
-                    st.caption(f"Rank {i+1} (Source: {src}, Rel: {score})")
-                    st.text(doc.page_content)
-                    st.divider()
-                    
-# --- TAB 3: PROMPT ARCHITECT ---
-with tab3:
-    st.header("🛠️ Mega-Prompt Factory")
-    st.info("Build standard templates for NIPRGPT.")
+                st.text(context_txt)
+
+# === TAB 2: KNOWLEDGE & TOOLS ===
+with tab2:
+    st.header("Document Processor")
     
-    c1, c2 = st.columns([1,1])
+    c1, c2 = st.columns([1, 1])
     with c1:
-        st.subheader("1. Parameters")
-        p = st.text_area("Persona", placeholder="Act as...", height=100)
-        c = st.text_area("Context", placeholder="Background...", height=100)
-        t = st.text_area("Task", placeholder="Action...", height=100)
-        v = st.text_input("Placeholder Name", value="PASTE_DATA_HERE")
-    
+        uploaded_file = st.file_uploader("Upload File (PDF, PPT, Doc, Text)", type=["pdf", "docx", "pptx", "txt", "md"])
     with c2:
-        st.subheader("2. Result")
-        final = f"### ROLE\n{p}\n### CONTEXT\n{c}\n### TASK\n{t}\n### INPUT DATA\n\"\"\"\n[{v}]\n\"\"\""
-        st.code(final, language="markdown")
-        st.download_button("💾 Download .txt", final, "template.txt")
+        use_vision = st.toggle("👁️ Enable Vision Mode", help="Use GPT-4o to read diagrams/tables. Requires API Key.")
+        if use_vision and "GPT-4o" not in opts:
+            st.warning("Vision requires OpenAI Access.")
 
-# --- TAB 4: KNOWLEDGE BASE ---
-with tab4:
-    st.header("🧠 Personal Knowledge Base")
-    st.info(f"Managing knowledge for: **{st.session_state.username}**")
-    
-    # We no longer check 'is_admin' for the whole tab
-    kb_tab1, kb_tab2 = st.tabs(["📤 Add Documents", "🗂️ Manage Database"])
-    
-    # --- SUB-TAB 1: UPLOAD (Unlocked for Everyone) ---
-    with kb_tab1:
-        st.subheader("Ingest New Knowledge")
-        uploaded_file = st.file_uploader("Upload Instructions, Manuals, or Logs", type=["pdf", "docx", "txt", "md"])
+    if uploaded_file:
+        # Save temp
+        temp_path = rag_engine.save_uploaded_file(uploaded_file)
         
-        col1, col2 = st.columns([1, 2])
-        with col1:
-            chunk_strategy = st.selectbox(
-                "Chunking Strategy", 
-                ["paragraph", "token", "page"],
-                help="Paragraph: Manuals. Token: Dense text. Page: Forms."
-            )
+        # ACTION BAR
+        col_a, col_b, col_c = st.columns(3)
         
-        if uploaded_file and st.button("Process & Add"):
-            with st.spinner("Analyzing and Indexing..."):
-                # 1. Save temp file
-                temp_path = rag_engine.save_uploaded_file(uploaded_file)
-                
-                # 2. Process locally
-                success, msg = rag_engine.process_and_add_document(
-                    temp_path, 
-                    st.session_state.username, 
-                    chunk_strategy
-                )
-                
-                if success:
-                    # 3. FIX: SYNC TO CLOUD IMMEDIATELY
-                    with st.spinner("Backing up to Cloud..."):
-                        tracker.upload_user_db(st.session_state.username)
+        # 1. ADD TO DB
+        with col_a:
+            if st.button("📥 Add to Knowledge Base", type="primary"):
+                with st.spinner("Ingesting..."):
+                    # Use Admin or User Key for Vision
+                    key = st.session_state.get("user_openai_key") or OPENAI_KEY
+                    
+                    ok, msg = rag_engine.process_and_add_document(
+                        temp_path, st.session_state.username, "paragraph", 
+                        use_vision=use_vision, api_key=key
+                    )
+                    
+                    if ok:
+                        tracker.upload_user_db(st.session_state.username) # Auto-Sync
+                        st.success(msg)
+                    else: st.error(msg)
+
+        # 2. SUMMARIZE
+        with col_b:
+            if st.button("📝 Summarize Document"):
+                with st.spinner("Reading & Summarizing..."):
+                    key = st.session_state.get("user_openai_key") or OPENAI_KEY
+                    # Extract raw text first
+                    class FileObj:
+                        def __init__(self, p, n): self.path=p; self.name=n
+                        def read(self): 
+                            with open(self.path, "rb") as f: return f.read()
+                    
+                    # Extraction
+                    raw = doc_loader.extract_text_from_file(
+                        FileObj(temp_path, uploaded_file.name), 
+                        use_vision=use_vision, api_key=key
+                    )
+                    
+                    # Call LLM
+                    prompt = f"Summarize this document into a key executive brief:\n\n{raw[:20000]}" # Truncate for safety
+                    msgs = [{"role":"user", "content": prompt}]
+                    summ, usage = query_model_universal(msgs, 1000, model_choice, st.session_state.get("user_openai_key"))
+                    
+                    st.subheader("Summary Result")
+                    st.markdown(summ)
+                    if usage:
+                        m_name = "GPT-4o" if "GPT-4o" in model_choice else model_choice.split()[0]
+                        tracker.log_usage(m_name, usage["input"], usage["output"])
+                        update_sidebar_metrics()
+
+        # 3. FLATTEN
+        with col_c:
+            if st.button("📄 Flatten Context"):
+                with st.spinner("Flattening..."):
+                    key = st.session_state.get("user_openai_key") or OPENAI_KEY
+                    # Extract
+                    with open(temp_path, "rb") as f:
+                        # Dummy object again for the loader
+                        class Wrapper:
+                            def __init__(self, data, n): self.data=data; self.name=n
+                            def read(self): return self.data
+                        raw = doc_loader.extract_text_from_file(
+                            Wrapper(f.read(), uploaded_file.name), use_vision=use_vision, api_key=key
+                        )
+
+                    # Parse
+                    proc = OutlineProcessor(raw)
+                    items = proc.parse()
+                    
+                    # Flatten
+                    out_txt = []
+                    bar = st.progress(0)
+                    for i, item in enumerate(items):
+                        # Use the Universal Router so it works with Granite too!
+                        p = f"Context: {item['context']}\nTarget: {item['target']}\nRewrite as one sentence."
+                        m = [{"role":"user", "content": p}]
+                        res, _ = query_model_universal(m, 300, model_choice, st.session_state.get("user_openai_key"))
+                        out_txt.append(res)
+                        bar.progress((i+1)/len(items))
                         
-                    st.success(msg)
-                    st.rerun() 
-                else:
-                    st.error(f"Failed: {msg}")
-
-        st.divider()
-        st.subheader("🔎 Quick Test")
-        test_query = st.text_input("Ask your brain something...")
-        if test_query:
-            results = rag_engine.search_knowledge_base(test_query, st.session_state.username)
-            if not results:
-                st.warning("No matches found.")
-            for i, doc in enumerate(results):
-                src_name = os.path.basename(doc.metadata.get('source', '?'))
-                score = doc.metadata.get('relevance_score', 'N/A')
-                with st.expander(f"Match {i+1}: {src_name} (Score: {score})"):
-                    st.write(doc.page_content)
+                    st.text_area("Result", "\n".join(out_txt), height=300)
 
-    # --- SUB-TAB 2: MANAGE (Unlocked for Everyone) ---
-    with kb_tab2:
-        st.subheader("🗄️ Database Inventory")
-        
-        docs = rag_engine.list_documents(st.session_state.username)
-        
-        if not docs:
-            st.info("Your Knowledge Base is empty.")
-        else:
-            st.markdown(f"**Total Documents:** {len(docs)}")
-            
-            for doc in docs:
-                c1, c2, c3, c4 = st.columns([3, 2, 1, 1])
-                with c1:
-                    st.text(f"📄 {doc['filename']}")
-                with c2:
-                    st.caption(f"⚙️ {doc.get('strategy', 'Unknown')}")
-                with c3:
-                    st.caption(f"{doc['chunks']}")
-                with c4:
-                    if st.button("🗑️", key=doc['source'], help="Delete Document"):
-                        with st.spinner("Deleting..."):
-                            success, msg = rag_engine.delete_document(st.session_state.username, doc['source'])
-                            if success:
-                                tracker.upload_user_db(st.session_state.username)
-                                st.success(msg)
-                                st.rerun()
-                            else:
-                                st.error(msg)
-
-            st.divider()
-            with st.expander("🚨 Danger Zone"):
-                # Allow ANY user to reset their OWN database
-                if st.button("☢️ RESET MY DATABASE", type="primary"):
-                    success, msg = rag_engine.reset_knowledge_base(st.session_state.username)
-                    if success:
-                        st.success(msg)
-                        st.rerun()
+    st.divider()
+    
+    # DB MANAGER
+    st.subheader("Database Management")
+    docs = rag_engine.list_documents(st.session_state.username)
+    if docs:
+        for d in docs:
+            c1, c2 = st.columns([4,1])
+            c1.text(f"📄 {d['filename']} ({d['chunks']} chunks)")
+            if c2.button("🗑️", key=d['source']):
+                rag_engine.delete_document(st.session_state.username, d['source'])
+                tracker.upload_user_db(st.session_state.username)
+                st.rerun()
+    else:
+        st.info("Database Empty.")