feat: initial web_platform offering

.gitignore

@@ -14,7 +14,6 @@ dist/
 downloads/
 eggs/
 .eggs/
-lib/
 lib64/
 parts/
 sdist/
@@ -61,6 +60,12 @@ local_settings.py
 db.sqlite3
 db.sqlite3-journal
 
+# SQLite database files (including WAL mode files)
+*.db
+*.db-wal
+*.db-shm
+*.db-journal
+
 # Flask stuff:
 instance/
 .webassets-cache

benchmarking/system_prompts.txt

@@ -5,6 +5,12 @@ You can make multiple tool calls in parallel or in sequence as needed for compre
 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.
 
+CRITICAL - FILE PATH HANDLING:
+- When you receive image paths in XML tags like <image_path id="1">path/to/file.jpg</image_path>, you MUST copy the EXACT path from between the tags
+- Do NOT type paths from memory, abbreviate, or modify the UUIDs in any way
+- Copy the ENTIRE path character-for-character including all hyphens and alphanumeric sequences
+- A single character error in a UUID will cause tool calls to fail
+
 CITATION REQUIREMENTS:
 - When referencing information from RAG and/or web search tools, ALWAYS include numbered citations [1], [2], [3], etc.
 - Use citations immediately after making claims or statements based on the above tool results.

medrax/agent/agent.py

@@ -115,5 +115,10 @@ def has_tool_calls(self, state: AgentState) -> bool:
         Returns:
             bool: True if tool calls exist, False otherwise.
         """
-        response = state["messages"][-1]
-        return len(response.tool_calls) > 0
+        messages = state["messages"]
+        if not messages:
+            return False
+
+        response = messages[-1]
+        # Safely check for tool_calls attribute
+        return hasattr(response, 'tool_calls') and len(response.tool_calls) > 0

medrax/docs/system_prompts.txt

@@ -5,6 +5,12 @@ You can make multiple tool calls in parallel or in sequence as needed for compre
 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.
 
+CRITICAL - FILE PATH HANDLING:
+- When you receive image paths in XML tags like <image_path id="1">path/to/file.jpg</image_path>, you MUST copy the EXACT path from between the tags
+- Do NOT type paths from memory, abbreviate, or modify the UUIDs in any way
+- Copy the ENTIRE path character-for-character including all hyphens and alphanumeric sequences
+- A single character error in a UUID will cause tool calls to fail
+
 CITATION REQUIREMENTS:
 - When referencing information from RAG and/or web search tools, ALWAYS include numbered citations [1], [2], [3], etc.
 - Use citations immediately after making claims or statements based on the above tool results.

medrax/llava/model/builder.py

@@ -1,3 +1,4 @@
+from typing import Optional
 from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig
 import torch
 from medrax.llava.model import LlavaMistralForCausalLM
@@ -15,17 +16,20 @@ def load_pretrained_model(
     load_in_8bit=False,
     load_in_4bit=True,
     device="cuda",
-    cache_dir: str = "/model-weights",
+    cache_dir: Optional[str] = None,
     low_cpu_mem_usage=True,
     torch_dtype=torch.bfloat16,
 ):
 
     kwargs = {}
 
-    if device != "cuda":
+    # Device mapping strategy:
+    # - CUDA: Use "auto" to let HuggingFace/Accelerate handle placement, avoiding meta-tensor issues
+    # - CPU/MPS: Use explicit device mapping for direct control
+    if device == "cuda":
+        kwargs["device_map"] = "auto"
+    else:
         kwargs["device_map"] = {"": device}
-    # else:
-    #     kwargs["device_map"] = "auto"
 
     if load_in_8bit:
         kwargs["load_in_8bit"] = True
@@ -44,12 +48,16 @@ def load_pretrained_model(
         # Load LLaVA model
         if "mistral" in model_name.lower():
             tokenizer = AutoTokenizer.from_pretrained(model_path, cache_dir=cache_dir)
+            # LLaVA-Med specific: Override to use explicit "cuda" mapping
+            # to ensure proper device placement for vision-language model
+            if device == "cuda":
+                kwargs["device_map"] = "cuda"
             model = LlavaMistralForCausalLM.from_pretrained(
                 model_path,
-                low_cpu_mem_usage=low_cpu_mem_usage,
-                use_flash_attention_2=False,
+                # Disable low_cpu_mem_usage to avoid meta tensor errors during model loading
+                low_cpu_mem_usage=False,
                 cache_dir=cache_dir,
-                torch_dtype=torch_dtype,
+                attn_implementation="eager",
                 **kwargs,
             )
 
@@ -110,11 +118,12 @@ def load_pretrained_model(
         if not vision_tower.is_loaded:
             vision_tower.load_model()
 
-        vision_tower.to(device=device, dtype=torch_dtype)
-        model.model.mm_projector.to(device=device, dtype=torch_dtype)
-
-        if not (load_in_4bit or load_in_8bit):
-            model.to(device=device, dtype=torch_dtype)
+        # Skip manual device placement when using auto device mapping to prevent conflicts
+        if kwargs.get("device_map") != "auto":
+            vision_tower.to(device=device, dtype=torch_dtype)
+            model.model.mm_projector.to(device=device, dtype=torch_dtype)
+            if not (load_in_4bit or load_in_8bit):
+                model.to(device=device, dtype=torch_dtype)
 
         image_processor = vision_tower.image_processor
 

medrax/llava/model/llava_arch.py

@@ -97,19 +97,28 @@ def initialize_vision_modules(self, model_args, fsdp=None, embed_tokens=None):
             def get_w(weights, keyword):
                 return {k.split(keyword + ".")[1]: v for k, v in weights.items() if keyword in k}
 
-            mm_projector_weights = torch.load(pretrain_mm_mlp_adapter, map_location="cpu")
-            self.mm_projector.load_state_dict(get_w(mm_projector_weights, "mm_projector"))
+            try:
+                mm_projector_weights = torch.load(pretrain_mm_mlp_adapter, map_location="cpu")
+                self.mm_projector.load_state_dict(get_w(mm_projector_weights, "mm_projector"))
+            except Exception as e:
+                print(
+                    f"Warning: failed to load mm_projector weights from '{pretrain_mm_mlp_adapter}' due to: {e}. Continuing without these optional weights."
+                )
 
             # also load additional learnable parameters during feature alignment
             checkpoint_folder = os.path.dirname(pretrain_mm_mlp_adapter)
             ckpts = glob(f"{checkpoint_folder}/checkpoint-*", recursive=False)
             if len(ckpts) > 0:
-                vision_module_weights = torch.load(f"{ckpts[-1]}/mm_projector.bin", map_location="cpu")
-                model_dict = get_w(vision_module_weights, "vision_tower")
-                print(f"Loading vision module weights from {ckpts[-1]}/mm_projector.bin")
-                # print keys in model_dict
-                print(f"Loaded keys: {model_dict.keys()}")
-                self.vision_tower.load_state_dict(model_dict, strict=False)
+                try:
+                    vision_module_weights = torch.load(f"{ckpts[-1]}/mm_projector.bin", map_location="cpu")
+                    model_dict = get_w(vision_module_weights, "vision_tower")
+                    print(f"Loading vision module weights from {ckpts[-1]}/mm_projector.bin")
+                    print(f"Loaded keys: {model_dict.keys()}")
+                    self.vision_tower.load_state_dict(model_dict, strict=False)
+                except Exception as e:
+                    print(
+                        f"Warning: failed to load optional vision module weights from '{ckpts[-1]}/mm_projector.bin' due to: {e}. Skipping."
+                    )
 
 
 class LlavaMetaForCausalLM(ABC):
@@ -351,17 +360,26 @@ def initialize_vision_tokenizer(self, model_args, tokenizer):
                     p.requires_grad = False
 
             if model_args.pretrain_mm_mlp_adapter:
-                mm_projector_weights = torch.load(model_args.pretrain_mm_mlp_adapter, map_location="cpu")
-                embed_tokens_weight = mm_projector_weights["model.embed_tokens.weight"]
-                assert num_new_tokens == 2
-                if input_embeddings.shape == embed_tokens_weight.shape:
-                    input_embeddings[-num_new_tokens:] = embed_tokens_weight[-num_new_tokens:]
-                elif embed_tokens_weight.shape[0] == num_new_tokens:
-                    input_embeddings[-num_new_tokens:] = embed_tokens_weight
-                else:
-                    raise ValueError(
-                        f"Unexpected embed_tokens_weight shape. Pretrained: {embed_tokens_weight.shape}. Current: {input_embeddings.shape}. Numer of new tokens: {num_new_tokens}."
+                embed_tokens_weight = None
+                try:
+                    mm_projector_weights = torch.load(model_args.pretrain_mm_mlp_adapter, map_location="cpu")
+                    embed_tokens_weight = mm_projector_weights["model.embed_tokens.weight"]
+                except Exception as e:
+                    print(
+                        f"Warning: failed to load embed token weights from '{model_args.pretrain_mm_mlp_adapter}' due to: {e}. "
+                        "Continuing with default-initialized tokens."
                     )
+
+                if embed_tokens_weight is not None:
+                    assert num_new_tokens == 2
+                    if input_embeddings.shape == embed_tokens_weight.shape:
+                        input_embeddings[-num_new_tokens:] = embed_tokens_weight[-num_new_tokens:]
+                    elif embed_tokens_weight.shape[0] == num_new_tokens:
+                        input_embeddings[-num_new_tokens:] = embed_tokens_weight
+                    else:
+                        raise ValueError(
+                            f"Unexpected embed_tokens_weight shape. Pretrained: {embed_tokens_weight.shape}. Current: {input_embeddings.shape}. Numer of new tokens: {num_new_tokens}."
+                        )
         elif model_args.mm_use_im_patch_token:
             if model_args.tune_mm_mlp_adapter:
                 for p in self.get_input_embeddings().parameters():

medrax/tools/__init__.py

@@ -5,7 +5,10 @@
 from .segmentation import *
 from .vqa import *
 from .grounding import *
-from .xray_generation import *
+# Avoid importing heavy generation pipeline at package import time to prevent
+# unnecessary diffusers/triton dependencies during unrelated tool loads.
+# Import `ChestXRayGeneratorTool` directly from `medrax.tools.xray_generation`
+# where needed instead of exposing it here.
 from .dicom import *
 from .utils import *
 from .rag import *

medrax/tools/browsing/duckduckgo.py

@@ -18,7 +18,7 @@
     CallbackManagerForToolRun,
 )
 from langchain_core.tools import BaseTool
-from pydantic import BaseModel, Field
+from pydantic import BaseModel, Field, ConfigDict
 
 try:
     from duckduckgo_search import DDGS
@@ -115,30 +115,32 @@ def _perform_search_sync(self, query: str, max_results: int = 5, region: str = "
         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)
+            ddgs = DDGS()
+
+            search_results = list(ddgs.text(
+                keywords=query,
+                region=region,
+                safesearch="moderate",
+                max_results=max_results
+            ))
+
+            if search_results is None:
+                logger.warning("DuckDuckGo returned None instead of results list")
+                search_results = []
+
+            formatted_results = []
+            for i, result in enumerate(search_results, 1):
+                if result is None or not isinstance(result, dict):
+                    continue
+
+                formatted_result = {
+                    "rank": i,
+                    "title": result.get("title", result.get("t", "No title")),
+                    "url": result.get("href", result.get("l", "No URL")),
+                    "snippet": result.get("body", result.get("a", "No description available")),
+                    "source": "DuckDuckGo",
+                }
+                formatted_results.append(formatted_result)
 
                 # Create summary for the agent
                 if formatted_results:
@@ -149,7 +151,6 @@ def _perform_search_sync(self, query: str, max_results: int = 5, region: str = "
                 else:
                     summary = f"No results found for '{query}'"
 
-                # Log successful completion
                 logger.info(f"Web search completed successfully: {len(formatted_results)} results")
 
                 return {
@@ -191,9 +192,7 @@ def _run(
             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
+            Tuple[Dict[str, Any], Dict[str, Any]]: Output dictionary and metadata dictionary
         """
         # Create metadata structure
         metadata = {
@@ -249,9 +248,7 @@ async def _arun(
             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
+            Tuple[Dict[str, Any], Dict[str, Any]]: Output dictionary and metadata dictionary
         """
         # Try to get LangGraph stream writer for progress updates
         writer = None
@@ -284,12 +281,10 @@ async def _arun(
                     }
                 )
 
-            # 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(
                     {

medrax/tools/browsing/web_browser.py

@@ -13,7 +13,7 @@
 import requests
 from bs4 import BeautifulSoup
 from langchain_core.tools import BaseTool
-from pydantic import BaseModel, Field
+from pydantic import BaseModel, Field, ConfigDict
 
 
 class WebBrowserSchema(BaseModel):
@@ -77,6 +77,7 @@ class WebBrowserTool(BaseTool):
     )
     max_results: int = 5
     args_schema: Type[BaseModel] = WebBrowserSchema
+    model_config = ConfigDict(arbitrary_types_allowed=True, protected_namespaces=())
 
     def __init__(self, search_api_key: Optional[str] = None, search_engine_id: Optional[str] = None, **kwargs):
         """Initialize the web browser tool with optional search API credentials.