from langchain_groq import ChatGroq
from langchain_core.tools import tool
from langchain.agents import create_agent 
from langchain_tavily import TavilySearch
import gradio as gr
import os
from dotenv import load_dotenv
load_dotenv

tavily_key = os.getenv("TAVILY_API_KEY") 
groq_key = os.getenv("GROQ_API_KEY")

@tool 
def add_tool(a,b):
    """Add two numbers a and b"""
    return a+b

web_search = TavilySearch(tavily_api_key=tavily_key,
                          max_results=5)
llm = ChatGroq(
    model="llama-3.3-70b-versatile",
    api_key=groq_key,

)


agent = create_agent(
    model=llm,
    tools=[add_tool,web_search],
    system_prompt="You are the helpful AI assistant, use tools if needed."
)
def chat_func(message, history):
    
    response = agent.invoke({
        "messages": [{"role": "user", "content": message}]
    })
    
    # Return the content of the response
    return response["messages"][-1].content

# This creates the ChatGPT-like layout instantly
demo = gr.ChatInterface(
    fn=chat_func, # Uses the modern bubble format
    title="Chat Agent",
    description="Ask me anything!",
   # You can use "soft", "glass", "monochrome", or "ocean"
)

demo.launch()


# print(response["messages"][-1].content)
# print(response["messages"][1].tool_calls)