Setup argparse

main.py

@@ -11,6 +11,9 @@ with different model weights, tools, and parameters.
 
 import warnings
 import os
+import argparse
+import threading
+import uvicorn
 from typing import Dict, List, Optional, Any
 from dotenv import load_dotenv
 from transformers import logging
@@ -19,6 +22,7 @@ from langgraph.checkpoint.memory import MemorySaver
 from medrax.models import ModelFactory
 
 from interface import create_demo
+from api import create_api
 from medrax.agent import *
 from medrax.tools import *
 from medrax.utils import *
@@ -37,7 +41,7 @@ def initialize_agent(
     model_dir: str = "/model-weights",
     temp_dir: str = "temp",
     device: str = "cuda",
-    model: str = "gemini-2.5-pro",
+    model: str = "gpt-4.1",
     temperature: float = 1.0,
     rag_config: Optional[RAGConfig] = None,
     model_kwargs: Dict[str, Any] = {},
@@ -137,56 +141,216 @@ def initialize_agent(
     return agent, tools_dict
 
 
-if __name__ == "__main__":
+def run_gradio_interface(agent, tools_dict, host="0.0.0.0", port=8686):
     """
-    This is the main entry point for the MedRAX application.
-    It initializes the agent with the selected tools and creates the demo.
+    Run the Gradio web interface.
+    
+    Args:
+        agent: The initialized MedRAX agent
+        tools_dict: Dictionary of available tools
+        host (str): Host to bind the server to
+        port (int): Port to run the server on
     """
-    print("Starting server...")
-
-    # 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
-
-        # Report Generation Tools
-        "ChestXRayReportGeneratorTool",  # For generating medical reports from X-rays
+    print(f"Starting Gradio interface on {host}:{port}")
+    demo = create_demo(agent, tools_dict)
+    demo.launch(server_name=host, server_port=port, share=True)
 
-        # 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
+def run_api_server(agent, tools_dict, host="0.0.0.0", port=8000):
+    """
+    Run the FastAPI server.
+    
+    Args:
+        agent: The initialized MedRAX agent
+        tools_dict: Dictionary of available tools
+        host (str): Host to bind the server to
+        port (int): Port to run the server on
+    """
+    print(f"Starting API server on {host}:{port}")
+    app = create_api(agent, tools_dict)
+    uvicorn.run(app, host=host, port=port)
 
-        # 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
+def parse_arguments():
+    """Parse command line arguments."""
+    parser = argparse.ArgumentParser(description="MedRAX - Medical Reasoning Agent for Chest X-ray")
+    
+    # Server configuration
+    parser.add_argument(
+        "--mode", 
+        choices=["gradio", "api", "both"], 
+        default="gradio",
+        help="Run mode: 'gradio' for web interface, 'api' for REST API, 'both' for both services"
+    )
+    parser.add_argument("--gradio-host", default="0.0.0.0", help="Gradio host address")
+    parser.add_argument("--gradio-port", type=int, default=8686, help="Gradio port")
+    parser.add_argument("--api-host", default="0.0.0.0", help="API host address")
+    parser.add_argument("--api-port", type=int, default=8000, help="API port")
+    
+    # Model and system configuration
+    parser.add_argument(
+        "--model-dir", 
+        default="/model-weights",
+        help="Directory containing model weights (default: uses MODEL_WEIGHTS_DIR env var or '/model-weights')"
+    )
+    parser.add_argument(
+        "--device", 
+        default="cuda",
+        help="Device to run models on (default: uses MEDRAX_DEVICE env var or 'cuda:1')"
+    )
+    parser.add_argument(
+        "--model", 
+        default="gpt-4.1",
+        help="Model to use (default: gpt-4.1). Examples: gpt-4.1-2025-04-14, gemini-2.5-pro, gpt-5"
+    )
+    parser.add_argument(
+        "--temperature", 
+        type=float, 
+        default=1.0,
+        help="Temperature for the model (default: 1.0)"
+    )
+    parser.add_argument(
+        "--temp-dir", 
+        default="temp2",
+        help="Directory for temporary files (default: temp2)"
+    )
+    parser.add_argument(
+        "--prompt-file", 
+        default="medrax/docs/system_prompts.txt",
+        help="Path to file containing system prompts (default: medrax/docs/system_prompts.txt)"
+    )
+    parser.add_argument(
+        "--system-prompt", 
+        default="MEDICAL_ASSISTANT",
+        help="System prompt to use (default: MEDICAL_ASSISTANT)"
+    )
+    
+    # RAG configuration
+    parser.add_argument(
+        "--rag-model", 
+        default="command-a-03-2025",
+        help="Chat model for RAG responses (default: command-a-03-2025)"
+    )
+    parser.add_argument(
+        "--rag-embedding-model", 
+        default="embed-v4.0",
+        help="Embedding model for RAG system (default: embed-v4.0)"
+    )
+    parser.add_argument(
+        "--rag-rerank-model", 
+        default="rerank-v3.5",
+        help="Reranking model for RAG system (default: rerank-v3.5)"
+    )
+    parser.add_argument(
+        "--rag-temperature", 
+        type=float, 
+        default=0.3,
+        help="Temperature for RAG model (default: 0.3)"
+    )
+    parser.add_argument(
+        "--pinecone-index", 
+        default="medrax2",
+        help="Pinecone index name (default: medrax2)"
+    )
+    parser.add_argument(
+        "--chunk-size", 
+        type=int, 
+        default=1500,
+        help="RAG chunk size (default: 1500)"
+    )
+    parser.add_argument(
+        "--chunk-overlap", 
+        type=int, 
+        default=300,
+        help="RAG chunk overlap (default: 300)"
+    )
+    parser.add_argument(
+        "--retriever-k", 
+        type=int, 
+        default=3,
+        help="Number of documents to retrieve (default: 3)"
+    )
+    parser.add_argument(
+        "--rag-docs-dir", 
+        default="rag_docs",
+        help="Directory for RAG documents (default: rag_docs)"
+    )
+    
+    # Tools configuration
+    parser.add_argument(
+        "--tools", 
+        nargs="*",
+        help="Specific tools to enable (if not provided, uses default set). Available tools: " +
+             "ImageVisualizerTool, DicomProcessorTool, MedSAM2Tool, ChestXRaySegmentationTool, " +
+             "ChestXRayGeneratorTool, TorchXRayVisionClassifierTool, ArcPlusClassifierTool, " +
+             "ChestXRayReportGeneratorTool, XRayPhraseGroundingTool, MedGemmaVQATool, " +
+             "XRayVQATool, LlavaMedTool, MedicalRAGTool, WebBrowserTool, DuckDuckGoSearchTool, " +
+             "PythonSandboxTool"
+    )
+    
+    return parser.parse_args()
 
-        # Development Tools
-        # "PythonSandboxTool",  # Add the Python sandbox tool
-    ]
 
-    # Share a single cache directory and device across tools
-    model_dir = os.getenv("MODEL_WEIGHTS_DIR", "/model-weights")
-    device = os.getenv("MEDRAX_DEVICE", "cuda:0")
+if __name__ == "__main__":
+    """
+    This is the main entry point for the MedRAX application.
+    It initializes the agent with the selected tools and creates the demo/API.
+    """
+    args = parse_arguments()
+    print(f"Starting MedRAX in {args.mode} mode...")
+
+    # Configure tools based on arguments
+    if args.tools is not None:
+        # Use tools specified via command line
+        selected_tools = args.tools
+    else:
+        # Use default tools selection
+        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
+
+            # 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
+
+            # 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
+        ]
+
+    # Configure model directory and device
+    model_dir = args.model_dir or os.getenv("MODEL_WEIGHTS_DIR", "/model-weights")
+    device = args.device or os.getenv("MEDRAX_DEVICE", "cuda:0")
+
+    print(f"Using model directory: {model_dir}")
+    print(f"Using device: {device}")
+    print(f"Using model: {args.model}")
+    print(f"Selected tools: {selected_tools}")
 
     # Setup the MedGemma environment if the MedGemmaVQATool is selected
     if "MedGemmaVQATool" in selected_tools:
@@ -195,15 +359,15 @@ if __name__ == "__main__":
     # Configure the Retrieval Augmented Generation (RAG) system
     # This allows the agent to access and use medical knowledge documents
     rag_config = RAGConfig(
-        model="command-a-03-2025",  # Chat model for generating responses
-        embedding_model="embed-v4.0",  # Embedding model for the RAG system
-        rerank_model="rerank-v3.5",  # Reranking model for the RAG system
-        temperature=0.3,
-        pinecone_index_name="medrax2",  # Name for the Pinecone index
-        chunk_size=1500,
-        chunk_overlap=300,
-        retriever_k=3,
-        local_docs_dir="rag_docs",  # Change this to the path of the documents for RAG
+        model=args.rag_model,
+        embedding_model=args.rag_embedding_model,
+        rerank_model=args.rag_rerank_model,
+        temperature=args.rag_temperature,
+        pinecone_index_name=args.pinecone_index,
+        chunk_size=args.chunk_size,
+        chunk_overlap=args.chunk_overlap,
+        retriever_k=args.retriever_k,
+        local_docs_dir=args.rag_docs_dir,
         huggingface_datasets=["VictorLJZ/medrax2"],  # List of HuggingFace datasets to load
         dataset_split="train",  # Which split of the datasets to use
     )
@@ -212,18 +376,33 @@ if __name__ == "__main__":
     model_kwargs = {}
 
     agent, tools_dict = initialize_agent(
-        prompt_file="medrax/docs/system_prompts.txt",
+        prompt_file=args.prompt_file,
         tools_to_use=selected_tools,
         model_dir=model_dir,
-        temp_dir="temp2",  # Change this to the path of the temporary directory
+        temp_dir=args.temp_dir,
         device=device,
-        model="gpt-4.1",  # 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,
+        model=args.model,
+        temperature=args.temperature,
         model_kwargs=model_kwargs,
         rag_config=rag_config,
-        system_prompt="MEDICAL_ASSISTANT",
+        system_prompt=args.system_prompt,
     )
 
-    # Create and launch the web interface
-    demo = create_demo(agent, tools_dict)
-    demo.launch(server_name="0.0.0.0", server_port=8686, share=True)
+    # Launch based on selected mode
+    if args.mode == "gradio":
+        run_gradio_interface(agent, tools_dict, args.gradio_host, args.gradio_port)
+
+    elif args.mode == "api":
+        run_api_server(agent, tools_dict, args.api_host, args.api_port)
+
+    elif args.mode == "both":
+        # Run both services in separate threads
+        api_thread = threading.Thread(
+            target=run_api_server, 
+            args=(agent, tools_dict, args.api_host, args.api_port)
+        )
+        api_thread.daemon = True
+        api_thread.start()
+        
+        # Run Gradio in main thread
+        run_gradio_interface(agent, tools_dict, args.gradio_host, args.gradio_port)

pyproject.toml

@@ -46,8 +46,9 @@ dependencies = [
     "gradio>=3.0.0",
     "gradio_client>=0.2.0",
     "httpx>=0.23.0",
-    "uvicorn>=0.15.0",
+    "uvicorn[standard]>=0.15.0",
     "fastapi>=0.68.0",
+    "python-multipart>=0.0.6",
     "einops>=0.3.0",
     "einops-exts>=0.0.4",
     "timm==0.5.4",