import os
import time
from typing import TypedDict
from langgraph.graph import StateGraph, START, END
from langchain_groq import ChatGroq
from langchain_community.tools import DuckDuckGoSearchRun
import gradio as gr

# ─────────────────────────────────────────
# 1. Config — set your key here or via HF Secrets
# ─────────────────────────────────────────
os.environ["GROQ_API_KEY"] = os.getenv("GROQ_API_KEY")

llm = ChatGroq(model="llama-3.1-8b-instant")
search = DuckDuckGoSearchRun()

# ─────────────────────────────────────────
# 2. State + Graph (same as research_agent.py)
# ─────────────────────────────────────────
class ResearchState(TypedDict):
    query: str
    search_results: str
    summary: str

def search_node(state: ResearchState):
    return {"search_results": search.run(state["query"])}

def summarize_node(state: ResearchState):
    prompt = f"""You are a research assistant. Answer the user's question clearly 
based on the search results below. Be factual and concise (3-5 sentences).

Question: {state["query"]}

Search Results:
{state["search_results"]}"""
    response = llm.invoke(prompt)
    return {"summary": response.content}

graph = StateGraph(ResearchState)
graph.add_node("search", search_node)
graph.add_node("summarize", summarize_node)
graph.add_edge(START, "search")
graph.add_edge("search", "summarize")
graph.add_edge("summarize", END)
agent = graph.compile()

# ─────────────────────────────────────────
# 3. Gradio handler
# ─────────────────────────────────────────
def run_agent(query):
    if not query.strip():
        return "", "⚠️ Please enter a question."
    
    start = time.time()
    result = agent.invoke({"query": query})
    elapsed = round(time.time() - start, 1)
    
    answer = result["summary"]
    status = f"✅ Done in {elapsed}s · LangGraph → DuckDuckGo → ChatGroq"
    return answer, status

def clear_all():
    return "", "", "Ready."

# ─────────────────────────────────────────
# 4. UI
# ─────────────────────────────────────────
EXAMPLES = [
    "What is LangGraph and how does it work?",
    "Latest AI language models released in 2025?",
    "How does Retrieval-Augmented Generation (RAG) work?",
    "What is the difference between LangChain and LangGraph?",
]

with gr.Blocks(theme=gr.themes.Soft(), title="AI Research Agent") as demo:

    gr.Markdown("""
    ## 🔍 AI Research Agent
    Ask any question. The agent searches the web and returns a clear, summarized answer.
    """)

    with gr.Row():
        # ── Left panel: Input ──
        with gr.Column(scale=2):
            query_input = gr.Textbox(
                label="Your Question",
                placeholder="e.g. What is LangGraph?",
                lines=3,
            )
            with gr.Row():
                submit_btn = gr.Button("🚀 Search", variant="primary", scale=3)
                clear_btn = gr.Button("Clear", scale=1)

            gr.Markdown("**Try an example:**")
            for example in EXAMPLES:
                gr.Button(example, size="sm").click(
                    fn=lambda e=example: e,
                    outputs=query_input
                )

        # ── Right panel: Output ──
        with gr.Column(scale=3):
            output_box = gr.Textbox(
                label="📄 Answer",
                lines=12,
                interactive=False,
            )
            status_md = gr.Markdown("Ready.")

    # ── Wire up events ──
    submit_btn.click(
        fn=run_agent,
        inputs=query_input,
        outputs=[output_box, status_md],
    )
    query_input.submit(          # also fires on Enter key
        fn=run_agent,
        inputs=query_input,
        outputs=[output_box, status_md],
    )
    clear_btn.click(
        fn=clear_all,
        outputs=[query_input, output_box, status_md],
    )

demo.launch()# app.py
import os
import time
from typing import TypedDict
from langgraph.graph import StateGraph, START, END
from langchain_groq import ChatGroq
from langchain_community.tools import DuckDuckGoSearchRun
import gradio as gr

# ─────────────────────────────────────────
# 1. Config — set your key here or via HF Secrets
# ─────────────────────────────────────────
os.environ["GROQ_API_KEY"] = "your_groq_api_key_here"

llm = ChatGroq(model="llama3-8b-8192", temperature=0.3)
search = DuckDuckGoSearchRun()

# ─────────────────────────────────────────
# 2. State + Graph (same as research_agent.py)
# ─────────────────────────────────────────
class ResearchState(TypedDict):
    query: str
    search_results: str
    summary: str

def search_node(state: ResearchState):
    return {"search_results": search.run(state["query"])}

def summarize_node(state: ResearchState):
    prompt = f"""You are a research assistant. Answer the user's question clearly 
based on the search results below. Be factual and concise (3-5 sentences).

Question: {state["query"]}

Search Results:
{state["search_results"]}"""
    response = llm.invoke(prompt)
    return {"summary": response.content}

graph = StateGraph(ResearchState)
graph.add_node("search", search_node)
graph.add_node("summarize", summarize_node)
graph.add_edge(START, "search")
graph.add_edge("search", "summarize")
graph.add_edge("summarize", END)
agent = graph.compile()

# ─────────────────────────────────────────
# 3. Gradio handler
# ─────────────────────────────────────────
def run_agent(query):
    if not query.strip():
        return "", "⚠️ Please enter a question."
    
    start = time.time()
    result = agent.invoke({"query": query})
    elapsed = round(time.time() - start, 1)
    
    answer = result["summary"]
    status = f"✅ Done in {elapsed}s · LangGraph → DuckDuckGo → ChatGroq"
    return answer, status

def clear_all():
    return "", "", "Ready."

# ─────────────────────────────────────────
# 4. UI
# ─────────────────────────────────────────
EXAMPLES = [
    "What is LangGraph and how does it work?",
    "Latest AI language models released in 2025?",
    "How does Retrieval-Augmented Generation (RAG) work?",
    "What is the difference between LangChain and LangGraph?",
]

with gr.Blocks(theme=gr.themes.Soft(), title="AI Research Agent") as demo:

    gr.Markdown("""
    ## 🔍 AI Research Agent
    Ask any question. The agent searches the web and returns a clear, summarized answer.
    """)

    with gr.Row():
        # ── Left panel: Input ──
        with gr.Column(scale=2):
            query_input = gr.Textbox(
                label="Your Question",
                placeholder="e.g. What is LangGraph?",
                lines=3,
            )
            with gr.Row():
                submit_btn = gr.Button("🚀 Search", variant="primary", scale=3)
                clear_btn = gr.Button("Clear", scale=1)

            gr.Markdown("**Try an example:**")
            for example in EXAMPLES:
                gr.Button(example, size="sm").click(
                    fn=lambda e=example: e,
                    outputs=query_input
                )

        # ── Right panel: Output ──
        with gr.Column(scale=3):
            output_box = gr.Textbox(
                label="📄 Answer",
                lines=12,
                interactive=False,
            )
            status_md = gr.Markdown("Ready.")

    # ── Wire up events ──
    submit_btn.click(
        fn=run_agent,
        inputs=query_input,
        outputs=[output_box, status_md],
    )
    query_input.submit(          # also fires on Enter key
        fn=run_agent,
        inputs=query_input,
        outputs=[output_box, status_md],
    )
    clear_btn.click(
        fn=clear_all,
        outputs=[query_input, output_box, status_md],
    )

demo.launch()