Update app.py

app.py

@@ -16,13 +16,15 @@ except Exception:
     HAS_PYTTSX3 = False
 
 # ============ CONFIG ============
-OPENROUTER_KEY = os.getenv("OPENROUTER_API_KEY")
+# Use st.secrets for Streamlit Cloud deployment, or environment variables for local
+OPENROUTER_KEY = os.getenv("OPENROUTER_API_KEY", st.secrets.get("OPENROUTER_API_KEY"))
 OPENROUTER_MODEL = os.getenv("OPENROUTER_MODEL", "gpt-4o-mini")
-ELEVEN_API_KEY = os.getenv("ELEVEN_API_KEY")
-HUGGINGFACE_KEY = os.getenv("HUGGINGFACE_API_KEY")
+ELEVEN_API_KEY = os.getenv("ELEVEN_API_KEY", st.secrets.get("ELEVEN_API_KEY"))
+HUGGINGFACE_KEY = os.getenv("HUGGINGFACE_API_KEY", st.secrets.get("HUGGINGFACE_API_KEY"))
 HF_MERMAID_MODEL = os.getenv("HF_MERMAID_MODEL", "TroyDoesAI/MermaidStable3B")
 
 # ============ HELPERS ============
+# (rest of the helper functions from your original code are here, unchanged)
 def clean_text(text: str) -> str:
     return re.sub(r"\s+", " ", text or "").strip()
 
@@ -58,11 +60,9 @@ def openrouter_chat(messages: List[dict], model: str = OPENROUTER_MODEL, max_tok
         resp = requests.post(url, json=payload, headers=headers, timeout=30)
         resp.raise_for_status()
         data = resp.json()
-        # robust parsing
         choices = data.get("choices", [])
         if choices:
             c = choices[0]
-            # handle variations
             if "message" in c and isinstance(c["message"], dict):
                 content = c["message"].get("content")
                 if isinstance(content, dict) and "content" in content:
@@ -71,7 +71,6 @@ def openrouter_chat(messages: List[dict], model: str = OPENROUTER_MODEL, max_tok
                     return True, content
             if "text" in c:
                 return True, c["text"]
-        # fallback: try top-level 'text' or 'output'
         if "text" in data:
             return True, data["text"]
         return False, "OpenRouter responded with unexpected shape"
@@ -80,17 +79,11 @@ def openrouter_chat(messages: List[dict], model: str = OPENROUTER_MODEL, max_tok
 
 # ============ Local extractive summarizer (offline) ============
 def extractive_summary(text: str, num_sentences: int = 6) -> str:
-    # Very simple frequency-based extractive summarizer (works offline)
     if not text:
         return ""
-    # split into sentences (naive)
     sentences = re.split(r'(?<=[.!?])\s+', text)
-    # build frequency table of words
     words = re.findall(r'\w+', text.lower())
-    stopwords = set([
-        # minimal stopwords; you can expand
-        "the","and","is","in","to","of","a","that","it","for","on","with","as","are","was","be","by","an","or"
-    ])
+    stopwords = set(["the","and","is","in","to","of","a","that","it","for","on","with","as","are","was","be","by","an","or"])
     freq = {}
     for w in words:
         if w in stopwords or len(w) < 2:
@@ -98,16 +91,13 @@ def extractive_summary(text: str, num_sentences: int = 6) -> str:
         freq[w] = freq.get(w, 0) + 1
     if not freq:
         return "Unable to summarize (text too short)."
-    # score sentences
     sent_scores = []
     for s in sentences:
         s_words = re.findall(r'\w+', s.lower())
         score = sum(freq.get(w, 0) for w in s_words)
         sent_scores.append((score, s))
-    # pick top sentences
     sent_scores.sort(reverse=True, key=lambda x: x[0])
     chosen = [s for _, s in sent_scores[:num_sentences]]
-    # preserve approximate original order
     chosen_sorted = sorted(chosen, key=lambda s: text.find(s))
     bullets = "\n".join(f"- {clean_text(s)}" for s in chosen_sorted if s.strip())
     return bullets if bullets else clean_text(" ".join(chosen_sorted))
@@ -133,13 +123,11 @@ def pyttsx3_tts_file(text: str):
         return False, "pyttsx3 not installed"
     try:
         engine = pyttsx3.init()
-        # create a temp wav file
         tf = tempfile.NamedTemporaryFile(delete=False, suffix=".wav")
         tf_name = tf.name
         tf.close()
         engine.save_to_file(text, tf_name)
         engine.runAndWait()
-        # read bytes
         with open(tf_name, "rb") as f:
             b = f.read()
         return True, b
@@ -158,7 +146,6 @@ def call_hf_mermaid(prompt: str, model: str = HF_MERMAID_MODEL):
         if not r.ok:
             return False, f"HuggingFace returned {r.status_code}: {r.text[:300]}"
         j = r.json()
-        # extract text
         if isinstance(j, list) and len(j) > 0 and isinstance(j[0], dict) and "generated_text" in j[0]:
             return True, j[0]["generated_text"]
         if isinstance(j, str):
@@ -176,10 +163,8 @@ def generate_mermaid_from_summary(summary: str):
               "Output only the Mermaid code block. Summary:\n\n" + summary)
     ok, hf_out = call_hf_mermaid(prompt)
     if ok:
-        # try to strip triple-backtick wrapper if present
         m = re.search(r"```(?:mermaid)?\n([\s\S]+?)```", hf_out, re.IGNORECASE)
         return hf_out if m is None else m.group(1).strip()
-    # fallback local
     lines = re.split(r"\n+|-{1,}\s*|•\s*", summary)
     nodes = [clean_text(l) for l in lines if clean_text(l)]
     nodes = nodes[:8]
@@ -188,7 +173,7 @@ def generate_mermaid_from_summary(summary: str):
     mermaid = "flowchart TD\n"
     for i, n in enumerate(nodes):
         node_text = n.replace('"', "'")[:80]
-        mermaid += '  A{idx}["{text}"]\n'.format(idx=i, text=node_text)
+        mermaid += f'  A{i}["{node_text}"]\n'
     for i in range(len(nodes) - 1):
         mermaid += f"  A{i} --> A{i+1}\n"
     return mermaid
@@ -211,139 +196,158 @@ def render_mermaid(mermaid_code: str, height: int = 420):
 st.set_page_config(page_title="PDF Q&A resilient", layout="wide")
 st.title("📄 PDF Q&A — resilient (OpenRouter → local fallback)")
 
+# Session state initialization
+if 'text_data' not in st.session_state:
+    st.session_state.text_data = None
+if 'summary' not in st.session_state:
+    st.session_state.summary = None
+if 'mermaid' not in st.session_state:
+    st.session_state.mermaid = None
+if 'diagnostics' not in st.session_state:
+    st.session_state.diagnostics = {"summary_source": None, "mermaid_source": None, "tts_source": None}
+
+def process_pdf():
+    uploaded_file = st.session_state.uploaded_file
+    if uploaded_file:
+        try:
+            with st.spinner("Extracting text from PDF..."):
+                raw_text = extract_text_from_pdf(uploaded_file)
+                st.session_state.text_data = raw_text
+                st.success(f"Extracted {len(raw_text)} characters")
+        except Exception as e:
+            st.error(f"PDF extraction failed: {e}")
+            st.session_state.text_data = None
+    else:
+        st.session_state.text_data = None
+
+def generate_outputs():
+    raw_text = st.session_state.text_data
+    if not raw_text:
+        st.error("No text available to process. Please upload a PDF.")
+        return
+
+    used = st.session_state.diagnostics
+    
+    # Summarize logic
+    with st.spinner("Trying OpenRouter summarization..."):
+        prompt = f"Summarize the following text in 5-8 concise bullets:\n\n{raw_text[:15000]}"
+        messages = [{"role": "system", "content": "You are a concise summarizer."},
+                    {"role": "user", "content": prompt}]
+        ok, out = openrouter_chat(messages, max_tokens=400, model=OPENROUTER_MODEL)
+        if ok:
+            used["summary_source"] = "openrouter"
+            summary = out
+        else:
+            used["summary_source"] = f"fallback_local (reason: {out})"
+            summary = extractive_summary(raw_text, num_sentences=6)
+    st.session_state.summary = summary
+    st.session_state.diagnostics = used
+
+    # Mermaid logic
+    with st.spinner("Generating Mermaid diagram (HF → local fallback)..."):
+        mermaid = generate_mermaid_from_summary(summary)
+        used["mermaid_source"] = "huggingface" if HUGGINGFACE_KEY and mermaid.strip().startswith(("flowchart","graph")) else "local"
+    st.session_state.mermaid = mermaid
+    st.session_state.diagnostics = used
+    
+    st.success("Summary and Diagram generated!")
+
+# UI layout
 c1, c2, c3 = st.columns(3)
 with c1:
     st.write("OpenRouter:")
-    if OPENROUTER_KEY:
-        st.success("Key present")
-    else:
-        st.error("Key missing — will use local summarizer/Q&A fallback")
+    st.success("Key present") if OPENROUTER_KEY else st.error("Key missing — will use local summarizer/Q&A fallback")
 with c2:
     st.write("Hugging Face:")
-    if HUGGINGFACE_KEY:
-        st.success("Key present (optional)")
-    else:
-        st.info("Key missing — using local Mermaid fallback")
+    st.success("Key present (optional)") if HUGGINGFACE_KEY else st.info("Key missing — using local Mermaid fallback")
 with c3:
     st.write("Audio:")
     if ELEVEN_API_KEY:
         st.success("ElevenLabs key present (preferred)")
+    elif HAS_PYTTSX3:
+        st.info("Using local pyttsx3 fallback TTS")
     else:
-        if HAS_PYTTSX3:
-            st.info("Using local pyttsx3 fallback TTS")
-        else:
-            st.info("No ElevenLabs key and pyttsx3 not available — audio will be disabled")
+        st.info("No ElevenLabs key and pyttsx3 not available")
+        
+st.file_uploader("Upload a PDF", type=["pdf"], key='uploaded_file', on_change=process_pdf)
 
-uploaded_file = st.file_uploader("Upload a PDF", type=["pdf"])
-if uploaded_file:
-    try:
-        with st.spinner("Extracting text from PDF..."):
-            raw_text = extract_text_from_pdf(uploaded_file)
-    except Exception as e:
-        st.error(f"PDF extraction failed: {e}")
-        raw_text = ""
+if st.session_state.text_data:
+    st.button("Summarize & Diagram", on_click=generate_outputs)
+    
+if st.session_state.summary:
+    st.subheader("📌 Summary")
+    st.write(st.session_state.summary)
+    st.markdown(f"**Summary source:** {st.session_state.diagnostics['summary_source']}")
+    
+    st.subheader("🗺️ Summary Diagram")
+    render_mermaid(st.session_state.mermaid, height=460)
+    st.code(st.session_state.mermaid, language="mermaid")
+    st.markdown(f"**Mermaid source:** {st.session_state.diagnostics['mermaid_source']}")
 
-    if not raw_text:
-        st.warning("No text extracted (maybe scanned PDF). Use OCR if needed.")
-    else:
-        st.success(f"Extracted {len(raw_text)} characters")
-        # Summarize & diagram
-        if st.button("Summarize & Diagram (tries OpenRouter → fallback)"):
-            used = {"summary": None, "summary_source": None, "tts_source": None, "mermaid_source": None}
-            # Try OpenRouter summary first
-            with st.spinner("Trying OpenRouter summarization..."):
-                prompt = f"Summarize the following text in 5-8 concise bullets:\n\n{raw_text[:15000]}"
-                messages = [{"role": "system", "content": "You are a concise summarizer."},
-                            {"role": "user", "content": prompt}]
-                ok, out = openrouter_chat(messages, max_tokens=400, model=OPENROUTER_MODEL)
+    st.write("### TTS Audio")
+    if st.checkbox("Generate audio for summary"):
+        with st.spinner("Attempting TTS..."):
+            audio_bytes = None
+            if ELEVEN_API_KEY:
+                ok, out = eleven_tts_bytes(st.session_state.summary)
                 if ok:
-                    used["summary_source"] = "openrouter"
-                    summary = out
+                    st.session_state.diagnostics["tts_source"] = "elevenlabs"
+                    audio_bytes = out
                 else:
-                    # fallback to local
-                    used["summary_source"] = f"fallback_local (reason: {out})"
-                    summary = extractive_summary(raw_text, num_sentences=6)
-
-            st.subheader("📌 Summary")
-            st.write(summary)
-            st.markdown(f"**Summary source:** {used['summary_source']}")
-
-            # Mermaid
-            with st.spinner("Generating Mermaid diagram (HF → local fallback)..."):
-                mermaid = generate_mermaid_from_summary(summary)
-                # determine mermaid source roughly
-                used["mermaid_source"] = "huggingface" if HUGGINGFACE_KEY and mermaid.strip().startswith(("flowchart","graph")) else "local"
-            st.subheader("🗺️ Summary Diagram")
-            render_mermaid(mermaid, height=460)
-            st.code(mermaid, language="mermaid")
-            st.markdown(f"**Mermaid source:** {used['mermaid_source']}")
-
-            # TTS: try ElevenLabs first, then pyttsx3
-            if st.checkbox("Enable audio for summary (try ElevenLabs → fallback)"):
-                with st.spinner("Attempting TTS..."):
-                    if ELEVEN_API_KEY:
-                        ok, out = eleven_tts_bytes(summary)
-                        if ok:
-                            used["tts_source"] = "elevenlabs"
-                            st.audio(out, format="audio/mp3")
-                        else:
-                            # record reason and try pyttsx3
-                            used["tts_source"] = f"elevenlabs_failed ({out})"
-                            if HAS_PYTTSX3:
-                                ok2, out2 = pyttsx3_tts_file(summary)
-                                if ok2:
-                                    used["tts_source"] = "pyttsx3"
-                                    st.audio(out2, format="audio/wav")
-                                else:
-                                    st.error(f"TTS fallback failed: {out2}")
-                            else:
-                                st.error("ElevenLabs TTS failed and pyttsx3 not available.")
-                    else:
-                        if HAS_PYTTSX3:
-                            ok2, out2 = pyttsx3_tts_file(summary)
-                            if ok2:
-                                used["tts_source"] = "pyttsx3"
-                                st.audio(out2, format="audio/wav")
-                            else:
-                                st.error(f"pyttsx3 TTS failed: {out2}")
-                        else:
-                            st.info("No TTS available (no ElevenLabs key and pyttsx3 missing).")
+                    st.session_state.diagnostics["tts_source"] = f"elevenlabs_failed ({out})"
+                    if HAS_PYTTSX3:
+                        ok2, out2 = pyttsx3_tts_file(st.session_state.summary)
+                        if ok2:
+                            st.session_state.diagnostics["tts_source"] = "pyttsx3"
+                            audio_bytes = out2
+            elif HAS_PYTTSX3:
+                ok2, out2 = pyttsx3_tts_file(st.session_state.summary)
+                if ok2:
+                    st.session_state.diagnostics["tts_source"] = "pyttsx3"
+                    audio_bytes = out2
+            
+            if audio_bytes:
+                if st.session_state.diagnostics["tts_source"] == "elevenlabs":
+                    st.audio(audio_bytes, format="audio/mp3")
+                else:
+                    st.audio(audio_bytes, format="audio/wav")
+            else:
+                st.error("Audio generation failed. Check your API keys and local setup.")
+            st.markdown(f"**TTS source:** {st.session_state.diagnostics['tts_source']}")
 
-            st.write("### Diagnostics")
-            st.json(used)
+    st.write("### Diagnostics")
+    st.json(st.session_state.diagnostics)
 
-        # Q&A box (tries OpenRouter, otherwise local naive search)
-        query = st.text_input("Ask a question about the PDF (press Enter):")
-        if query:
+    st.markdown("---")
+    st.subheader("❓ Q&A")
+    query = st.text_input("Ask a question about the PDF:")
+    if query:
+        with st.spinner("Processing your question..."):
             if OPENROUTER_KEY:
-                with st.spinner("Asking OpenRouter..."):
-                    prompt = f"Context:\n{raw_text[:15000]}\n\nQuestion: {query}\nAnswer concisely."
-                    messages = [{"role": "system", "content": "You are a helpful assistant."},
-                                {"role": "user", "content": prompt}]
-                    ok, out = openrouter_chat(messages, max_tokens=600, model=OPENROUTER_MODEL)
-                    if ok:
-                        st.subheader("💡 Answer (OpenRouter)")
-                        st.write(out)
+                prompt = f"Context:\n{st.session_state.text_data[:15000]}\n\nQuestion: {query}\nAnswer concisely."
+                messages = [{"role": "system", "content": "You are a helpful assistant."},
+                            {"role": "user", "content": prompt}]
+                ok, out = openrouter_chat(messages, max_tokens=600, model=OPENROUTER_MODEL)
+                if ok:
+                    st.info("Answer from OpenRouter:")
+                    st.write(out)
+                else:
+                    st.warning(f"OpenRouter failed: {out}\nFalling back to local Q&A.")
+                    sentences = re.split(r'(?<=[.!?])\s+', st.session_state.text_data)
+                    q_words = re.findall(r'\w+', query.lower())[:3]
+                    matches = [s for s in sentences if all(w in s.lower() for w in q_words)]
+                    if matches:
+                        st.info("Answer from local fallback:")
+                        st.write(matches[:3])
                     else:
-                        st.warning(f"OpenRouter failed: {out}\nFalling back to local Q&A.")
-                        # fallback to very naive local answer: search for query words in text and return matching sentences
-                        q = query.lower()
-                        sentences = re.split(r'(?<=[.!?])\s+', raw_text)
-                        matches = [s for s in sentences if all(w in s.lower() for w in re.findall(r'\w+', q)[:3])]
-                        if matches:
-                            st.subheader("💡 Answer (local fallback)")
-                            st.write(matches[:3])
-                        else:
-                            st.info("No good local match found.")
+                        st.info("No good local match found.")
             else:
-                st.info("OpenRouter key missing — using local Q&A fallback.")
-                q = query.lower()
-                sentences = re.split(r'(?<=[.!?])\s+', raw_text)
-                matches = [s for s in sentences if all(w in s.lower() for w in re.findall(r'\w+', q)[:3])]
+                st.info("OpenRouter key missing. Using local Q&A fallback.")
+                sentences = re.split(r'(?<=[.!?])\s+', st.session_state.text_data)
+                q_words = re.findall(r'\w+', query.lower())[:3]
+                matches = [s for s in sentences if all(w in s.lower() for w in q_words)]
                 if matches:
-                    st.subheader("💡 Answer (local fallback)")
+                    st.info("Answer from local fallback:")
                     st.write(matches[:3])
                 else:
-                    st.info("No good local match found.")
-else:
-    st.info("Upload a PDF to begin.")
+                    st.info("No good local match found.")