agent.py

import os
import shutil
from langchain_groq import ChatGroq
from langchain.prompts import PromptTemplate
from langgraph.graph import START, StateGraph, MessagesState
from langgraph.prebuilt import ToolNode, tools_condition
from langchain_community.tools.tavily_search import TavilySearchResults
from langchain_community.document_loaders import WikipediaLoader
from langchain_core.messages import  SystemMessage, HumanMessage
from langchain.tools import tool
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.runnables import Runnable
from dotenv import load_dotenv
from langchain.vectorstores import Chroma
from langchain_huggingface import ChatHuggingFace, HuggingFaceEndpoint, HuggingFaceEmbeddings
from langchain.text_splitter import CharacterTextSplitter
from langchain.tools.retriever import create_retriever_tool
from typing import TypedDict, Annotated, List
from langchain_community.tools import DuckDuckGoSearchRun, WikipediaQueryRun, ArxivQueryRun
from langchain_community.utilities import WikipediaAPIWrapper, ArxivAPIWrapper
from langchain.tools import Tool

# Load environment variables from .env
load_dotenv()

# Custom Agent Prompt Template
Agent_prompt_template = '''You are a helpful assistant tasked with answering questions using a set of tools. 
Now, I will ask you a question. Report your thoughts, and finish your answer with the following template: 
FINAL ANSWER: [YOUR FINAL ANSWER]. 
YOUR FINAL ANSWER should be a number OR as few words as possible OR a comma separated list of numbers and/or strings. If you are asked for a number, don't use comma to write your number neither use units such as $ or percent sign unless specified otherwise. If you are asked for a string, don't use articles, neither abbreviations (e.g. for cities), and write the digits in plain text unless specified otherwise. If you are asked for a comma separated list, apply the above rules depending of whether the element to be put in the list is a number or a string.
Your answer should only start with "FINAL ANSWER: ", then follows with the answer.  '''

sys_msg = SystemMessage(content=Agent_prompt_template)


# Initialize LLM
def initialize_llm():
    """Initializes the ChatGroq LLM."""
    llm = ChatGroq(
        temperature=0,
        model_name="qwen-qwq-32b",
        groq_api_key=os.getenv("GROQ_API_KEY")
    )
    return llm

# Initialize Tavily Search Tool
def initialize_search_tool():
    """Initializes the TavilySearchResults tool."""
    return TavilySearchResults()

# Weather tool
def get_weather(location: str, search_tool: TavilySearchResults = None) -> str:
    """
    Fetches the current weather information for a given location using Tavily search.

    Args:
        location (str): The name of the location to search for.
        search_tool (TavilySearchResults, optional):  Defaults to None.

    Returns:
        str: The weather information for the specified location.
    """
    if search_tool is None:
        search_tool = initialize_search_tool()
    query = f"current weather in {location}"
    return search_tool.run(query)

# Recommendation chain
def initialize_recommendation_chain(llm: ChatGroq) -> Runnable:
    """
    Initializes the recommendation chain.

    Args:
      llm(ChatGroq):The LLM to use

    Returns:
        Runnable: A runnable sequence to generate recommendations.
    """
    recommendation_prompt = ChatPromptTemplate.from_template("""
    You are a helpful assistant that gives weather-based advice.
    
    Given the current weather condition: "{weather_condition}", provide:
    1. Clothing or activity recommendations suited for this weather.
    2. At least one health tip to stay safe or comfortable in this condition.

    Be concise and clear.
    """)
    return recommendation_prompt | llm

def get_recommendation(weather_condition: str, recommendation_chain: Runnable = None) -> str:
    """
    Gives activity/clothing recommendations and health tips based on the weather condition.

    Args:
        weather_condition (str): The current weather condition.
        recommendation_chain (Runnable, optional): The recommendation chain to use. Defaults to None.

    Returns:
        str:  Recommendations and health tips for the given weather condition.
    """
    if recommendation_chain is None:
        llm = initialize_llm()
        recommendation_chain = initialize_recommendation_chain(llm)
    return recommendation_chain.invoke({"weather_condition": weather_condition})

# Math tools
@tool
def add(x: int, y: int) -> int:
    """
    Adds two integers.

    Args:
        x (int): The first integer.
        y (int): The second integer.

    Returns:
        int: The sum of x and y.
    """
    return x + y

@tool
def subtract(x: int, y: int) -> int:
    """
    Subtracts two integers.

    Args:
        x (int): The first integer.
        y (int): The second integer.

    Returns:
        int: The difference between x and y.
    """
    return x - y

@tool
def multiply(x: int, y: int) -> int:
    """
    Multiplies two integers.

    Args:
        x (int): The first integer.
        y (int): The second integer.

    Returns:
        int: The product of x and y.
    """
    return x * y

@tool
def divide(x: int, y: int) -> float:
    """
    Divides two numbers.

    Args:
        x (int): The numerator.
        y (int): The denominator.

    Returns:
        float: The result of the division.

    Raises:
        ValueError: If y is zero.
    """
    if y == 0:
        raise ValueError("Cannot divide by zero.")
    return x / y

@tool
def square(x: int) -> int:
    """
    Calculates the square of a number.

    Args:
        x (int): The number to square.

    Returns:
        int: The square of x.
    """
    return x * x

@tool
def cube(x: int) -> int:
    """
    Calculates the cube of a number.

    Args:
        x (int): The number to cube.

    Returns:
        int: The cube of x.
    """
    return x * x * x

@tool
def power(x: int, y: int) -> int:
    """
    Raises a number to the power of another number.

    Args:
        x (int): The base number.
        y (int): The exponent.

    Returns:
        int: x raised to the power of y.
    """
    return x ** y

@tool
def factorial(n: int) -> int:
    """
    Calculates the factorial of a non-negative integer.

    Args:
        n (int): The non-negative integer.

    Returns:
        int: The factorial of n.

    Raises:
        ValueError: If n is negative.
    """
    if n < 0:
        raise ValueError("Factorial is not defined for negative numbers.")
    if n == 0 or n == 1:
        return 1
    result = 1
    for i in range(2, n + 1):
        result *= i
    return result

@tool
def mean(numbers: list) -> float:
    """
    Calculates the mean of a list of numbers.

    Args:
        numbers (list): A list of numbers.

    Returns:
        float: The mean of the numbers.

    Raises:
        ValueError: If the list is empty.
    """
    if not numbers:
        raise ValueError("The list is empty.")
    return sum(numbers) / len(numbers)

@tool
def standard_deviation(numbers: list) -> float:
    """
    Calculates the standard deviation of a list of numbers.

    Args:
        numbers (list): A list of numbers.

    Returns:
        float: The standard deviation of the numbers.

    Raises:
        ValueError: If the list is empty.
    """
    if not numbers:
        raise ValueError("The list is empty.")
    mean_value = mean(numbers)
    variance = sum((x - mean_value) ** 2 for x in numbers) / len(numbers)
    return variance ** 0.5

# --- Vector Store + Retriever ---
# State schema
class MessagesState(TypedDict):
    messages: Annotated[List[HumanMessage], "Messages in the conversation"]

# === VECTOR STORE SETUP ===
PERSIST_DIR = "./chroma_store"

def initialize_chroma_store():
    # Optional: clear existing store if desired
    if os.path.exists(PERSIST_DIR):
        shutil.rmtree(PERSIST_DIR)
        os.makedirs(PERSIST_DIR)

    # Initialize embeddings
    embeddings = HuggingFaceEmbeddings(model_name="sentence-transformers/all-mpnet-base-v2")

    # Load existing or empty vector store
    vectorstore = Chroma(
        embedding_function=embeddings,
        persist_directory=PERSIST_DIR
    )
    return vectorstore

vector_store = initialize_chroma_store()

# Create retriever tool
retriever_tool = create_retriever_tool(
    retriever=vector_store.as_retriever(),
    name="Question Search",
    description="A tool to retrieve similar questions from a vector store."
)



@tool
def weather_tool(location: str) -> str:
    """
    Fetches the weather for a location.

    Args:
        location (str): The location to fetch weather for.

    Returns:
        str: The weather information.
    """
    return get_weather(location, search_tool)


from langchain_community.tools.tavily_search import TavilySearchResults
from langchain_community.tools.ddg_search import DuckDuckGoSearchRun
from langchain_community.tools.wikipedia.tool import WikipediaQueryRun
from langchain_community.utilities.wikipedia import WikipediaAPIWrapper
from langchain_community.tools.arxiv.tool import ArxivQueryRun
from langchain_community.utilities.arxiv import ArxivAPIWrapper
from langchain.tools import tool

# 1. Tavily Web Search Tool (already in correct format)
@tool
def web_search(query: str) -> str:
    """Search the web for a given query and return the summary."""
    search_tool = TavilySearchResults()
    result = search_tool.run(query)
    return result[0]['content']

# 2. DuckDuckGo Search Tool
@tool
def duckduckgo_search(query: str) -> str:
    """Search the web using DuckDuckGo for a given query and return the result."""
    search_tool = DuckDuckGoSearchRun(verbose=False)
    return search_tool.run(query)

# 3. Wikipedia Search Tool
@tool
def wikipedia_search(query: str) -> str:
    """Search Wikipedia for a given query and return the top 3 results."""
    wrapper = WikipediaAPIWrapper(top_k_results=3)
    wikipedia = WikipediaQueryRun(api_wrapper=wrapper, verbose=False)
    return wikipedia.run(query)

# 4. Arxiv Search Tool
@tool
def arxiv_search(query: str) -> str:
    """Search arXiv for academic papers based on a query and return the top 3 results."""
    wrapper = ArxivAPIWrapper(
        top_k_results=3,
        ARXIV_MAX_QUERY_LENGTH=300,
        load_max_docs=3,
        load_all_available_meta=False,
        doc_content_chars_max=40000
    )
    arxiv = ArxivQueryRun(api_wrapper=wrapper, verbose=False)
    return arxiv.run(query)

tools = [arxiv_search, duckduckgo_search, web_search,wikipedia_search,
            add, subtract, multiply, divide, square, cube, power, factorial, mean, standard_deviation]

# === LLM with Tools ===

llm = ChatGroq(
    temperature=0,
    model_name="qwen-qwq-32b",
    groq_api_key=os.getenv("GROQ_API_KEY")
)

# tools = [weather_tool, wiki_search, web_search,
#          add, subtract, multiply, divide, square, cube,
#          power, factorial, mean, standard_deviation, arxiv_tool,wikisearch_tool, search_tool ]

llm_with_tools = llm.bind_tools(tools)

# === LangGraph State ===

class ToolAgentState(TypedDict):
    messages: Annotated[List[HumanMessage], "Messages in the conversation"]

def assistant(state: ToolAgentState):
    return {"messages": [llm_with_tools.invoke([sys_msg] + state["messages"])]}

# === Build Graph ===

def build_graph():
    builder = StateGraph(ToolAgentState)
    builder.add_node("assistant", assistant)
    builder.add_node("tools", ToolNode(tools))

    builder.set_entry_point("assistant")
    builder.add_conditional_edges("assistant", tools_condition)
    builder.add_edge("tools", "assistant")
    return builder.compile()

# === Run ===

if __name__ == "__main__":
    question = "When did India won a world cup in cricket before 2000?"
    graph = build_graph()
    messages = [HumanMessage(content=question)]
    result = graph.invoke({"messages": messages})

    for msg in result["messages"]:
        msg.pretty_print()