Update app.py

app.py

@@ -3,134 +3,126 @@ import os
 
 from config import (
     OPENAI_API_KEY,
-    GEMINI_API_KEY,
-    DEFAULT_CHUNK_SIZE,
+    OPENAI_DEFAULT_MODEL,
+    MAX_PUBMED_RESULTS
 )
-from models import configure_llms, openai_chat, gemini_chat
-from pubmed_utils import (
-    search_pubmed,
-    fetch_pubmed_abstracts,
-    chunk_and_summarize
+from pubmed_rag import (
+    search_pubmed, fetch_pubmed_abstracts, chunk_and_summarize,
+    upsert_documents, semantic_search
 )
-from image_pipeline import load_image_model, analyze_image
+from models import chat_with_openai
+from image_pipeline import analyze_medical_image
 
 ###############################################################################
-#                             STREAMLIT PAGE CONFIG                           #
+#                               STREAMLIT SETUP                               #
 ###############################################################################
-st.set_page_config(page_title="RAG + Image Captioning Demo", layout="wide")
+st.set_page_config(page_title="Advanced Medical AI", layout="wide")
 
-###############################################################################
-#                           INITIALIZE & LOAD MODELS                          #
-###############################################################################
-@st.cache_resource
-def initialize_app():
-    """
-    Configure LLMs (OpenAI/Gemini) and load the image captioning model once.
-    """
-    configure_llms()
-    model = load_image_model()
-    return model
+def main():
+    st.title("Advanced Medical AI: Multi-Modal RAG & Image Diagnostics")
 
-image_model = initialize_app()
+    st.markdown("""
+    **Features**:
+    1. **PubMed RAG**: Retrieve and summarize medical literature, store in a vector DB, 
+       and use advanced semantic search for context. 
+    2. **LLM Q&A**: Leverage OpenAI for final question-answering with RAG context.
+    3. **Medical Image Analysis**: Use `HuggingFaceTB/SmolVLM-500M-Instruct` for diagnostic insights.
+    4. **Multi-Lingual & Extended Triage**: Placeholder expansions for real-time translation or advanced triage logic.
+    5. **Production-Ready**: Modular, concurrent, disclaimers, and synergy across tasks.
+    """)
 
-###############################################################################
-#                    HELPER: BUILD SYSTEM PROMPT WITH REFS                    #
-###############################################################################
-def build_system_prompt_with_refs(pmids, summaries):
-    """
-    Creates a system prompt for the LLM that includes references [Ref1], [Ref2], etc.
-    """
-    system_context = "You have access to the following summarized PubMed articles:\n\n"
-    for idx, pmid in enumerate(pmids, start=1):
-        ref_label = f"[Ref{idx}]"
-        system_context += f"{ref_label} (PMID {pmid}): {summaries[pmid]}\n\n"
-    system_context += "Use this info to answer the user's question, citing references if needed."
-    return system_context
+    menu = ["PubMed RAG Q&A", "Medical Image Analysis", "Semantic Search (Vector DB)"]
+    choice = st.sidebar.selectbox("Select Task", menu)
 
-###############################################################################
-#                                MAIN APP                                     #
-###############################################################################
-def main():
-    st.title("RAG + Image Captioning: Production Demo")
+    if choice == "PubMed RAG Q&A":
+        pubmed_rag_qna()
+    elif choice == "Medical Image Analysis":
+        medical_image_analysis()
+    else:
+        vector_db_search_ui()
 
+    st.markdown("---")
     st.markdown("""
-    This demonstration shows:
-    1. **PubMed RAG**: Retrieve abstracts, summarize, and feed them into an LLM.
-    2. **Image Captioning**: Upload an image for analysis using a known stable model.
+    **Disclaimer**: This is an **advanced demonstration** for educational or research purposes only. 
+    Always consult a qualified healthcare professional for personal medical decisions.
     """)
 
-    # Section A: Image Upload / Caption
-    st.subheader("Image Captioning")
-    uploaded_img = st.file_uploader("Upload an image (optional)", type=["png", "jpg", "jpeg"])
-    if uploaded_img:
-        with st.spinner("Analyzing image..."):
-            caption = analyze_image(uploaded_img, image_model)
-        st.image(uploaded_img, use_column_width=True)
-        st.write("**Caption**:", caption)
-        st.write("---")
-
-    # Section B: PubMed-based RAG
-    st.subheader("PubMed RAG Pipeline")
-    user_query = st.text_input("Enter a medical question:", "What are the latest treatments for type 2 diabetes?")
-    
-    c1, c2, c3 = st.columns([2,1,1])
-    with c1:
-        st.markdown("**Parameters**:")
-        max_papers = st.slider("Number of Articles", 1, 10, 3)
-        chunk_size = st.slider("Chunk Size", 128, 1024, DEFAULT_CHUNK_SIZE)
-    with c2:
-        llm_choice = st.selectbox("Choose LLM", ["OpenAI: GPT-3.5", "Gemini: PaLM2"])
-    with c3:
-        temperature = st.slider("LLM Temperature", 0.0, 1.0, 0.3, step=0.1)
-
-    if st.button("Run RAG Pipeline"):
-        if not user_query.strip():
+def pubmed_rag_qna():
+    st.subheader("PubMed Retrieval-Augmented Q&A")
+    query = st.text_area(
+        "Ask a medical question (e.g., 'What are the latest treatments for type 2 diabetes?'):",
+        height=100
+    )
+    max_art = st.slider("Number of PubMed Articles to Retrieve", 1, 10, 5)
+
+    if st.button("Search & Summarize"):
+        if not query.strip():
             st.warning("Please enter a query.")
             return
-        
+
         with st.spinner("Searching PubMed..."):
-            pmids = search_pubmed(user_query, max_papers)
+            pmids = search_pubmed(query, max_art)
         
         if not pmids:
-            st.error("No PubMed results. Try a different query.")
+            st.error("No articles found. Try another query.")
             return
-        
-        with st.spinner("Fetching & Summarizing..."):
-            abstracts_map = fetch_pubmed_abstracts(pmids)
-            summarized_map = {}
-            for pmid, text in abstracts_map.items():
-                if text.startswith("Error:"):
-                    summarized_map[pmid] = text
+
+        with st.spinner("Fetching and Summarizing..."):
+            raw_abstracts = fetch_pubmed_abstracts(pmids)
+            # Summarize each
+            summarized = {}
+            for pmid, text in raw_abstracts.items():
+                if text.startswith("Error"):
+                    summarized[pmid] = text
                 else:
-                    summarized_map[pmid] = chunk_and_summarize(text, chunk_size=chunk_size)
+                    summary = chunk_and_summarize(text)
+                    summarized[pmid] = summary
+        
+        st.subheader("Summaries")
+        for i, (pmid, summary) in enumerate(summarized.items(), start=1):
+            st.markdown(f"**[Ref{i}] PMID {pmid}**")
+            st.write(summary)
 
-        st.subheader("Retrieved & Summarized PubMed Articles")
-        for idx, pmid in enumerate(pmids, start=1):
-            st.markdown(f"**[Ref{idx}] PMID {pmid}**")
-            st.write(summarized_map[pmid])
-            st.write("---")
+        # Upsert into vector DB
+        upsert_documents(summarized)  # store raw or summarized texts
 
-        st.subheader("RAG-Enhanced Final Answer")
-        system_prompt = build_system_prompt_with_refs(pmids, summarized_map)
-        with st.spinner("Generating LLM response..."):
-            if llm_choice == "OpenAI: GPT-3.5":
-                answer = openai_chat(system_prompt, user_query, temperature=temperature)
-            else:
-                answer = gemini_chat(system_prompt, user_query, temperature=temperature)
+        # Build system prompt
+        system_prompt = "You are an advanced medical assistant with the following references:\n"
+        for i, (pmid, summary) in enumerate(summarized.items(), start=1):
+            system_prompt += f"[Ref{i}] PMID {pmid}: {summary}\n"
+        system_prompt += "\nUsing these references, provide an evidence-based answer."
 
-        st.write(answer)
-        st.success("Pipeline Complete.")
+        with st.spinner("Generating final answer..."):
+            final_answer = chat_with_openai(system_prompt, query)
+        st.subheader("Final Answer")
+        st.write(final_answer)
 
-    st.markdown("---")
-    st.markdown("""
-    **Production Tips**:
-    - Vector DB for advanced retrieval
-    - Precise citation parsing
-    - Rate limiting on PubMed
-    - Multi-lingual expansions
-    - Logging & monitoring
-    - Security & privacy compliance
-    """)
+def medical_image_analysis():
+    st.subheader("Medical Image Analysis")
+    uploaded_file = st.file_uploader("Upload a Medical Image (PNG/JPG)", type=["png", "jpg", "jpeg"])
+    if uploaded_file is not None:
+        st.image(uploaded_file, caption="Uploaded Image", use_column_width=True)
+        if st.button("Analyze Image"):
+            with st.spinner("Analyzing..."):
+                result = analyze_medical_image(uploaded_file)
+            st.subheader("Diagnostic Insight")
+            st.write(result)
+
+def vector_db_search_ui():
+    st.subheader("Semantic Search in Vector DB")
+    user_query = st.text_input("Enter a query to find relevant documents", "")
+    top_k = st.slider("Number of results", 1, 10, 3)
+    if st.button("Search"):
+        if not user_query.strip():
+            st.warning("Please enter a query.")
+            return
+        with st.spinner("Performing semantic search..."):
+            results = semantic_search(user_query, top_k=top_k)
+        st.subheader("Search Results")
+        for i, doc in enumerate(results, start=1):
+            st.markdown(f"**Result {i}** - PMID {doc['pmid']} (Distance: {doc['score']:.4f})")
+            st.write(doc["text"])
+            st.write("---")
 
 if __name__ == "__main__":
     main()