Update src/app.py

src/app.py

@@ -13,7 +13,7 @@ from openai import OpenAI
 from datetime import datetime
 from test_integration import run_tests
 from core.QuizEngine import QuizEngine
-from core.PineconeManager import PineconeManager # FIXED: Added missing import
+from core.PineconeManager import PineconeManager
 
 # --- CONFIGURATION ---
 st.set_page_config(page_title="Navy AI Toolkit", page_icon="⚓", layout="wide")
@@ -35,6 +35,10 @@ if "quiz_state" not in st.session_state:
         "generated_question_text": ""
     }
 
+# NEW: Quiz History for Study Guide
+if "quiz_history" not in st.session_state:
+    st.session_state.quiz_history = []
+
 if "active_index" not in st.session_state:
     st.session_state.active_index = None
 
@@ -82,7 +86,6 @@ class OutlineProcessor:
 # --- 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
@@ -96,8 +99,6 @@ def query_model_universal(messages, max_tokens, model_choice, user_key=None):
             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",
@@ -109,7 +110,6 @@ def query_model_universal(messages, max_tokens, model_choice, user_key=None):
         
         url = f"{API_URL_ROOT}/generate"
         
-        # Flatten history for Ollama
         hist = ""
         sys_msg = "You are a helpful assistant."
         for m in messages:
@@ -128,18 +128,28 @@ def query_model_universal(messages, max_tokens, model_choice, user_key=None):
             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"])
 
+def generate_study_guide_md(history):
+    """Converts quiz history to a Markdown string."""
+    md = "# ⚓ Study Guide\n\n"
+    md += f"Generated: {datetime.now().strftime('%Y-%m-%d %H:%M')}\n\n"
+    for item in history:
+        md += f"## Q: {item['question']}\n"
+        md += f"**Your Answer:** {item['user_answer']}\n\n"
+        md += f"**Grade:** {item['grade']}\n\n"
+        md += f"**Context/Correct Info:**\n> {item['context']}\n\n"
+        md += "---\n\n"
+    return md
+
 # --- LOGIN ---
 if "authentication_status" not in st.session_state or st.session_state["authentication_status"] is None:
     login_tab, register_tab = st.tabs(["🔑 Login", "📝 Register"])
     with login_tab:
         if tracker.check_login():
-            # Session Isolation Logic
             if "last_user" in st.session_state and st.session_state.last_user != st.session_state.username:
                 st.session_state.messages = []
                 st.session_state.user_openai_key = None
@@ -154,14 +164,10 @@ if "authentication_status" not in st.session_state or st.session_state["authenti
             new_email = st.text_input("Email")
             new_pwd = st.text_input("Password", type="password")
             invite = st.text_input("Invitation Passcode") 
-            
             if st.form_submit_button("Register"):
                 success, msg = tracker.register_user(new_email, new_user, new_name, new_pwd, invite)
-                if success:
-                    st.success(msg)
-                else:
-                    st.error(msg)
-                    
+                if success: st.success(msg)
+                else: st.error(msg)
     if not st.session_state.get("authentication_status"): st.stop()
 
 # --- SIDEBAR ---
@@ -169,47 +175,33 @@ metric_placeholder = None
 with st.sidebar:
     st.header("👤 User Profile")
     st.write(f"Welcome, **{st.session_state.name}**")
-    
     st.header("📊 Usage Tracker")
     metric_placeholder = st.empty()
-    
-    # Admin Tools
     if "admin" in st.session_state.roles:
         st.divider()
         st.header("🛡️ Admin Tools")
         log_path = tracker.get_log_path()
         if log_path.exists():
-            with open(log_path, "r") as f:
-                log_data = f.read()
-            st.download_button(
-                label="📥 Download Usage Logs",
-                data=log_data,
-                file_name=f"usage_log_{datetime.now().strftime('%Y-%m-%d')}.json",
-                mime="application/json"
-            )
-    
+            with open(log_path, "r") as f: log_data = f.read()
+            st.download_button("📥 Download Usage Logs", log_data, f"usage_log_{datetime.now().strftime('%Y-%m-%d')}.json", "application/json")
     st.divider()
 
     st.header("🌲 Pinecone Settings")
-    # Initialize Manager
     pc_key = os.getenv("PINECONE_API_KEY")
     if pc_key:
         pm = PineconeManager(pc_key)
         indexes = pm.list_indexes()
-        
-        # 1. INDEX SELECTOR
         selected_index = st.selectbox("Active Index", indexes)
         st.session_state.active_index = selected_index
         
-        # 2. SAFETY CHECK VISUAL
+        # 2. SAFETY CHECK VISUAL (FIXED)
         if selected_index:
-            # Get current model dimension dynamically
+            # Check if the user has already selected a model; default to MiniLM if not
             current_model = st.session_state.get("active_embed_model", "sentence-transformers/all-MiniLM-L6-v2")
             try:
                 emb_fn = rag_engine.get_embedding_func(current_model)
-                test_vec = emb_fn.embed_query("this is a test")
+                test_vec = emb_fn.embed_query("test")
                 active_model_dim = len(test_vec)
-                
                 is_compatible = pm.check_dimension_compatibility(selected_index, active_model_dim)
                 if is_compatible:
                     st.caption(f"✅ Compatible with Model ({active_model_dim}d)")
@@ -218,47 +210,30 @@ with st.sidebar:
             except Exception as e:
                 st.caption(f"⚠️ Could not verify dims: {e}")
         
-        # 3. CREATE NEW INDEX
         with st.expander("Create New Index"):
             new_idx_name = st.text_input("Index Name")
-            
-            # NEW: Dimension Selector
-            new_idx_dim = st.selectbox(
-                "Vector Dimension", 
-                [384, 768, 1024, 1536, 3072],
-                index=0, # Defaults to 384
-                help="384=All-MiniLM, 768=MPNet/Nomic, 1536=OpenAI-Small, 3072=OpenAI-Large"
-            )
-            
+            new_idx_dim = st.selectbox("Vector Dimension", [384, 768, 1024, 1536, 3072], index=0)
             if st.button("Create"):
                 with st.spinner("Provisioning Cloud Index..."):
-                    # We pass the selected dimension to the manager
                     ok, msg = pm.create_index(new_idx_name, dimension=new_idx_dim)
                     if ok: 
                         st.success(msg)
-                        time.sleep(2) # Give Pinecone a moment to propagate
+                        time.sleep(2)
                         st.rerun()
-                    else: 
-                        st.error(msg)
+                    else: st.error(msg)
     else:
         st.warning("No Pinecone Key Found")
     
-    # Model Selector
     st.header("🧠 Intelligence")
-
-    # 1. EMBEDDING MODEL SELECTOR (New!)
     st.subheader("1. Embeddings (The Memory)")
     embed_options = {
         "Standard (All-MiniLM, 384d)": "sentence-transformers/all-MiniLM-L6-v2",
         "High-Perf (MPNet, 768d)": "sentence-transformers/all-mpnet-base-v2",
         "OpenAI Small (1536d)": "text-embedding-3-small"
     }
-    
     embed_choice_label = st.selectbox("Select Embedding Model", list(embed_options.keys()))
-    # Store the actual API string in session state
     st.session_state.active_embed_model = embed_options[embed_choice_label]
     
-    # 2. LLM SELECTOR (The Brain)
     st.subheader("2. Chat Model (The Brain)")
     model_map = {
         "Granite 4 (IBM)": "granite4:latest", 
@@ -267,51 +242,22 @@ with st.sidebar:
     }
     opts = list(model_map.keys())
     model_captions = ["Slower, free, private" for _ in opts]
-    
-    # Vision Key Input (User or Admin)
     is_admin = "admin" in st.session_state.roles
     user_key = None
     if not is_admin:
-        user_key = st.text_input(
-            "🔓 Unlock GPT-4o (Enter API Key)", 
-            type="password", 
-            key=f"key_{st.session_state.username}",
-            help="Required for Vision Mode and GPT-4o."
-        )
+        user_key = st.text_input("🔓 Unlock GPT-4o (Enter API Key)", type="password", key=f"key_{st.session_state.username}")
         if user_key: 
             st.session_state.user_openai_key = user_key
             st.caption("✅ Key Active")
-        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
-    
-    # Unlock GPT-4o option
+        else: st.session_state.user_openai_key = None
+    else: st.session_state.user_openai_key = None
     if is_admin or st.session_state.get("user_openai_key"):
         opts.append("GPT-4o (Omni)")
         model_captions.append("Fast, smart, sends data to OpenAI")
-        
     model_choice = st.radio("Select Model:", opts, captions=model_captions, key="model_selector_radio")
     st.info(f"Connected to: **{model_choice}**")
-    
     st.divider()
-    if st.session_state.authenticator:
-        st.session_state.authenticator.logout(location='sidebar')
-
-    st.divider()
-    st.subheader("🔧 System Diagnostics")
-    
-    if st.button("Run Integration Test"):
-        with st.spinner("Running diagnostics..."):
-            f = io.StringIO()
-            try:
-                with contextlib.redirect_stdout(f):
-                    run_tests()
-                st.success("Tests Completed")
-                st.code(f.getvalue(), language="text")
-            except Exception as e:
-                st.error(f"Test Execution Failed: {e}")
+    if st.session_state.authenticator: st.session_state.authenticator.logout(location='sidebar')
 
 update_sidebar_metrics()
 
@@ -323,28 +269,20 @@ tab1, tab2, tab3 = st.tabs(["💬 Chat Playground", "📂 Knowledge & Tools", "
 with tab1:
     st.header("Discussion & Analysis")
     if "messages" not in st.session_state: st.session_state.messages = []
-    
     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)
-    
     for msg in st.session_state.messages:
         with st.chat_message(msg["role"]): st.markdown(msg["content"])
-        
     if prompt := st.chat_input("Input command..."):
         st.session_state.messages.append({"role": "user", "content": prompt})
         with st.chat_message("user"): st.markdown(prompt)
-        
-        # RAG Search
         context_txt = ""
         sys_p = "You are a helpful AI assistant."
-        
         if use_rag:
-            if not st.session_state.active_index:
-                st.error("⚠️ Please select an Active Index in the sidebar first.")
+            if not st.session_state.active_index: st.error("⚠️ Please select an Active Index in the sidebar first.")
             else:
                 with st.spinner("Searching Knowledge Base..."):
-                    # FIXED: Added index_name parameter
                     docs = rag_engine.search_knowledge_base(
                         query=prompt, 
                         username=st.session_state.username,
@@ -358,74 +296,44 @@ with tab1:
                             "If the Context contains the answer, output it clearly. "
                             "If the Context does NOT contain the answer, simply state: "
                             "'I cannot find that specific information in the documents provided.'"
-                        )
+                        )"
                         for i, d in enumerate(docs):
                             src = d.metadata.get('source', 'Unknown')
                             context_txt += f"<document index='{i+1}' source='{src}'>\n{d.page_content}\n</document>\n"
-                    
-        # Construct Payload
         if context_txt:
-            final_prompt = (
-                f"User Question: {prompt}\n\n"
-                f"<context>\n{context_txt}\n</context>\n\n"
-                "Instruction: Answer the question using the context above."
-            )
-        else:
-            final_prompt = prompt
-        
-        # Generation
+            final_prompt = f"User Question: {prompt}\n\n<context>\n{context_txt}\n</context>\n\nInstruction: Answer using the context above."
+        else: final_prompt = prompt
         with st.chat_message("assistant"):
             with st.spinner("Thinking..."):
                 hist = [{"role":"system", "content":sys_p}] + st.session_state.messages[-6:-1] + [{"role":"user", "content":final_prompt}]
-                
                 resp, usage = query_model_universal(hist, 2000, model_choice, st.session_state.get("user_openai_key"))
                 st.markdown(resp)
-                
                 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})
-        
         if use_rag and context_txt:
-            with st.expander("📚 View Context Used"):
-                st.text(context_txt)
+            with st.expander("📚 View Context Used"): st.text(context_txt)
 
 # === TAB 2: KNOWLEDGE & TOOLS ===
 with tab2:
     st.header("Document Processor")
-    
     c1, c2 = st.columns([1, 1])
-    with c1:
-        uploaded_file = st.file_uploader("Upload File (PDF, PPT, Doc, Text)", type=["pdf", "docx", "pptx", "txt", "md"])
+    with c1: uploaded_file = st.file_uploader("Upload File (PDF, PPT, Doc, Text)", type=["pdf", "docx", "pptx", "txt", "md"])
     with c2:
-        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.")
+        use_vision = st.toggle("👁️ Enable Vision Mode", help="Use GPT-4o to read diagrams/tables.")
+        if use_vision and "GPT-4o" not in opts: st.warning("Vision requires OpenAI Access.")
 
     if uploaded_file:
-        # Save temp
         temp_path = rag_engine.save_uploaded_file(uploaded_file, st.session_state.username)
-        
-        # ACTION BAR
         col_a, col_b, col_c = st.columns(3)
-        
-        # 1. ADD TO DB
         with col_a:
-            chunk_strategy = st.selectbox(
-                "Chunking Strategy", 
-                ["paragraph", "token"], 
-                help="Paragraph: Standard. Token: Dense text.",
-                key="chunk_selector"
-            )
-            
+            chunk_strategy = st.selectbox("Chunking Strategy", ["paragraph", "token"], key="chunk_selector")
             if st.button("📥 Add to Knowledge Base", type="primary"):
-                if not st.session_state.active_index:
-                    st.error("Please select an Active Index in the sidebar.")
+                if not st.session_state.active_index: st.error("Please select an Active Index.")
                 else:
                     with st.spinner("Ingesting..."):
-                        # FIXED: Added index_name parameter
                         ok, msg = rag_engine.ingest_file(
                             file_path=temp_path, 
                             username=st.session_state.username, 
@@ -433,64 +341,37 @@ with tab2:
                             strategy=chunk_strategy,
                             embed_model_name=st.session_state.active_embed_model
                         )
-                        
-                        if ok:
-                            tracker.upload_user_db(st.session_state.username) # Auto-Sync
+                        if ok: 
+                            tracker.upload_user_db(st.session_state.username)
                             st.success(msg)
-                        else:
-                            st.error(msg)
-
-        # 2. SUMMARIZE
+                        else: st.error(msg)
         with col_b:
-            st.write("") 
-            st.write("") 
+            st.write(""); st.write("")
             if st.button("📝 Summarize Document"):
-                with st.spinner("Reading & Summarizing..."):
+                with st.spinner("Reading..."):
                     key = st.session_state.get("user_openai_key") or OPENAI_KEY
                     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()
-                    
-                    raw = doc_loader.extract_text_from_file(
-                        FileObj(temp_path, uploaded_file.name), 
-                        use_vision=use_vision, api_key=key
-                    )
-                    
-                    prompt = f"Summarize this document into a key executive brief:\n\n{raw[:20000]}"
+                    raw = doc_loader.extract_text_from_file(FileObj(temp_path, uploaded_file.name), use_vision=use_vision, api_key=key)
+                    prompt = f"Summarize this document:\n\n{raw[:20000]}"
                     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
+                    st.subheader("Summary Result"); st.markdown(summ)
         with col_c:
-            st.write("") 
-            st.write("") 
-            
-            if "flattened_result" not in st.session_state:
-                st.session_state.flattened_result = None
-            
+            st.write(""); st.write("")
+            if "flattened_result" not in st.session_state: st.session_state.flattened_result = None
             if st.button("📄 Flatten Context"):
                 with st.spinner("Flattening..."):
                     key = st.session_state.get("user_openai_key") or OPENAI_KEY
-                    
                     with open(temp_path, "rb") as f:
                         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
-                        )
-
+                        raw = doc_loader.extract_text_from_file(Wrapper(f.read(), uploaded_file.name), use_vision=use_vision, api_key=key)
                     proc = OutlineProcessor(raw)
                     items = proc.parse()
-                    
                     out_txt = []
                     bar = st.progress(0)
                     for i, item in enumerate(items):
@@ -499,76 +380,52 @@ with tab2:
                         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))
-                    
                     final_flattened_text = "\n".join(out_txt)
-                    st.session_state.flattened_result = {
-                        "text": final_flattened_text,
-                        "source": f"{uploaded_file.name}_flat"
-                    }
+                    st.session_state.flattened_result = {"text": final_flattened_text, "source": f"{uploaded_file.name}_flat"}
                     st.rerun()
-
             if st.session_state.flattened_result:
                 res = st.session_state.flattened_result
                 st.success("Flattening Complete!")
                 st.text_area("Result", res["text"], height=200)
-                
                 if st.button("📥 Index This Flattened Version"):
-                    if not st.session_state.active_index:
-                        st.error("Please select an Active Index in the sidebar.")
+                    if not st.session_state.active_index: st.error("Please select an Active Index.")
                     else:
-                        with st.spinner("Indexing Flattened Text..."):
-                            # FIXED: Added index_name parameter
+                        with st.spinner("Indexing..."):
                             ok, msg = rag_engine.process_and_add_text(
                                 text=res["text"], 
                                 source_name=res["source"], 
                                 username=st.session_state.username,
                                 index_name=st.session_state.active_index
                             )
-                            if ok:
-                                tracker.upload_user_db(st.session_state.username) 
+                            if ok: 
+                                tracker.upload_user_db(st.session_state.username)
                                 st.success(msg)
-                            else:
-                                st.error(msg)
-
+                            else: st.error(msg)
     st.divider()
-    
-    # DB MANAGER
     st.subheader("Database Management")
-    # 1. RESYNC BUTTON (The Fix)
     col_db_1, col_db_2 = st.columns([2, 1])
-    with col_db_1:
-        st.info("If Quiz Mode is failing, your local files might be missing (due to restart).")
+    with col_db_1: st.info("If Quiz Mode is failing, your local files might be missing.")
     with col_db_2:
         if st.button("🔄 Resync from Pinecone"):
-            if not st.session_state.active_index:
-                st.error("Select Index first.")
+            if not st.session_state.active_index: st.error("Select Index first.")
             else:
-                with st.spinner("Downloading memories from Pinecone..."):
-                    ok, msg = rag_engine.rebuild_cache_from_pinecone(
-                        st.session_state.username, 
-                        st.session_state.active_index
-                    )
+                with st.spinner("Downloading memories..."):
+                    ok, msg = rag_engine.rebuild_cache_from_pinecone(st.session_state.username, st.session_state.active_index)
                     if ok: st.success(msg); time.sleep(1); st.rerun()
                     else: st.error(msg)
     st.divider()
-
-    # 2. FILE LIST
-    # This reads from local cache so no index needed
     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']} (Cached)")
             if c2.button("🗑️", key=d['source']):
-                if not st.session_state.active_index:
-                    st.error("Select Index first.")
+                if not st.session_state.active_index: st.error("Select Index first.")
                 else:
                     rag_engine.delete_document(st.session_state.username, d['source'], st.session_state.active_index)
                     tracker.upload_user_db(st.session_state.username)
                     st.rerun()
-    else:
-        st.warning("Local Cache Empty. Click 'Resync' above if you have data in Pinecone.")
+    else: st.warning("Local Cache Empty. Click 'Resync' above if you have data in Pinecone.")
 
 # === TAB 3: QUIZ MODE ===
 with tab3:
@@ -577,104 +434,82 @@ with tab3:
     # 1. MODE SELECTION & RESET LOGIC
     col_mode, col_streak = st.columns([3, 1])
     with col_mode:
-        quiz_mode = st.radio(
-            "Select Quiz Mode:", 
-            ["⚡ Acronym Lightning Round", "📖 Document Deep Dive"],
-            horizontal=True
-        )
-
-    # Initialize Session State Variables if missing
-    if "last_quiz_mode" not in st.session_state:
-        st.session_state.last_quiz_mode = quiz_mode
-        
-    if "quiz_trigger" not in st.session_state:
-        st.session_state.quiz_trigger = False
+        quiz_mode = st.radio("Select Quiz Mode:", ["⚡ Acronym Lightning Round", "📖 Document Deep Dive"], horizontal=True)
+    
+    # New: Focus Topic Input
+    if "Document" in quiz_mode:
+        focus_topic = st.text_input("🎯 Focus Topic (Optional)", placeholder="e.g., PPBE, Shipyards, Radar...", help="Leave empty for random questions.")
+    else:
+        focus_topic = None
+
+    if "last_quiz_mode" not in st.session_state: st.session_state.last_quiz_mode = quiz_mode
+    if "quiz_trigger" not in st.session_state: st.session_state.quiz_trigger = False
 
-    # GHOST IMAGE FIX: Detect Mode Switch
-    # If the user toggled the radio button since the last run, wipe the state.
     if st.session_state.last_quiz_mode != quiz_mode:
         st.session_state.quiz_state["active"] = False
         st.session_state.quiz_state["question_data"] = None
         st.session_state.quiz_state["feedback"] = None
         st.session_state.quiz_state["generated_question_text"] = ""
         st.session_state.last_quiz_mode = quiz_mode
-        st.rerun() # Force a clean refresh immediately
+        st.rerun()
     
-    # Initialize Engine & Shortcut to State
     quiz = QuizEngine()
     qs = st.session_state.quiz_state
-
-    # Display Streak
+    
     with col_streak:
-        st.metric("Current Streak", qs["streak"])
+        st.metric("Streak", qs["streak"])
         if st.button("Reset"): qs["streak"] = 0
 
+    # New: Study Guide Download
+    if st.session_state.quiz_history:
+        with st.expander(f"📚 Review Study Guide ({len(st.session_state.quiz_history)} items)"):
+            st.download_button(
+                "📥 Download Markdown", 
+                generate_study_guide_md(st.session_state.quiz_history), 
+                f"StudyGuide_{datetime.now().strftime('%Y%m%d')}.md"
+            )
+
     st.divider()
 
-    # --- GENERATION FUNCTION ---
     def generate_question():
         with st.spinner("Consulting the Board..."):
-            # MODE A: ACRONYMS
             if "Acronym" in quiz_mode:
                 q_data = quiz.get_random_acronym()
                 if q_data:
-                    qs["active"] = True
-                    qs["question_data"] = q_data
-                    qs["feedback"] = None
-                    qs["generated_question_text"] = q_data["question"]
-                else:
-                    st.error("No acronyms found! Run the extractor first.")
-
-            # MODE B: DOCUMENTS
+                    qs["active"] = True; qs["question_data"] = q_data; qs["feedback"] = None; qs["generated_question_text"] = q_data["question"]
+                else: st.error("No acronyms found! Run the extractor first.")
             else:
                 valid_question_found = False
                 attempts = 0
-                
-                # RETRY LOOP: Increased to 5 attempts to find a good chunk
                 while not valid_question_found and attempts < 5:
                     attempts += 1
-                    q_ctx = quiz.get_document_context(st.session_state.username)
+                    # Pass the focus topic here!
+                    q_ctx = quiz.get_document_context(st.session_state.username, topic_filter=focus_topic)
                     
                     if q_ctx:
                         prompt = quiz.construct_question_generation_prompt(q_ctx["context_text"])
-                        question_text, usage = query_model_universal(
-                            [{"role": "user", "content": prompt}], 
-                            300, model_choice, st.session_state.get("user_openai_key")
-                        )
-                        
-                        # LOGIC UPDATE: Check for 'UNABLE' instead of 'SKIP'
-                        # We also check length to ensure we didn't get a blank response
+                        question_text, usage = query_model_universal([{"role": "user", "content": prompt}], 300, model_choice, st.session_state.get("user_openai_key"))
                         if "UNABLE" not in question_text and len(question_text) > 10:
-                            valid_question_found = True
-                            qs["active"] = True
-                            qs["question_data"] = q_ctx
-                            qs["generated_question_text"] = question_text
-                            qs["feedback"] = None
+                            valid_question_found = True; qs["active"] = True; qs["question_data"] = q_ctx; qs["generated_question_text"] = question_text; qs["feedback"] = None
                 
                 if not valid_question_found:
-                    st.warning("Could not generate a question after 5 attempts. The selected documents might be too sparse or formatted as pure data tables.")
+                    if focus_topic: st.warning(f"No documents found containing '{focus_topic}'. Try a different keyword.")
+                    else: st.warning("Could not generate a question. Documents may be too sparse.")
 
-    # 2. AUTO-TRIGGER (Chained Question Logic)
     if st.session_state.quiz_trigger:
         st.session_state.quiz_trigger = False
         generate_question()
         st.rerun()
 
-    # 3. MANUAL START BUTTON
     if not qs["active"]:
         if st.button("🚀 Generate New Question", type="primary"):
             generate_question()
             st.rerun()
 
-    # 4. QUIZ INTERFACE
     if qs["active"]:
         st.markdown(f"### {qs['generated_question_text']}")
+        if "document" in qs.get("question_data", {}).get("type", ""): st.caption(f"Source: *{qs['question_data']['source_file']}*")
         
-        # Context Hint
-        if "document" in qs.get("question_data", {}).get("type", ""):
-            st.caption(f"Source: *{qs['question_data']['source_file']}*")
-        
-        # Answer Form
         with st.form(key="quiz_response"):
             user_ans = st.text_area("Your Answer:")
             sub = st.form_submit_button("Submit Answer")
@@ -682,57 +517,43 @@ with tab3:
         if sub and user_ans:
             with st.spinner("Grading..."):
                 data = qs["question_data"]
-                
-                # Grading Logic Branch
-                if data["type"] == "acronym":
-                    prompt = quiz.construct_acronym_grading_prompt(
-                        data["term"], data["correct_definition"], user_ans
-                    )
-                else:
-                    prompt = quiz.construct_grading_prompt(
-                        qs["generated_question_text"], user_ans, data["context_text"]
-                    )
+                if data["type"] == "acronym": prompt = quiz.construct_acronym_grading_prompt(data["term"], data["correct_definition"], user_ans)
+                else: prompt = quiz.construct_grading_prompt(qs["generated_question_text"], user_ans, data["context_text"])
 
                 msgs = [{"role": "user", "content": prompt}]
-                grade, _ = query_model_universal(
-                    msgs, 500, model_choice, st.session_state.get("user_openai_key")
-                )
-                
+                grade, _ = query_model_universal(msgs, 500, model_choice, st.session_state.get("user_openai_key"))
                 qs["feedback"] = grade
                 
-                # Streak Logic
-                if "GRADE:** PASS" in grade or "GRADE:** Pass" in grade:
-                    qs["streak"] += 1
-                elif "GRADE:** FAIL" in grade:
-                    qs["streak"] = 0
-                    
+                # Update Streak
+                is_pass = "PASS" in grade
+                if is_pass: qs["streak"] += 1
+                elif "FAIL" in grade: qs["streak"] = 0
+
+                # Save to History
+                correct_info = data['correct_definition'] if data['type'] == 'acronym' else data['context_text']
+                st.session_state.quiz_history.append({
+                    "question": qs["generated_question_text"],
+                    "user_answer": user_ans,
+                    "grade": "PASS" if is_pass else "FAIL",
+                    "context": correct_info
+                })
+                
                 st.rerun()
 
-    # 5. FEEDBACK AREA (Deduplicated)
     if qs["feedback"]:
         st.divider()
-        if "PASS" in qs["feedback"]:
-            st.success("✅ CORRECT")
+        if "PASS" in qs["feedback"]: st.success("✅ CORRECT")
         else:
-            if "FAIL" in qs["feedback"]:
-                st.error("❌ INCORRECT")
-            else:
-                st.warning("⚠️ PARTIAL / COMMENTARY")
-            
+            if "FAIL" in qs["feedback"]: st.error("❌ INCORRECT")
+            else: st.warning("⚠️ PARTIAL / COMMENTARY")
         st.markdown(qs["feedback"])
         
-        # Display Answer Key
         data = qs["question_data"]
-        if data["type"] == "acronym":
-             st.info(f"**Official Definition:** {data['correct_definition']}")
+        if data["type"] == "acronym": st.info(f"**Official Definition:** {data['correct_definition']}")
         elif data["type"] == "document":
-            with st.expander("Show Source Text (Answer Key)"):
-                st.info(data["context_text"])
+            with st.expander("Show Source Text (Answer Key)"): st.info(data["context_text"])
         
-        # Next Question Button
         if st.button("Next Question ➡️"):
             st.session_state.quiz_trigger = True
-            qs["active"] = False
-            qs["question_data"] = None
-            qs["feedback"] = None
+            qs["active"] = False; qs["question_data"] = None; qs["feedback"] = None
             st.rerun()