Update app.py

app.py

@@ -350,122 +350,100 @@ def get_current_time_in_timezone(timezone: str) -> str:
 
 final_answer = FinalAnswerTool()
 
-# Create Perplexity R1 model implementation directly without referencing an undefined variable
-
-# Import necessary modules (already imported above)
-# from huggingface_hub import InferenceClient
+# Remove the huggingface_api_key parameter - it's not supported
+model = HfApiModel(
+    max_tokens=2096,
+    temperature=0.5,
+    model_id='https://pflgm2locj2t89co.us-east-1.aws.endpoints.huggingface.cloud',  # Using the backup endpoint
+    custom_role_conversions=None
+)
 
-# Create a new model implementation that uses the larger context window model through InferenceClient
-class PerplexityR1Model:
-    def __init__(self, temperature=0.5, max_tokens=1500):
-        """Initialize Perplexity R1-1776 model with 128K context window."""
-        self.temperature = temperature
-        self.max_tokens = max_tokens
-        self.model_name = "perplexity-ai/r1-1776"
-        self.provider = "fireworks-ai"
-        self.last_input_token_count = 0
-        self.last_output_token_count = 0  # Added attribute for output tokens
-        # Get the API key
-        self.api_key = os.getenv("HF_API_KEY") or os.getenv("HUGGINGFACE_API_KEY")
-        if not self.api_key:
+# Add fallback logic that only activates if the primary model fails
+def try_model_call_with_fallbacks(prompt):
+    """Try to use the primary model first, fall back to alternatives if it fails."""
+    # First attempt with primary model
+    try:
+        return model(prompt)
+    except Exception as primary_error:
+        print(f"Primary model call failed: {str(primary_error)}")
+        print("Trying fallback models...")
+        
+        # List of fallback models
+        fallbacks = [
+            {
+                "provider": "sambanova",
+                "model_name": "Qwen/Qwen2.5-Coder-32B-Instruct",
+                "display_name": "Qwen 2.5 Coder 32B"
+            },
+            {
+                "provider": "hf-inference",
+                "model_name": "deepseek-ai/DeepSeek-R1-Distill-Qwen-32B",
+                "display_name": "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")
-        # Create the inference client
-        self.client = InferenceClient(provider=self.provider, api_key=self.api_key)
-        print("Initialized Perplexity R1-1776 model with 128K context window")
-    
-    def __call__(self, prompt):
-        """Call the model with the prompt."""
-        # Determine message format and count tokens
-        if isinstance(prompt, list):
-            # Convert each message's content to a string to avoid nested lists
-            combined_prompt = " ".join(str(msg.get("content", "")) for msg in prompt)
-            self.last_input_token_count = len(combined_prompt.split())
-            messages = prompt  # Already in message format
-        elif isinstance(prompt, str):
-            self.last_input_token_count = len(prompt.split())
-            messages = [{"role": "user", "content": prompt}]
-        else:
-            prompt_str = str(prompt)
-            self.last_input_token_count = len(prompt_str.split())
-            messages = [{"role": "user", "content": prompt_str}]
-            
-        print(f"Sending approximately {self.last_input_token_count} tokens to Perplexity R1-1776")
         
-        try:
-            completion = self.client.chat.completions.create(
-                model=self.model_name,
-                messages=messages,
-                temperature=self.temperature,
-                max_tokens=self.max_tokens
-            )
-            output = completion.choices[0].message.content
-            self.last_output_token_count = len(output.split())
-            return output
-        except Exception as e:
-            print(f"Error calling Perplexity R1-1776: {str(e)}")
-            # For context length errors, try simple truncation
-            if "context length" in str(e).lower() or "token limit" in str(e).lower():
-                print("Context length error with R1-1776 - truncating prompt and retrying")
-                if isinstance(prompt, str):
-                    truncated_prompt = prompt[-80000:] if len(prompt) > 80000 else prompt
-                    messages = [{"role": "user", "content": truncated_prompt}]
-                else:
-                    combined_prompt = " ".join(str(msg.get("content", "")) for msg in prompt)
-                    truncated_prompt = combined_prompt[-80000:] if len(combined_prompt) > 80000 else combined_prompt
-                    messages = [{"role": "user", "content": truncated_prompt}]
-                    
-                try:
-                    completion = self.client.chat.completions.create(
-                        model=self.model_name,
-                        messages=messages,
-                        temperature=self.temperature,
-                        max_tokens=self.max_tokens
-                    )
-                    output = completion.choices[0].message.content
-                    self.last_output_token_count = len(output.split())
-                    return output
-                except Exception as retry_error:
-                    print(f"Error on retry: {str(retry_error)}")
-                    return f"ERROR: Model call failed even with reduced context. Please try a shorter query."
-            else:
-                return f"ERROR: {str(e)}"
+        # 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
+        
+        # If all fallbacks fail, re-raise the original error
+        raise primary_error
 
-# Initialize our model with Perplexity R1-1776
-model = PerplexityR1Model(temperature=0.5, max_tokens=1500)
+# Monkey patch the model's __call__ method to use our fallback logic
+original_call = model.__call__
+model.__call__ = try_model_call_with_fallbacks
 
-# Import tool from Hub - do this before using the tool in the agent
+# Import tool from Hub
 image_generation_tool = load_tool("agents-course/text-to-image", trust_remote_code=True)
 
-# Load prompt templates before using them in the agent
 with open("prompts.yaml", 'r') as stream:
     prompt_templates = yaml.safe_load(stream)
-
-# Initialize the agent with all required components already defined
+    
+# Initialize the agent using standard smolagents patterns
 agent = CodeAgent(
     model=model,
     tools=[
         final_answer,
         Sonar_Web_Search_Tool,
-        primary_search_tool,
+        primary_search_tool,  # This is already set to either DuckDuckGo, Google, or fallback
         get_current_time_in_timezone,
         image_generation_tool,
         Dataset_Creator_Tool,
         Check_Dataset_Validity,
-        visit_webpage_tool,
+        visit_webpage_tool,  # This is correctly initialized as VisitWebpageTool()
     ],
-    max_steps=12,
+    max_steps=6,
     verbosity_level=1,
     grammar=None,
-    planning_interval=2,
+    planning_interval=3,
     name="Research Assistant",
-    description="""An AI assistant that can search the web, create datasets, and answer questions.
-                Using Perplexity R1-1776 model with 128K token context window.""",
+    description="""An AI assistant that can search the web, create datasets, and answer questions # Note about working within token limits
+                # When using with queries that might exceed token limits, consider:
+                # 1. Breaking tasks into smaller sub-tasks
+                # 2. Limiting the amount of data returned by search tools
+                # 3. Using the planning_interval to enable more effective reasoning""",
     prompt_templates=prompt_templates
 )
 
-# Add informative message about the model
-print("Using Perplexity R1-1776 model with 128K token context window")
-
 # Add informative message about which search tool is being used
 print(f"Agent initialized with {search_tool_name} as primary search tool")
 print(f"Available tools: final_answer, Sonar_Web_Search_Tool, {search_tool_name}, get_current_time_in_timezone, image_generation_tool, Dataset_Creator_Tool, Check_Dataset_Validity, visit_webpage_tool")