Update server.py

server.py

@@ -2,7 +2,8 @@
 import sys, json, asyncio
 import yfinance as yf
 from statsmodels.tsa.holtwinters import ExponentialSmoothing
-from langchain_core.tools import tool, BaseTool
+from langchain_core.tools import StructuredTool
+from pydantic import BaseModel
 
 # --- Compact MCP Server Logic ---
 class MCPToolServer:
@@ -39,7 +40,10 @@ class MCPToolServer:
 # --- Compact Tool Definitions ---
 COMMODITY_TICKERS = {"gold": "GC=F", "silver": "SI=F"}
 
-@tool
+class PriceInput(BaseModel):
+    commodity_name: str
+
+@tool(args_schema=PriceInput)
 async def get_current_price(commodity_name: str) -> str:
     """Gets the most recent 'live' price for gold or silver."""
     ticker = COMMODITY_TICKERS.get(commodity_name.lower())
@@ -49,19 +53,28 @@ async def get_current_price(commodity_name: str) -> str:
         return f"The current price of {commodity_name} is approx. ${price:.2f} USD."
     except Exception as e: return f"Error fetching price: {e}"
 
-@tool
+
+
+class ForecastInput(BaseModel):
+    commodity_name: str
+    forecast_days: int
+
+@tool(args_schema=ForecastInput)
 async def get_price_forecast(commodity_name: str, forecast_days: int) -> str:
     """Generates a 3 to 5 day price forecast for gold or silver."""
     ticker = COMMODITY_TICKERS.get(commodity_name.lower())
-    if not ticker: return f"Error: '{commodity_name}' is not supported. Use 'gold' or 'silver'."
-    if not 3 <= forecast_days <= 5: return "Error: Forecast must be for 3, 4, or 5 days."
+    if not ticker:
+        return f"Error: '{commodity_name}' is not supported. Use 'gold' or 'silver'."
+    if not 3 <= forecast_days <= 5:
+        return "Error: Forecast must be for 3, 4, or 5 days."
     try:
         data = yf.download(ticker, period="6mo", progress=False)['Close']
-        if data.empty: return "Not enough data to forecast."
+        if data.empty:
+            return "Not enough data to forecast."
         forecast = ExponentialSmoothing(data, trend='add').fit().forecast(steps=forecast_days)
-        forecast_lines = [f"Day {i}: ${val:.2f} USD" for i, val in enumerate(forecast, 1)]
-        return f"The {forecast_days}-day forecast for {commodity_name} is:\n" + "\n".join(forecast_lines)
-    except Exception as e: return f"Error during forecast: {e}"
+        return "\n".join([f"Day {i+1}: ${val:.2f} USD" for i, val in enumerate(forecast)])
+    except Exception as e:
+        return f"Error during forecast: {e}"
 
 # --- Main Server Entrypoint ---
 if __name__ == "__main__":