Update server.py

server.py

@@ -2,23 +2,24 @@
 import sys, json, asyncio
 import yfinance as yf
 from statsmodels.tsa.holtwinters import ExponentialSmoothing
-from langchain_core.tools import StructuredTool
 from pydantic import BaseModel
-from langchain_core.tools import BaseTool
-from langchain_core.tools import tool
+from langchain_core.tools import tool, BaseTool
 
-
-
-# --- Compact MCP Server Logic ---
+# --- MCP Server Logic ---
 class MCPToolServer:
     def __init__(self, tools: list[BaseTool]):
         self.tools = {t.name: t for t in tools}
 
     async def _handle_request(self, req: dict):
         method, params, req_id = req.get("method"), req.get("params", {}), req.get("id")
+        
         if method == "discover":
-            result = [{"name": t.name, "description": t.description, "args_schema": t.args} for t in self.tools.values()]
+            result = [
+                {"name": t.name, "description": t.description, "args_schema": t.args}
+                for t in self.tools.values()
+            ]
             return {"jsonrpc": "2.0", "result": result, "id": req_id}
+
         if method == "execute":
             tool_name, tool_args = params.get("tool_name"), params.get("tool_args", {})
             if tool_to_exec := self.tools.get(tool_name):
@@ -26,7 +27,12 @@ class MCPToolServer:
                     result = await tool_to_exec.ainvoke(tool_args)
                     return {"jsonrpc": "2.0", "result": result, "id": req_id}
                 except Exception as e:
-                    return {"jsonrpc": "2.0", "error": {"code": -32603, "message": str(e)}, "id": req_id}
+                    return {
+                        "jsonrpc": "2.0",
+                        "error": {"code": -32603, "message": str(e)},
+                        "id": req_id,
+                    }
+
         return {"jsonrpc": "2.0", "error": {"code": -32601, "message": "Method not found"}, "id": req_id}
 
     async def serve(self):
@@ -34,14 +40,17 @@ class MCPToolServer:
         await asyncio.get_event_loop().connect_read_pipe(lambda: asyncio.StreamReaderProtocol(reader), sys.stdin)
         writer_transport, _ = await asyncio.get_event_loop().connect_write_pipe(asyncio.streams.FlowControlMixin, sys.stdout)
         writer = asyncio.StreamWriter(writer_transport, _, None, asyncio.get_event_loop())
+
         while line := await reader.readline():
             try:
                 response = await self._handle_request(json.loads(line))
-                writer.write(json.dumps(response).encode() + b'\n')
+                writer.write(json.dumps(response).encode() + b"\n")
                 await writer.drain()
-            except json.JSONDecodeError: continue
+            except json.JSONDecodeError:
+                continue
 
-# --- Compact Tool Definitions ---
+
+# --- Tool Definitions ---
 COMMODITY_TICKERS = {"gold": "GC=F", "silver": "SI=F"}
 
 class PriceInput(BaseModel):
@@ -49,15 +58,15 @@ class PriceInput(BaseModel):
 
 @tool(args_schema=PriceInput)
 async def get_current_price(commodity_name: str) -> str:
-    """Gets the most recent 'live' price for gold or silver."""
+    """Gets the most recent price 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 ticker:
+        return f"Error: '{commodity_name}' is not supported."
     try:
         price = yf.Ticker(ticker).history(period="1d")['Close'].iloc[-1]
         return f"The current price of {commodity_name} is approx. ${price:.2f} USD."
-    except Exception as e: return f"Error fetching price: {e}"
-
-
+    except Exception as e:
+        return f"Error fetching price: {e}"
 
 class ForecastInput(BaseModel):
     commodity_name: str
@@ -65,12 +74,12 @@ class ForecastInput(BaseModel):
 
 @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."""
+    """Generates a 3 to 5 day price forecast."""
     ticker = COMMODITY_TICKERS.get(commodity_name.lower())
     if not ticker:
-        return f"Error: '{commodity_name}' is not supported. Use 'gold' or 'silver'."
+        return "Error: Unknown commodity. Use 'gold' or 'silver'."
     if not 3 <= forecast_days <= 5:
-        return "Error: Forecast must be for 3, 4, or 5 days."
+        return "Error: Forecast must be 3, 4, or 5 days."
     try:
         data = yf.download(ticker, period="6mo", progress=False)['Close']
         if data.empty:
@@ -80,7 +89,7 @@ async def get_price_forecast(commodity_name: str, forecast_days: int) -> str:
     except Exception as e:
         return f"Error during forecast: {e}"
 
-# --- Main Server Entrypoint ---
+# --- EntryPoint ---
 if __name__ == "__main__":
     server = MCPToolServer(tools=[get_current_price, get_price_forecast])
-    asyncio.run(server.serve())
+    asyncio.run(server.serve())