Update app.py

app.py

@@ -5,50 +5,49 @@ import pytz
 import yaml
 from tools.final_answer import FinalAnswerTool
 from Gradio_UI import GradioUI
-# from tavily import TavilyClient
+from tavily import TavilyClient
 import torch
 from diffusers import FluxPipeline
 # Below is an example of a tool that does nothing. Amaze us with your creativity !
-
-
-# فقط یک بار لود می‌شه (برای سرعت)
-pipe = FluxPipeline.from_pretrained(
-    "black-forest-labs/FLUX.1-schnell",
-    torch_dtype=torch.bfloat16
-)
-pipe.enable_model_cpu_offload()  # خیلی مهم! باعث می‌شه روی CPU هم کار کنه
+from tavily import TavilyClient
 
 @tool
-def generate_image(prompt: str, width: int = 1024, height: int = 1024, steps: int = 4) -> str:
-    """Generates an image using FLUX.1-schnell based on the given prompt.
+def search_internet(query: str) -> str:
+    """Search the internet using Tavily API to retrieve relevant results for a given query.
     
     Args:
-        prompt: The text description of the image you want to generate.
-        width: Image width in pixels (default: 1024)
-        height: Image height in pixels (default: 1024)
-        steps: Number of inference steps (4 is enough for schnell, more = better quality but slower)
-    
-    Returns:
-        URL of the generated image hosted on the Space
+        query: The search query string to look up information about.
     """
-    print(f"Generating image with prompt: {prompt}")
-    
-    # تولید عکس
-    image = pipe(
-        prompt=prompt,
-        height=height,
-        width=width,
-        num_inference_steps=steps,
-        guidance_scale=0.0,  # برای schnell باید 0 باشه
-        max_sequence_length=256
-    ).images[0]
-    
-    # ذخیره موقت و برگرداندن URL
-    image_path = "/tmp/generated_image.png"
-    image.save(image_path)
-    
-    # این ترفند باعث می‌شه عکس در گرادیو نمایش داده بشه
-    return image_path
+    print(f"🔍 Searching with Tavily: {query}")
+
+    # Get API key from tavily.com (1000 free searches/month)
+    api_key = "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"
+
+    try:
+        client = TavilyClient(api_key=api_key)
+
+        response = client.search(
+            query=query,
+            max_results=3
+        )
+
+        results = response.get('results', [])
+
+        if not results:
+            return "No results found."
+
+        output = []
+        for i, result in enumerate(results, 1):
+            title = result.get('title', 'No title')
+            content = result.get('content', 'No description')
+            url = result.get('url', '')
+            output.append(f"{i}. {title}\n{content[:200]}...\nSource: {url}\n")
+
+        return "\n".join(output)
+
+    except Exception as e:
+        return f"Search error: {str(e)}"
+
     
 
 @tool