ai_agent.py

# Step 1: Setup the API keys of OPENAI, Tavily and GROK

from dotenv import load_dotenv
import os

load_dotenv()

OPENAI_API_KEY = os.getenv("OPENAI_API_KEY")
GROQ_API_KEY = os.getenv("GROQ_API_KEY")
TAVILY_API_KEY = os.getenv("TAVILY_API_KEY")


# Step 2: Setup LLM and tools


from langchain_groq import ChatGroq
from langchain_openai import ChatOpenAI
from langchain_community.tools.tavily_search import TavilySearchResults
from langchain_core.messages.ai import AIMessage


# Step 3: Setup AI agent with Search tool functionality

from langgraph.prebuilt import create_react_agent



def get_response_from_ai_agent(llm_id, query, allow_search, system_prompt, provider):

    if provider.lower()=="groq": 
        llm=ChatGroq(model=llm_id, api_key=GROQ_API_KEY)
    elif provider.lower()=="openai":
        llm=ChatOpenAI(model=llm_id, api_key=OPENAI_API_KEY)
    else:
        llm=ChatOpenAI(model=llm_id, api_key=OPENAI_API_KEY)

    tools = [TavilySearchResults(max_results=3, api_key=TAVILY_API_KEY)] if allow_search else []


    agent = create_react_agent(
        model=llm, 
        tools=tools,
        state_modifier=system_prompt,
    )


    state= {"messages" :query}

    response= agent.invoke(state)

    messages= response["messages"]

    ai_messages= [message.content for message in messages if isinstance(message, AIMessage)]

    return ai_messages[-1]