import os
from langchain_groq import ChatGroq
from langchain.agents import initialize_agent, AgentType
from memory import memory
from tools import tools  # Import Flight Booking Tool

# Load API Key
API_KEY = os.getenv("API_KEY")

# Ensure API Key is set
if not API_KEY:
    raise ValueError("API_KEY is not set. Please define it in your environment variables.")

# Initialize the LLM (Groq's Mixtral)
llm = ChatGroq(
    groq_api_key=API_KEY,
    model_name="mixtral-8x7b-32768",
    temperature=0.7,
    max_tokens=512,
)

# Initialize the conversational agent with Flight Booking Tool
agent = initialize_agent(
    tools=tools,
    llm=llm,
    agent=AgentType.CONVERSATIONAL_REACT_DESCRIPTION,
    verbose=True,  # Hide debug output
    memory=memory,
)

# 🛠 Custom Logging Function to Improve Execution
def log_agent_action(prompt, response):
    print(f"\n🟢 **User Query:** {prompt}")
    print(f"🔵 **Agent Thought Process:**")
    print(f"✅ **Final Response:** {response}")