updates

benchmarking/cli.py

@@ -60,14 +60,6 @@ def run_benchmark_command(args) -> None:
     
     # Create LLM provider
     provider_kwargs = {}
-    if args.provider == "medrax":
-        provider_kwargs = {
-            "tools_to_use": args.medrax_tools.split(",") if args.medrax_tools else None,
-            "model_dir": args.model_dir,
-            "temp_dir": args.temp_dir,
-            "device": args.device,
-            "rag_config": None,  # You might want to add RAG config options
-        }
     
     llm_provider = create_llm_provider(args.model, args.provider, **provider_kwargs)
     
@@ -82,12 +74,8 @@ def run_benchmark_command(args) -> None:
         benchmark_name=args.benchmark,
         output_dir=args.output_dir,
         max_questions=args.max_questions,
-        start_index=args.start_index,
         temperature=args.temperature,
-        max_tokens=args.max_tokens,
-        system_prompt=args.system_prompt,
-        save_frequency=args.save_frequency,
-        log_level=args.log_level,
+        max_tokens=args.max_tokens
     )
     
     # Run benchmark
@@ -126,39 +114,11 @@ def main():
     run_parser.add_argument("--data-dir", required=True, help="Directory containing benchmark data")
     run_parser.add_argument("--output-dir", default="benchmark_results", help="Output directory for results")
     run_parser.add_argument("--max-questions", type=int, help="Maximum number of questions to process")
-    run_parser.add_argument("--start-index", type=int, default=0, help="Starting index for questions")
     run_parser.add_argument("--temperature", type=float, default=0.7, help="Model temperature")
     run_parser.add_argument("--max-tokens", type=int, default=1500, help="Maximum tokens per response")
-    run_parser.add_argument("--system-prompt", help="System prompt for the model")
-    run_parser.add_argument("--save-frequency", type=int, default=10, help="Save results every N questions")
-    run_parser.add_argument("--log-level", default="INFO", choices=["DEBUG", "INFO", "WARNING", "ERROR"])
-    
-    # MedRAX-specific arguments
-    run_parser.add_argument("--medrax-tools", help="Comma-separated list of tools for MedRAX (e.g., WebBrowserTool,MedicalRAGTool)")
-    run_parser.add_argument("--model-dir", default="/model-weights", help="Directory containing model weights for MedRAX")
-    run_parser.add_argument("--temp-dir", default="temp", help="Temporary directory for MedRAX")
-    run_parser.add_argument("--device", default="cuda", help="Device for MedRAX models")
-    
-
     
     run_parser.set_defaults(func=run_benchmark_command)
     
-    # Evaluate results command
-    eval_parser = subparsers.add_parser("evaluate", help="Evaluate benchmark results")
-    eval_parser.add_argument("results_files", nargs="+", help="Path(s) to results files")
-    eval_parser.add_argument("--output-dir", default="evaluation_results", help="Output directory for evaluation")
-    eval_parser.add_argument("--report-name", default="evaluation_report", help="Name for the evaluation report")
-    eval_parser.add_argument("--statistical-test", action="store_true", help="Run statistical significance tests")
-    eval_parser.set_defaults(func=evaluate_results_command)
-    
-    # List providers command
-    list_providers_parser = subparsers.add_parser("list-providers", help="List available LLM providers")
-    list_providers_parser.set_defaults(func=list_providers_command)
-    
-    # List benchmarks command
-    list_benchmarks_parser = subparsers.add_parser("list-benchmarks", help="List available benchmarks")
-    list_benchmarks_parser.set_defaults(func=list_benchmarks_command)
-    
     args = parser.parse_args()
     
     if args.command is None:

benchmarking/llm_providers/base.py

@@ -1,10 +1,9 @@
 """Base class for LLM providers."""
 
 from abc import ABC, abstractmethod
-from typing import Dict, List, Optional, Any, Union
+from typing import Dict, List, Optional, Any
 from dataclasses import dataclass
 import base64
-import time
 from pathlib import Path
 
 

benchmarking/llm_providers/google_provider.py

@@ -2,7 +2,6 @@
 
 import os
 import time
-from typing import Dict, Any
 from tenacity import retry, wait_exponential, stop_after_attempt
 from langchain_google_genai import ChatGoogleGenerativeAI
 from langchain_core.messages import HumanMessage, SystemMessage

benchmarking/llm_providers/medrax_provider.py

@@ -1,23 +1,14 @@
 """MedRAX LLM provider implementation."""
 
-import os
 import time
 import tempfile
 import shutil
-from typing import Dict, Any, List, Optional
 from pathlib import Path
-import json
 
 from .base import LLMProvider, LLMRequest, LLMResponse
 
-# Import MedRAX components
-from medrax.agent import Agent
-from medrax.tools import *
-from medrax.utils import load_prompts_from_file
 from medrax.rag.rag import RAGConfig
-from medrax.models import ModelFactory
-from langgraph.checkpoint.memory import MemorySaver
-from langchain_core.messages import HumanMessage
+from main import initialize_agent
 
 
 class MedRAXProvider(LLMProvider):
@@ -30,21 +21,7 @@ class MedRAXProvider(LLMProvider):
             model_name (str): Base LLM model name (e.g., "gpt-4.1-2025-04-14")
             **kwargs: Additional configuration parameters
         """
-        # MedRAX-specific configuration
-        self.tools_to_use = kwargs.get("tools_to_use", [
-            "WebBrowserTool",
-            "MedicalRAGTool", 
-            "PythonSandboxTool"
-        ])
-        self.model_dir = kwargs.get("model_dir", "/model-weights")
-        self.temp_dir = kwargs.get("temp_dir", "temp")
-        self.device = kwargs.get("device", "cuda")
-        self.temperature = kwargs.get("temperature", 0.7)
-        self.top_p = kwargs.get("top_p", 0.95)
-        self.rag_config = kwargs.get("rag_config")
-        self.prompt_file = kwargs.get("prompt_file", "medrax/docs/system_prompts.txt")
-        
-        # Initialize agent as None, will be created in _setup
+        self.model_name = model_name
         self.agent = None
         self.tools_dict = None
         
@@ -53,71 +30,60 @@ class MedRAXProvider(LLMProvider):
     def _setup(self) -> None:
         """Set up MedRAX agent system."""
         try:
-            # Load system prompts
-            prompts = load_prompts_from_file(self.prompt_file)
-            prompt = prompts["MEDICAL_ASSISTANT"]
-            
-            # Initialize tools
-            all_tools = {
-                "TorchXRayVisionClassifierTool": lambda: TorchXRayVisionClassifierTool(device=self.device),
-                "ArcPlusClassifierTool": lambda: ArcPlusClassifierTool(cache_dir=self.model_dir, device=self.device),
-                "ChestXRaySegmentationTool": lambda: ChestXRaySegmentationTool(device=self.device),
-                "LlavaMedTool": lambda: LlavaMedTool(cache_dir=self.model_dir, device=self.device, load_in_8bit=True),
-                "XRayVQATool": lambda: XRayVQATool(cache_dir=self.model_dir, device=self.device),
-                "ChestXRayReportGeneratorTool": lambda: ChestXRayReportGeneratorTool(
-                    cache_dir=self.model_dir, device=self.device
-                ),
-                "XRayPhraseGroundingTool": lambda: XRayPhraseGroundingTool(
-                    cache_dir=self.model_dir, temp_dir=self.temp_dir, load_in_8bit=True, device=self.device
-                ),
-                "ChestXRayGeneratorTool": lambda: ChestXRayGeneratorTool(
-                    model_path=f"{self.model_dir}/roentgen", temp_dir=self.temp_dir, device=self.device
-                ),
-                "ImageVisualizerTool": lambda: ImageVisualizerTool(),
-                "DicomProcessorTool": lambda: DicomProcessorTool(temp_dir=self.temp_dir),
-                "MedicalRAGTool": lambda: RAGTool(config=self.rag_config) if self.rag_config else None,
-                "WebBrowserTool": lambda: WebBrowserTool(),
-            }
-            
-            # Add PythonSandboxTool if available
-            try:
-                all_tools["PythonSandboxTool"] = lambda: create_python_sandbox()
-            except Exception as e:
-                print(f"Warning: PythonSandboxTool not available: {e}")
-            
-            # Initialize selected tools
-            self.tools_dict = {}
-            for tool_name in self.tools_to_use:
-                if tool_name in all_tools:
-                    try:
-                        tool_instance = all_tools[tool_name]()
-                        if tool_instance is not None:
-                            self.tools_dict[tool_name] = tool_instance
-                    except Exception as e:
-                        print(f"Warning: Failed to initialize {tool_name}: {e}")
-            
-            # Set up checkpointing
-            checkpointer = MemorySaver()
-            
-            # Create the language model
-            llm = ModelFactory.create_model(
-                model_name=self.model_name,
-                temperature=self.temperature,
-                top_p=self.top_p
-            )
-            
-            # Create the agent
-            self.agent = Agent(
-                llm,
-                tools=list(self.tools_dict.values()),
-                log_tools=False,  # Disable logging for benchmarking
-                system_prompt=prompt,
-                checkpointer=checkpointer,
-                debug=False,
+            print("Starting server...")
+
+            selected_tools = [
+                "ImageVisualizerTool",  # For displaying images in the UI
+                # "DicomProcessorTool",  # For processing DICOM medical image files
+                # "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
+                # "ChestXRayReportGeneratorTool",  # For generating medical reports from X-rays
+                # "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
+                "WebBrowserTool",  # For web browsing and search capabilities
+                "MedicalRAGTool",  # For retrieval-augmented generation with medical knowledge
+                "PythonSandboxTool",  # Add the Python sandbox tool
+            ]
+
+            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=7,
+                local_docs_dir="rag_docs",  # Change this to the path of the documents for RAG
+                huggingface_datasets=["VictorLJZ/medrax2"],  # List of HuggingFace datasets to load
+                dataset_split="train",  # Which split of the datasets to use
             )
-            
+
+            # Prepare any additional model-specific kwargs
+            model_kwargs = {}
+
             # Create temporary directory for this session
             self.session_temp_dir = Path(tempfile.mkdtemp(prefix="medrax_bench_"))
+
+            agent, tools_dict = initialize_agent(
+                prompt_file="medrax/docs/system_prompts.txt",
+                tools_to_use=selected_tools,
+                model_dir="/model-weights",
+                temp_dir=self.session_temp_dir,  # Change this to the path of the temporary directory
+                device="cuda",
+                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.7,
+                top_p=0.95,
+                model_kwargs=model_kwargs,
+                rag_config=rag_config,
+                debug=True,
+            )
+            
+            self.agent = agent
+            self.tools_dict = tools_dict
             
             print(f"MedRAX agent initialized with tools: {list(self.tools_dict.keys())}")
             

main.py

@@ -9,14 +9,12 @@ The system uses OpenAI's language models for reasoning and can be configured
 with different model weights, tools, and parameters.
 """
 
-import os
 import warnings
-from typing import Dict, List, Optional, Tuple, Any
+from typing import Dict, List, Optional, Any
 from dotenv import load_dotenv
 from transformers import logging
 
 from langgraph.checkpoint.memory import MemorySaver
-from langchain_openai import ChatOpenAI
 from medrax.models import ModelFactory
 
 from interface import create_demo
@@ -138,7 +136,7 @@ if __name__ == "__main__":
     # Example: initialize with only specific tools
     # Here three tools are commented out, you can uncomment them to use them
     selected_tools = [
-        # "ImageVisualizerTool",  # For displaying images in the UI
+        "ImageVisualizerTool",  # For displaying images in the UI
         # "DicomProcessorTool",  # For processing DICOM medical image files
         # "TorchXRayVisionClassifierTool",  # For classifying chest X-ray images using TorchXRayVision
         # "ArcPlusClassifierTool",  # For advanced chest X-ray classification using ArcPlus
@@ -172,14 +170,6 @@ if __name__ == "__main__":
     # Prepare any additional model-specific kwargs
     model_kwargs = {}
 
-    # Set up API keys for the web browser tool
-    # You'll need to set these environment variables:
-    # - GOOGLE_SEARCH_API_KEY: Your Google Custom Search API key
-    # - GOOGLE_SEARCH_ENGINE_ID: Your Google Custom Search Engine ID
-    # - COHERE_API_KEY: Your Cohere API key
-    # - OPENAI_API_KEY: Your OpenAI API key
-    # - PINECONE_API_KEY: Your Pinecone API key
-
     agent, tools_dict = initialize_agent(
         prompt_file="medrax/docs/system_prompts.txt",
         tools_to_use=selected_tools,