Update app.py

app.py

@@ -46,15 +46,39 @@ def generate_image(prompt:str) -> list:
     return image_generation_tool(prompt)    
 
 @tool
-def duckduckgo_search(query: str, max_results: int=5) -> str:
-    """Search DuckDuckGo for a query and return the top results.
+def duckduckgo_search(query: str, max_results: int = 5) -> str:
+    """Search DuckDuckGo for a query and return the top N results."""
+    searcher = DuckDuckGoSearchTool(max_results=max_results)
+    return searcher(query) 
+
+
+@tool
+def is_prime(number: int) -> bool:
+    """
+    Check if a number is prime using an optimized 6k±1 algorithm.
+
     Args:
-        query: the search terms
-        max_results: how many results to fetch
+        number: The integer to check for primality.
+
+    Returns:
+        True if `number` is prime, False otherwise.
     """
-    searcher = DuckDuckGoSearchTool()
-    results = searcher(query, max_results=max_results)
-    return results    
+    # Numbers less than 2 are not prime
+    if number <= 1:
+        return False
+    # 2 and 3 are prime
+    if number <= 3:
+        return True
+    # Eliminate multiples of 2 and 3
+    if number % 2 == 0 or number % 3 == 0:
+        return False
+    # Test 6k ± 1 factors up to sqrt(number)
+    i = 5
+    while i * i <= number:
+        if number % i == 0 or number % (i + 2) == 0:
+            return False
+        i += 6
+    return True
 
 
 
@@ -78,7 +102,8 @@ with open("prompts.yaml", 'r') as stream:
     
 agent = CodeAgent(
     model=model,
-    tools=[final_answer, duckduckgo_search, get_current_time_in_timezone, my_custom_tool, generate_image], ## add your tools here (don't remove final answer)
+    tools=[final_answer, duckduckgo_search, get_current_time_in_timezone, my_custom_tool, generate_image,
+          is_prime], ## add your tools here (don't remove final answer)
     max_steps=6,
     verbosity_level=1,
     grammar=None,