Update app.py

app.py

@@ -39,10 +39,10 @@ def call_llm(messages, model="deepseek/deepseek-chat-v3-0324:free", max_tokens=3
     return result["choices"][0]["message"]["content"]
 
 def get_related_images(topic, num_images=3):
-    prompt = f"Give me {num_images} direct image URLs from the web that are relevant to the topic: '{topic}'. Only return raw URLs on separate lines."
+    prompt = f"You are an AI assistant helping to find visuals. Please return {num_images} direct image URLs (one per line) that are relevant to the topic: '{topic}'. Ensure each line is a raw direct link to an image file (ending in .jpg, .png, etc.) and nothing else."
     try:
         content = call_llm([{"role": "user", "content": prompt}])
-        urls = [line.strip() for line in content.split('\n') if line.strip().startswith("http")]
+        urls = [line.strip() for line in content.split('\n') if line.strip().startswith("http") and any(ext in line for ext in ['.jpg', '.jpeg', '.png', '.webp'])]
         return urls[:num_images]
     except Exception as e:
         print("Error fetching images:", e)