added doc strings to tool

agent.py

@@ -15,6 +15,7 @@ load_dotenv()
 
 # Initialize LLM
 def initialize_llm():
+    """Initializes the ChatGroq LLM."""
     llm = ChatGroq(
         temperature=0,
         model_name="qwen-qwq-32b",
@@ -24,10 +25,21 @@ def initialize_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}"
@@ -35,6 +47,15 @@ def get_weather(location: str, search_tool: TavilySearchResults = None) -> str:
 
 # 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.
     
@@ -47,6 +68,16 @@ def initialize_recommendation_chain(llm: ChatGroq) -> Runnable:
     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)
@@ -55,36 +86,119 @@ def get_recommendation(weather_condition: str, recommendation_chain: Runnable =
 # 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:
@@ -96,12 +210,36 @@ def factorial(n: int) -> int:
 
 @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)
@@ -110,16 +248,40 @@ def standard_deviation(numbers: list) -> float:
 
 # Build the LangGraph
 def build_graph():
+    """
+    Builds the LangGraph with the defined tools and assistant node.
+
+    Returns:
+        RunnableGraph: The compiled LangGraph.
+    """
     llm = initialize_llm()
     search_tool = initialize_search_tool()
     recommendation_chain = initialize_recommendation_chain(llm)
 
     @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)
 
     @tool
     def recommendation_tool(weather_condition: str) -> str:
+        """
+        Provides recommendations based on weather conditions.
+
+        Args:
+            weather_condition (str): The weather condition.
+
+        Returns:
+            str: The recommendations.
+        """
         return get_recommendation(weather_condition, recommendation_chain)
 
     tools = [weather_tool, recommendation_tool, add, subtract, multiply, divide, square, cube, power, factorial, mean, standard_deviation]
@@ -127,6 +289,15 @@ def build_graph():
     llm_with_tools = llm.bind_tools(tools)
 
     def assistant(state: MessagesState):
+        """
+        Assistant node in the LangGraph.
+
+        Args:
+            state (MessagesState): The current state of the conversation.
+
+        Returns:
+            dict: The next state of the conversation.
+        """
         print("Entering assistant node...")
         response = llm_with_tools.invoke(state["messages"])
         print(f"Assistant says: {response.content}")