Update app.py

app.py

@@ -10,9 +10,6 @@ Tabs:
     Results | Genes | Trials | Variants | Graph | Metrics | Visuals
 """
 
-##############################################################################
-# Std-lib / third-party
-##############################################################################
 import os
 import pathlib
 import asyncio
@@ -24,18 +21,13 @@ import plotly.express as px
 from fpdf import FPDF
 from streamlit_agraph import agraph
 
-##############################################################################
-# Internal helpers
-##############################################################################
 from mcp.orchestrator import orchestrate_search, answer_ai_question
 from mcp.workspace import get_workspace, save_query
 from mcp.knowledge_graph import build_agraph
 from mcp.graph_metrics import build_nx, get_top_hubs, get_density
 from mcp.alerts import check_alerts
 
-# ---------------------------------------------------------------------------
-#  Streamlit telemetry directory → /tmp
-# ---------------------------------------------------------------------------
+# Streamlit telemetry directory → /tmp
 os.environ.update({
     "STREAMLIT_DATA_DIR": "/tmp/.streamlit",
     "XDG_STATE_HOME"    : "/tmp",
@@ -46,9 +38,6 @@ pathlib.Path("/tmp/.streamlit").mkdir(parents=True, exist_ok=True)
 ROOT = Path(__file__).parent
 LOGO = ROOT / "assets" / "logo.png"
 
-##############################################################################
-# Utility helpers
-##############################################################################
 
 def _latin1_safe(txt: str) -> str:
     """Coerce UTF-8 → Latin-1 with replacement (for FPDF)."""
@@ -66,7 +55,12 @@ def _pdf(papers: list[dict]) -> bytes:
         pdf.set_font("Helvetica", "B", 11)
         pdf.multi_cell(0, 7, _latin1_safe(f"{i}. {p.get('title','')}"))
         pdf.set_font("Helvetica", "", 9)
-        body = f"{p.get('authors','')}\n{p.get('summary','')}\n{p.get('link','')}\n"
+        # fixed string literal with explicit newlines
+        body = (
+            f"{p.get('authors','')}\n"
+            f"{p.get('summary','')}\n"
+            f"{p.get('link','')}\n"
+        )
         pdf.multi_cell(0, 6, _latin1_safe(body))
         pdf.ln(1)
     return pdf.output(dest="S").encode("latin-1", "replace")
@@ -83,9 +77,6 @@ def _workspace_sidebar():
             with st.expander(f"{i}. {item['query']}"):
                 st.write(item['result']['ai_summary'])
 
-##############################################################################
-# Main UI renderer
-##############################################################################
 
 def render_ui():
     st.set_page_config("MedGenesis AI", layout="wide")
@@ -176,10 +167,9 @@ def render_ui():
         st.subheader("AI summary")
         st.info(res['ai_summary'])
 
-        # --- Genes tab
+    # --- Genes tab
     with tabs[1]:
         st.header("Gene / Variant signals")
-        # Filter out non-dict entries (e.g., errors)
         valid_genes = [g for g in res['genes'] if isinstance(g, dict)]
         if valid_genes:
             for g in valid_genes:
@@ -188,14 +178,12 @@ def render_ui():
         else:
             st.info("No gene signals returned.")
 
-        # MeSH definitions (skip errors/non-strings)
         mesh_list = [d for d in res['mesh_defs'] if isinstance(d, str) and d]
         if mesh_list:
             st.markdown("### MeSH definitions")
             for d in mesh_list:
                 st.write(f"- {d}")
 
-        # DisGeNET links (skip non-dict entries)
         gene_disease = [d for d in res['gene_disease'] if isinstance(d, dict)]
         if gene_disease:
             st.markdown("### DisGeNET links")
@@ -207,8 +195,12 @@ def render_ui():
         trials = res['clinical_trials']
         if not trials:
             st.info("No trials returned (rate-limited or none found).")
-        for t in trials:
-            st.markdown(f"**{t.get('nctId','')}** – {t.get('briefTitle','')}  Phase {t.get('phase','?')} | Status {t.get('status','?')}")
+        else:
+            for t in trials:
+                st.markdown(
+                    f"**{t.get('nctId','')}** – {t.get('briefTitle','')}  "
+                    f"Phase {t.get('phase','?')} | Status {t.get('status','?')}"
+                )
 
     # --- Variants tab
     with tabs[3]:
@@ -239,25 +231,28 @@ def render_ui():
         if years:
             st.plotly_chart(px.histogram(years, nbins=12, title="Publication Year"))
 
-    # Follow-up QA
+    # Follow-up QA (outside any tab)
     st.markdown("---")
-    st.text_input("Ask follow-up question:", key="followup_input")
-
-    def _on_ask():
-        q = st.session_state.followup_input
-        if not q.strip():
-            st.warning("Please type a question first.")
-            return
-        with st.spinner("Querying LLM …"):
-            ans = asyncio.run(answer_ai_question(q, context=st.session_state.last_query, llm=st.session_state.last_llm))
-            st.session_state.followup_response = ans['answer']
-
-    st.button("Ask AI", on_click=_on_ask)
+    input_col, button_col = st.columns([4, 1])
+    with input_col:
+        followup = st.text_input("Ask follow-up question:", key="followup_input")
+    with button_col:
+        if st.button("Ask AI"):
+            if followup.strip():
+                with st.spinner("Querying LLM …"):
+                    ans = asyncio.run(
+                        answer_ai_question(
+                            question=followup,
+                            context=st.session_state.last_query,
+                            llm=st.session_state.last_llm,
+                        )
+                    )
+                    st.session_state.followup_response = ans.get('answer', '')
+            else:
+                st.warning("Please type a question first.")
     if st.session_state.followup_response:
         st.write(st.session_state.followup_response)
 
-##############################################################################
-# Entrypoint
-##############################################################################
+
 if __name__ == "__main__":
     render_ui()