Update app.py

app.py

@@ -350,100 +350,181 @@ def get_current_time_in_timezone(timezone: str) -> str:
 
 final_answer = FinalAnswerTool()
 
-# Define the primary DeepSeek model using InferenceClient
-def primary_deepseek_model(prompt):
-    """Use the DeepSeek-R1-Distill-Qwen-32B model as the primary endpoint."""
-    try:
-        print("Using primary model: DeepSeek-R1-Distill-Qwen-32B via HF Inference API")
-        # Get API key
-        api_key = os.getenv("HF_API_KEY") or os.getenv("HUGGINGFACE_API_KEY")
-        if not api_key:
-            raise ValueError("No Hugging Face API key found in environment variables")
-            
-        client = InferenceClient(
-            provider="hf-inference",
-            api_key=api_key
-        )
+# Create a custom model class that implements the HfApiModel interface
+class FallbackChainModel:
+    """
+    A custom model class that implements the interface expected by smolagents,
+    with an automatic fallback chain from primary to backup models.
+    
+    Following smolagents best practices:
+    - Simple implementation with robust logging
+    - Thorough error handling with detailed messages
+    - Support for all parameters used by CodeAgent
+    """
+    def __init__(self):
+        # Initialize token tracking for compatibility with smolagents
+        self.last_input_token_count = 0
+        self.last_output_token_count = 0
         
-        messages = [{"role": "user", "content": prompt}]
-        completion = client.chat.completions.create(
-            model="deepseek-ai/DeepSeek-R1-Distill-Qwen-32B", 
-            messages=messages, 
+        # Set up the backup model as a proper HfApiModel instance
+        self.backup_model = HfApiModel(
             max_tokens=2096,
-            temperature=0.5
+            temperature=0.5,
+            model_id='https://pflgm2locj2t89co.us-east-1.aws.endpoints.huggingface.cloud',
         )
-        print("Primary model call successful")
-        return completion.choices[0].message.content
-    except Exception as e:
-        print(f"Primary DeepSeek model failed: {str(e)}")
-        raise e
-
-# Keep the original endpoint as a backup
-backup_model = HfApiModel(
-    max_tokens=2096,
-    temperature=0.5,
-    model_id='https://pflgm2locj2t89co.us-east-1.aws.endpoints.huggingface.cloud',
-    custom_role_conversions=None
-)
-
-# Updated fallback mechanism that tries primary first, then original model, then additional fallbacks
-def try_model_call_with_fallbacks(prompt):
-    """Try to use the primary model first, fall back to original model, then try other fallbacks if both fail."""
-    # First attempt with primary DeepSeek model
-    try:
-        return primary_deepseek_model(prompt)
-    except Exception as primary_error:
-        print(f"Primary model call failed: {str(primary_error)}")
-        print("Trying backup model...")
         
-        # Second attempt with original backup model
-        try:
-            print("Using backup model: HfApiModel endpoint")
-            result = backup_model(prompt)
-            print("Backup model call successful")
-            return result
-        except Exception as backup_error:
-            print(f"Backup model call failed: {str(backup_error)}")
-            print("Trying additional fallback models...")
+        print("Initialized FallbackChainModel with DeepSeek as primary and HfApiModel as backup")
+    
+    def __call__(self, prompt, stop_sequences=None, temperature=0.5, max_tokens=2096):
+        """
+        Call method that matches the HfApiModel interface required by smolagents.
+        
+        Args:
+            prompt: The input prompt (can be string or structured format)
+            stop_sequences: Optional list of sequences to stop generation
+            temperature: Controls randomness in generation
+            max_tokens: Maximum tokens to generate
             
-            # List of additional fallback models
-            fallbacks = [
-                {
-                    "provider": "sambanova",
-                    "model_name": "Qwen/Qwen2.5-Coder-32B-Instruct",
-                    "display_name": "Qwen 2.5 Coder 32B"
-                }
-                # DeepSeek is now the primary model, so we removed it from fallbacks
-            ]
+        Returns:
+            Generated text from the model
+        """
+        # Track all calls and parameters for debugging
+        print(f"Model call with temp={temperature}, max_tokens={max_tokens}")
+        if stop_sequences:
+            print(f"Using stop sequences: {stop_sequences}")
+        
+        # Try the primary model (DeepSeek)
+        try:
+            print("🚀 Attempting to use primary model: DeepSeek-R1-Distill-Qwen-32B")
             
             # Get API key
             api_key = os.getenv("HF_API_KEY") or os.getenv("HUGGINGFACE_API_KEY")
             if not api_key:
-                raise ValueError("No Hugging Face API key found in environment variables")
+                raise ValueError("No Hugging Face API key found")
             
-            # Try each fallback model in sequence
-            for fallback in fallbacks:
-                try:
-                    print(f"Trying fallback model: {fallback['display_name']}")
-                    client = InferenceClient(provider=fallback["provider"], api_key=api_key)
-                    messages = [{"role": "user", "content": prompt}]
-                    completion = client.chat.completions.create(
-                        model=fallback["model_name"],
-                        messages=messages,
-                        max_tokens=2096,
-                        temperature=0.5
-                    )
-                    print(f"Successfully used fallback model: {fallback['display_name']}")
-                    return completion.choices[0].message.content
-                except Exception as e:
-                    print(f"Fallback model {fallback['display_name']} failed: {str(e)}")
-                    continue
+            # Format the prompt appropriately for DeepSeek
+            # For structured prompts (needed by some agents), convert to proper format
+            if isinstance(prompt, list) and all(isinstance(item, dict) for item in prompt):
+                # Handle chat format with roles
+                messages = prompt
+                print(f"Using structured chat format with {len(messages)} messages")
+            else:
+                # Simple string prompt
+                if isinstance(prompt, (dict, list)):
+                    # If it's a dict or list but not in expected chat format, convert to string
+                    import json
+                    prompt_str = json.dumps(prompt)
+                    print("Converting complex prompt to JSON string")
+                else:
+                    prompt_str = str(prompt)
+                
+                # Create a single user message
+                messages = [{"role": "user", "content": prompt_str}]
+            
+            # Create the InferenceClient instance
+            client = InferenceClient(
+                provider="hf-inference",
+                api_key=api_key
+            )
+            
+            # Call the DeepSeek model
+            completion = client.chat.completions.create(
+                model="deepseek-ai/DeepSeek-R1-Distill-Qwen-32B", 
+                messages=messages, 
+                max_tokens=max_tokens,
+                temperature=temperature,
+                stop=stop_sequences  # Pass stop_sequences as stop parameter
+            )
+            
+            # Extract the response text
+            response = completion.choices[0].message.content
             
-            # If all fallbacks fail, re-raise the original error
-            raise primary_error
+            # Track token usage (approximate for now)
+            prompt_text = str(prompt)
+            self.last_input_token_count = len(prompt_text.split()) * 1.3  # Rough approximation
+            self.last_output_token_count = len(response.split()) * 1.3    # Rough approximation
+            
+            print("✅ Primary model call successful")
+            return response
+            
+        except Exception as primary_error:
+            # Detailed error logging for the primary model
+            print(f"❌ Primary model error: {str(primary_error)}")
+            
+            # Try the backup model (original endpoint)
+            try:
+                print("🔄 Falling back to backup model (HfApiModel endpoint)")
+                # Direct call to backup with all parameters
+                result = self.backup_model(
+                    prompt, 
+                    stop_sequences=stop_sequences,
+                    temperature=temperature,
+                    max_tokens=max_tokens
+                )
+                print("✅ Backup model call successful")
+                return result
+                
+            except Exception as backup_error:
+                # Log backup error and try final fallbacks
+                print(f"❌ Backup model error: {str(backup_error)}")
+                print("🔄 Trying additional fallback models...")
+                
+                # Additional fallback options (last resort)
+                fallbacks = [
+                    {
+                        "provider": "sambanova",
+                        "model_name": "Qwen/Qwen2.5-Coder-32B-Instruct",
+                        "display_name": "Qwen 2.5 Coder 32B"
+                    }
+                ]
+                
+                # Get API key again to ensure it's available
+                api_key = os.getenv("HF_API_KEY") or os.getenv("HUGGINGFACE_API_KEY")
+                if not api_key:
+                    raise ValueError("No Hugging Face API key found")
+                
+                # Try each fallback in sequence
+                for fallback in fallbacks:
+                    try:
+                        print(f"🔄 Trying fallback model: {fallback['display_name']}")
+                        client = InferenceClient(provider=fallback["provider"], api_key=api_key)
+                        
+                        # Format prompt for this model
+                        if isinstance(prompt, (dict, list)):
+                            import json
+                            prompt_str = json.dumps(prompt)
+                        else:
+                            prompt_str = str(prompt)
+                            
+                        messages = [{"role": "user", "content": prompt_str}]
+                        
+                        # Call the fallback model
+                        completion = client.chat.completions.create(
+                            model=fallback["model_name"],
+                            messages=messages,
+                            max_tokens=max_tokens,
+                            temperature=temperature,
+                            stop=stop_sequences
+                        )
+                        
+                        print(f"✅ Successfully used fallback model: {fallback['display_name']}")
+                        return completion.choices[0].message.content
+                        
+                    except Exception as e:
+                        print(f"❌ Fallback model {fallback['display_name']} failed: {str(e)}")
+                        continue
+                
+                # If all fallbacks failed, raise a comprehensive error
+                error_msg = (
+                    f"All models in fallback chain failed. "
+                    f"Primary error: {str(primary_error)}. "
+                    f"Backup error: {str(backup_error)}"
+                )
+                print(f"❌ FATAL ERROR: {error_msg}")
+                raise RuntimeError(error_msg)
 
-# Create a model wrapper that uses our fallback logic
-model = lambda prompt: try_model_call_with_fallbacks(prompt)
+# Create an instance of our fallback chain model
+model = FallbackChainModel()
 
 # Import tool from Hub
 image_generation_tool = load_tool("agents-course/text-to-image", trust_remote_code=True)