|
|
|
|
|
import gradio as gr |
|
|
import json |
|
|
|
|
|
question_examples = [ |
|
|
["Given a patient with WHIM syndrome on prophylactic antibiotics, is it advisable to co-administer Xolremdi with fluconazole?"], |
|
|
["What treatment options exist for HER2+ breast cancer resistant to trastuzumab?"] |
|
|
] |
|
|
|
|
|
def extract_tool_name_and_clean_content(msg): |
|
|
tool_name = "Tool Result" |
|
|
content = msg.get("content") if isinstance(msg, dict) else getattr(msg, "content", "") |
|
|
try: |
|
|
parsed = json.loads(content) |
|
|
if isinstance(parsed, dict): |
|
|
tool_name = parsed.get("tool_name", tool_name) |
|
|
content = parsed.get("content", content) |
|
|
except Exception: |
|
|
pass |
|
|
if isinstance(content, (dict, list)): |
|
|
content = json.dumps(content, indent=2) |
|
|
return f"Tool: {tool_name}", content |
|
|
|
|
|
def format_collapsible(content, title="Answer"): |
|
|
return ( |
|
|
f"<details style='border: 1px solid #ccc; border-radius: 8px; padding: 10px; margin-top: 10px;'>" |
|
|
f"<summary style='font-size: 16px; font-weight: bold; color: #3B82F6;'>{title}</summary>" |
|
|
f"<div style='margin-top: 8px; font-size: 15px; line-height: 1.6; white-space: pre-wrap;'>{content}</div></details>" |
|
|
) |
|
|
|
|
|
def create_ui(agent): |
|
|
with gr.Blocks(theme=gr.themes.Soft()) as demo: |
|
|
gr.Markdown("<h1 style='text-align: center;'>💊 TxAgent: Therapeutic Reasoning</h1>") |
|
|
chatbot = gr.Chatbot(label="TxAgent", height=600, type="messages") |
|
|
message_input = gr.Textbox(placeholder="Ask a biomedical question...", show_label=False) |
|
|
send_button = gr.Button("Send", variant="primary") |
|
|
conversation_state = gr.State([]) |
|
|
|
|
|
def handle_chat(message, history, conversation): |
|
|
generator = agent.run_gradio_chat( |
|
|
message=message, |
|
|
history=history, |
|
|
temperature=0.3, |
|
|
max_new_tokens=1024, |
|
|
max_token=8192, |
|
|
call_agent=False, |
|
|
conversation=conversation, |
|
|
max_round=30 |
|
|
) |
|
|
for update in generator: |
|
|
formatted = [] |
|
|
for m in update: |
|
|
role = m.get("role") if isinstance(m, dict) else getattr(m, "role", "assistant") |
|
|
if role == "assistant": |
|
|
title, clean = extract_tool_name_and_clean_content(m) |
|
|
content = format_collapsible(clean, title) |
|
|
else: |
|
|
content = m.get("content") if isinstance(m, dict) else getattr(m, "content", "") |
|
|
formatted.append({"role": role, "content": content}) |
|
|
yield formatted |
|
|
|
|
|
send_button.click(fn=handle_chat, inputs=[message_input, chatbot, conversation_state], outputs=chatbot) |
|
|
message_input.submit(fn=handle_chat, inputs=[message_input, chatbot, conversation_state], outputs=chatbot) |
|
|
|
|
|
gr.Examples(examples=question_examples, inputs=message_input) |
|
|
gr.Markdown("<small style='color: gray;'>DISCLAIMER: This demo is for research purposes only and does not provide medical advice.</small>") |
|
|
|
|
|
return demo |
|
|
|
