Update src/streamlit_app.py

src/streamlit_app.py

@@ -1,40 +1,209 @@
-import altair as alt
-import numpy as np
-import pandas as pd
+import os
+import logging
+from dotenv import load_dotenv
 import streamlit as st
+from langchain_chroma import Chroma
+from langchain_huggingface import HuggingFaceEmbeddings
+from langchain_openai import ChatOpenAI
 
+# Get a logger for this module
+logger = logging.getLogger(__name__)
+
+logger.info("Design Page...")
+# -------------------------------
+# PAGE CONFIG (MUST BE FIRST)
+# -------------------------------
+PORT = int(os.environ.get("PORT", 8501))
+
+st.markdown("""
+<style>
+.main-title {
+    font-size: 52px;
+    font-weight: 800;
+    text-align: center;
+    color: #0B5ED7;
+    margin-bottom: 5px;
+}
+
+.sub-title {
+    font-size: 20px;
+    text-align: center;
+    color: #555555;
+    margin-bottom: 30px;
+}
+</style>
+""", unsafe_allow_html=True)
+
+st.markdown(
+    '<div class="main-title">💊 AI Medical Labelling System</div>',
+    unsafe_allow_html=True
+)
+
+st.markdown(
+    '<div class="sub-title">Simplifying FDA Drug Safety Information using Generative AI & RAG</div>',
+    unsafe_allow_html=True
+)
+
+# -------------------------------
+# CUSTOM CSS (FANCY DESIGN)
+# -------------------------------
+st.markdown("""
+<style>
+.main {
+    background-color: #f7f9fc;
+}
+
+.big-title {
+    font-size:40px;
+    font-weight:700;
+    color:#1f4e79;
+}
+
+.subtitle {
+    font-size:18px;
+    color:#555;
+}
+
+.result-card {
+    background-color:white;
+    padding:20px;
+    border-radius:12px;
+    box-shadow:0px 2px 10px rgba(0,0,0,0.08);
+    margin-top:15px;
+}
+</style>
+""", unsafe_allow_html=True)
+
+# -------------------------------
+# HEADER
+# -------------------------------
+
+st.divider()
+
+# -------------------------------
+# SIDEBAR CONTROLS
+# -------------------------------
+with st.sidebar:
+    st.header("⚙️ Search Options")
+
+    drug_name = st.text_input(
+        "Drug Name",
+        placeholder="PHENYTOIN SODIUM"
+    )
+
+    selected_results = st.radio(
+        "Information Type",
+        ["Side Effects", "Warnings", "Both"]
+    )
+
+    run_button = st.button("🔍 Generate Explanation")
+
+# -------------------------------
+# LOAD ENV + MODELS
+# -------------------------------
+
+logger.info("Loading HuggingFace embedding model...")
+
+load_dotenv()
+
+working_dir = os.path.dirname(os.path.abspath(__file__))
+
+embeddings = HuggingFaceEmbeddings(
+    model_name="sentence-transformers/all-MiniLM-L6-v2"
+)
+
+vectordb = Chroma(
+    persist_directory=os.path.join(working_dir, "Chroma_db"),
+    embedding_function=embeddings
+)
+
+logger.info("Calling OpenAI model gpt-4o-mini...")
+
+llm = ChatOpenAI(
+    model="gpt-4o-mini",
+    temperature=0
+)
+
+# -------------------------------
+# RAG FUNCTION
+# -------------------------------
+def generate_section(drug_name, section, rules):
+
+    results = vectordb.get(
+        where={
+            "$and": [
+                {"generic_name": drug_name},
+                {"section": section}
+            ]
+        }
+    )
+
+    documents = results.get("documents", [])
+
+    if not documents:
+        st.warning(f"No data found for {section}")
+        return
+
+    context = "\n".join(set(documents))
+
+    prompt = f"""
+You are a medical assistant.
+
+Rewrite the FDA drug information into simplified,
+easy-to-understand language.
+
+Rules:
+{rules}
+
+Drug: {drug_name}
+
+FDA TEXT:
+{context}
 """
-# Welcome to Streamlit!
 
-Edit `/streamlit_app.py` to customize this app to your heart's desire :heart:.
-If you have any questions, checkout our [documentation](https://docs.streamlit.io) and [community
-forums](https://discuss.streamlit.io).
+    with st.spinner("🧠 AI is analysing FDA data..."):
+        response = llm.invoke(prompt)
+
+    st.markdown(
+        f'<div class="result-card">{response.content}</div>',
+        unsafe_allow_html=True
+    )
+
+logger.info("Configuring prompt..")
+# -------------------------------
+# RULES
+# -------------------------------
+SIDE_EFFECT_RULES = """
+- Use simple English
+- Bullet points (max 7)
+- Group similar side effects
+- Separate common vs serious
+"""
 
-In the meantime, below is an example of what you can do with just a few lines of code:
+WARNING_RULES = """
+- Use simple English
+- Bullet points (max 7)
+- Group warnings clearly
 """
 
-num_points = st.slider("Number of points in spiral", 1, 10000, 1100)
-num_turns = st.slider("Number of turns in spiral", 1, 300, 31)
-
-indices = np.linspace(0, 1, num_points)
-theta = 2 * np.pi * num_turns * indices
-radius = indices
-
-x = radius * np.cos(theta)
-y = radius * np.sin(theta)
-
-df = pd.DataFrame({
-    "x": x,
-    "y": y,
-    "idx": indices,
-    "rand": np.random.randn(num_points),
-})
-
-st.altair_chart(alt.Chart(df, height=700, width=700)
-    .mark_point(filled=True)
-    .encode(
-        x=alt.X("x", axis=None),
-        y=alt.Y("y", axis=None),
-        color=alt.Color("idx", legend=None, scale=alt.Scale()),
-        size=alt.Size("rand", legend=None, scale=alt.Scale(range=[1, 150])),
-    ))
+SECTION_MAP = {
+    "Side Effects": [("adverse_reactions", SIDE_EFFECT_RULES)],
+    "Warnings": [("warnings_and_cautions", WARNING_RULES)],
+    "Both": [
+        ("adverse_reactions", SIDE_EFFECT_RULES),
+        ("warnings_and_cautions", WARNING_RULES),
+    ],
+}
+
+# -------------------------------
+# MAIN ACTION
+# -------------------------------
+if run_button and drug_name:
+
+    st.subheader(f"Results for: {drug_name.upper()}")
+
+    for section, rules in SECTION_MAP[selected_results]:
+        generate_section(drug_name, section, rules)
+
+elif run_button:
+    st.warning("Please enter a drug name.")