Update app.py

app.py

@@ -1,4 +1,5 @@
 from smolagents import CodeAgent,DuckDuckGoSearchTool, HfApiModel,load_tool,tool
+import os 
 import datetime
 import requests
 import pytz
@@ -7,22 +8,74 @@ from tools.final_answer import FinalAnswerTool
 
 from Gradio_UI import GradioUI
 
-# Below is my news aggregator tool! 
+from transformers import pipeline
+from diffusers import StableDiffusionPipeline
+import torch
+
+
+pipe = StableDiffusionPipeline.from_pretrained("stabilityai/stable-diffusion-2-1", torch_dtype=torch.float16)
+pipe = pipe.to("cuda")  # Use GPU if available
+
+@tool
+def generate_image(prompt: str) -> str:
+    """
+    A tool that generates an image based on a text prompt.
+    
+    Args:
+        prompt: The text prompt to generate an image from.
+    
+    Returns:
+        The file path or URL of the generated image.
+    """
+    try:
+        # Generate the image
+        image = pipe(prompt).images[0]
+        
+        # Save the image to a file
+        image_path = f"generated_images/{prompt[:20]}.png"  # Save with a unique name
+        image.save(image_path)
+        
+        return image_path
+    except Exception as e:
+        return f"Error generating image: {str(e)}"
+
+
+summarizer = pipeline("summarization")
+
+@tool
+def summarize_article(article_text: str) -> str:
+    """
+    A tool that summarizes the content of a news article.
+    
+    Args:
+        article_text: The text content of the article to summarize.
+    
+    Returns:
+        A summarized version of the article.
+    """
+    try:
+        # Summarize the article text
+        summary = summarizer(article_text, max_length=50, min_length=25, do_sample=False)
+        return summary[0]['summary_text']
+    except Exception as e:
+        return f"Error summarizing article: {str(e)}"
+
 @tool
 def news_aggregator(keyword: str, max_results: int = 5) -> str:
     """
-    A tool that fetches the latest news articles related to a specific keyword.
+    A tool that fetches the latest news articles, summarizes them and generates images.
     
     Args:
         keyword: The topic or keyword to search for (e.g., "AI", "technology").
         max_results: The maximum number of articles to return (default is 5).
     
     Returns:
-        A string containing the titles and URLs of the latest news articles.
+        A string containing the titles, summaries, and image paths of the latest news articles.
     """
     try:
         
         api_key = 'd43565ac0e674145a22c5f0de4dcd69c'
+        
         # Construct the API URL
         url = f"https://newsapi.org/v2/everything?q={keyword}&pageSize={max_results}&apiKey={api_key}"
         
@@ -38,12 +91,24 @@ def news_aggregator(keyword: str, max_results: int = 5) -> str:
         if not articles:
             return f"No news articles found for the keyword '{keyword}'."
         
-        # Format the results
+        # Process each article
         result = f"Here are the latest news articles about '{keyword}':\n"
         for i, article in enumerate(articles, 1):
             title = article.get("title", "No title")
             url = article.get("url", "#")
-            result += f"{i}. {title}\n   {url}\n"
+            content = article.get("content", "")  # Get the article content
+            
+            # Summarize the article
+            summary = summarize_article(content)
+            
+            # Generate an image based on the title
+            image_path = generate_image(title)
+            
+            # Format the results
+            result += f"{i}. {title}\n"
+            result += f"   Summary: {summary}\n"
+            result += f"   Image: {image_path}\n"
+            result += f"   URL: {url}\n"
         
         return result
     
@@ -51,8 +116,6 @@ def news_aggregator(keyword: str, max_results: int = 5) -> str:
         return f"Error fetching news articles: {str(e)}"
     except Exception as e:
         return f"An unexpected error occurred: {str(e)}"
-
-
 final_answer = FinalAnswerTool()
 
 # If the agent does not answer, the model is overloaded, please use another model or the following Hugging Face Endpoint that also contains qwen2.5 coder:
@@ -73,7 +136,7 @@ with open("prompts.yaml", 'r') as stream:
 # Agent and tools
 agent = CodeAgent(
     model=model,
-    tools=[final_answer, news_aggregator], ## add your tools here (don't remove final answer)
+    tools=[final_answer, news_aggregator, summarize_article, generate_image], ## add your tools here (don't remove final answer)
     max_steps=6,
     verbosity_level=1,
     grammar=None,