Fix merge conflicts

.env.example

@@ -6,4 +6,5 @@ GOOGLE_SEARCH_ENGINE_ID=
 OPENROUTER_API_KEY=
 OPENROUTER_BASE_URL=
 COHERE_API_KEY=
-PINECONE_API_KEY=
+PINECONE_API_KEY=
+MEDGEMMA_API_URL=

.gitignore

@@ -179,4 +179,6 @@ model-weights/
 
 .DS_Store
 
-benchmarking/data/
+benchmarking/data/
+model_cache/
+medgemma/

README.md

@@ -22,12 +22,14 @@ MedRAX is built on a robust technical foundation:
 
 ### Integrated Tools
 - **Visual QA**: Utilizes CheXagent and LLaVA-Med for complex visual understanding and medical reasoning
+- **MedGemma VQA**: Advanced medical visual question answering using Google's MedGemma 4B model for comprehensive medical image analysis across multiple modalities
 - **Segmentation**: Employs MedSAM2 (advanced medical image segmentation) and PSPNet model trained on ChestX-Det for precise anatomical structure identification
 - **Grounding**: Uses Maira-2 for localizing specific findings in medical images
 - **Report Generation**: Implements SwinV2 Transformer trained on CheXpert Plus for detailed medical reporting
 - **Disease Classification**: Leverages DenseNet-121 from TorchXRayVision for detecting 18 pathology classes
 - **X-ray Generation**: Utilizes RoentGen for synthetic CXR generation
 - **Web Browser**: Provides web search capabilities and URL content retrieval using Google Custom Search API
+- **DuckDuckGo Search**: Offers privacy-focused web search capabilities using DuckDuckGo search engine for medical research, fact-checking, and accessing current medical information without API keys
 - **Python Sandbox**: Executes Python code in a secure, stateful sandbox environment using `langchain-sandbox` and Pyodide. Supports custom data analysis, calculations, and dynamic package installations. Pre-configured with medical analysis packages including pandas, numpy, pydicom, SimpleITK, scikit-image, Pillow, scikit-learn, matplotlib, seaborn, and openpyxl. **Requires Deno runtime.**
 - **Utilities**: Includes DICOM processing, visualization tools, and custom plotting capabilities
 <br><br>
@@ -130,6 +132,10 @@ PINECONE_API_KEY=
 # Requires Google Custom Search API credentials.
 GOOGLE_SEARCH_API_KEY=
 GOOGLE_SEARCH_ENGINE_ID=
+
+# MedGemma VQA Tool (Optional)
+# URL for the MedGemma FastAPI service
+MEDGEMMA_API_URL=
 ```
 
 ### Getting Started
@@ -159,6 +165,7 @@ selected_tools = [
     "ChestXRaySegmentationTool",
     "PythonSandboxTool",              # Python code execution
     "WebBrowserTool",                 # Web search and URL access
+    "DuckDuckGoSearchTool",           # Privacy-focused web search
     # Add or remove tools as needed
 ]
 
@@ -174,17 +181,10 @@ agent, tools_dict = initialize_agent(
 
 The following tools will automatically download their model weights when initialized:
 
-### Classification Tools
+### Classification Tool
 ```python
 # TorchXRayVision-based classifier (original)
 TorchXRayVisionClassifierTool(device=device)
-
-# ArcPlus SwinTransformer-based classifier (new)
-ArcPlusClassifierTool(
-    model_path="/path/to/Ark6_swinLarge768_ep50.pth.tar",  # Optional
-    num_classes=18,  # Default
-    device=device
-)
 ```
 
 ### Segmentation Tool
@@ -232,6 +232,21 @@ XRayVQATool(
 ```
 - CheXagent weights download automatically
 
+### MedGemma VQA Tool
+```python
+MedGemmaAPIClientTool(
+    device=device,
+    cache_dir=model_dir, 
+    api_url=MEDGEMMA_API_URL)
+)
+```
+- Uses Google's MedGemma 4B instruction-tuned model for comprehensive medical image analysis
+- Specialized for chest X-rays, dermatology, ophthalmology, and pathology images
+- Provides radiologist-level medical reasoning and diagnosis assistance
+- Supports up to 128K context length and 896x896 image resolution
+- 4-bit quantization available (~4GB VRAM) with full precision option (~8GB VRAM)
+- Model weights download automatically when the service starts
+
 ### MedSAM2 Tool
 ```python
 MedSAM2Tool(
@@ -263,6 +278,7 @@ No additional model weights required:
 ImageVisualizerTool()
 DicomProcessorTool(temp_dir=temp_dir)
 WebBrowserTool()  # Requires Google Search API credentials
+DuckDuckGoSearchTool()  # No API key required, privacy-focused search
 ```
 <br>
 
@@ -281,6 +297,25 @@ ChestXRayGeneratorTool(
   2. Place weights in `{model_dir}/roentgen`
   3. Optional tool, can be excluded if not needed
 
+### ArcPlus SwinTransformer-based Classifier
+```python
+ArcPlusClassifierTool(
+    model_path="/path/to/Ark6_swinLarge768_ep50.pth.tar",  # Optional
+    num_classes=18,  # Default
+    device=device
+)
+```
+
+The ArcPlus classifier requires manual setup as the pre-trained model is not publicly available for automatic download:
+
+1. **Request Access**: Visit [https://github.com/jlianglab/Ark](https://github.com/jlianglab/Ark) and request the pretrained model through their Google Forms
+2. **Download Model**: Once approved, download the `Ark6_swinLarge768_ep50.pth.tar` file
+3. **Place in Directory**: Drag the downloaded file into your `model-weights` directory
+4. **Initialize Tool**: The tool will automatically look for the model file in the specified `cache_dir`
+
+The ArcPlus model provides advanced chest X-ray classification across 6 medical datasets (MIMIC, CheXpert, NIH, RSNA, VinDr, Shenzhen) with 52+ pathology categories.
+```
+
 ### Knowledge Base Setup (MedicalRAGTool)
 
 The `MedicalRAGTool` uses a Pinecone vector database to store and retrieve medical knowledge. To use this tool, you need to set up a Pinecone account and a Cohere account.
@@ -383,6 +418,8 @@ If you are running a local LLM using frameworks like [Ollama](https://ollama.com
 
 **WebBrowserTool**: Requires Google Custom Search API credentials, which can be set in the `.env` file.
 
+**DuckDuckGoSearchTool**: No API key required. Uses DuckDuckGo's privacy-focused search engine for medical research and fact-checking.
+
 **PythonSandboxTool**: Requires Deno runtime installation:
 ```bash
 # Verify Deno is installed

benchmarking/benchmarks/base.py

@@ -4,6 +4,7 @@ from abc import ABC, abstractmethod
 from typing import Dict, List, Optional, Any, Iterator, Tuple
 from dataclasses import dataclass
 from pathlib import Path
+import random
 
 
 @dataclass
@@ -31,17 +32,31 @@ class Benchmark(ABC):
         Args:
             data_dir (str): Directory containing benchmark data
             **kwargs: Additional configuration parameters
+                random_seed (int): Random seed for shuffling data (default: None, no shuffling)
         """
         self.data_dir = Path(data_dir)
         self.config = kwargs
         self.data_points = []
         self._load_data()
+        self._shuffle_data()
 
     @abstractmethod
     def _load_data(self) -> None:
         """Load benchmark data from the data directory."""
         pass
 
+    def _shuffle_data(self) -> None:
+        """Shuffle the data points if a random seed is provided.
+        
+        This method is called automatically after data loading to ensure
+        reproducible benchmark runs when a random seed is specified.
+        """
+        random_seed = self.config.get("random_seed", None)
+        if random_seed is not None:
+            random.seed(random_seed)
+            random.shuffle(self.data_points)
+            print(f"Shuffled {len(self.data_points)} data points with seed {random_seed}")
+
     def get_data_point(self, index: int) -> BenchmarkDataPoint:
         """Get a specific data point by index.
         

benchmarking/benchmarks/chestagentbench_benchmark.py

@@ -1,5 +1,4 @@
 import json
-import random
 from pathlib import Path
 from typing import Dict, Optional, Any
 from .base import Benchmark, BenchmarkDataPoint
@@ -31,10 +30,6 @@ class ChestAgentBenchBenchmark(Benchmark):
                 except Exception as e:
                     print(f"Error loading item {i}: {e}")
                     continue
-        
-        # Shuffle the final data
-        random.seed(42)
-        random.shuffle(self.data_points)
 
     def _parse_item(self, item: Dict[str, Any], index: int) -> Optional[BenchmarkDataPoint]:
         # Use full_question_id or question_id if available, else fallback

benchmarking/benchmarks/rexvqa_benchmark.py

@@ -34,20 +34,20 @@ class ReXVQABenchmark(Benchmark):
             data_dir (str): Directory to store/cache downloaded data
             **kwargs: Additional configuration parameters
                 split (str): Dataset split to use (default: 'test')
-                cache_dir (str): Directory for caching HuggingFace datasets
                 trust_remote_code (bool): Whether to trust remote code (default: False)
                 max_questions (int): Maximum number of questions to load (default: None, load all)
                 images_dir (str): Directory containing extracted PNG images (default: None)
         """
         self.split = kwargs.get("split", "test")
-        self.cache_dir = kwargs.get("cache_dir", None)
         self.trust_remote_code = kwargs.get("trust_remote_code", False)
         self.max_questions = kwargs.get("max_questions", None)
-        self.images_dir = "benchmarking/data/rexvqa/images/deid_png"
         self.image_dataset = None
         self.image_mapping = {}  # Maps study_id to image data
         
         super().__init__(data_dir, **kwargs)
+        
+        # Set images_dir after parent initialization
+        self.images_dir = f"{self.data_dir}/images/deid_png"
 
     @staticmethod
     def download_rexgradient_images(output_dir: str = "benchmarking/data/rexvqa", repo_id: str = "rajpurkarlab/ReXGradient-160K"):
@@ -166,8 +166,8 @@ class ReXVQABenchmark(Benchmark):
         """Load ReXVQA data from local JSON file."""
         try:
             # Check for images and test_vqa_data.json, download if missing
-            self.download_test_vqa_data_json()
-            self.download_rexgradient_images()
+            self.download_test_vqa_data_json(self.data_dir)
+            self.download_rexgradient_images(self.data_dir)
             
             # Construct path to the JSON file
             json_file_path = os.path.join("benchmarking", "data", "rexvqa", "metadata", "test_vqa_data.json")
@@ -197,7 +197,7 @@ class ReXVQABenchmark(Benchmark):
                 self.image_dataset = load_dataset(
                     "rajpurkarlab/ReXGradient-160K",
                     split="test",
-                    cache_dir=self.cache_dir,
+                    cache_dir=self.data_dir,
                     trust_remote_code=self.trust_remote_code
                 )
                 print(f"Loaded {len(self.image_dataset)} image metadata entries from ReXGradient-160K")

benchmarking/cli.py

@@ -73,6 +73,8 @@ def run_benchmark_command(args) -> None:
     
     # Create benchmark
     benchmark_kwargs = {}
+    if args.random_seed is not None:
+        benchmark_kwargs["random_seed"] = args.random_seed
     
     benchmark = create_benchmark(benchmark_name=args.benchmark, data_dir=args.data_dir, **benchmark_kwargs)
     
@@ -135,12 +137,14 @@ def main():
                            help="Output directory for results (default: benchmark_results)")
     run_parser.add_argument("--max-questions", type=int, 
                            help="Maximum number of questions to process (default: all)")
-    run_parser.add_argument("--temperature", type=float, default=0.7, 
+    run_parser.add_argument("--temperature", type=float, default=1, 
                            help="Model temperature for response generation (default: 0.7)")
     run_parser.add_argument("--top-p", type=float, default=0.95, 
                            help="Top-p nucleus sampling parameter (default: 0.95)")
     run_parser.add_argument("--max-tokens", type=int, default=5000, 
                            help="Maximum tokens per model response (default: 5000)")
+    run_parser.add_argument("--random-seed", type=int, default=42, 
+                           help="Random seed for shuffling benchmark data (enables reproducible runs, default: None)")
     
     run_parser.set_defaults(func=run_benchmark_command)
     

benchmarking/llm_providers/medrax_provider.py

@@ -33,20 +33,36 @@ class MedRAXProvider(LLMProvider):
             print("Starting server...")
 
             selected_tools = [
+                # Image Processing Tools
                 # "ImageVisualizerTool",  # For displaying images in the UI
                 # "DicomProcessorTool",  # For processing DICOM medical image files
                 # "ChestXRaySegmentationTool",  # For segmenting anatomical regions in chest X-rays
-                # "LlavaMedTool",  # For multimodal medical image understanding
                 # "ChestXRayGeneratorTool",  # For generating synthetic chest X-rays
-                # "PythonSandboxTool",  # Add the Python sandbox tool
-                
+
+                # Classification Tools
+                "TorchXRayVisionClassifierTool",  # For classifying chest X-ray images using TorchXRayVision
+                "ArcPlusClassifierTool",  # For advanced chest X-ray classification using ArcPlus
+
+                # Report Generation Tools
                 "ChestXRayReportGeneratorTool",  # For generating medical reports from X-rays
+
+                # Grounding Tools
+                "XRayPhraseGroundingTool",  # For locating described features in X-rays
+
+                # VQA Tools
+                # "LlavaMedTool",  # For multimodal medical image understanding
+                # "XRayVQATool",  # For visual question answering on X-rays
+                "MedGemmaVQATool",
+
+                # RAG Tools
                 "MedicalRAGTool",  # For retrieval-augmented generation with medical knowledge
+
+                # Search Tools
                 # "WebBrowserTool",  # For web browsing and search capabilities
-                # "XRayVQATool",  # For visual question answering on X-rays
-                "TorchXRayVisionClassifierTool",  # For classifying chest X-ray images using TorchXRayVision
-                "ArcPlusClassifierTool",  # For advanced chest X-ray classification using ArcPlus
-                # "XRayPhraseGroundingTool",  # For locating described features in X-rays
+                # "DuckDuckGoSearchTool",  # For privacy-focused web search using DuckDuckGo
+
+                # Development Tools
+                # "PythonSandboxTool",  # Add the Python sandbox tool
             ]
 
             rag_config = RAGConfig(
@@ -69,11 +85,11 @@ class MedRAXProvider(LLMProvider):
             agent, tools_dict = initialize_agent(
                 prompt_file="medrax/docs/system_prompts.txt",
                 tools_to_use=selected_tools,
-                model_dir="/model-weights",
+                model_dir="/scratch/ssd004/scratch/victorli/model-weights",
                 temp_dir="temp",  # Change this to the path of the temporary directory
                 device="cuda:1",
                 model=self.model_name,  # Change this to the model you want to use, e.g. gpt-4.1-2025-04-14, gemini-2.5-pro
-                temperature=0.3,
+                temperature=1.0,
                 top_p=0.95,
                 model_kwargs=model_kwargs,
                 rag_config=rag_config,

benchmarking/runner.py

@@ -268,12 +268,6 @@ class BenchmarkRunner:
         if match:
             return match.group(1).upper()
         
-        # Fallback: look for the '<|A|>' format (legacy code, will remove later on)
-        legacy_pattern = r'\s*<\|([A-F])\|>'
-        match = re.search(legacy_pattern, response_text)
-        if match:
-            return match.group(1).upper()
-        
         # If no pattern matches, return the full response
         return response_text.strip()
 

interface.py

@@ -193,7 +193,11 @@ class ChatInterface:
 
                                 # First, display the tool usage card
                                 try:
-                                    tool_output_json = json.loads(msg.content)
+                                    # Handle case where tool returns tuple (output, metadata)
+                                    content = msg.content
+                                    content_tuple = ast.literal_eval(content)
+                                    content = json.dumps(content_tuple[0])
+                                    tool_output_json = json.loads(content)
                                     tool_output_str = json.dumps(tool_output_json, indent=2)
                                 except (json.JSONDecodeError, TypeError):
                                     tool_output_str = str(msg.content)
@@ -216,19 +220,19 @@ class ChatInterface:
                                 # Special handling for image_visualizer
                                 if tool_name == "image_visualizer":
                                     try:
-                                        # Tool returns (output, metadata) tuple
-                                        # msg.content should be the serialized version of this
-                                        result = eval(msg.content)  # Safe here as it's from our tool
-                                        if isinstance(result, tuple) and len(result) >= 1:
-                                            output_dict = result[0]
-                                            if isinstance(output_dict, dict) and "image_path" in output_dict:
-                                                self.display_file_path = output_dict["image_path"]
-                                                chat_history.append(
-                                                    ChatMessage(
-                                                        role="assistant",
-                                                        content={"path": self.display_file_path},
-                                                    )
+                                        # Handle case where tool returns tuple (output, metadata)
+                                        content = msg.content
+                                        content_tuple = ast.literal_eval(content)
+                                        result = content_tuple[0]
+                                        
+                                        if isinstance(result, dict) and "image_path" in result:
+                                            self.display_file_path = result["image_path"]
+                                            chat_history.append(
+                                                ChatMessage(
+                                                    role="assistant",
+                                                    content={"path": self.display_file_path},
                                                 )
+                                            )
                                     except Exception:
                                         pass
                                 

main.py

@@ -10,6 +10,7 @@ with different model weights, tools, and parameters.
 """
 
 import warnings
+import os
 from typing import Dict, List, Optional, Any
 from dotenv import load_dotenv
 from transformers import logging
@@ -33,11 +34,11 @@ _ = load_dotenv()
 def initialize_agent(
     prompt_file: str,
     tools_to_use: Optional[List[str]] = None,
-    model_dir: str = "/model-weights",
+    model_dir: str = "model-weights",
     temp_dir: str = "temp",
     device: str = "cpu",
-    model: str = "gpt-4.1-2025-04-14",
-    temperature: float = 0.7,
+    model: str = "gemini-2.5-pro",
+    temperature: float = 1.0,
     top_p: float = 0.95,
     rag_config: Optional[RAGConfig] = None,
     model_kwargs: Dict[str, Any] = {},
@@ -66,12 +67,15 @@ def initialize_agent(
     prompts = load_prompts_from_file(prompt_file)
     prompt = prompts[system_prompt]
 
+    # Define the URL of the MedGemma FastAPI service.
+    MEDGEMMA_API_URL = os.getenv("MEDGEMMA_API_URL", "http://172.17.8.141:8002")
+
     all_tools = {
         "TorchXRayVisionClassifierTool": lambda: TorchXRayVisionClassifierTool(device=device),
         "ArcPlusClassifierTool": lambda: ArcPlusClassifierTool(cache_dir=model_dir, device=device),
         "ChestXRaySegmentationTool": lambda: ChestXRaySegmentationTool(device=device),
         "LlavaMedTool": lambda: LlavaMedTool(cache_dir=model_dir, device=device, load_in_8bit=True),
-        "XRayVQATool": lambda: XRayVQATool(cache_dir=model_dir, device=device),
+        "CheXagentXRayVQATool": lambda: CheXagentXRayVQATool(cache_dir=model_dir, device=device),
         "ChestXRayReportGeneratorTool": lambda: ChestXRayReportGeneratorTool(
             cache_dir=model_dir, device=device
         ),
@@ -85,23 +89,29 @@ def initialize_agent(
         "DicomProcessorTool": lambda: DicomProcessorTool(temp_dir=temp_dir),
         "MedicalRAGTool": lambda: RAGTool(config=rag_config),
         "WebBrowserTool": lambda: WebBrowserTool(),
+        "DuckDuckGoSearchTool": lambda: DuckDuckGoSearchTool(),
         "MedSAM2Tool": lambda: MedSAM2Tool(
             device=device, cache_dir=model_dir, temp_dir=temp_dir
         ),
-    }
-
-    try:
-        tools_dict["PythonSandboxTool"] = create_python_sandbox()
-    except Exception as e:
-        print(f"Error creating PythonSandboxTool: {e}")
-        print("Skipping PythonSandboxTool")
+        "MedGemmaVQATool": lambda: MedGemmaAPIClientTool(cache_dir=model_dir, device=device, api_url=MEDGEMMA_API_URL)
+    }    
 
     # Initialize only selected tools or all if none specified
     tools_dict: Dict[str, BaseTool] = {}
-    tools_to_use = tools_to_use or all_tools.keys()
+
+    if tools_to_use is None:
+        tools_to_use = []
+        
     for tool_name in tools_to_use:
+        if tool_name == "PythonSandboxTool":
+            try:
+                tools_dict["PythonSandboxTool"] = create_python_sandbox()
+            except Exception as e:
+                print(f"Error creating PythonSandboxTool: {e}")
+                print("Skipping PythonSandboxTool")
         if tool_name in all_tools:
             tools_dict[tool_name] = all_tools[tool_name]()
+    
 
     # Set up checkpointing for conversation state
     checkpointer = MemorySaver()
@@ -139,22 +149,47 @@ if __name__ == "__main__":
     # Example: initialize with only specific tools
     # Here three tools are commented out, you can uncomment them to use them
     selected_tools = [
+        # Image Processing Tools
         "ImageVisualizerTool",  # For displaying images in the UI
         # "DicomProcessorTool",  # For processing DICOM medical image files
+
+        # Segmentation Tools
+        "MedSAM2Tool",  # For advanced medical image segmentation using MedSAM2
+        "ChestXRaySegmentationTool",  # For segmenting anatomical regions in chest X-rays
+
+        # Generation Tools
+        # "ChestXRayGeneratorTool",  # For generating synthetic chest X-rays
+
+        # Classification Tools
         "TorchXRayVisionClassifierTool",  # For classifying chest X-ray images using TorchXRayVision
         "ArcPlusClassifierTool",  # For advanced chest X-ray classification using ArcPlus
-        "ChestXRaySegmentationTool",  # For segmenting anatomical regions in chest X-rays
+
+        # Report Generation Tools
         "ChestXRayReportGeneratorTool",  # For generating medical reports from X-rays
+
+        # Grounding Tools
+        "XRayPhraseGroundingTool",  # For locating described features in X-rays
+
+        # VQA Tools
+        "MedGemmaVQATool",  # Google MedGemma VQA tool
         "XRayVQATool",  # For visual question answering on X-rays
         # "LlavaMedTool",  # For multimodal medical image understanding
-        "XRayPhraseGroundingTool",  # For locating described features in X-rays
-        # "ChestXRayGeneratorTool",  # For generating synthetic chest X-rays
-        "MedSAM2Tool",  # For advanced medical image segmentation using MedSAM2
-        "WebBrowserTool",  # For web browsing and search capabilities
+
+        # RAG Tools
         "MedicalRAGTool",  # For retrieval-augmented generation with medical knowledge
+
+        # Search Tools
+        "WebBrowserTool",  # For web browsing and search capabilities
+        "DuckDuckGoSearchTool",  # For privacy-focused web search using DuckDuckGo
+
+        # Development Tools
         # "PythonSandboxTool",  # Add the Python sandbox tool
     ]
 
+    # Setup the MedGemma environment if the MedGemmaVQATool is selected
+    if "MedGemmaVQATool" in selected_tools:
+        setup_medgemma_env()
+
     # Configure the Retrieval Augmented Generation (RAG) system
     # This allows the agent to access and use medical knowledge documents
     rag_config = RAGConfig(
@@ -177,11 +212,11 @@ if __name__ == "__main__":
     agent, tools_dict = initialize_agent(
         prompt_file="medrax/docs/system_prompts.txt",
         tools_to_use=selected_tools,
-        model_dir="/model-weights",
+        model_dir="model-weights",
         temp_dir="temp",  # Change this to the path of the temporary directory
         device="cuda:1",
-        model="gpt-4.1-2025-04-14",  # Change this to the model you want to use, e.g. gpt-4.1-2025-04-14, gemini-2.5-pro
-        temperature=0.7,
+        model="gpt-4.1-2025-04-14",  # Change this to the model you want to use, e.g. gpt-4.1-2025-04-14, gemini-2.5-pro, gpt-5
+        temperature=1.0,
         top_p=0.95,
         model_kwargs=model_kwargs,
         rag_config=rag_config,

medrax/docs/system_prompts.txt

@@ -17,10 +17,9 @@ Examples:
 - "Based on clinical guidelines [3], the recommended treatment approach is..."
 
 [CHESTAGENTBENCH_PROMPT]
-You are an expert medical AI assistant who can answer any medical questions and analyze medical images similar to a doctor.
-Solve using your own vision and reasoning and use tools (if available) to complement your reasoning.
-You can make multiple tool calls in parallel or in sequence as needed for comprehensive answers.
-Think critically about and criticize the tool outputs.
-If you need to look up some information before asking a follow up question, you are allowed to do that.
+You are an expert medical assistant who can answer medical questions and analyze medical images with world-class accuracy.
+Use your state-of-the art reasoning and critical thinking skills to answer the questions that you are asked.
+You may use tools (if available) to complement your reasoning and you are allowed to make multiple tool calls in parallel or in sequence as needed for comprehensive answers.
+Think critically about how to best use the tools available to you and scrutinize the tool outputs.
 When encountering a multiple-choice question, your final response should end with "Final answer: \boxed{A}" from list of possible choices A, B, C, D, E, F.
-It is extremely important that you strictly answer in the format mentioned above.
+It is extremely important that you answer strictly in the format described above.

medrax/tools/__init__.py

@@ -3,14 +3,11 @@
 from .classification import *
 from .report_generation import *
 from .segmentation import *
-from .xray_vqa import *
-from .llava_med import *
+from .vqa import *
 from .grounding import *
-from .generation import *
+from .xray_generation import *
 from .dicom import *
 from .utils import *
 from .rag import *
-from .web_browser import *
+from .browsing import *
 from .python_tool import *
-from .medsam2 import *
-

medrax/tools/browsing/__init__.py

@@ -0,0 +1,13 @@
+"""Web browsing tools for MedRAX2 medical agents."""
+
+from .duckduckgo import DuckDuckGoSearchTool, WebSearchInput
+from .web_browser import WebBrowserTool, WebBrowserSchema, SearchQuerySchema, VisitUrlSchema
+
+__all__ = [
+    "DuckDuckGoSearchTool",
+    "WebSearchInput",
+    "WebBrowserTool", 
+    "WebBrowserSchema",
+    "SearchQuerySchema",
+    "VisitUrlSchema"
+] 

medrax/tools/browsing/duckduckgo.py

@@ -0,0 +1,403 @@
+"""
+Web search tool for MedRAX2 medical agents.
+
+Provides DuckDuckGo search capabilities for medical agents to retrieve
+real-time information from the web with proper error handling
+and result formatting. Designed specifically for medical research,
+fact-checking, and accessing current medical information.
+"""
+
+import asyncio
+import logging
+import time
+from datetime import datetime
+from typing import Dict, Any, Tuple
+
+from langchain_core.callbacks import (
+    AsyncCallbackManagerForToolRun,
+    CallbackManagerForToolRun,
+)
+from langchain_core.tools import BaseTool
+from pydantic import BaseModel, Field
+
+try:
+    from duckduckgo_search import DDGS
+except ImportError:
+    DDGS = None
+
+logger = logging.getLogger(__name__)
+
+
+class WebSearchInput(BaseModel):
+    """Input schema for web search tool."""
+
+    query: str = Field(
+        ...,
+        description="The search query to look up on the web. Be specific and include relevant medical keywords for better results.",
+        min_length=1,
+        max_length=500,
+    )
+    max_results: int = Field(
+        default=5,
+        description="Maximum number of search results to return (1-10)",
+        ge=1,
+        le=10,
+    )
+    region: str = Field(
+        default="us-en",
+        description="Region for search results (e.g., 'us-en', 'uk-en', 'ca-en')",
+    )
+
+
+class DuckDuckGoSearchTool(BaseTool):
+    """
+    Tool that performs web searches using DuckDuckGo search engine for medical research.
+
+    This tool provides access to real-time web information through DuckDuckGo's
+    search API, specifically designed for medical agents that need to retrieve current
+    medical information, verify facts, or find resources on medical topics.
+
+    Features:
+        - Real-time web search capability for medical information
+        - Configurable number of results (1-10)
+        - Regional search support for localized medical results
+        - Robust error handling for network issues
+        - Structured result formatting for easy parsing
+        - Privacy-focused (DuckDuckGo doesn't track users)
+        - Medical-focused search optimization
+
+    Use Cases:
+        - Medical fact checking and verification
+        - Finding current medical news and updates
+        - Researching specific medical topics or questions
+        - Gathering multiple perspectives on medical issues
+        - Locating official medical resources and documentation
+        - Accessing current clinical guidelines and research
+
+    Rate Limiting:
+        DuckDuckGo has rate limits. Avoid making too many rapid requests
+        to prevent temporary blocking.
+    """
+
+    name: str = "duckduckgo_search"
+    description: str = (
+        "Search the web using DuckDuckGo to find current medical information, research, and resources. "
+        "Input should be a clear search query with relevant medical keywords. The tool returns a list of relevant web results "
+        "with titles, URLs, and brief snippets. Useful for medical fact-checking, finding current medical events, "
+        "researching medical topics, and gathering information from reliable medical sources. "
+        "Results are privacy-focused and don't track user searches. Optimized for medical research and clinical information."
+    )
+    args_schema: type[BaseModel] = WebSearchInput
+    return_direct: bool = False
+
+    def __init__(self, **kwargs):
+        """Initialize the DuckDuckGo search tool."""
+        super().__init__(**kwargs)
+
+        if DDGS is None:
+            logger.error(
+                "duckduckgo-search package not installed. Install with: pip install duckduckgo-search"
+            )
+            raise ImportError(
+                "duckduckgo-search package is required for web search functionality"
+            )
+
+        logger.info("DuckDuckGo search tool initialized successfully")
+
+    def _perform_search_sync(
+        self, query: str, max_results: int = 5, region: str = "us-en"
+    ) -> Dict[str, Any]:
+        """
+        Perform the actual web search using DuckDuckGo synchronously.
+
+        Args:
+            query (str): The search query.
+            max_results (int): Maximum number of results to return.
+            region (str): Region for localized results.
+
+        Returns:
+            Dict[str, Any]: Structured search results.
+        """
+        logger.info(
+            f"Performing web search: '{query}' (max_results={max_results}, region={region})"
+        )
+
+        try:
+            # Initialize DDGS with error handling
+            with DDGS() as ddgs:
+                # Perform the search
+                search_results = list(
+                    ddgs.text(
+                        keywords=query,
+                        region=region,
+                        safesearch="moderate",
+                        timelimit=None,
+                        max_results=max_results,
+                    )
+                )
+
+                # Format results for the agent
+                formatted_results = []
+                for i, result in enumerate(search_results, 1):
+                    formatted_result = {
+                        "rank": i,
+                        "title": result.get("title", "No title"),
+                        "url": result.get("href", "No URL"),
+                        "snippet": result.get("body", "No description available"),
+                        "source": "DuckDuckGo",
+                    }
+                    formatted_results.append(formatted_result)
+
+                # Create summary for the agent
+                if formatted_results:
+                    summary = (
+                        f"Found {len(formatted_results)} results for '{query}'. Top results include: "
+                        + ", ".join([f"{r['title']}" for r in formatted_results[:3]])
+                    )
+                else:
+                    summary = f"No results found for '{query}'"
+
+                # Log successful completion
+                logger.info(
+                    f"Web search completed successfully: {len(formatted_results)} results"
+                )
+
+                return {
+                    "query": query,
+                    "results_count": len(formatted_results),
+                    "results": formatted_results,
+                    "summary": summary,
+                    "search_engine": "DuckDuckGo",
+                    "timestamp": datetime.now().isoformat(),
+                }
+
+        except Exception as e:
+            error_msg = f"Web search failed for query '{query}': {str(e)}"
+            logger.error(f"{error_msg}")
+
+            return {
+                "query": query,
+                "results_count": 0,
+                "results": [],
+                "error": error_msg,
+                "search_engine": "DuckDuckGo",
+                "timestamp": datetime.now().isoformat(),
+            }
+
+    def _run(
+        self,
+        query: str,
+        max_results: int = 5,
+        region: str = "us-en",
+        run_manager: CallbackManagerForToolRun | None = None,
+    ) -> Tuple[Dict[str, Any], Dict[str, Any]]:
+        """
+        Execute the web search synchronously.
+
+        Args:
+            query (str): Search query
+            max_results (int): Maximum number of results
+            region (str): Search region
+            run_manager: Callback manager (unused)
+
+        Returns:
+            Tuple[Dict[str, Any], Dict[str, Any]]: A tuple containing:
+                - output: Dictionary with search results
+                - metadata: Dictionary with execution metadata
+        """
+        # Create metadata structure
+        metadata = {
+            "query": query,
+            "max_results": max_results,
+            "region": region,
+            "timestamp": time.time(),
+            "tool": "duckduckgo_search",
+            "operation": "search",
+        }
+
+        try:
+            result = self._perform_search_sync(query, max_results, region)
+            
+            # Check if search was successful
+            if "error" in result:
+                metadata["analysis_status"] = "failed"
+                metadata["error_details"] = result["error"]
+            else:
+                metadata["analysis_status"] = "completed"
+                metadata["results_count"] = result.get("results_count", 0)
+
+            return result, metadata
+
+        except Exception as e:
+            error_result = {
+                "query": query,
+                "results_count": 0,
+                "results": [],
+                "error": str(e),
+                "search_engine": "DuckDuckGo",
+                "timestamp": datetime.now().isoformat(),
+            }
+            metadata["analysis_status"] = "failed"
+            metadata["error_details"] = str(e)
+            
+            return error_result, metadata
+
+    async def _arun(
+        self,
+        query: str,
+        max_results: int = 5,
+        region: str = "us-en",
+        run_manager: AsyncCallbackManagerForToolRun | None = None,
+    ) -> Tuple[Dict[str, Any], Dict[str, Any]]:
+        """
+        Execute the web search asynchronously.
+
+        Args:
+            query (str): Search query
+            max_results (int): Maximum number of results
+            region (str): Search region
+            run_manager: Callback manager (unused)
+
+        Returns:
+            Tuple[Dict[str, Any], Dict[str, Any]]: A tuple containing:
+                - output: Dictionary with search results
+                - metadata: Dictionary with execution metadata
+        """
+        # Try to get LangGraph stream writer for progress updates
+        writer = None
+        try:
+            from langgraph.config import get_stream_writer
+
+            writer = get_stream_writer()
+        except Exception:
+            # Stream writer not available (outside LangGraph context)
+            pass
+
+        if writer:
+            writer(
+                {
+                    "tool_name": "DuckDuckGoSearchTool",
+                    "status": "started",
+                    "query": query,
+                    "max_results": max_results,
+                    "step": "Initiating web search",
+                }
+            )
+
+        try:
+            if writer:
+                writer(
+                    {
+                        "tool_name": "DuckDuckGoSearchTool",
+                        "status": "searching",
+                        "step": "Fetching results from DuckDuckGo API",
+                    }
+                )
+
+            # Use asyncio to run sync search in executor
+            loop = asyncio.get_event_loop()
+            result, metadata = await loop.run_in_executor(
+                None, self._run, query, max_results, region
+            )
+
+            if writer:
+                # Parse result to get count for progress update
+                results_count = result.get("results_count", 0)
+                writer(
+                    {
+                        "tool_name": "DuckDuckGoSearchTool",
+                        "status": "completed",
+                        "step": f"Search completed with {results_count} results",
+                        "results_count": results_count,
+                    }
+                )
+
+            return result, metadata
+
+        except Exception as e:
+            if writer:
+                writer(
+                    {
+                        "tool_name": "DuckDuckGoSearchTool",
+                        "status": "error",
+                        "step": f"Search failed: {str(e)}",
+                        "error": str(e),
+                    }
+                )
+
+            error_result = {
+                "query": query,
+                "results_count": 0,
+                "results": [],
+                "error": str(e),
+                "search_engine": "DuckDuckGo",
+                "timestamp": datetime.now().isoformat(),
+            }
+            
+            metadata = {
+                "query": query,
+                "max_results": max_results,
+                "region": region,
+                "timestamp": time.time(),
+                "tool": "duckduckgo_search",
+                "operation": "search",
+                "analysis_status": "failed",
+                "error_details": str(e),
+            }
+            
+            return error_result, metadata
+
+    def get_search_summary(
+        self, query: str, max_results: int = 3
+    ) -> dict[str, str | list[str]]:
+        """
+        Get a quick summary of search results for a given query.
+
+        Args:
+            query (str): The search query.
+            max_results (int): Maximum number of results to summarize.
+
+        Returns:
+            Dict[str, Union[str, List[str]]]: Summary of search results.
+        """
+        try:
+            result, _ = self._run(query, max_results)
+
+            if "error" in result:
+                return {
+                    "query": query,
+                    "status": "error",
+                    "error": result["error"],
+                    "results": [],
+                }
+
+            # Extract key information
+            results = result.get("results", [])
+            titles = [r["title"] for r in results]
+            urls = [r["url"] for r in results]
+            snippets = [
+                (
+                    r["snippet"][:100] + "..."
+                    if len(r["snippet"]) > 100
+                    else r["snippet"]
+                )
+                for r in results
+            ]
+
+            return {
+                "query": query,
+                "status": "success",
+                "total_results": result.get("results_count", 0),
+                "titles": titles,
+                "urls": urls,
+                "snippets": snippets,
+            }
+
+        except Exception as e:
+            logger.error(f"Error getting search summary: {e}")
+            return {
+                "query": query,
+                "status": "error",
+                "error": str(e),
+                "results": [],
+            }

medrax/tools/classification/arcplus.py

@@ -345,7 +345,8 @@ class ArcPlusClassifierTool(BaseTool):
                     predictions = predictions[: len(self.disease_list)]
 
             # Create output dictionary mapping disease names to probabilities
-            output = dict(zip(self.disease_list, predictions.astype(float)))
+            # Convert numpy floats to native Python floats for proper serialization
+            output = dict(zip(self.disease_list, [float(pred) for pred in predictions]))
 
             metadata = {
                 "image_path": image_path,

medrax/tools/segmentation/__init__.py

@@ -0,0 +1,12 @@
+"""Medical image segmentation tools for MedRAX2."""
+
+from .segmentation import ChestXRaySegmentationTool, ChestXRaySegmentationInput, OrganMetrics
+from .medsam2 import MedSAM2Tool, MedSAM2Input
+
+__all__ = [
+    "ChestXRaySegmentationTool",
+    "ChestXRaySegmentationInput",
+    "OrganMetrics",
+    "MedSAM2Tool",
+    "MedSAM2Input"
+] 

medrax/tools/{medsam2.py → segmentation/medsam2.py}

@@ -15,7 +15,7 @@ from langchain_core.callbacks import (
 from langchain_core.tools import BaseTool
 
 # Add MedSAM2 to Python path for proper module resolution
-medsam2_path = str(Path(__file__).parent.parent.parent / "MedSAM2")
+medsam2_path = str(Path(__file__).parent.parent.parent.parent / "MedSAM2")
 if medsam2_path not in sys.path:
     sys.path.append(medsam2_path)
 
@@ -93,7 +93,7 @@ class MedSAM2Tool(BaseTool):
             if GlobalHydra.instance().is_initialized():
                 GlobalHydra.instance().clear()
             
-            config_dir = Path(__file__).parent.parent.parent / "MedSAM2" / "sam2" / "configs"
+            config_dir = Path(__file__).parent.parent.parent.parent / "MedSAM2" / "sam2" / "configs"
             initialize_config_dir(config_dir=str(config_dir), version_base="1.2")
             
             hf_hub_download(

medrax/tools/vqa/__init__.py

@@ -0,0 +1,16 @@
+"""Visual Question Answering tools for medical images."""
+
+from .llava_med import LlavaMedTool, LlavaMedInput
+from .xray_vqa import CheXagentXRayVQATool, XRayVQAToolInput  
+from .medgemma.medgemma_client import MedGemmaAPIClientTool, MedGemmaVQAInput
+from .medgemma.medgemma_setup import setup_medgemma_env
+
+__all__ = [
+    "LlavaMedTool",
+    "LlavaMedInput",
+    "CheXagentXRayVQATool", 
+    "XRayVQAToolInput",
+    "MedGemmaAPIClientTool",
+    "MedGemmaVQAInput",
+    "setup_medgemma_env"
+] 

medrax/tools/{llava_med.py → vqa/llava_med.py}

@@ -151,7 +151,7 @@ class LlavaMedTool(BaseTool):
             output = {
                 "answer": answer,
             }
-            
+           
             metadata = {
                 "question": question,
                 "image_path": image_path,

medrax/tools/vqa/medgemma/medgemma.py

@@ -0,0 +1,425 @@
+import asyncio
+import os
+from pathlib import Path
+import sys
+import traceback
+from typing import Any, Dict, List, Optional, Tuple
+import uuid
+
+from PIL import Image
+
+from fastapi import FastAPI, File, Form, HTTPException, UploadFile
+from pydantic import BaseModel, Field
+import torch
+import transformers
+from transformers import BitsAndBytesConfig, pipeline
+import uvicorn
+
+# Configuration
+UPLOAD_DIR = "./medgemma_images"
+
+# Create directories if they don't exist
+os.makedirs(UPLOAD_DIR, exist_ok=True)
+
+# Pydantic Models for API
+class VQAInput(BaseModel):
+    """Input schema for the MedGemma VQA API endpoint.
+    
+    Defines the structure for requests to the /analyze-images/ endpoint.
+    Used for validating incoming API requests and generating OpenAPI documentation.
+    """
+    prompt: str = Field(..., description="Question or instruction about the medical images")
+    system_prompt: Optional[str] = Field(
+        "You are an expert radiologist.",
+        description="System prompt to set the context for the model",
+    )
+    max_new_tokens: int = Field(
+        300, description="Maximum number of tokens to generate in the response"
+    )
+
+class VQAResponse(BaseModel):
+    """Response schema for successful MedGemma VQA API requests.
+    
+    Defines the structure of successful responses from the /analyze-images/ endpoint.
+    Used for response validation and OpenAPI documentation.
+    """
+    response: str = Field(..., description="Generated medical analysis response from MedGemma model")
+    metadata: Dict[str, Any] = Field(..., description="Additional metadata about the analysis request and results")
+
+class ErrorResponse(BaseModel):
+    """Error response schema for failed MedGemma VQA API requests.
+    
+    Defines the structure of error responses from the /analyze-images/ endpoint.
+    Used for error response validation and OpenAPI documentation.
+    """
+    error: str = Field(..., description="Human-readable error message describing what went wrong")
+    metadata: Dict[str, Any] = Field(..., description="Additional metadata about the error and request context")
+
+# MedGemma Model Handling
+class MedGemmaModel:
+    """Medical visual question answering model using Google's MedGemma 4B model.
+
+    MedGemma is a specialized multimodal AI model trained on medical images and text.
+    It provides expert-level analysis for chest X-rays, dermatology images,
+    ophthalmology images, and histopathology slides.
+
+    Key capabilities:
+    - Medical image classification and analysis across multiple modalities
+    - Visual question answering for radiology, dermatology, pathology, ophthalmology
+    - Clinical reasoning and medical knowledge integration
+    - Multi-modal medical understanding (text + images)
+    - Support for up to 128K context length
+
+    Performance:
+    - Full precision (bfloat16): ~8GB VRAM, recommended for medical applications
+    - 4-bit quantization (default): Available but may affect quality on some systems
+
+    This class implements a singleton pattern to ensure only one model instance
+    is loaded in memory, optimizing resource usage for the FastAPI service.
+    """
+    
+    _instance = None
+
+    def __new__(cls, *args, **kwargs):
+        """Create or return the singleton instance of MedGemmaModel.
+        
+        Ensures only one model instance exists in memory, preventing
+        multiple model loads and conserving GPU memory.
+        
+        Returns:
+            MedGemmaModel: The singleton instance
+        """
+        if not cls._instance:
+            cls._instance = super(MedGemmaModel, cls).__new__(cls)
+        return cls._instance
+
+    def __init__(
+        self,
+        model_name: str = "google/medgemma-4b-it",
+        device: Optional[str] = "cuda",
+        dtype: torch.dtype = torch.bfloat16,
+        cache_dir: Optional[str] = None,
+        load_in_4bit: bool = True,
+        **kwargs: Any,
+    ) -> None:
+        """Initialize the MedGemmaModel.
+
+        Args:
+            model_name: Name of the MedGemma model to use (default: "google/medgemma-4b-it")
+            device: Device to run model on - "cuda" or "cpu" (default: "cuda")
+            dtype: Data type for model weights - bfloat16 recommended for efficiency (default: torch.bfloat16)
+            cache_dir: Directory to cache downloaded models (default: None)
+            load_in_4bit: Whether to load model in 4-bit quantization for memory efficiency (default: True)
+            **kwargs: Additional arguments passed to the model pipeline
+
+        Raises:
+            RuntimeError: If model initialization fails (e.g., insufficient GPU memory)
+        """
+        # Re-initialization guard
+        if hasattr(self, 'pipe') and self.pipe is not None:
+            return
+
+        self.device = device if device and torch.cuda.is_available() else "cpu"
+        self.dtype = dtype
+        self.cache_dir = cache_dir
+
+        # Setup model configuration
+        model_kwargs = {
+            "torch_dtype": self.dtype,
+        }
+
+        if cache_dir:
+            model_kwargs["cache_dir"] = cache_dir
+
+        # Handle device mapping and quantization
+        pipeline_kwargs = {
+            "model": model_name,
+            "model_kwargs": model_kwargs,
+            "trust_remote_code": True,
+            "use_cache": True,
+        }
+
+        if load_in_4bit:
+            model_kwargs["quantization_config"] = BitsAndBytesConfig(load_in_4bit=True)
+        model_kwargs["device_map"] = {"": self.device}
+
+        try:
+            self.pipe = pipeline("image-text-to-text", **pipeline_kwargs)
+        except Exception as e:
+            raise RuntimeError(f"Failed to initialize MedGemma pipeline: {str(e)}")
+
+    def _prepare_messages(
+        self, image_paths: List[str], prompt: str, system_prompt: str
+    ) -> Tuple[List[Dict[str, Any]], List[Image.Image]]:
+        """Prepare chat messages in the format expected by MedGemma.
+
+        Converts image paths to PIL Image objects and formats them into the
+        chat message structure that MedGemma expects for multimodal input.
+
+        Args:
+            image_paths: List of file paths to medical images
+            prompt: User's question or instruction about the images
+            system_prompt: System context message to set the model's role
+
+        Returns:
+            Tuple containing:
+                - List of formatted chat messages for MedGemma
+                - List of loaded PIL Image objects
+
+        Raises:
+            FileNotFoundError: If any image file cannot be found
+        """
+        images = []
+        for path in image_paths:
+            if not Path(path).is_file():
+                raise FileNotFoundError(f"Image file not found: {path}")
+
+            image = Image.open(path)
+            if image.mode != "RGB":
+                image = image.convert("RGB")
+            images.append(image)
+
+        # Create messages in chat format
+        messages = [
+            {"role": "system", "content": [{"type": "text", "text": system_prompt}]},
+            {
+                "role": "user",
+                "content": [{"type": "text", "text": prompt}]
+                + [{"type": "image", "image": img} for img in images],
+            },
+        ]
+
+        return messages, images
+
+    def _generate_response(self, messages: List[Dict[str, Any]], max_new_tokens: int) -> str:
+        """Generate response using MedGemma pipeline.
+
+        Processes the formatted messages through the MedGemma model to generate
+        a medical analysis response.
+
+        Args:
+            messages: Formatted chat messages with images and text
+            max_new_tokens: Maximum number of tokens to generate in response
+
+        Returns:
+            Generated response text from MedGemma model
+        """
+        # Generate using pipeline
+        output = self.pipe(
+            text=messages,
+            max_new_tokens=max_new_tokens,
+            do_sample=False,
+        )
+
+        # Extract generated text from pipeline output
+        if (
+            isinstance(output, list)
+            and output
+            and isinstance(output[0].get("generated_text"), list)
+        ):
+            generated_text = output[0]["generated_text"]
+            if generated_text:
+                return generated_text[-1].get("content", "").strip()
+
+        return "No response generated"
+
+    def _create_error_response(
+        self,
+        image_paths: List[str],
+        prompt: str,
+        error_message: str,
+        error_type: str,
+        error_details: str,
+    ) -> Dict[str, Any]:
+        """Create standardized error response metadata.
+
+        Generates consistent error metadata structure for logging and debugging
+        purposes across different error scenarios.
+
+        Args:
+            image_paths: List of image paths that were being processed
+            prompt: User prompt that was being processed
+            error_message: Human-readable error message
+            error_type: Categorization of the error (e.g., "memory_error", "file_not_found")
+            error_details: Detailed technical error information
+
+        Returns:
+            Dictionary containing standardized error metadata
+        """
+        return {
+            "image_paths": image_paths,
+            "prompt": prompt,
+            "analysis_status": "failed",
+            "error_type": error_type,
+            "error_details": error_details,
+        }
+
+    async def aget_response(self, image_paths: List[str], prompt: str, system_prompt: str, max_new_tokens: int) -> str:
+        """Async method to get response from MedGemma model.
+
+        Main entry point for generating medical analysis responses. Handles
+        the complete pipeline from image loading to response generation
+        in an asynchronous manner.
+
+        Args:
+            image_paths: List of file paths to medical images
+            prompt: User's question or instruction about the images
+            system_prompt: System context message to set the model's role
+            max_new_tokens: Maximum number of tokens to generate in response
+
+        Returns:
+            Generated medical analysis response as a string
+
+        Raises:
+            FileNotFoundError: If any image file cannot be found
+            RuntimeError: If model inference fails
+        """
+        loop = asyncio.get_event_loop()
+        messages, _ = await loop.run_in_executor(None, self._prepare_messages, image_paths, prompt, system_prompt)
+
+        def _generate():
+            return self._generate_response(messages, max_new_tokens)
+
+        return await loop.run_in_executor(None, _generate)
+
+# FastAPI Application
+app = FastAPI(
+    title="MedGemma VQA API",
+    description="API for medical visual question answering using Google's MedGemma model."
+)
+
+medgemma_model: Optional[MedGemmaModel] = None
+
+@app.on_event("startup")
+async def startup_event():
+    """Load the MedGemma model at application startup.
+    
+    This function is called when the FastAPI application starts up.
+    It initializes the MedGemma model as a global singleton instance,
+    ensuring the model is loaded and ready to handle requests.
+    
+    The model is loaded with default settings optimized for medical
+    image analysis, including 4-bit quantization for memory efficiency.
+    
+    Raises:
+        SystemExit: If model loading fails, the application will exit
+                   to prevent serving requests with an unavailable model.
+    """
+    global medgemma_model
+    try:
+        medgemma_model = MedGemmaModel()
+        print("MedGemma model loaded successfully.")
+    except RuntimeError as e:
+        print(f"Error loading MedGemma model: {e}")
+        exit(1)
+
+@app.post("/analyze-images/",
+            response_model=VQAResponse,
+            responses={
+                500: {"model": ErrorResponse, "description": "Internal server error or model inference failure"},
+                404: {"model": ErrorResponse, "description": "Image file not found"},
+                400: {"description": "Invalid request format or unsupported image type"},
+                503: {"description": "Model not available or not loaded"}
+            },
+            summary="Analyze one or more medical images",
+            description="Upload medical images and receive AI-powered analysis using Google's MedGemma model.")
+async def analyze_images(
+    images: List[UploadFile] = File(..., description="List of medical image files to analyze (JPG or PNG)."),
+    prompt: str = Form(..., description="Question or instruction about the medical images."),
+    system_prompt: Optional[str] = Form("You are an expert radiologist.", description="System prompt to set the context for the model."),
+    max_new_tokens: int = Form(100, description="Maximum number of tokens to generate in the response.")
+):
+    """Analyze medical images using MedGemma AI model.
+    
+    This endpoint accepts one or more medical images along with a prompt
+    and returns AI-generated medical analysis.
+    
+    The endpoint handles the complete pipeline:
+    1. Validates uploaded image files
+    2. Saves images temporarily to disk
+    3. Processes images through MedGemma model
+    4. Returns structured analysis with metadata
+    5. Cleans up temporary files
+    
+    Args:
+        images: List of uploaded image files (JPG/PNG format)
+        prompt: Medical question or instruction about the images
+        system_prompt: Context setting for the AI model (default: radiologist role)
+        max_new_tokens: Maximum response length (default: 100)
+    
+    Returns:
+        VQAResponse: Contains the AI-generated analysis and request metadata
+        
+    Raises:
+        HTTPException 400: Invalid image format or request structure
+        HTTPException 404: Image file not found during processing
+        HTTPException 500: Model inference error or memory issues
+        HTTPException 503: Model not available for processing
+    """
+    # Check if model is available
+    if medgemma_model is None or medgemma_model.pipe is None:
+        raise HTTPException(status_code=503, detail="Model is not available. Please try again later.")
+
+    # Process uploaded images
+    image_paths = []
+    for image in images:
+        # Validate image format
+        if image.content_type not in ["image/jpeg", "image/png"]:
+            raise HTTPException(status_code=400, detail=f"Unsupported image format: {image.content_type}. Only JPG and PNG are supported.")
+        
+        # Generate unique filename to avoid conflicts
+        unique_filename = f"{uuid.uuid4()}_{image.filename}"
+        file_path = os.path.join(UPLOAD_DIR, unique_filename)
+        
+        try:
+            # Save uploaded image to disk
+            with open(file_path, "wb") as buffer:
+                buffer.write(await image.read())
+            image_paths.append(file_path)
+        except Exception as e:
+            raise HTTPException(status_code=500, detail=f"Failed to save uploaded image: {str(e)}")
+
+    try:
+        # Generate AI analysis
+        response_text = await medgemma_model.aget_response(image_paths, prompt, system_prompt, max_new_tokens)
+        
+        # Prepare success response
+        metadata = {
+            "image_paths": image_paths,
+            "prompt": prompt,
+            "system_prompt": system_prompt,
+            "max_new_tokens": max_new_tokens,
+            "num_images": len(image_paths),
+            "analysis_status": "completed",
+        }
+        return VQAResponse(response=response_text, metadata=metadata)
+        
+    except FileNotFoundError as e:
+        raise HTTPException(status_code=404, detail=f"Image file not found: {str(e)}")
+    except torch.cuda.OutOfMemoryError as e:
+        error_message = "GPU memory exhausted. Try reducing image resolution or max_new_tokens."
+        metadata = medgemma_model._create_error_response(
+            image_paths, prompt, error_message, "memory_error", str(e)
+        )
+        raise HTTPException(status_code=500, detail=error_message)
+    except Exception as e:
+        traceback.print_exc()
+        metadata = medgemma_model._create_error_response(
+            image_paths, prompt, f"Analysis failed: {str(e)}", "general_error", str(e)
+        )
+        raise HTTPException(status_code=500, detail=f"Analysis failed: {str(e)}")
+    finally:
+        # Clean up temporary image files
+        for path in image_paths:
+            try:
+                os.remove(path)
+            except OSError:
+                pass
+
+if __name__ == "__main__":
+    """Launch the MedGemma VQA API server.
+    
+    Starts the FastAPI application with uvicorn server, binding to all
+    network interfaces on port 8002.
+    """
+    uvicorn.run(app, host="0.0.0.0", port=8002)

medrax/tools/vqa/medgemma/medgemma_client.py

@@ -0,0 +1,278 @@
+import os
+from typing import Any, Dict, List, Optional, Tuple, Type
+
+import httpx
+from langchain_core.callbacks import (
+    AsyncCallbackManagerForToolRun,
+    CallbackManagerForToolRun,
+)
+from langchain_core.tools import BaseTool
+from pydantic import BaseModel, Field
+
+class MedGemmaVQAInput(BaseModel):
+    """Input schema for the MedGemma VQA Tool. Only supports JPG or PNG images."""
+    image_paths: List[str] = Field(
+        ...,
+        description="List of paths to medical image files to analyze, only supports JPG or PNG images",
+    )
+    prompt: str = Field(..., description="Question or instruction about the medical images")
+    system_prompt: Optional[str] = Field(
+        "You are an expert radiologist.",
+        description="System prompt to set the context for the model",
+    )
+    max_new_tokens: int = Field(
+        300, description="Maximum number of tokens to generate in the response"
+    )
+
+class MedGemmaAPIClientTool(BaseTool):
+    """Medical visual question answering tool using Google's MedGemma 4B model via API.
+
+    MedGemma is a specialized multimodal AI model trained on medical images and text.
+    It provides expert-level analysis for chest X-rays, dermatology images,
+    ophthalmology images, and histopathology slides.
+
+    Key capabilities:
+    - Medical image classification and analysis across multiple modalities
+    - Visual question answering for radiology, dermatology, pathology, ophthalmology
+    - Clinical reasoning and medical knowledge integration
+    - Multi-modal medical understanding (text + images)
+    - Support for up to 128K context length
+
+    Performance:
+    - Full precision (bfloat16): ~8GB VRAM, recommended for medical applications
+    - 4-bit quantization (default): Available but may affect quality on some systems
+    """
+
+    name: str = "medgemma_medical_vqa"
+    description: str = (
+        "Advanced medical visual question answering tool using Google's MedGemma 4B instruction-tuned model via API. "
+        "Specialized for comprehensive medical image analysis across multiple modalities including chest X-rays, "
+        "dermatology images, ophthalmology images, and histopathology slides. Provides expert-level medical "
+        "reasoning, diagnosis assistance, and detailed image interpretation with radiologist-level expertise. "
+        "Input: List of medical image paths and medical question/prompt with optional custom system prompt. "
+        "Output: Comprehensive medical analysis and answers based on visual content with detailed reasoning. "
+        "Supports multi-image analysis, comparative studies, and complex medical reasoning tasks. "
+        "Model handles images up to 896x896 resolution and supports context up to 128K tokens."
+    )
+    args_schema: Type[BaseModel] = MedGemmaVQAInput
+    return_direct: bool = True
+
+    # API configuration
+    api_url: str  # The URL of the running FastAPI service
+
+    def __init__(self, api_url: str, **kwargs: Any):
+        """Initialize the MedGemmaAPIClientTool.
+
+        Args:
+            api_url: The URL of the running MedGemma FastAPI service
+            **kwargs: Additional arguments passed to BaseTool
+        """
+        super().__init__(api_url=api_url, **kwargs)
+
+    def _prepare_request_data(
+        self, image_paths: List[str], prompt: str, system_prompt: str, max_new_tokens: int
+    ) -> Tuple[List, Dict]:
+        """Prepare multipart form data for API request.
+
+        Args:
+            image_paths: List of paths to medical images
+            prompt: Question or instruction about the images
+            system_prompt: System context for the model
+            max_new_tokens: Maximum number of tokens to generate
+
+        Returns:
+            Tuple of files list and data dictionary
+        """
+        files_to_send = []
+        opened_files = []
+        
+        for path in image_paths:
+            with open(path, "rb") as f:
+                files_to_send.append(("images", (os.path.basename(path), f.read(), "image/jpeg")))
+
+        data = {
+            "prompt": prompt,
+            "system_prompt": system_prompt,
+            "max_new_tokens": max_new_tokens,
+        }
+        
+        return files_to_send, data, opened_files
+
+    def _create_error_response(
+        self,
+        image_paths: List[str],
+        prompt: str,
+        error_message: str,
+        error_type: str,
+        error_details: str,
+    ) -> Tuple[Dict[str, Any], Dict]:
+        """Create standardized error response.
+
+        Args:
+            image_paths: List of image paths
+            prompt: User prompt
+            error_message: Human-readable error message
+            error_type: Type of error
+            error_details: Detailed error information
+
+        Returns:
+            Tuple of error output and metadata
+        """
+        output = {"error": error_message}
+        metadata = {
+            "image_paths": image_paths,
+            "prompt": prompt,
+            "analysis_status": "failed",
+            "error_type": error_type,
+            "error_details": error_details,
+        }
+        return output, metadata
+
+    def _run(
+        self,
+        image_paths: List[str],
+        prompt: str,
+        system_prompt: str = "You are an expert radiologist.",
+        max_new_tokens: int = 300,
+        run_manager: Optional[CallbackManagerForToolRun] = None,
+    ) -> Tuple[Dict[str, Any], Dict]:
+        """Execute medical visual question answering via API.
+
+        Args:
+            image_paths: List of paths to medical images
+            prompt: Question or instruction about the images
+            system_prompt: System context for the model
+            max_new_tokens: Maximum number of tokens to generate
+            run_manager: Optional callback manager
+
+        Returns:
+            Tuple of output dictionary and metadata
+        """
+        # httpx is a modern HTTP client that supports sync and async
+        timeout_config = httpx.Timeout(300.0, connect=10.0)
+        client = httpx.Client(timeout=timeout_config)
+        
+        try:
+            # Prepare the multipart form data
+            files_to_send, data, opened_files = self._prepare_request_data(
+                image_paths, prompt, system_prompt, max_new_tokens
+            )
+            
+            response = client.post(
+                f"{self.api_url}/analyze-images/",
+                data=data,
+                files=files_to_send,
+            )
+            response.raise_for_status()  # Raise an exception for bad status codes (4xx or 5xx)
+            
+            response_data = response.json()
+            output = {"response": response_data["response"]}
+            
+            metadata = {
+                "image_paths": image_paths,
+                "prompt": prompt,
+                "system_prompt": system_prompt,
+                "max_new_tokens": max_new_tokens,
+                "num_images": len(image_paths),
+                "analysis_status": "completed",
+            }
+
+            return output, metadata
+
+        except httpx.TimeoutException as e:
+            return self._create_error_response(
+                image_paths,
+                prompt,
+                f"Error: The request to the MedGemma API timed out after {timeout_config.read} seconds. The server might be overloaded or the model is taking too long to load. Try again later.",
+                "timeout_error",
+                str(e)
+            )
+        except httpx.ConnectError as e:
+            return self._create_error_response(
+                image_paths,
+                prompt,
+                f"Error: Could not connect to the MedGemma API. Check if the server address '{self.api_url}' is correct and running.",
+                "connection_error",
+                str(e)
+            )
+        except httpx.HTTPStatusError as e:
+            return self._create_error_response(
+                image_paths,
+                prompt,
+                f"Error: The MedGemma API returned an error (Status {e.response.status_code}): {e.response.text}",
+                "http_error",
+                f"Status {e.response.status_code}: {e.response.text}"
+            )
+        except Exception as e:
+            return self._create_error_response(
+                image_paths,
+                prompt,
+                f"An unexpected error occurred in the MedGemma client tool: {str(e)}",
+                "general_error",
+                str(e)
+            )
+        finally:
+            # Ensure all opened files are closed
+            if 'opened_files' in locals():
+                for f in opened_files:
+                    f.close()
+
+    async def _arun(
+        self,
+        image_paths: List[str],
+        prompt: str,
+        system_prompt: str = "You are an expert radiologist.",
+        max_new_tokens: int = 300,
+        run_manager: Optional[AsyncCallbackManagerForToolRun] = None,
+    ) -> Tuple[Dict[str, Any], Dict]:
+        """Execute the tool asynchronously."""
+        async with httpx.AsyncClient() as client:
+            try:
+                # Prepare the multipart form data
+                files_to_send, data, opened_files = self._prepare_request_data(
+                    image_paths, prompt, system_prompt, max_new_tokens
+                )
+                
+                response = await client.post(
+                    f"{self.api_url}/analyze-images/",
+                    data=data,
+                    files=files_to_send,
+                    timeout=120.0
+                )
+                response.raise_for_status()
+                
+                response_data = response.json()
+                output = {"response": response_data["response"]}
+                
+                metadata = {
+                    "image_paths": image_paths,
+                    "prompt": prompt,
+                    "system_prompt": system_prompt,
+                    "max_new_tokens": max_new_tokens,
+                    "num_images": len(image_paths),
+                    "analysis_status": "completed",
+                }
+
+                return output, metadata
+
+            except httpx.HTTPStatusError as e:
+                return self._create_error_response(
+                    image_paths,
+                    prompt,
+                    f"Error calling MedGemma API: {e.response.status_code} - {e.response.text}",
+                    "http_error",
+                    f"Status {e.response.status_code}: {e.response.text}"
+                )
+            except Exception as e:
+                return self._create_error_response(
+                    image_paths,
+                    prompt,
+                    f"An unexpected error occurred: {str(e)}",
+                    "general_error",
+                    str(e)
+                )
+            finally:
+                # Ensure all opened files are closed
+                if 'opened_files' in locals():
+                    for f in opened_files:
+                        f.close()

medrax/tools/vqa/medgemma/medgemma_requirements_standard.txt

@@ -0,0 +1,55 @@
+accelerate==1.9.0
+annotated_types==0.7.0
+anyio==4.9.0
+bitsandbytes==0.46.0
+certifi==2025.7.14
+charset_normalizer==3.4.2
+click==8.2.1
+fastapi==0.116.1
+filelock==3.18.0
+fsspec==2025.7.0
+h11==0.16.0
+hf_xet==1.1.3
+httpcore==1.0.9
+httpx==0.28.1
+huggingface-hub==0.34.3
+idna==3.10
+inquirerpy==0.3.4
+jinja2==3.1.6
+jsonpatch==1.33
+jsonpointer==3.0.0
+langchain-core==0.3.72
+langsmith==0.4.8
+MarkupSafe==2.1.5
+mpmath==1.3.0
+networkx==3.5
+numpy==2.2.2
+orjson==3.10.5
+packaging==25.0
+pfzy==0.3.4
+pillow==11.1.0
+prompt_toolkit==3.0.51
+psutil==6.1.1
+pydantic==2.11.7
+pydantic_core==2.33.2
+python_multipart==0.0.20
+PyYAML==6.0.2
+regex==2024.11.6
+requests==2.32.4
+requests_toolbelt==1.0.0
+safetensors==0.5.3
+sniffio==1.3.1
+sshuttle==1.3.1
+starlette==0.47.2
+sympy==1.14.0
+tenacity==9.1.2
+tokenizers==0.21.1
+torch==2.7.1
+tqdm==4.67.1
+transformers==4.54.1
+typing_extensions==4.14.1
+typing_inspection==0.4.1
+urllib3==2.5.0
+uvicorn==0.35.0
+wcwidth==0.2.13
+zstandard==0.23.0 

medrax/tools/vqa/medgemma/medgemma_setup.py

@@ -0,0 +1,64 @@
+import os
+from pathlib import Path
+import subprocess
+import venv
+
+def setup_medgemma_env():
+    """Set up MedGemma virtual environment and launch the FastAPI service.
+    
+    This function performs the following steps:
+    1. Creates a virtual environment for MedGemma if it doesn't exist
+    2. Installs MedGemma-specific dependencies from requirements.txt
+    3. Launches the MedGemma FastAPI service in the isolated environment
+    
+    Returns:
+        None: Launches MedGemma service as a background process
+        
+    Raises:
+        subprocess.CalledProcessError: If pip installation fails
+        FileNotFoundError: If required files are missing
+        OSError: If virtual environment creation fails
+    """
+    # Get the directory containing this script
+    current_dir = Path(__file__).resolve().parent
+    
+    # Define paths for MedGemma components
+    medgemma_path = current_dir / "medgemma.py"
+    requirements_path = current_dir / "medgemma_requirements_standard.txt"
+    env_dir = current_dir / "medgemma_env"
+
+    # Determine executable paths based on operating system
+    if os.name == "nt":  # Windows
+        pip_executable = env_dir / "Scripts" / "pip"
+        python_executable = env_dir / "Scripts" / "python"
+    else:  # Unix/Linux/macOS
+        pip_executable = env_dir / "bin" / "pip"
+        python_executable = env_dir / "bin" / "python"
+
+    # Create virtual environment if it doesn't exist
+    if not env_dir.exists():
+        print("Creating MedGemma virtual environment...")
+        venv.create(env_dir, with_pip=True)
+        
+        # Install MedGemma dependencies
+        print("Installing MedGemma dependencies...")
+        subprocess.check_call([
+            str(pip_executable), 
+            "install", 
+            "-r", 
+            str(requirements_path)
+        ])
+
+    # Ensure environment exists before accessing executables
+    if not env_dir.exists():
+        raise RuntimeError("Failed to create MedGemma virtual environment")
+
+    # Launch MedGemma FastAPI service
+    print("Launching MedGemma FastAPI service...")
+    subprocess.Popen([
+        str(python_executable), 
+        str(medgemma_path)
+    ])
+    # Note: stdout and stderr redirection commented out for debugging
+    # stdout=subprocess.DEVNULL,
+    # stderr=subprocess.DEVNULL,

medrax/tools/{xray_vqa.py → vqa/xray_vqa.py}

@@ -24,10 +24,10 @@ class XRayVQAToolInput(BaseModel):
     )
 
 
-class XRayVQATool(BaseTool):
+class CheXagentXRayVQATool(BaseTool):
     """Tool that leverages CheXagent for comprehensive chest X-ray analysis."""
 
-    name: str = "chest_xray_expert"
+    name: str = "chexagent_xray_vqa"
     description: str = (
         "A versatile tool for analyzing chest X-rays. "
         "Can perform multiple tasks including: visual question answering, report generation, "
@@ -51,7 +51,7 @@ class XRayVQATool(BaseTool):
         cache_dir: Optional[str] = None,
         **kwargs: Any,
     ) -> None:
-        """Initialize the XRayVQATool.
+        """Initialize the CheXagentXRayVQATool.
 
         Args:
             model_name: Name of the CheXagent model to use

pyproject.toml

@@ -57,7 +57,6 @@ dependencies = [
     "torch>=2.2.0",
     "torchvision>=0.10.0",
     "scikit-image>=0.18.0",
-    "gradio>=5.0.0",
     "opencv-python>=4.8.0",
     "matplotlib>=3.8.0",
     "diffusers>=0.20.0",
@@ -65,16 +64,15 @@ dependencies = [
     "pylibjpeg>=1.0.0",
     "jupyter>=1.0.0",
     "albumentations>=1.0.0",
-    "pyarrow>=10.0.0",
     "chromadb>=0.0.10",
     "pinecone-client>=3.2.2",
     "langchain-pinecone>=0.0.1",
     "langchain-google-genai>=0.1.0",
     "ray>=2.9.0",
-    "langchain-sandbox>=0.0.6",
     "seaborn>=0.12.0",
     "huggingface_hub>=0.17.0",
     "iopath>=0.1.10",
+    "duckduckgo-search>=4.0.0",
 ]
 
 [project.optional-dependencies]