Fix: Add conditional import for MLX with CPU fallback

src/generation/mlx_wrapper.py

@@ -1,29 +1,52 @@
-# src/generation/mlx_wrapper.py
 import os
 from typing import Any, List, Optional
 from langchain_core.callbacks.manager import CallbackManagerForLLMRun
 from langchain_core.language_models.llms import LLM
-from mlx_lm import load, generate
 from dotenv import load_dotenv
 
 load_dotenv()
 
+# --- CRITICAL FIX: Handle Import Error ---
+try:
+    from mlx_lm import load, generate
+    HAS_MLX = True
+except ImportError:
+    HAS_MLX = False
+# ----------------------------------------
+
 class MLXLLM(LLM):
-    """Custom LangChain Wrapper for MLX Models"""
+    """Custom LangChain Wrapper for MLX Models (with Cloud Fallback)"""
     
     model_id: str = os.getenv("MODEL_ID", "mlx-community/Llama-3.2-3B-Instruct-4bit")
     model: Any = None
     tokenizer: Any = None
     max_tokens: int = int(os.getenv("MAX_TOKENS", 512))
+    pipeline: Any = None # For Cloud Fallback
 
     def __init__(self, **kwargs):
         super().__init__(**kwargs)
-        print(f"🚀 Loading MLX Model: {self.model_id}")
-        self.model, self.tokenizer = load(self.model_id)
+        
+        if HAS_MLX:
+            print(f"🚀 Loading MLX Model: {self.model_id}")
+            self.model, self.tokenizer = load(self.model_id)
+        else:
+            print(f"⚠️ MLX not found. Falling back to HuggingFace Transformers (CPU/Cloud).")
+            # Fallback: Use standard Transformers
+            from transformers import pipeline, AutoTokenizer, AutoModelForCausalLM
+            
+            # Use the MODEL_ID env var (set to 'gpt2' or 'facebook/opt-125m' in HF Secrets)
+            # Do NOT use the MLX model ID here, as it requires MLX format.
+            cloud_model_id = os.getenv("MODEL_ID", "gpt2") 
+            
+            self.pipeline = pipeline(
+                "text-generation", 
+                model=cloud_model_id, 
+                max_new_tokens=self.max_tokens
+            )
 
     @property
     def _llm_type(self) -> str:
-        return "mlx_llama"
+        return "mlx_llama" if HAS_MLX else "transformers_fallback"
 
     def _call(
         self,
@@ -35,16 +58,23 @@ class MLXLLM(LLM):
         if stop is not None:
             raise ValueError("stop kwargs are not permitted.")
 
-        messages = [{"role": "user", "content": prompt}]
-        formatted_prompt = self.tokenizer.apply_chat_template(
-            messages, tokenize=False, add_generation_prompt=True
-        )
-
-        response = generate(
-            self.model, 
-            self.tokenizer, 
-            prompt=formatted_prompt, 
-            verbose=False, 
-            max_tokens=self.max_tokens
-        )
-        return response
+        if HAS_MLX:
+            # MLX Generation Logic
+            messages = [{"role": "user", "content": prompt}]
+            formatted_prompt = self.tokenizer.apply_chat_template(
+                messages, tokenize=False, add_generation_prompt=True
+            )
+            response = generate(
+                self.model, 
+                self.tokenizer, 
+                prompt=formatted_prompt, 
+                verbose=False, 
+                max_tokens=self.max_tokens
+            )
+            return response
+        else:
+            # Cloud/CPU Fallback Logic
+            # Simple text generation for MVP
+            response = self.pipeline(prompt)[0]['generated_text']
+            # Remove the prompt from the response if needed
+            return response[len(prompt):]