Upload agents/priority_router.py with huggingface_hub

agents/priority_router.py

@@ -8,13 +8,12 @@ from typing import Any, Dict, Optional, List
 
 from .base_agent import BaseAgent, AgentResult
 
-# Try to import torch and transformers
+# Import the unified MedGemma engine
 try:
-    import torch
-    from transformers import AutoModelForCausalLM, AutoTokenizer
-    TORCH_AVAILABLE = True
+    from .medgemma_engine import get_engine, MedGemmaEngine
+    ENGINE_AVAILABLE = True
 except ImportError:
-    TORCH_AVAILABLE = False
+    ENGINE_AVAILABLE = False
 
 
 class PriorityRouterAgent(BaseAgent):
@@ -22,7 +21,7 @@ class PriorityRouterAgent(BaseAgent):
     Agent 4: MedGemma Priority Router
     
     Assesses case urgency and determines appropriate routing
-    based on radiology report and findings.
+    based on radiology report and findings using MedGemma.
     """
     
     # Priority level definitions
@@ -68,39 +67,20 @@ class PriorityRouterAgent(BaseAgent):
             model_name="google/medgemma-4b-it"
         )
         self.demo_mode = demo_mode
-        self.tokenizer = None
+        self.engine = None
     
     def load_model(self) -> bool:
-        """Load MedGemma model for priority assessment."""
-        if self.demo_mode:
-            self.is_loaded = True
-            return True
-        
-        if not TORCH_AVAILABLE:
-            self.demo_mode = True
+        """Load MedGemma model via unified engine."""
+        if self.demo_mode or not ENGINE_AVAILABLE:
             self.is_loaded = True
             return True
         
         try:
-            self.tokenizer = AutoTokenizer.from_pretrained(
-                self.model_name,
-                trust_remote_code=True
-            )
-            
-            self.model = AutoModelForCausalLM.from_pretrained(
-                self.model_name,
-                trust_remote_code=True,
-                torch_dtype=torch.bfloat16 if torch.cuda.is_available() else torch.float32,
-                device_map="auto" if torch.cuda.is_available() else None,
-                low_cpu_mem_usage=True
-            )
-            
-            self.model.eval()
-            self.is_loaded = True
+            self.engine = get_engine(force_demo=self.demo_mode)
+            self.is_loaded = self.engine.is_loaded
             return True
-            
         except Exception as e:
-            print(f"Failed to load model for priority routing: {e}")
+            print(f"Failed to load MedGemma engine: {e}")
             self.demo_mode = True
             self.is_loaded = True
             return True
@@ -135,13 +115,13 @@ class PriorityRouterAgent(BaseAgent):
         # Get original findings if passed through context
         original_findings = context.get("original_findings", []) if context else []
         
-        # Process based on mode
-        if self.demo_mode:
-            routing = self._simulate_priority_assessment(
+        # Process - always try to use real model if available
+        if self.engine and self.engine.is_loaded and self.engine.backend != "demo":
+            routing = self._run_model_inference(
                 report_sections, full_report, findings_count, original_findings, context
             )
         else:
-            routing = self._run_model_inference(
+            routing = self._simulate_priority_assessment(
                 report_sections, full_report, findings_count, original_findings, context
             )
         
@@ -162,25 +142,12 @@ class PriorityRouterAgent(BaseAgent):
         original_findings: List[Dict],
         context: Optional[Dict]
     ) -> Dict:
-        """Use MedGemma to assess priority."""
+        """Use MedGemma to assess priority via unified engine."""
         try:
             prompt = self._build_priority_prompt(full_report, original_findings)
             
-            inputs = self.tokenizer(prompt, return_tensors="pt")
-            if torch.cuda.is_available():
-                inputs = {k: v.cuda() for k, v in inputs.items()}
-            
-            with torch.no_grad():
-                outputs = self.model.generate(
-                    **inputs,
-                    max_new_tokens=512,
-                    temperature=0.2,  # Lower temperature for more deterministic output
-                    top_p=0.9,
-                    do_sample=True,
-                    pad_token_id=self.tokenizer.eos_token_id
-                )
-            
-            response = self.tokenizer.decode(outputs[0], skip_special_tokens=True)
+            # Use the unified engine to assess priority
+            response = self.engine.generate(prompt, max_tokens=256)
             
             return self._parse_priority_response(response, original_findings)