app.py

import gradio as gr
import interface
import utils
import treegraph as tg

system_prompt = """
You are a critical AI radiology assistant.
You are helping a radiologist correctly fill out a radiology report.
The report is regarding a Knee MRI.
"""

graph, nodes_dict = tg.build_tree_from_file("knee_template.json")
report_interface = interface.ReportChecklistInterface(
    llm=utils.LLM(model="gpt-3.5-turbo"),
    system_prompt=system_prompt,
    graph=graph,
    nodes_dict=nodes_dict,
)

if report_interface.prime_model() is False:
    print("Model priming failed. Please try again.")
    exit()
else:
    print("Model priming successful.")

with gr.Blocks(theme="finlaymacklon/smooth_slate") as demo:
    gr.Markdown("## Radiology Report Assistant")
    gr.Markdown(report_interface.help_message)

    running = gr.components.Variable(True)
    report_textbox = gr.TextArea(label="Report", lines=20, max_lines=50)
    check_btn = gr.Button(
        value="Check Report",
    )
    clear_btn = gr.ClearButton(
        value="Clear Messages",
    )
    quit_btn = gr.Button(
        value="Quit",
    )
    results_textbox = gr.TextArea(label="Results", lines=20, max_lines=50)
    clear_btn.add([results_textbox, report_textbox])

    def check_report(report):
        if running:
            results = report_interface.process_input(report)
            if results == "quit":
                quit_fn()
            elif results == "help":
                return report_interface.help_message
            elif results == "exception":
                return "An exception occurred. Please try again."
            else:
                return results
        else:
            return "Model has been stopped."

    def quit_fn():
        running.value = False
        results_textbox.value = "Model has been stopped."

    check_btn.click(fn=check_report, inputs=[report_textbox], outputs=[results_textbox])
    quit_btn.click(fn=quit_fn)

demo.launch(auth=("admin", "radiologyreport"), share=True)