Update app.py

app.py

@@ -4,7 +4,7 @@ import requests
 import pandas as pd
 import yaml
 import gradio as gr
-
+from bs4 import BeautifulSoup
 from datetime import datetime
 from typing import Optional
 
@@ -35,19 +35,52 @@ def get_current_time_in_timezone(timezone: str) -> str:
         return f"Error fetching time for timezone '{timezone}': {e}"
 
 @tool
-def wiki(query: str) -> str:
+def wiki(query: str, section: Optional[str] = None) -> str:
     """
-    Searches Wikipedia for a query and returns the page content (English).
+    Searches Wikipedia for `query` and returns either:
+      - the full page text (if section is None), or
+      - the text of a specific section (if section is given).
 
     Args:
         query: The Wikipedia page title or search term.
+        section: (optional) The section heading to extract (e.g. "Discography").
     """
     try:
+        # 1) resolve the page
         wikipedia.set_lang("en")
         page = wikipedia.page(query, auto_suggest=False)
+    except Exception as e:
+        return f"Error finding Wikipedia page for '{query}': {e}"
+
+    # 2) if no section requested, return raw content
+    if not section:
         return page.content
+
+    # 3) otherwise fetch the HTML and parse
+    try:
+        resp = requests.get(page.url, timeout=10)
+        resp.raise_for_status()
+        soup = BeautifulSoup(resp.text, "lxml")
     except Exception as e:
-        return f"Error fetching Wikipedia page for '{query}': {e}"
+        return f"Error fetching HTML for '{page.url}': {e}"
+
+    # 4) locate the heading matching `section`
+    header = None
+    for h in soup.find_all(["h2", "h3"]):
+        if section.lower() in h.get_text().lower():
+            header = h
+            break
+    if not header:
+        return f"Section '{section}' not found on page '{query}'."
+
+    # 5) collect text until the next same-level heading
+    texts = []
+    for sib in header.next_siblings:
+        if getattr(sib, "name", None) in ["h2", "h3"]:
+            break
+        texts.append(sib.get_text(separator="\n").strip())
+
+    return "\n\n".join([t for t in texts if t])
 
 # -----------------------------------------------------------------------------
 # 2) LOAD PROMPTS & RAG MAP