Convert to simiple tools that return one scalar value, along with a helper function to improve efficiency

app.py

@@ -4,7 +4,7 @@ import requests
 import pytz
 import yaml
 import yfinance as yf
-from typing import Dict, Any
+from typing import Tuple, Dict, Any
 from tools.final_answer import FinalAnswerTool
 
 from Gradio_UI import GradioUI
@@ -40,132 +40,113 @@ def get_current_time_in_timezone(timezone: str) -> str:
         return f"Error fetching time for timezone '{timezone}': {str(e)}"
 
 @tool
-def yf_get_ticker_price(ticker: str) -> float:
+def yf_underlying_price(ticker: str, dte: int = 0) -> float:
     """
-    Return the latest closing price for a stock ticker.
+    Get the current underlying stock price for the given ticker.
 
     Args:
-        ticker: A string representing a stock ticker.
-    """
-    stock = yf.Ticker(ticker)
-
-    # Get current underlying price
-    hist = stock.history(period="1d")
-    if hist.empty:
-        return {"error": f"No price data found for {ticker}."}
+        ticker (str): The stock ticker symbol, e.g., AAPL.
+        dte (int): Days-to-expiration used to select the option chain.
+            Defaults to 0 (nearest expiration).
 
-    return float(hist["Close"].iloc[-1])
+    Returns:
+        float: The current underlying stock price.
+    """
+    underlying_price, _, _, _, _ = _get_atm_option_data(ticker, dte)
+    return underlying_price
 
 @tool
-def yf_get_atm_straddle_price(ticker: str, dte: int = 0) -> Dict[str, Any]:
+def yf_atm_strike(ticker: str, dte: int = 0) -> float:
     """
-    Calculate the price of the at-the-money (ATM) straddle for a given stock ticker,
-    including the implied volatility (IV) of the ATM call and put options. The ATM
-    strike is defined as the strike closest to the current underlying price. The
-    expiration date is selected based on the specified days-to-expiration (DTE).
+    Get the at-the-money (ATM) strike price for the given ticker.
 
     Args:
-        ticker (str): The stock ticker symbol, e.g., AAPL.
-        dte (int): Desired days-to-expiration. The tool selects the expiration date
-            closest to this value. A value of 0 selects the nearest available
-            expiration. Defaults to 0.
+        ticker (str): The stock ticker symbol.
+        dte (int): Desired days-to-expiration. Defaults to 0.
 
     Returns:
-        Dict[str, Any]: A dictionary containing:
-            - 'underlying_price' (float): The current stock price.
-            - 'expiration' (str): The selected option expiration date.
-            - 'strike' (float): The ATM strike price selected.
-            - 'call_price' (float): The last traded price of the ATM call option.
-            - 'put_price' (float): The last traded price of the ATM put option.
-            - 'call_iv' (float): The implied volatility of the ATM call option.
-            - 'put_iv' (float): The implied volatility of the ATM put option.
-            - 'straddle_price' (float): The total ATM straddle price (call + put).
+        float: The ATM strike price.
     """
-    stock = yf.Ticker(ticker)
-
-    # Get current underlying price
-    underlying_price = yf_get_ticker_price(stock)
+    _, _, atm_strike, _, _ = _get_atm_option_data(ticker, dte)
+    return atm_strike
 
-    # Get option expiration dates
-    expirations = stock.options
-    if not expirations:
-        return {"error": f"No options data available for {ticker}."}
-
-    # Convert expiration strings to datetime objects
-    exp_dates = [datetime.strptime(exp, "%Y-%m-%d") for exp in expirations]
-    today = datetime.now()
+@tool
+def yf_atm_call_price(ticker: str, dte: int = 0) -> float:
+    """
+    Get the last traded price of the ATM call option.
 
-    # Compute DTE for each expiration
-    dtes = [(exp - today).days for exp in exp_dates]
+    Args:
+        ticker (str): The stock ticker symbol.
+        dte (int): Desired days-to-expiration. Defaults to 0.
 
-    # Select expiration closest to requested DTE
-    target_idx = min(range(len(dtes)), key=lambda i: abs(dtes[i] - dte))
-    expiration = expirations[target_idx]
+    Returns:
+        float: The ATM call option price.
+    """
+    _, _, _, call_row, _ = _get_atm_option_data(ticker, dte)
+    return float(call_row.get("lastPrice", 0.0))
 
-    # Load option chain for selected expiration
-    opt_chain = stock.option_chain(expiration)
-    calls = opt_chain.calls
-    puts = opt_chain.puts
+@tool
+def yf_atm_put_price(ticker: str, dte: int = 0) -> float:
+    """
+    Get the last traded price of the ATM put option.
 
-    if calls.empty or puts.empty:
-        return {"error": f"Options chain incomplete for {ticker} on {expiration}."}
+    Args:
+        ticker (str): The stock ticker symbol.
+        dte (int): Desired days-to-expiration. Defaults to 0.
 
-    # Find ATM strike
-    calls["diff"] = (calls["strike"] - underlying_price).abs()
-    atm_strike = float(calls.sort_values("diff").iloc[0]["strike"])
+    Returns:
+        float: The ATM put option price.
+    """
+    _, _, _, _, put_row = _get_atm_option_data(ticker, dte)
+    return float(put_row.get("lastPrice", 0.0))
 
-    # Extract ATM call and put rows
-    call_row = calls[calls["strike"] == atm_strike].iloc[0]
-    put_row = puts[puts["strike"] == atm_strike].iloc[0]
+@tool
+def yf_atm_call_iv(ticker: str, dte: int = 0) -> float:
+    """
+    Get the implied volatility (IV) of the ATM call option.
 
-    call_price = float(call_row.get("lastPrice", 0.0))
-    put_price = float(put_row.get("lastPrice", 0.0))
+    Args:
+        ticker (str): The stock ticker symbol.
+        dte (int): Desired days-to-expiration. Defaults to 0.
 
-    call_iv = float(call_row.get("impliedVolatility", 0.0))
-    put_iv = float(put_row.get("impliedVolatility", 0.0))
+    Returns:
+        float: The implied volatility of the ATM call option.
+    """
+    _, _, _, call_row, _ = _get_atm_option_data(ticker, dte)
+    return float(call_row.get("impliedVolatility", 0.0))
 
-    straddle_price = call_price + put_price
+@tool
+def yf_atm_put_iv(ticker: str, dte: int = 0) -> float:
+    """
+    Get the implied volatility (IV) of the ATM put option.
 
-    return {
-        "underlying_price": underlying_price,
-        "expiration": expiration,
-        "strike": atm_strike,
-        "call_price": call_price,
-        "put_price": put_price,
-        "call_iv": call_iv,
-        "put_iv": put_iv,
-        "straddle_price": straddle_price,
-    }
+    Args:
+        ticker (str): The stock ticker symbol.
+        dte (int): Desired days-to-expiration. Defaults to 0.
 
+    Returns:
+        float: The implied volatility of the ATM put option.
+    """
+    _, _, _, _, put_row = _get_atm_option_data(ticker, dte)
+    return float(put_row.get("impliedVolatility", 0.0))
 
 @tool
-def estimate_expected_move(ticker:str, dte:int = 0) -> Dict[str, Any]:
+def yf_atm_straddle_price(ticker: str, dte: int = 0) -> float:
     """
-    Roughly estimates the expected move for a given stock ticker by adding/subtracting the price of the ATM straddle for
-        a given expiration to the underlying price.
+    Calculate the price of the ATM straddle (call + put).
 
     Args:
-        ticker (str): The stock ticker to estimate the expected move for.
-        dte (int): Days to expiration.
+        ticker (str): The stock ticker symbol.
+        dte (int): Desired days-to-expiration. Defaults to 0.
 
     Returns:
-        Dict[str, Any]: A dictionary containing:
-            - 'underlying_price' (float): The current stock price.
-            - 'em' (float): The expected move (price of the ATM straddle).
-            - 'em_lower' (float): The lower expected move.
-            - 'em_upper' (float): The upper expected move.
-    """
-    atm_straddle = yf_get_atm_straddle_price(ticker, dte)
-    underlying_price = yf_get_ticker_price(ticker)
-    em_upper = underlying_price + atm_straddle
-    em_lower = underlying_price - atm_straddle
-
-    return {
-        "underlying_price": underlying_price,
-        "em": atm_straddle,
-        "em_upper": em_upper,
-        "em_lower": em_lower,
-    }
+        float: The total ATM straddle price.
+    """
+    _, _, _, call_row, put_row = _get_atm_option_data(ticker, dte)
+    call_price = float(call_row.get("lastPrice", 0.0))
+    put_price = float(put_row.get("lastPrice", 0.0))
+    return call_price + put_price
+
 
 ddgs = DuckDuckGoSearchTool();
 
@@ -190,7 +171,7 @@ with open("prompts.yaml", 'r') as stream:
     
 agent = CodeAgent(
     model=model,
-    tools=[final_answer, ddgs, yf_get_ticker_price, yf_get_atm_straddle_price, estimate_expected_move], ## add your tools here (don't remove final answer)
+    tools=[final_answer, ddgs, yf_underlying_price, yf_atm_strike, yf_atm_call_price, yf_atm_put_price, yf_atm_call_iv, yf_atm_put_iv, yf_atm_straddle_price], ## add your tools here (don't remove final answer)
     max_steps=6,
     verbosity_level=1,
     grammar=None,