Merge pull request #2 from bowang-lab/victor

README.md

@@ -15,7 +15,7 @@ Chest X-rays (CXRs) play an integral role in driving critical decisions in disea
 ## MedRAX
 MedRAX is built on a robust technical foundation:
 - **Core Architecture**: Built on LangChain and LangGraph frameworks
-- **Language Model**: Uses GPT-4o with vision capabilities as the backbone LLM
+- **Language Models**: Supports multiple LLM providers including OpenAI (GPT-4o) and Google (Gemini) models
 - **Deployment**: Supports both local and cloud-based deployments
 - **Interface**: Production-ready interface built with Gradio
 - **Modular Design**: Tool-agnostic architecture allowing easy integration of new capabilities
@@ -27,6 +27,7 @@ MedRAX is built on a robust technical foundation:
 - **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
 - **Utilities**: Includes DICOM processing, visualization tools, and custom plotting capabilities
 <br><br>
 
@@ -180,6 +181,7 @@ No additional model weights required:
 ```python
 ImageVisualizerTool()
 DicomProcessorTool(temp_dir=temp_dir)
+WebBrowserTool()  # Requires Google Search API credentials
 ```
 <br>
 
@@ -212,12 +214,45 @@ ChestXRayGeneratorTool(
 - Some tools (LLaVA-Med, Grounding) are more resource-intensive
 <br>
 
-### Local LLMs
+### Language Model Options
+MedRAX supports multiple language model providers:
+
+#### OpenAI Models
+Supported prefixes: `gpt-` and `chatgpt-`
+```
+export OPENAI_API_KEY="your-openai-api-key"
+export OPENAI_BASE_URL="https://api.openai.com/v1"  # Optional for custom endpoints
+```
+
+#### Google Gemini Models
+Supported prefix: `gemini-`
+```
+export GOOGLE_API_KEY="your-google-api-key"
+```
+
+#### OpenRouter Models (Open Source & Proprietary)
+Supported prefix: `openrouter-`
+
+Access many open source and proprietary models via [OpenRouter](https://openrouter.ai/):
+```
+export OPENROUTER_API_KEY="your-openrouter-api-key"
+```
+
+**Note:** Tool compatibility may vary with open-source models. For best results with tools, we recommend using OpenAI or Google Gemini models.
+
+#### Local LLMs
 If you are running a local LLM using frameworks like [Ollama](https://ollama.com/) or [LM Studio](https://lmstudio.ai/), you need to configure your environment variables accordingly. For example:
 ```
 export OPENAI_BASE_URL="http://localhost:11434/v1"
 export OPENAI_API_KEY="ollama"
 ```
+
+#### WebBrowserTool Configuration
+If you're using the WebBrowserTool, you'll need to set these environment variables:
+```
+export GOOGLE_SEARCH_API_KEY="your-google-search-api-key"
+export GOOGLE_SEARCH_ENGINE_ID="your-google-search-engine-id"
+```
 <br>
 
 ## Star History

main.py

@@ -23,7 +23,7 @@ def initialize_agent(
     model_dir="/model-weights",
     temp_dir="temp",
     device="cuda",
-    model="gpt-4o",
+    model="chatgpt-4o-latest",
     temperature=0.7,
     top_p=0.95,
     model_kwargs={}
@@ -137,7 +137,7 @@ if __name__ == "__main__":
         model_dir="/m_weights",  # Change this to the path of the model weights
         temp_dir="temp",  # Change this to the path of the temporary directory
         device="cpu",  # Change this to the device you want to use
-        model="gemini-2.5-pro",  # Change this to the model you want to use, e.g. gpt-4o-mini, gemini-2.5-pro
+        model="gpt-4o-mini",  # Change this to the model you want to use, e.g. gpt-4o-mini, gemini-2.5-pro
         temperature=0.7,
         top_p=0.95,
         model_kwargs=model_kwargs

medrax/models/model_factory.py

@@ -22,10 +22,21 @@ class ModelFactory:
             "env_key": "OPENAI_API_KEY",
             "base_url_key": "OPENAI_BASE_URL"
         },
+        "chatgpt": {
+            "class": ChatOpenAI,
+            "env_key": "OPENAI_API_KEY",
+            "base_url_key": "OPENAI_BASE_URL"
+        },
         "gemini": {
             "class": ChatGoogleGenerativeAI,
             "env_key": "GOOGLE_API_KEY"
         },
+        "openrouter": {
+            "class": ChatOpenAI,  # OpenRouter uses OpenAI-compatible interface
+            "env_key": "OPENROUTER_API_KEY",
+            "base_url_key": "OPENROUTER_BASE_URL",
+            "default_base_url": "https://openrouter.ai/api/v1"
+        },
         # Add more providers with default configurations here
     }
     
@@ -91,17 +102,26 @@ class ModelFactory:
             print(f"Warning: Environment variable {env_key} not found. Authentication may fail.")
         
         # Check for base_url if applicable
-        if "base_url_key" in provider and provider["base_url_key"] in os.environ:
-            provider_kwargs["base_url"] = os.environ[provider["base_url_key"]]
+        if "base_url_key" in provider:
+            if provider["base_url_key"] in os.environ:
+                provider_kwargs["base_url"] = os.environ[provider["base_url_key"]]
+            elif "default_base_url" in provider:
+                provider_kwargs["base_url"] = provider["default_base_url"]
         
         # Merge with any additional provider-specific settings from the registry
         for k, v in provider.items():
-            if k not in ["class", "env_key", "base_url_key"]:
+            if k not in ["class", "env_key", "base_url_key", "default_base_url"]:
                 provider_kwargs[k] = v
         
+        # Strip the provider prefix from the model name
+        # For example, 'openrouter-anthropic/claude-sonnet-4' becomes 'anthropic/claude-sonnet-4'
+        actual_model_name = model_name
+        if model_name.startswith(f"{provider_prefix}-"):
+            actual_model_name = model_name[len(provider_prefix)+1:]
+            
         # Create and return the model instance
         return model_class(
-            model=model_name,
+            model=actual_model_name,
             temperature=temperature,
             top_p=top_p,
             **provider_kwargs,