feat: add context memory to reuse outputs across modules

app.py

@@ -1,17 +1,84 @@
 import streamlit as st
 import os
 import glob
+import re
+import base64
 
+# ==============================
 # MUST BE FIRST
-st.set_page_config(page_title="LexTransition AI", page_icon="⚖️", layout="wide", initial_sidebar_state="expanded")
+# ==============================
+st.set_page_config(
+    page_title="LexTransition AI",
+    page_icon="⚖️",
+    layout="wide",
+    initial_sidebar_state="expanded"
+)
 
-# --- Component Imports ---
+# ==============================
+# COMPONENT IMPORTS
+# ==============================
 from app.components.theme import init_theme, load_css, apply_theme_css
 from app.components.header import render_header
 from app.components.sidebar import render_sidebar
 from app.components.footer import render_footer
 
-# --- Global Definitions ---
+# ==============================
+# ENGINE LOADING
+# ==============================
+IMPORT_ERROR = None
+try:
+    from engine.ocr_processor import extract_text, available_engines
+    from engine.mapping_logic import map_ipc_to_bns, add_mapping
+    from engine.rag_engine import search_pdfs, add_pdf, index_pdfs
+    from engine.db import (
+        import_mappings_from_csv,
+        import_mappings_from_excel,
+        export_mappings_to_json,
+        export_mappings_to_csv,
+    )
+    from engine.comparator import compare_ipc_bns
+    from engine import glossary as glossary_engine
+    ENGINES_AVAILABLE = True
+except Exception as e:
+    IMPORT_ERROR = str(e)
+    ENGINES_AVAILABLE = False
+
+# LLM fallback
+try:
+    from engine.llm import summarize as llm_summarize
+except Exception:
+    def llm_summarize(text, question=None):
+        return None
+
+# ==============================
+# CONTEXT MEMORY
+# ==============================
+if "context_memory" not in st.session_state:
+    st.session_state.context_memory = ""
+
+def store_context(text: str):
+    if text:
+        st.session_state.context_memory = text
+
+def clear_context():
+    st.session_state.context_memory = ""
+
+# ==============================
+# CLEANUP TEMP AUDIO
+# ==============================
+def cleanup_temp_audio():
+    if os.path.exists("temp_audio"):
+        for f in glob.glob("temp_audio/*.wav"):
+            try:
+                os.remove(f)
+            except Exception:
+                pass
+
+cleanup_temp_audio()
+
+# ==============================
+# NAVIGATION SETUP
+# ==============================
 NAV_ITEMS = [
     ("Home", "Home"),
     ("Mapper", "IPC -> BNS Mapper"),
@@ -24,12 +91,8 @@
     ("Privacy", "Privacy Policy"),
 ]
 
-# --- Initialization & Cleanup ---
-def cleanup_temp_audio():
-    if os.path.exists("temp_audio"):
-        for audio_file in glob.glob("temp_audio/*.wav"):
-            try: os.remove(audio_file)
-            except Exception: pass
+if "current_page" not in st.session_state:
+    st.session_state.current_page = "Home"
 
 def get_url_page():
     try:
@@ -38,74 +101,87 @@ def get_url_page():
     except Exception:
         return None
 
-# --- Session State Setup ---
-cleanup_temp_audio()
-init_theme()
-
-if "current_page" not in st.session_state:
-    st.session_state.current_page = "Home"
-
 url_page = get_url_page()
+
 if "pending_page" in st.session_state:
     st.session_state.current_page = st.session_state.pop("pending_page")
-elif url_page in dict(NAV_ITEMS).keys():
+elif url_page in dict(NAV_ITEMS):
     st.session_state.current_page = url_page
 
 current_page = st.session_state.current_page
 
-# --- Engine Pre-load (Silent) ---
-try:
-    from engine.rag_engine import index_pdfs
-    if not st.session_state.get("pdf_indexed"):
+# ==============================
+# INITIAL PDF INDEXING
+# ==============================
+if ENGINES_AVAILABLE and not st.session_state.get("pdf_indexed"):
+    try:
         index_pdfs("law_pdfs")
         st.session_state.pdf_indexed = True
-except Exception:
-    pass # Degrade gracefully if engine fails
+    except Exception:
+        pass
 
-# --- UI Setup ---
+# ==============================
+# UI SETUP
+# ==============================
 load_css("assets/styles.css")
+init_theme()
 apply_theme_css()
 render_sidebar(NAV_ITEMS)
 render_header(NAV_ITEMS, current_page)
 
-# ============================================================================
+if IMPORT_ERROR:
+    st.error(f"⚠️ Engines failed to load.\n\nError: `{IMPORT_ERROR}`")
+
+# ==============================
 # PAGE ROUTER
-# ============================================================================
+# ==============================
 try:
+
     if current_page == "Home":
         from app.pages.home import render_home_page
         render_home_page()
+
     elif current_page == "Mapper":
         from app.pages.mapper import render_mapper_page
         render_mapper_page()
+
     elif current_page == "OCR":
         from app.pages.ocr import render_ocr_page
         render_ocr_page()
+
     elif current_page == "Glossary":
         from app.pages.glossary import render_glossary_page
         render_glossary_page()
+
     elif current_page == "Fact":
         from app.pages.fact_checker import render_fact_checker_page
         render_fact_checker_page()
+
     elif current_page == "Community":
         from app.pages.community import render_community_page
         render_community_page()
+
     elif current_page == "Settings":
         from app.pages.settings import render_settings_page
         render_settings_page()
+
     elif current_page == "FAQ":
         from app.pages.faq import render_faq_page
         render_faq_page()
+
     elif current_page == "Privacy":
         from app.pages.privacy import render_privacy_page
         render_privacy_page()
+
     else:
         st.error("Page not found.")
 
 except Exception as e:
-    st.error("🚨 An unexpected error occurred.")
+    st.error("🚨 Unexpected error occurred.")
     st.exception(e)
 
-# --- Footer ---
+# ==============================
+# FOOTER
+# ==============================
 st.divider()
 render_footer()

app/pages/fact_checker.py

@@ -54,9 +54,14 @@ def render_fact_checker_page():
 
     with col1:
         # Bind the value to our session state so Voice Input auto-fills this box
+        default_question = (
+            st.session_state.get("context_memory", "")
+            or st.session_state.get("fact_search_val", "")
+        )
+
         user_question = st.text_input(
-            "Question", 
-            value=st.session_state.get('fact_search_val', ''),
+            "Question",
+            value=default_question,
             placeholder="e.g., penalty for cheating?",
             label_visibility="collapsed"
         )
@@ -104,7 +109,20 @@ def render_fact_checker_page():
             st.session_state['fact_search_val'] = voice_query
             st.session_state['fact_auto_search'] = True 
             st.rerun()
-
+    # ===== Show Context Memory =====
+    if st.session_state.get("context_memory"):
+        col_m1, col_m2 = st.columns(2)
+
+        with col_m1:
+            st.info("📌 Using stored context")
+
+        with col_m2:
+            if st.button("🧹 Clear Memory"):
+                st.session_state.context_memory = ""
+                st.rerun()
+
+        with st.expander("📌 Stored Context"):
+            st.write(st.session_state.context_memory)
     # --- Auto-Search Trigger ---
     if st.session_state.get('fact_auto_search'):
         verify_btn = True

app/pages/mapper.py

@@ -190,6 +190,19 @@ def render_mapper_page():
 
         st.write("###")
 
+        # ===== Context Memory Buttons =====
+        col_ctx1, col_ctx2 = st.columns(2)
+
+        with col_ctx1:
+            if st.button("📥 Use in Fact Checker"):
+                st.session_state.context_memory = f"IPC {ipc} mapped to {bns}. Notes: {notes}"
+                st.success("Stored in memory. Go to Fact Checker.")
+
+        with col_ctx2:
+            if st.button("🧹 Clear Memory"):
+                st.session_state.context_memory = ""
+                st.info("Context cleared.")
+
         # --- STEP 3: Action Buttons ---
         col_a, col_b, col_c, col_d = st.columns(4)
 

app/pages/ocr.py

@@ -69,6 +69,9 @@ def render_ocr_page():
 
                 st.success("✅ Text extraction completed!")
                 st.text_area("Extracted Text", extracted, height=300)
+                if st.button("📥 Use in Fact Checker"):
+                st.session_state.context_memory = extracted
+                st.success("OCR text stored for reuse.")
 
                 copy_to_clipboard(extracted, "Copy OCR Text")
 
@@ -174,4 +177,4 @@ def render_ocr_page():
                     st.warning("⚠ AI Engine failed to generate summary.")
 
             except Exception as e:
-                st.error(f"❌ Error during processing: {str(e)}")
+                st.error(f"❌ Error during processing: {str(e)}")