app.py

import streamlit
import streamlit as st

# st.set_page_config(layout="wide")

from ring import LLM_GPT35
from ring import LLM_GPT4
from ring import LLM_GPT4o
from ring import LLM_LLAMA2
from ring import LLM_GEMINI
from ring import LLM_AI21
from ring import LLM_Claude2
from ring import LLM_Cohere
from ring import PREF_ORDER_LLMS


from ring import requires_openai_key
from ring import get_openai_api_key
from ring import generate_response
from ring import generate_panel_response

__version__ = "0.2.2"

col1, col2 = st.columns([3, 1])

with col1:
    st.title("OpenWormLLM")
    st.markdown("**v%s** - Note: work in progress!" % __version__)

with col2:
    st.image("images/OpenWormLogo.png")


tab_free, tab_panel, tab_pubs, tab_data, tab_corpus, tab_model, tab_about = st.tabs(
    [
        "Individual LLMs",
        "Panel discussion",
        "Publications",
        "Structured data",
        "OpenWorm Corpus",
        "Run model",
        "About",
    ]
)


with tab_free:

    st.markdown("**Ask individual LLMs questions about _C. elegans_**")

    with st.form("form_free"):

        text = st.text_area(
            "Ask a question related to _C. elegans_:",
            "What is the primary role of the C. elegans neuron AVBL?",
        )

        llm_ver = st.selectbox("Which LLM version should I use?", PREF_ORDER_LLMS)

        temperature = st.text_input("Temperature", value=0.1)

        only_celegans = st.checkbox(
            "Only answer questions related to _C. elegans_", value=True, disabled=True
        )

        submitted = st.form_submit_button("Submit")

        if requires_openai_key(llm_ver) and not get_openai_api_key():
            st.info("Please add your OpenAI API key to continue.")
        elif submitted:
            response = generate_response(text, llm_ver, temperature, only_celegans)

            st.info(response)

with tab_panel:

    st.markdown("**Get a consensus answer across multiple LLMs**")

    with st.form("form_panel"):

        text = st.text_area(
            "Ask a question related to _C. elegans_:",
            "What is the typical length of the worm C. elegans?",
        )

        temperature = st.text_input("Temperature", value=0.1)

        ###Start of Kaan Changes
        # Selectbox to determine which LLM becomes panel lead
        st.markdown(
            "<style> .stRadio> div {gap:30px;} </style>", 
            unsafe_allow_html=True,
        )
        panel_lead = st.radio(
            "Which LLM would you like to lead this panel?",
            PREF_ORDER_LLMS,
            horizontal=True,
        )

        # Create checkboxes to determine which LLMs are included in the panel
        st.markdown(
            "<div style='text-align: left; font-size: 14px;'>Which LLMs would you like to attend the panel?</div>",
            unsafe_allow_html=True,
        )
        options = PREF_ORDER_LLMS  # [llm for llm in PREF_ORDER_LLMS if llm !=panel_lead] creates a list of options excluding the chosen panel lead
        selected_options = []
        
        #Put checkboxes into table arranged 4x2
        for row in range(2):  # 2 rows
            cols = st.columns(4)  # 4 columns in each row
            for col in range(4):  # Iterate over each column
                index = row * 4 + col
                if index < len(options):
                    with cols[col]:
                        selected = st.checkbox(options[index], key=options[index])
                        if selected:
                            selected_options.append(options[index])

        submitted = st.form_submit_button("Submit")

        if len(selected_options) == 1 and panel_lead in selected_options:
            st.markdown("Please choose a panelist that is not the panel lead.")
        elif len(selected_options) == 0:
            st.markdown("Please choose a panelist to join the panel lead.")
        elif requires_openai_key(llm_ver) and not get_openai_api_key():
            st.info("Please add your OpenAI API key to continue.")
        elif submitted:
            # Create new list of LLMs that does not include panel lead
            llm_panelists = [llm for llm in selected_options if llm != panel_lead]
            response = generate_panel_response(
                text,
                llm_panelists=llm_panelists,
                llm_panel_chair=panel_lead,
                temperature=temperature,
            )

            st.info(response)

with tab_pubs:

    from publications import find_basis_paper

    st.markdown(
        "**Find literature related to _C. elegans._** Uses: [Semantic Scholar](https://www.semanticscholar.org)"
    )

    with st.form("form_pubs"):

        text = st.text_input("Enter a topic:", "C. elegans locomotion")

        num = st.text_input("Number of results", value=10)

        submitted = st.form_submit_button("Submit")

        if submitted:
            response = find_basis_paper(query=text, result_limit=num)

            st.info(response)


with tab_data:

    st.markdown("**Query structured datasets**")

    with st.form("form_data"):

        from datasources import DS_WORMNEUROATLAS
        from datasources import FORMATS
        from datasources import query_data_source

        text = st.text_area("Which neuron would you like to know about?", "AVBL")

        source = st.selectbox("Select data source to use:", (DS_WORMNEUROATLAS,))
        format = st.selectbox("Return format:", FORMATS)

        submitted = st.form_submit_button("Submit")

        if submitted:
            response = query_data_source(text, source, format)

            st.info(response)

with tab_corpus:

    from corpus_query import run_query

    st.markdown(
        "**Search a (small) corpus of _C. elegans_ literature** Uses: [paper-qa](https://github.com/whitead/paper-qa)"
    )

    with st.form("form_corpus"):

        text = st.text_input(
            "Enter a query:",
            "What types of ion channels are present in C. elegans neurons?",
        )

        llm_ver = st.selectbox(
            "Which LLM version should I use?", (LLM_LLAMA2, LLM_GPT4o)
        )

        submitted = st.form_submit_button("Submit")

        if submitted:
            response = run_query(query=text, llm_ver=llm_ver)

            st.info(response)

with tab_model:

    st.markdown("**Run a _C. elegans_ cell model**")

    with st.form("form_model"):

        from model import run_model
        from model import MODELS

        text = st.text_area("Current injection level:", "4.1 pA")
        model_to_sim = st.selectbox("Model to simulate", list(MODELS.keys()))

        submitted = st.form_submit_button("Submit")

        if submitted:
            response, traces, events = run_model(text, model_to_sim)
            st.info(response)
            try:

                import numpy as np
                import matplotlib.pyplot as plt

                fig, ax = plt.subplots()

                for key in sorted(traces.keys()):
                    if key != "t":
                        ts = traces["t"]
                        vs = traces[key]
                        ax.plot(ts, vs, label=key)
                        ax.legend()

                        plt.xlabel("Time (s)")
                        plt.ylabel("(SI units)")

                st.pyplot(fig)
            except Exception as e:
                st.info("There was a problem...\n\n%s" % e)

with tab_about:

    st.markdown('### About "Project Sydney"')
    st.markdown(
        "This is an initiative by the [OpenWorm project](https://openworm.org) to investigate the use of LLMs for interacting with scientific literature and structured datasets related to _C. elegans_."
    )
    st.markdown("**Work in progress!!** Subject to change/removal without notice!")