medication_copilot.py

from functools import partial
import gradio as gr

from src.gradio_utils import (
    extract_table_from_chat,
    upload_file,
    redo,
    undo,
    edit_or_save_changes,
    update_llm_selection,
)
from src.llm_calls import query_llm
from src.data_handler import generate_excel_base64

SYSTEM_PROMPT = """You are a pharmacology assistant specialized in analyzing and structuring medical data.

Inputs You Receive:
A JSON dataset representing medications for Retinitis Pigmentosa

A user query requesting additional details to be added to the dataset

Your Task:
Analyze the dataset and determine what new information is needed

Research and generate new details based on the user’s request

Enhance the dataset by adding the requested information

Ensure completeness: The updated dataset must always include all medications

Your Output:
A succinct response explaining your findings and how the dataset was extended

A fully updated JSON dataset, strictly following this format:

json
Copy
Edit
{
  "Medications": [
    {"Name": "Medication Name", "key1": "value1", "key2": "value2", ...},
    {"Name": "Medication Name", "key1": "value1", "key2": "value2", ...}
  ]
}
Key Requirements:
- JSON output is mandatory in every response
- All medications must be present in the JSON, even if unchanged
- Extend the dataset with newly generated information—do not just retrieve existing data
- No repetition of the example JSON—only return the updated data
- Verify the JSON before responding to ensure it is well-formed and complete

Always structure your response clearly:

- Text Summary: Explanation of findings and dataset extensions

- Updated JSON Dataset: Full dataset with all medications, including new information

- References & Sources (if applicable)
"""


with gr.Blocks(theme=gr.themes.Glass()) as app:
    df_before = gr.State([])  # Undo history
    df_state = gr.State(None)  # Current DataFrame
    df_after = gr.State([])  # Redo history
    last_response = gr.State("")  # Store last LLM response
    edit_mode = gr.State("Edit")  # Track edit mode
    base64data = gr.State(None)

    with gr.Sidebar():
        gr.Markdown("### Configuration")
        llm_type = gr.Radio(
            choices=["Perplexity", "OpenAI"], label="LLM Type", value="Perplexity"
        )

        api_key = gr.Textbox(
            label="OpenAI API Key",
            placeholder=f"Enter {llm_type.value} Key",
            interactive=True,
            type="password"
        )

        gr.Markdown("### Upload existing data")
        file_upload = gr.File(label="Upload Excel File", file_types=[".xlsx"])
        gr.Markdown("### Download table to Excel")
        excel_output = gr.State(None)

        excel_data = gr.Textbox(visible=False)
        download_button = gr.DownloadButton(label="Download dataset")

        download_button.click(
            generate_excel_base64, inputs=[df_state], outputs=[excel_data]
        )

        excel_data.change(
            None,
            [excel_data],
            None,
            js="""
        (base64Data) => {
            const binaryString = atob(base64Data);
            const len = binaryString.length;
            const bytes = new Uint8Array(len);
            for (let i = 0; i < len; i++) {
                bytes[i] = binaryString.charCodeAt(i);
            }
            const blob = new Blob([bytes], { type: "application/vnd.openxmlformats-officedocument.spreadsheetml.sheet" });

            const link = document.createElement("a");
            link.href = URL.createObjectURL(blob);
            link.download = "dataset.xlsx";
            document.body.appendChild(link);
            link.click();
            document.body.removeChild(link);
        }
        """,
        )

        llm_type.change(update_llm_selection, inputs=[llm_type], outputs=[api_key])
        with gr.Accordion("System Prompt", open=False):
            system_prompt_box = gr.Textbox(
                value=SYSTEM_PROMPT, interactive=True, lines=10, label="System Prompt"
            )

    gr.Markdown("## Medications Data CoPilot")
    # Chat Interface
    chat = gr.ChatInterface(
        fn=query_llm,
        type="messages",
        description="Chat with an LLM to create a data representation of medications.",
        stop_btn=False,
        save_history=False,
        additional_inputs=[df_state, llm_type, api_key, system_prompt_box],
        examples=[
            [
                "List 10 medications that are known to be effective for Retinitis Pigmentosa"
            ],
            [
                "Add a column specifying if the medication passes the Retinal Blood Barrier"
            ],
            ["Add the safety profile for each medication"],
            ["Create four columns, each specifying each of the ADME profile factors"],
            [
                "Categorize each column into up to five categories, for simple classification."
            ],
        ],
    )
    with gr.Row():
        gr.Markdown("### Medications Table")
    with gr.Row():
        update_button = gr.Button(
            "Update table using the chat information", scale=8, interactive=True
        )
    with gr.Row():
        dataframe_display = gr.DataFrame(interactive=False)
    with gr.Row():
        prev_button = gr.Button("<-", interactive=False, scale=1)
        edit_save_button = gr.Button("Edit", interactive=True, scale=2)
        next_button = gr.Button("->", interactive=False, scale=1)

    # Save user changes
    edit_save_button.click(
        edit_or_save_changes,
        inputs=[dataframe_display, df_before, df_state, df_after, edit_mode],
        outputs=[
            dataframe_display,  # Updated DataFrame
            df_before,  # Undo history
            df_state,  # Current state
            df_after,  # Redo history
            prev_button,  # Update prev button
            next_button,  # Update next button
            dataframe_display,  # Update DataFrame interactivity
            edit_save_button,  # Update button label
            edit_mode,  # Update edit mode
        ],
    )
    # Undo button
    prev_button.click(
        undo,
        inputs=[df_before, df_state, df_after],
        outputs=[
            dataframe_display,
            df_before,
            df_state,
            df_after,
            prev_button,
            next_button,
        ],
    )
    # Redo button
    next_button.click(
        redo,
        inputs=[df_before, df_state, df_after],
        outputs=[
            dataframe_display,
            df_before,
            df_state,
            df_after,
            prev_button,
            next_button,
        ],
    )
    # File upload event
    file_upload.change(
        upload_file,
        inputs=[file_upload, df_before, df_state, df_after],
        outputs=[
            dataframe_display,
            df_before,
            df_state,
            df_after,
            prev_button,
            next_button,
        ],
    )
    # Update button copies chat history to text box
    update_button.click(
        partial(extract_table_from_chat, key="Medications"),
        inputs=[chat.chatbot, df_before, df_state, df_after, llm_type, api_key],
        outputs=[
            dataframe_display,
            df_before,
            df_state,
            df_after,
            prev_button,
            next_button,
        ],
    )

# Launch App
app.launch()