Spaces:
Runtime error
Runtime error
| import re | |
| import gradio as gr | |
| from pathlib import Path | |
| import time | |
| import shutil | |
| from typing import AsyncGenerator, List, Optional, Tuple | |
| from gradio import ChatMessage | |
| class ChatInterface: | |
| """ | |
| A chat interface for interacting with a medical AI agent through Gradio. | |
| Handles file uploads, message processing, and chat history management. | |
| Supports both regular image files and DICOM medical imaging files. | |
| """ | |
| def __init__(self, agent, tools_dict): | |
| """ | |
| Initialize the chat interface. | |
| Args: | |
| agent: The medical AI agent to handle requests | |
| tools_dict (dict): Dictionary of available tools for image processing | |
| """ | |
| self.agent = agent | |
| self.tools_dict = tools_dict | |
| self.upload_dir = Path("temp") | |
| self.upload_dir.mkdir(exist_ok=True) | |
| self.current_thread_id = None | |
| # Separate storage for original and display paths | |
| self.original_file_path = None # For LLM (.dcm or other) | |
| self.display_file_path = None # For UI (always viewable format) | |
| def handle_upload(self, file_path: str) -> str: | |
| """ | |
| Handle new file upload and set appropriate paths. | |
| Args: | |
| file_path (str): Path to the uploaded file | |
| Returns: | |
| str: Display path for UI, or None if no file uploaded | |
| """ | |
| if not file_path: | |
| return None | |
| source = Path(file_path) | |
| timestamp = int(time.time()) | |
| # Save original file with proper suffix | |
| suffix = source.suffix.lower() | |
| saved_path = self.upload_dir / f"upload_{timestamp}{suffix}" | |
| shutil.copy2(file_path, saved_path) # Use file_path directly instead of source | |
| self.original_file_path = str(saved_path) | |
| # Handle DICOM conversion for display only | |
| if suffix == ".dcm": | |
| output, _ = self.tools_dict["DicomProcessorTool"]._run(str(saved_path)) | |
| self.display_file_path = output["image_path"] | |
| else: | |
| self.display_file_path = str(saved_path) | |
| return self.display_file_path | |
| def add_message( | |
| self, message: str, display_image: str, history: List[dict] | |
| ) -> Tuple[List[dict], gr.Textbox]: | |
| """ | |
| Add a new message to the chat history. | |
| Args: | |
| message (str): Text message to add | |
| display_image (str): Path to image being displayed | |
| history (List[dict]): Current chat history | |
| Returns: | |
| Tuple[List[dict], gr.Textbox]: Updated history and textbox component | |
| """ | |
| image_path = self.original_file_path or display_image | |
| if image_path is not None: | |
| history.append({"role": "user", "content": {"path": image_path}}) | |
| if message is not None: | |
| history.append({"role": "user", "content": message}) | |
| return history, gr.Textbox(value=message, interactive=False) | |
| async def process_message( | |
| self, message: str, display_image: Optional[str], chat_history: List[ChatMessage] | |
| ) -> AsyncGenerator[Tuple[List[ChatMessage], Optional[str], str], None]: | |
| """ | |
| Process a message and generate responses. | |
| Args: | |
| message (str): User message to process | |
| display_image (Optional[str]): Path to currently displayed image | |
| chat_history (List[ChatMessage]): Current chat history | |
| Yields: | |
| Tuple[List[ChatMessage], Optional[str], str]: Updated chat history, display path, and empty string | |
| """ | |
| chat_history = chat_history or [] | |
| # Initialize thread if needed | |
| if not self.current_thread_id: | |
| self.current_thread_id = str(time.time()) | |
| messages = [] | |
| image_path = self.original_file_path or display_image | |
| if image_path is not None: | |
| messages.append({"role": "user", "content": f"path: {image_path}"}) | |
| if message is not None: | |
| messages.append({"role": "user", "content": message}) | |
| try: | |
| for event in self.agent.workflow.stream( | |
| {"messages": messages}, {"configurable": {"thread_id": self.current_thread_id}} | |
| ): | |
| if isinstance(event, dict): | |
| if "process" in event: | |
| content = event["process"]["messages"][-1].content | |
| if content: | |
| content = re.sub(r"temp/[^\s]*", "", content) | |
| chat_history.append(ChatMessage(role="assistant", content=content)) | |
| yield chat_history, self.display_file_path, "" | |
| elif "execute" in event: | |
| for message in event["execute"]["messages"]: | |
| tool_name = message.name | |
| tool_result = eval(message.content)[0] | |
| if tool_result: | |
| metadata = {"title": f"🖼️ Image from tool: {tool_name}"} | |
| formatted_result = " ".join( | |
| line.strip() for line in str(tool_result).splitlines() | |
| ).strip() | |
| metadata["description"] = formatted_result | |
| chat_history.append( | |
| ChatMessage( | |
| role="assistant", | |
| content=formatted_result, | |
| metadata=metadata, | |
| ) | |
| ) | |
| # For image_visualizer, use display path | |
| if tool_name == "image_visualizer": | |
| self.display_file_path = tool_result["image_path"] | |
| chat_history.append( | |
| ChatMessage( | |
| role="assistant", | |
| # content=gr.Image(value=self.display_file_path), | |
| content={"path": self.display_file_path}, | |
| ) | |
| ) | |
| yield chat_history, self.display_file_path, "" | |
| except Exception as e: | |
| chat_history.append( | |
| ChatMessage( | |
| role="assistant", content=f"❌ Error: {str(e)}", metadata={"title": "Error"} | |
| ) | |
| ) | |
| yield chat_history, self.display_file_path | |
| def create_demo(agent, tools_dict): | |
| """ | |
| Create a Gradio demo interface for the medical AI agent. | |
| Args: | |
| agent: The medical AI agent to handle requests | |
| tools_dict (dict): Dictionary of available tools for image processing | |
| Returns: | |
| gr.Blocks: Gradio Blocks interface | |
| """ | |
| interface = ChatInterface(agent, tools_dict) | |
| with gr.Blocks(theme=gr.themes.Soft()) as demo: | |
| with gr.Column(): | |
| gr.Markdown( | |
| """ | |
| # 🏥 MedRAX | |
| Medical Reasoning Agent for Chest X-ray | |
| """ | |
| ) | |
| with gr.Row(): | |
| with gr.Column(scale=3): | |
| chatbot = gr.Chatbot( | |
| [], | |
| height=800, | |
| container=True, | |
| show_label=True, | |
| elem_classes="chat-box", | |
| type="messages", | |
| label="Agent", | |
| avatar_images=( | |
| None, | |
| "assets/medrax_logo.jpg", | |
| ), | |
| ) | |
| with gr.Row(): | |
| with gr.Column(scale=3): | |
| txt = gr.Textbox( | |
| show_label=False, | |
| placeholder="Ask about the X-ray...", | |
| container=False, | |
| ) | |
| with gr.Column(scale=3): | |
| image_display = gr.Image( | |
| label="Image", type="filepath", height=700, container=True | |
| ) | |
| with gr.Row(): | |
| upload_button = gr.UploadButton( | |
| "📎 Upload X-Ray", | |
| file_types=["image"], | |
| ) | |
| dicom_upload = gr.UploadButton( | |
| "📄 Upload DICOM", | |
| file_types=["file"], | |
| ) | |
| with gr.Row(): | |
| clear_btn = gr.Button("Clear Chat") | |
| new_thread_btn = gr.Button("New Thread") | |
| # Event handlers | |
| def clear_chat(): | |
| interface.original_file_path = None | |
| interface.display_file_path = None | |
| return [], None | |
| def new_thread(): | |
| interface.current_thread_id = str(time.time()) | |
| return [], interface.display_file_path | |
| def handle_file_upload(file): | |
| return interface.handle_upload(file.name) | |
| chat_msg = txt.submit( | |
| interface.add_message, inputs=[txt, image_display, chatbot], outputs=[chatbot, txt] | |
| ) | |
| bot_msg = chat_msg.then( | |
| interface.process_message, | |
| inputs=[txt, image_display, chatbot], | |
| outputs=[chatbot, image_display, txt], | |
| ) | |
| bot_msg.then(lambda: gr.Textbox(interactive=True), None, [txt]) | |
| upload_button.upload(handle_file_upload, inputs=upload_button, outputs=image_display) | |
| dicom_upload.upload(handle_file_upload, inputs=dicom_upload, outputs=image_display) | |
| clear_btn.click(clear_chat, outputs=[chatbot, image_display]) | |
| new_thread_btn.click(new_thread, outputs=[chatbot, image_display]) | |
| return demo | |