updates

benchmarking/cli.py

@@ -139,8 +139,8 @@ def main():
                            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=1000, 
-                           help="Maximum tokens per model response (default: 1000)")
+    run_parser.add_argument("--max-tokens", type=int, default=5000, 
+                           help="Maximum tokens per model response (default: 5000)")
     
     run_parser.set_defaults(func=run_benchmark_command)
     

benchmarking/llm_providers/medrax_provider.py

@@ -33,19 +33,19 @@ class MedRAXProvider(LLMProvider):
             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
+                "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
+                "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
+                "PythonSandboxTool",  # Add the Python sandbox tool
             ]
 
             rag_config = RAGConfig(
@@ -68,7 +68,7 @@ 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="model-weights",
                 temp_dir="temp",  # Change this to the path of the temporary directory
                 device="cpu",
                 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

medrax/agent/agent.py

@@ -1,12 +1,8 @@
-import json
-import operator
 from pathlib import Path
 from dotenv import load_dotenv
-from datetime import datetime
-from typing import List, Dict, Any, TypedDict, Annotated, Optional
+from typing import List, Any, TypedDict, Optional
 
 from langgraph.prebuilt import create_react_agent
-from langchain_core.messages import AnyMessage
 from langgraph.prebuilt.chat_agent_executor import AgentState
 from langchain_core.language_models import BaseLanguageModel
 from langchain_core.tools import BaseTool

medrax/tools/rag.py

@@ -48,14 +48,14 @@ class RAGTool(BaseTool):
         self.rag = CohereRAG(config)
         self.chain = self.rag.initialize_rag(with_memory=True)
 
-    def _run(self, query: str) -> Tuple[Dict[str, Any], Dict[str, Any]]:
+    def _run(self, query: str) -> Tuple[Dict[str, Any], Dict]:
         """Execute the RAG tool with the given query.
 
         Args:
             query (str): Medical question to answer
 
         Returns:
-            Tuple[Dict[str, Any], Dict[str, Any]]: Output dictionary and metadata dictionary
+            Tuple[Dict[str, Any], Dict]: Output dictionary and metadata dictionary
         """
         try:
             result = self.chain.invoke({"query": query})
@@ -87,14 +87,14 @@ class RAGTool(BaseTool):
             }
             return output, metadata
 
-    async def _arun(self, query: str) -> Tuple[Dict[str, Any], Dict[str, Any]]:
+    async def _arun(self, query: str) -> Tuple[Dict[str, Any], Dict]:
         """Async version of _run.
 
         Args:
             query (str): Medical question to answer
 
         Returns:
-            Tuple[Dict[str, Any], Dict[str, Any]]: Output dictionary and metadata dictionary
+            Tuple[Dict[str, Any], Dict]: Output dictionary and metadata dictionary
 
         Raises:
             NotImplementedError: Async not implemented yet