Update app.py

app.py

@@ -9,10 +9,10 @@ from typing import List, Optional
 
 # ============ CONFIG =============
 OPENROUTER_KEY = os.getenv("OPENROUTER_API_KEY")
-OPENROUTER_MODEL = os.getenv("OPENROUTER_MODEL", "gpt-4o-mini")  # change if you prefer
+OPENROUTER_MODEL = os.getenv("OPENROUTER_MODEL", "gpt-4o-mini")
 ELEVEN_API_KEY = os.getenv("ELEVEN_API_KEY")
-HUGGINGFACE_KEY = os.getenv("HUGGINGFACE_API_KEY")  # optional: if set, we'll call a HF mermaid model
-HF_MERMAID_MODEL = os.getenv("HF_MERMAID_MODEL", "TroyDoesAI/MermaidStable3B")  # example community model
+HUGGINGFACE_KEY = os.getenv("HUGGINGFACE_API_KEY")
+HF_MERMAID_MODEL = os.getenv("HF_MERMAID_MODEL", "TroyDoesAI/MermaidStable3B")
 
 # ============ HELPERS ============
 def clean_text(text: str) -> str:
@@ -38,59 +38,43 @@ def chunk_text_by_chars(text: str, chunk_size: int = 3000, overlap: int = 200) -
         start = max(end - overlap, end)
     return chunks
 
-# ---------- OpenRouter chat (replacement for openai.ChatCompletion) ----------
+# ---------- OpenRouter chat ----------
 def openrouter_chat(messages: List[dict], model: str = OPENROUTER_MODEL, max_tokens: int = 800, temperature: float = 0.2) -> str:
-    """
-    Send messages (OpenAI-style) to OpenRouter's chat completions endpoint.
-    Requires OPENROUTER_API_KEY in ENV.
-    """
     if not OPENROUTER_KEY:
         raise RuntimeError("OPENROUTER_API_KEY not set")
-
     url = "https://api.openrouter.ai/v1/chat/completions"
     headers = {"Authorization": f"Bearer {OPENROUTER_KEY}", "Content-Type": "application/json"}
-    payload = {
-        "model": model,
-        "messages": messages,
-        "max_tokens": max_tokens,
-        "temperature": temperature,
-    }
+    payload = {"model": model, "messages": messages, "max_tokens": max_tokens, "temperature": temperature}
     resp = requests.post(url, json=payload, headers=headers, timeout=60)
     try:
         resp.raise_for_status()
     except Exception as e:
         raise RuntimeError(f"OpenRouter API error: {resp.status_code} {resp.text}") from e
-
     data = resp.json()
-    # robustly extract text
     text = ""
-    try:
-        choices = data.get("choices", [])
-        if choices:
-            c = choices[0]
-            # OpenRouter returns similar shape to OpenAI
-            if "message" in c and "content" in c["message"]:
-                text = c["message"]["content"]
-            elif "text" in c:
-                text = c["text"]
-    except Exception:
-        text = ""
+    choices = data.get("choices", [])
+    if choices:
+        c = choices[0]
+        if "message" in c and isinstance(c["message"], dict) and "content" in c["message"]:
+            # some OpenRouter shapes put content directly
+            content = c["message"]["content"]
+            # content might be dict or string; handle both
+            if isinstance(content, dict) and "content" in content:
+                text = content["content"]
+            elif isinstance(content, str):
+                text = content
+        elif "text" in c:
+            text = c["text"]
     return text or ""
 
 def ask_model_for_summary(text: str) -> str:
     prompt = f"Summarize the following text clearly and concisely (bullet points, 5-8 bullets max):\n\n{text}"
-    messages = [
-        {"role": "system", "content": "You are a concise summarizer."},
-        {"role": "user", "content": prompt},
-    ]
+    messages = [{"role": "system", "content": "You are a concise summarizer."}, {"role": "user", "content": prompt}]
     return openrouter_chat(messages, max_tokens=400)
 
 def ask_model_question(question: str, context: str) -> str:
     prompt = f"Context:\n{context}\n\nQuestion: {question}\nAnswer in a concise helpful way."
-    messages = [
-        {"role": "system", "content": "You are a helpful assistant."},
-        {"role": "user", "content": prompt},
-    ]
+    messages = [{"role": "system", "content": "You are a helpful assistant."}, {"role": "user", "content": prompt}]
     return openrouter_chat(messages, max_tokens=600)
 
 # ---------- ElevenLabs TTS ----------
@@ -107,12 +91,8 @@ def text_to_speech_eleven(text: str, voice_id: str = "pnCWbS8Aqipqqr5wzjuy") ->
         st.warning(f"ElevenLabs TTS failed: {r.status_code} {r.text[:300]}")
         return None
 
-# ---------- Mermaid generation (Hugging Face model optional) ----------
+# ---------- Hugging Face mermaid (optional) ----------
 def call_hf_mermaid(prompt: str, model: str = HF_MERMAID_MODEL) -> Optional[str]:
-    """
-    If HUGGINGFACE_KEY is set, call Hugging Face Inference API for model that outputs Mermaid or Mermaid-like code.
-    Many community models/Spaces are simple text-output LLMs that can be prompted to return mermaid code.
-    """
     if not HUGGINGFACE_KEY:
         return None
     url = f"https://api-inference.huggingface.co/models/{model}"
@@ -123,16 +103,13 @@ def call_hf_mermaid(prompt: str, model: str = HF_MERMAID_MODEL) -> Optional[str]
         st.warning(f"Hugging Face model call failed: {r.status_code} {r.text[:300]}")
         return None
     j = r.json()
-    # shape varies by model; try to extract text
     if isinstance(j, dict) and "error" in j:
         st.warning(f"Hugging Face error: {j['error']}")
         return None
     if isinstance(j, list) and len(j) > 0 and isinstance(j[0], dict) and "generated_text" in j[0]:
         return j[0]["generated_text"]
-    # some models return plain text in str
     if isinstance(j, str):
         return j
-    # fallback: try to get 'output' key
     if isinstance(j, dict):
         for k in ("generated_text", "output", "text"):
             if k in j:
@@ -140,47 +117,32 @@ def call_hf_mermaid(prompt: str, model: str = HF_MERMAID_MODEL) -> Optional[str]
     return None
 
 def generate_mermaid_from_summary(summary: str) -> str:
-    """
-    Try HF model first (if key provided). If not available or fails, produce a clean Mermaid flowchart locally.
-    We'll create a simple flow: split summary into sentences / bullets and link them sequentially.
-    """
-    # first try HF
-    prompt = (
-        "Given the following concise summary, produce a Mermaid flowchart (use 'graph TD' or 'flowchart TD' syntax). "
-        "Output only the Mermaid code block (no extra explanation). Summary:\n\n" + summary
-    )
+    prompt = ("Given the following concise summary, produce a Mermaid flowchart (use 'graph TD' or 'flowchart TD' syntax). "
+              "Output only the Mermaid code block (no extra explanation). Summary:\n\n" + summary)
     hf_output = call_hf_mermaid(prompt)
     if hf_output:
-        # try to extract just the mermaid text
-        # if the model wrapped in ```mermaid ... ``` try to strip
         m = re.search(r"```(?:mermaid)?\n([\s\S]+?)```", hf_output, re.IGNORECASE)
         if m:
             return m.group(1).strip()
         return hf_output.strip()
 
-    # fallback local generator
-    # split by bullet/newline or sentences
+    # fallback: create simple sequential flowchart
     lines = re.split(r"\n+|-{1,}\s*|•\s*", summary)
     nodes = [clean_text(l) for l in lines if clean_text(l)]
-    # keep a reasonable number
     nodes = nodes[:8]
     if not nodes:
         nodes = ["Summary empty"]
     mermaid = "flowchart TD\n"
-    # create nodes with safe ids
     for i, n in enumerate(nodes):
-        # short id
-        mermaid += f'  A{i}["{n.replace(\'"\', "\\\'")[:80]}"]\n'
+        # sanitize node text: replace double quotes with single quotes to avoid breaking mermaid quotes
+        node_text = n.replace('"', "'")[:80]
+        # use .format to avoid backslashes inside f-string expressions
+        mermaid += '  A{idx}["{text}"]\n'.format(idx=i, text=node_text)
     for i in range(len(nodes) - 1):
         mermaid += f"  A{i} --> A{i+1}\n"
     return mermaid
 
-# ---------- Render mermaid in browser ----------
 def render_mermaid(mermaid_code: str, height: int = 400):
-    """
-    Render Mermaid chart client-side using mermaid.js in an HTML component.
-    """
-    # wrap in HTML that loads mermaid CDN
     html = f"""
     <div id="mermaid-target">
       <pre class="mermaid">
@@ -192,13 +154,13 @@ def render_mermaid(mermaid_code: str, height: int = 400):
       mermaid.initialize({{startOnLoad:true}});
     </script>
     """
+    # render
     st.components.v1.html(html, height=height, scrolling=True)
 
 # ============ STREAMLIT UI ============
 st.set_page_config(page_title="PDF Q&A + Summary Diagram", layout="wide")
 st.title("📄 PDF Q&A + Summary Diagram + Audio")
 
-# API status
 c1, c2, c3 = st.columns(3)
 with c1:
     if OPENROUTER_KEY:
@@ -232,31 +194,26 @@ if uploaded_file:
         if st.button("Summarize and generate diagram"):
             try:
                 with st.spinner("Summarizing with OpenRouter..."):
-                    # limit to avoid huge inputs
                     to_sum = raw_text[:15000]
                     summary = ask_model_for_summary(to_sum)
                 st.subheader("📌 Summary")
                 st.write(summary)
 
-                # TTS summary
                 audio = text_to_speech_eleven(summary)
                 if audio:
                     st.audio(audio, format="audio/mp3")
                 elif not ELEVEN_API_KEY:
                     st.info("TTS not available (ELEVEN_API_KEY missing).")
 
-                # produce mermaid
                 mermaid_code = generate_mermaid_from_summary(summary)
                 st.subheader("🗺️ Summary Diagram (Mermaid)")
                 render_mermaid(mermaid_code, height=480)
-                # also show the raw mermaid code for copy/paste
                 st.markdown("**Mermaid code (copy/paste):**")
                 st.code(mermaid_code, language="mermaid")
 
             except Exception as e:
                 st.error(f"Summarize/diagram generation failed: {e}")
 
-        # Q&A box
         query = st.text_input("Ask a question about the PDF (use Enter):")
         if query:
             if not OPENROUTER_KEY:
@@ -266,7 +223,7 @@ if uploaded_file:
                     with st.spinner("Answering via OpenRouter..."):
                         chunks = chunk_text_by_chars(raw_text, chunk_size=3000, overlap=200)
                         answers = []
-                        for c in chunks[:3]:  # limit to 3 chunks
+                        for c in chunks[:3]:
                             a = ask_model_question(query, c)
                             if a:
                                 answers.append(a)