Update app.py

app.py

@@ -1,4 +1,3 @@
-# app.py
 import gradio as gr
 from search import search_google
 from scraper import scrape_url
@@ -8,42 +7,38 @@ from llm import generate_answer
 vs = VectorStore()
 
 def ask_agent(question):
-    urls = search_google(question, num_results=3)
-    texts = [scrape_url(url) for url in urls]
+    urls = [u for u in search_google(question, num_results=3) if u.startswith("http")]
+
+    texts_images = [scrape_url(url) for url in urls]
+    texts = [ti[0] for ti in texts_images if not ti[0].startswith("[Error")]
+    images = [ti[1] for ti in texts_images]  # list of list of images
+
+    # add to vector store
     vs.add_texts(texts)
     relevant = vs.retrieve(question, top_k=2)
     context = "\n\n".join(relevant)
+
+    # generate answer
     answer = generate_answer(context, question)
-    return f"## 🧠 Answer\n\n{answer}\n\n---\n### 🔗 Sources\n" + "\n".join(f"- [{url}]({url})" for url in urls)
-
-with gr.Blocks(theme=gr.themes.Soft(primary_hue="violet", secondary_hue="blue")) as demo:
-    gr.Markdown("""
-    # 🔍 **AI Web RAG Agent**  
-    _Ask anything, I'll search, scrape & answer in real time!_
-    """)
-
-    with gr.Column(scale=1):
-        question = gr.Textbox(
-            label="💡 Your question",
-            placeholder="e.g., Find cheapest flights Kanpur to Mumbai on 30 July",
-            show_label=True,
-            scale=1
-        )
-
-        btn = gr.Button("🚀 Ask")
-
-        output = gr.Markdown("🤖 *Your answer will appear here...*")
-
-        # Examples help new users
-        gr.Examples(
-            examples=[
-                "Best laptop under 50,000 INR",
-                "Latest news about ISRO moon mission",
-                "What are some tourist places near Mumbai"
-            ],
-            inputs=[question]
-        )
-
-    btn.click(fn=ask_agent, inputs=question, outputs=output)
-
-demo.launch()
+
+    # build image markdown with source
+    image_markdown = ""
+    for url, imgs in zip(urls, images):
+        if imgs:
+            # show first image as thumbnail
+            img_url = imgs[0]
+            image_markdown += f"![image]({img_url})\n"
+        image_markdown += f"[Source]({url})\n\n"
+
+    final_output = f"## 🧠 Answer\n\n{answer}\n\n---\n## 📸 Images & Sources\n\n{image_markdown}"
+    return final_output
+
+with gr.Blocks(theme=gr.themes.Soft(primary_hue="violet")) as demo:
+    gr.Markdown("# 🔍 **AI Web RAG Agent**\nAsk me anything, I'll search, scrape text & images, and answer!")
+    inp = gr.Textbox(label="Your question", placeholder="e.g., Best laptop under 50,000 INR")
+    btn = gr.Button("Ask")
+    out = gr.Markdown()
+
+    btn.click(fn=ask_agent, inputs=inp, outputs=out)
+
+demo.launch()