Create shared_models.py

shared_models.py

@@ -0,0 +1,285 @@
+# shared_models.py
+"""
+Shared model manager for Mimir agents.
+Uses Llama-3.2-3B-Instruct with transformers for all agents.
+"""
+import torch
+import threading
+import logging
+import os
+from typing import Optional, List
+from transformers import pipeline, AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig
+
+logger = logging.getLogger(__name__)
+
+# ZeroGPU support
+try:
+    import spaces
+    HF_SPACES_AVAILABLE = True
+    logger.info("✅ ZeroGPU (spaces) available")
+except ImportError:
+    HF_SPACES_AVAILABLE = False
+    class DummySpaces:
+        @staticmethod
+        def GPU(duration=90):
+            def decorator(func):
+                return func
+            return decorator
+    spaces = DummySpaces()
+    logger.warning("⚠️ ZeroGPU not available - running without GPU allocation")
+
+HF_TOKEN = os.getenv("HF_TOKEN") or os.getenv("HUGGINGFACEHUB_API_TOKEN")
+
+# Model configuration
+LLAMA_MODEL_ID = "meta-llama/Llama-3.2-3B-Instruct"
+
+
+class LlamaSharedAgent:
+    """
+    Singleton agent using Llama-3.2-3B-Instruct for all Mimir operations.
+    Thread-safe with ZeroGPU allocation management.
+    
+    Used by:
+    - ToolDecisionAgent
+    - PromptRoutingAgents (all 4 agents)
+    - ThinkingAgents (all reasoning agents)
+    - ResponseAgent
+    """
+    
+    _instance = None
+    _lock = threading.Lock()
+    
+    def __new__(cls):
+        """Ensure only one instance exists (singleton pattern)"""
+        if cls._instance is None:
+            with cls._lock:
+                if cls._instance is None:
+                    cls._instance = super().__new__(cls)
+        return cls._instance
+    
+    def __init__(self):
+        """Initialize only once"""
+        if hasattr(self, '_initialized'):
+            return
+        
+        self.pipe = None
+        self.tokenizer = None
+        self.model = None
+        self.model_loaded = False
+        self._initialized = True
+        logger.info("LlamaSharedAgent instance created (singleton)")
+    
+    @spaces.GPU(duration=120)
+    def _ensure_loaded(self):
+        """
+        Load model with GPU allocation (ZeroGPU).
+        Only ONE @spaces.GPU decorator for Llama across entire app!
+        """
+        if self.model_loaded:
+            logger.info("✅ Llama-3.2-3B already loaded, reusing existing instance")
+            return
+        
+        logger.info("="*60)
+        logger.info("LOADING SHARED LLAMA-3.2-3B-INSTRUCT")
+        logger.info("="*60)
+        
+        try:
+            # 4-bit quantization config for memory efficiency
+            quantization_config = BitsAndBytesConfig(
+                load_in_4bit=True,
+                bnb_4bit_quant_type="nf4",
+                bnb_4bit_compute_dtype=torch.bfloat16,
+                bnb_4bit_use_double_quant=True,
+            )
+            
+            logger.info(f"Loading model: {LLAMA_MODEL_ID}")
+            logger.info("Configuration: 4-bit NF4 quantization")
+            
+            # Load tokenizer
+            self.tokenizer = AutoTokenizer.from_pretrained(
+                LLAMA_MODEL_ID,
+                token=HF_TOKEN,
+                trust_remote_code=True,
+            )
+            
+            # Load model with quantization
+            self.model = AutoModelForCausalLM.from_pretrained(
+                LLAMA_MODEL_ID,
+                quantization_config=quantization_config,
+                device_map="auto",
+                token=HF_TOKEN,
+                trust_remote_code=True,
+                torch_dtype=torch.bfloat16,
+            )
+            
+            # Create pipeline
+            self.pipe = pipeline(
+                "text-generation",
+                model=self.model,
+                tokenizer=self.tokenizer,
+                torch_dtype=torch.bfloat16,
+                device_map="auto",
+            )
+            
+            self.model_loaded = True
+            
+            logger.info("="*60)
+            logger.info("✅ SHARED LLAMA-3.2-3B LOADED SUCCESSFULLY")
+            logger.info(f"  Model: {LLAMA_MODEL_ID}")
+            logger.info(f"  Quantization: 4-bit NF4")
+            logger.info(f"  Memory: ~1GB VRAM (vs 3.3GB GGUF)")
+            logger.info("  Context: 128K tokens")
+            logger.info("  This model will be reused by:")
+            logger.info("  - ToolDecisionAgent")
+            logger.info("  - PromptRoutingAgents (all 4 agents)")
+            logger.info("  - ThinkingAgents (all reasoning)")
+            logger.info("  - ResponseAgent (final responses)")
+            logger.info("="*60)
+            
+        except Exception as e:
+            logger.error(f"Failed to load Llama-3.2-3B: {e}")
+            raise
+    
+    def generate(
+        self, 
+        system_prompt: str,
+        user_message: str,
+        max_tokens: int = 100,
+        temperature: float = 0.7,
+        stop_sequences: Optional[List[str]] = None
+    ) -> str:
+        """
+        Generate response using shared Llama-3.2-3B model.
+        
+        Args:
+            system_prompt: System instruction
+            user_message: User query
+            max_tokens: Max tokens to generate
+            temperature: Sampling temperature
+            stop_sequences: Optional list of stop sequences (not used with pipeline)
+            
+        Returns:
+            Generated text
+        """
+        # Ensure model is loaded (triggers @spaces.GPU only once)
+        self._ensure_loaded()
+        
+        # Format messages using Llama 3.2 chat template (handled automatically)
+        messages = [
+            {"role": "system", "content": system_prompt},
+            {"role": "user", "content": user_message},
+        ]
+        
+        try:
+            # Generate using pipeline
+            outputs = self.pipe(
+                messages,
+                max_new_tokens=max_tokens,
+                temperature=temperature,
+                do_sample=True,
+                top_p=0.9,
+                top_k=40,
+                repetition_penalty=1.15,
+            )
+            
+            # Extract generated text (pipeline returns full conversation)
+            result = outputs[0]["generated_text"][-1]["content"]
+            
+            return result.strip()
+            
+        except Exception as e:
+            logger.error(f"Generation error: {e}")
+            return ""
+    
+    def generate_streaming(
+        self,
+        system_prompt: str,
+        user_message: str,
+        max_tokens: int = 512,
+        temperature: float = 0.7,
+    ):
+        """
+        Generate response with streaming (for ResponseAgent).
+        
+        Yields:
+            str: Generated text chunks
+        """
+        self._ensure_loaded()
+        
+        messages = [
+            {"role": "system", "content": system_prompt},
+            {"role": "user", "content": user_message},
+        ]
+        
+        try:
+            # Use TextIteratorStreamer for streaming
+            from transformers import TextIteratorStreamer
+            from threading import Thread
+            
+            # Apply chat template
+            input_ids = self.tokenizer.apply_chat_template(
+                messages,
+                add_generation_prompt=True,
+                return_tensors="pt"
+            ).to(self.model.device)
+            
+            streamer = TextIteratorStreamer(
+                self.tokenizer,
+                skip_prompt=True,
+                skip_special_tokens=True
+            )
+            
+            generation_kwargs = dict(
+                input_ids=input_ids,
+                streamer=streamer,
+                max_new_tokens=max_tokens,
+                temperature=temperature,
+                do_sample=True,
+                top_p=0.9,
+                top_k=40,
+                repetition_penalty=1.15,
+            )
+            
+            # Generate in separate thread
+            thread = Thread(target=self.model.generate, kwargs=generation_kwargs)
+            thread.start()
+            
+            # Yield generated text
+            for text in streamer:
+                yield text
+                
+        except Exception as e:
+            logger.error(f"Streaming generation error: {e}")
+            yield ""
+    
+    def get_model_info(self) -> dict:
+        """Get model information for diagnostics"""
+        return {
+            "status": "loaded" if self.model_loaded else "not_loaded",
+            "model_id": LLAMA_MODEL_ID,
+            "model_type": "llama-3.2-3b-instruct",
+            "quantization": "4-bit NF4",
+            "size_gb": 1.0,
+            "context_length": 128000,
+            "zerogpu_ready": True,
+            "transformers_pipeline": True,
+            "shared_instance": True,
+        }
+
+
+# Global singleton instance
+_shared_llama = None
+
+def get_shared_llama() -> LlamaSharedAgent:
+    """Get or create the shared Llama-3.2-3B agent instance"""
+    global _shared_llama
+    if _shared_llama is None:
+        _shared_llama = LlamaSharedAgent()
+    return _shared_llama
+
+
+# Backwards compatibility aliases
+Qwen3SharedAgent = LlamaSharedAgent
+MistralSharedAgent = LlamaSharedAgent
+get_shared_qwen3 = get_shared_llama
+get_shared_mistral = get_shared_llama