Add estimate_expected_move

app.py

@@ -47,24 +47,31 @@ def yf_get_ticker_price(ticker: str) -> float:
     Args:
         ticker: A string representing a stock ticker.
     """
-    data = yf.Ticker(ticker).history(period="1d")
-    return float(data["Close"].iloc[-1])
+    # Get current underlying price
+    hist = stock.history(period="1d")
+    if hist.empty:
+        return {"error": f"No price data found for {ticker}."}
+
+    return float(hist["Close"].iloc[-1])
 
 @tool
-def yf_get_atm_straddle_price(ticker: str) -> Dict[str, Any]:
+def yf_get_atm_straddle_price(ticker: str, dte: int = 0) -> Dict[str, Any]:
     """
     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, using
-    the nearest available expiration date.
+    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).
 
     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.
 
     Returns:
         Dict[str, Any]: A dictionary containing:
             - 'underlying_price' (float): The current stock price.
-            - 'expiration' (str): The nearest option expiration date used.
+            - '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.
@@ -75,21 +82,26 @@ def yf_get_atm_straddle_price(ticker: str) -> Dict[str, Any]:
     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}."}
-
-    underlying_price = float(hist["Close"].iloc[-1])
+    underlying_price = yf_get_ticker_price(ticker)
 
     # Get option expiration dates
     expirations = stock.options
     if not expirations:
         return {"error": f"No options data available for {ticker}."}
 
-    # Use the nearest expiration
-    expiration = expirations[0]
-    opt_chain = stock.option_chain(expiration)
+    # Convert expiration strings to datetime objects
+    exp_dates = [datetime.strptime(exp, "%Y-%m-%d") for exp in expirations]
+    today = datetime.now()
+
+    # Compute DTE for each expiration
+    dtes = [(exp - today).days for exp in exp_dates]
 
+    # Select expiration closest to requested DTE
+    target_idx = min(range(len(dtes)), key=lambda i: abs(dtes[i] - dte))
+    expiration = expirations[target_idx]
+
+    # Load option chain for selected expiration
+    opt_chain = stock.option_chain(expiration)
     calls = opt_chain.calls
     puts = opt_chain.puts
 
@@ -123,6 +135,36 @@ def yf_get_atm_straddle_price(ticker: str) -> Dict[str, Any]:
         "straddle_price": straddle_price,
     }
 
+
+@tool
+def estimate_expected_move(ticker:str, dte:int = 0) -> Dict[str, Any]:
+    """
+    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.
+
+    Args:
+        ticker (str): The stock ticker to estimate the expected move for.
+        dte (int): Days to expiration.
+
+    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,
+    }
+
 ddgs = DuckDuckGoSearchTool();
 
 final_answer = FinalAnswerTool()
@@ -146,7 +188,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], ## add your tools here (don't remove final answer)
+    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)
     max_steps=6,
     verbosity_level=1,
     grammar=None,