app.py

# app.py
import os, asyncio, aiohttp, nest_asyncio
from llama_index.tools.duckduckgo import DuckDuckGoSearchToolSpec
from llama_index.tools.weather import OpenWeatherMapToolSpec
from llama_index.tools.playwright import PlaywrightToolSpec
from llama_index.core.tools import FunctionTool
from llama_index.core.agent.workflow import ReActAgent, FunctionAgent, AgentWorkflow
from llama_index.llms.huggingface_api import HuggingFaceInferenceAPI
from llama_index.llms.openai import OpenAI
from llama_index.core.memory import ChatMemoryBuffer
from llama_index.readers.web import RssReader
from llama_index.core.workflow import Context
import gradio as gr
import subprocess
subprocess.run(["playwright", "install"])

# allow nested loops in Spaces
nest_asyncio.apply()

# --- Secrets via env vars ---
HF_TOKEN            = os.getenv("HF_TOKEN")
OPENAI_API_KEY      = os.getenv("OPENAI_API_KEY")
OPENWEATHERMAP_KEY  = os.getenv("OPENWEATHERMAP_API_KEY")
SERPER_API_KEY      = os.getenv("SERPER_API_KEY")

# --- LLMs ---
llm = HuggingFaceInferenceAPI(
    model_name="Qwen/Qwen2.5-Coder-32B-Instruct",
    token=HF_TOKEN, 
    task="conversational"
)
# OpenAI for pure function-calling
openai_llm = OpenAI(
    model="gpt-4o",
    api_key=OPENAI_API_KEY,
    temperature=0.0,
    streaming=False,
)

# --- Memory ---
memory = ChatMemoryBuffer.from_defaults(token_limit=4096)

# --- Tools Setup ---
# DuckDuckGo
duck_spec = DuckDuckGoSearchToolSpec()
search_tool = FunctionTool.from_defaults(duck_spec.duckduckgo_full_search)

# Weather
openweather_api_key=OPENWEATHERMAP_KEY
weather_tool_spec = OpenWeatherMapToolSpec(key=openweather_api_key)
weather_tool_spec = OpenWeatherMapToolSpec(key=openweather_api_key)
weather_tool = FunctionTool.from_defaults(
    weather_tool_spec.weather_at_location,
    name="current_weather",
    description="Get the current weather at a specific location (city, country)."
)
forecast_tool = FunctionTool.from_defaults(
    weather_tool_spec.forecast_tommorrow_at_location,
    name="weather_forecast",
    description="Get tomorrow's weather forecast for a specific location (city, country)."
)

# Playwright (synchronous start)
async def _start_browser():
    return await PlaywrightToolSpec.create_async_playwright_browser(headless=True)
browser = asyncio.get_event_loop().run_until_complete(_start_browser())
playwright_tool_spec = PlaywrightToolSpec.from_async_browser(browser)
navigate_tool = FunctionTool.from_defaults(
    playwright_tool_spec.navigate_to,
    name="web_navigate",
    description="Navigate to a specific URL."
)
extract_text_tool = FunctionTool.from_defaults(
    playwright_tool_spec.extract_text,
    name="web_extract_text",
    description="Extract all text from the current page."
)
extract_links_tool = FunctionTool.from_defaults(
    playwright_tool_spec.extract_hyperlinks,
    name="web_extract_links",
    description="Extract all hyperlinks from the current page."
)

# Google News RSS
def fetch_google_news_rss():
    docs = RssReader(html_to_text=True).load_data(["https://news.google.com/rss"])
    return [{"title":d.metadata.get("title",""), "url":d.metadata.get("link","")} for d in docs]
google_rss_tool = FunctionTool.from_defaults(
    fn=fetch_google_news_rss,
    name="fetch_google_news_rss",
    description="Get headlines & URLs from Google News RSS."
)

# Serper
async def fetch_serper(ctx, query):
    if not SERPER_API_KEY:
        raise ValueError("SERPER_API_KEY missing")
    url = f"https://google.serper.dev/news?q={query}&tbs=qdr%3Ad"
    hdr = {"X-API-KEY": SERPER_API_KEY, "Content-Type":"application/json"}
    async with aiohttp.ClientSession() as s:
        r = await s.get(url, headers=hdr)
        r.raise_for_status()
        return await r.json()
serper_news_tool = FunctionTool.from_defaults(
    fetch_serper, name="fetch_news_from_serper", 
    description="Search today’s news via Serper."
)

# --- Agents ---
# 1. Google News RSS Agent (replaces old google_news_agent)
google_rss_agent = FunctionAgent(
    name="google_rss_agent",
    description="Fetches latest headlines and URLs from Google News RSS feed.",
    system_prompt="You are an agent that fetches the latest headlines and URLs from the Google News RSS feed.",
    tools=[google_rss_tool],
    llm=openai_llm,
    memory=memory,
)

# 2. Web Browsing Agent
web_browsing_agent = ReActAgent(
    name="web_browsing_agent",
    description="Fetches Serper URLs, navigates to each link, extracts the text and builds a summary",
    system_prompt=(
        "You are a news-content agent. When asked for details on a headline:\n"
        "1. Call `fetch_news_from_serper(query)` to get JSON with article URLs.\n"
        "2. For each top URL, call `web_navigate(url)` then `web_extract_text()`.\n"
        "3. Synthesize those texts into a concise summary."
    ),
    tools=[serper_news_tool, navigate_tool, extract_text_tool, extract_links_tool],
    llm=llm,
    memory=memory,
)

# 3. Weather Agent
weather_agent = ReActAgent(
    name="weather_agent",
    description="Answers weather-related questions.",
    system_prompt="You are a weather agent that provides current weather and forecasts.",
    tools=[weather_tool, forecast_tool],
    llm=openai_llm,
)

# 4. DuckDuckGo Search Agent
search_agent = ReActAgent(
    name="search_agent",
    description="Searches general info using DuckDuckGo.",
    system_prompt="You are a search agent that uses DuckDuckGo to answer questions.",
    tools=[search_tool],
    llm=openai_llm,
)

router_agent = ReActAgent(
    name="router_agent",
    description="Routes queries to the correct specialist agent.",
    system_prompt=(
        "You are RouterAgent. "
        "Given the user query, reply with exactly one name from: "
        "['google_rss_agent','weather_agent','search_agent','web_browsing_agent']."
    ),
    llm=llm,
    tools=[
        FunctionTool.from_defaults(
            fn=lambda ctx, choice: choice,
            name="choose_agent",
            description="Return the chosen agent name."
        )
    ],
    can_handoff_to=[
        "google_rss_agent",
        "weather_agent",
        "search_agent",
        "web_browsing_agent",
    ],
)

workflow = AgentWorkflow(
    agents=[router_agent, google_rss_agent, web_browsing_agent, weather_agent, search_agent],
    root_agent="router_agent"
)
ctx = Context(workflow)

# # Sync wrapper
# async def respond(query: str) -> str:
#     out = await workflow.run(user_msg=query, ctx=ctx, memory=memory)
#     return out.response.blocks[0].text
    
# # Async response handler for Gradio
# async def respond_gradio(query, chat_history):
#     answer = await respond(query)
#     return chat_history + [[query, answer]]

# New unified respond() function
async def respond(message, history):
    out = await workflow.run(user_msg=message, ctx=ctx, memory=memory)
    answer = out.response.blocks[0].text
    # Return the full updated history
    return history + [[message, answer]]

# --- Gradio UI ---
# with gr.Blocks() as demo:
#     gr.Markdown("## 🤖 Perspicacity")
#     gr.Markdown(
#     "This bot can check the news, tell you the weather, and even browse websites to answer follow-up questions — all powered by a team of tiny AI agents working behind the scenes.  \n\n"
#     "🧪 Built for fun during the [AI Agents course](https://huggingface.co/learn/agents-course/unit0/introduction) at Hugging Face — it's just a demo to show what agents can do.  \n"
#     "🙌 Got ideas or improvements? PRs welcome!"
#     )
#     chatbot = gr.Chatbot()
#     txt = gr.Textbox(placeholder="Ask me about news, weather, etc…")

#     txt.submit(respond_gradio, inputs=[txt, chatbot], outputs=chatbot)

# demo.launch()
with gr.Blocks() as demo:
    gr.Markdown("## 🗞️ Multi‐Agent News & Weather Chatbot")
    gr.Markdown(
        "This bot can check the news, tell you the weather, and even browse websites to answer follow-up questions — all powered by a team of tiny AI agents working behind the scenes.\n\n"
        "🧪 Built for fun during the [AI Agents course](https://huggingface.co/learn/agents-course/unit0/introduction) — it's just a demo to show what agents can do.  \n"
        "🙌 Got ideas or improvements? PRs welcome!  \n\n"
        "👉 _Try asking “What’s the weather in Montreal?” or “What’s in the news today?”_"
    )

    chatbot = gr.Chatbot()
    txt = gr.Textbox(placeholder="Ask me about news, weather or anything…")

    txt.submit(respond, inputs=[txt, chatbot], outputs=chatbot)

demo.launch()