import gradio as gr
from langchain_groq import ChatGroq
from langchain_community.utilities import ArxivAPIWrapper, WikipediaAPIWrapper
from langchain_community.tools import ArxivQueryRun, WikipediaQueryRun, DuckDuckGoSearchRun
from langchain.agents import initialize_agent, AgentType
from langchain.callbacks import StdOutCallbackHandler
import os

# Tools Setup (identical to original)
arxiv_wrapper = ArxivAPIWrapper(top_k_results=1, doc_content_chars_max=200)
arxiv = ArxivQueryRun(api_wrapper=arxiv_wrapper)
wiki_wrapper = WikipediaAPIWrapper(top_k_results=1, doc_content_chars_max=200)
wiki = WikipediaQueryRun(api_wrapper=wiki_wrapper)
search = DuckDuckGoSearchRun(name="Search")

# Gradio Chat Interface - UPDATED VERSION
def respond(message, chat_history, api_key):
    if not api_key:
        chat_history.append((message, "⚠️ Please enter your Groq API key"))
        return "", chat_history
    
    try:
        # Initialize agent with error handling
        llm = ChatGroq(groq_api_key=api_key, model_name="Llama3-8b-8192")
        tools = [search, arxiv, wiki]
        
        agent = initialize_agent(
            tools, 
            llm, 
            agent=AgentType.ZERO_SHOT_REACT_DESCRIPTION,
            handle_parsing_errors=True
        )
        
        # Using the recommended .invoke() method instead of deprecated .run()
        callback = StdOutCallbackHandler()
        response = agent.invoke(
            {"input": message},
            {"callbacks": [callback]}
        )
        
        chat_history.append((message, response["output"]))
        
    except Exception as e:
        error_msg = f"❌ Error: {str(e)}"
        if "API key" in str(e):
            error_msg = "🔑 Invalid or missing API key"
        chat_history.append((message, error_msg))
    
    return "", chat_history

# UI Layout (unchanged)
with gr.Blocks(title="🔎 LangChain - Chat with Search") as demo:
    gr.Markdown("""
    ## 🔎 LangChain - Chat with Search
    Chatbot that can search the web, arXiv, and Wikipedia.
    """)
    
    with gr.Row():
        with gr.Column(scale=1):
            api_key = gr.Textbox(
                label="Groq API Key", 
                type="password",
                placeholder="Enter your API key"
            )
        with gr.Column(scale=4):
            chatbot = gr.Chatbot(height=500)
            msg = gr.Textbox(
                placeholder="What is machine learning?",
                label="Your Question"
            )
            clear = gr.ClearButton([msg, chatbot])
    
    msg.submit(
        respond, 
        [msg, chatbot, api_key], 
        [msg, chatbot]
    )

if __name__ == "__main__":
    demo.launch(debug=True)  # Added debug mode for development