Update server.py

server.py

@@ -1,120 +1,69 @@
 # server.py
-import sys
-import json
-import asyncio
+import sys, json, asyncio
 import yfinance as yf
-import pandas as pd
 from statsmodels.tsa.holtwinters import ExponentialSmoothing
 from langchain_core.tools import tool, BaseTool
 
-# --- MCP Server Logic (Handles communication) ---
-class MCPMultiToolServer:
-    """A server that exposes multiple LangChain tools over JSON-RPC."""
+# --- Compact MCP Server Logic ---
+class MCPToolServer:
     def __init__(self, tools: list[BaseTool]):
-        self.tools = {tool.name: tool for tool in tools}
-
-    async def _handle_request(self, request: dict):
-        method = request.get("method")
-        params = request.get("params", {})
-        request_id = request.get("id")
+        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":
-            # Send back the schema for all available tools
-            tool_specs = [
-                {"name": t.name, "description": t.description, "args_schema": t.args}
-                for t in self.tools.values()
-            ]
-            return {"jsonrpc": "2.0", "result": tool_specs, "id": request_id}
-        
+            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 = params.get("tool_name")
-            tool_args = params.get("tool_args", {})
-            tool_to_execute = self.tools.get(tool_name)
-            if not tool_to_execute:
-                return {"jsonrpc": "2.0", "error": {"code": -32601, "message": "Method not found"}, "id": request_id}
-            
-            try:
-                # Execute the tool with the provided arguments
-                result = await tool_to_execute.ainvoke(tool_args)
-                return {"jsonrpc": "2.0", "result": result, "id": request_id}
-            except Exception as e:
-                return {"jsonrpc": "2.0", "error": {"code": -32603, "message": f"Internal error: {e}"}, "id": request_id}
-
-        return {"jsonrpc": "2.0", "error": {"code": -32601, "message": "Method not found"}, "id": request_id}
+            tool_name, tool_args = params.get("tool_name"), params.get("tool_args", {})
+            if tool_to_exec := self.tools.get(tool_name):
+                try:
+                    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": -32601, "message": "Method not found"}, "id": req_id}
 
     async def serve(self):
-        """Listens to stdin for JSON-RPC requests and sends responses to stdout."""
         reader = asyncio.StreamReader()
         await asyncio.get_event_loop().connect_read_pipe(lambda: asyncio.StreamReaderProtocol(reader), sys.stdin)
-        writer_transport, writer_protocol = await asyncio.get_event_loop().connect_write_pipe(asyncio.streams.FlowControlMixin, sys.stdout)
-        writer = asyncio.StreamWriter(writer_transport, writer_protocol, None, asyncio.get_event_loop())
-
-        while not reader.at_eof():
-            line = await reader.readline()
-            if not line: break
+        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:
-                request = json.loads(line)
-                response = await self._handle_request(request)
+                response = await self._handle_request(json.loads(line))
                 writer.write(json.dumps(response).encode() + b'\n')
                 await writer.drain()
-            except json.JSONDecodeError:
-                continue
+            except json.JSONDecodeError: continue
 
-# --- Tool Definitions (The actual work is done here) ---
+# --- Compact Tool Definitions ---
 COMMODITY_TICKERS = {"gold": "GC=F", "silver": "SI=F"}
 
 @tool
 async def get_current_price(commodity_name: str) -> str:
-    """
-    Gets the most recent 'live' price for a given commodity.
-    Use this for questions about the current or today's price of gold or silver.
-    """
-    commodity_name = commodity_name.lower()
-    ticker = COMMODITY_TICKERS.get(commodity_name)
-    if not ticker:
-        return f"Error: '{commodity_name}' is not supported. Please use 'gold' or 'silver'."
-    
+    """Gets the most recent 'live' 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'."
     try:
-        data = yf.Ticker(ticker).history(period="1d")
-        if data.empty: return "Could not retrieve price data."
-        price = data['Close'].iloc[-1]
-        return f"The current price of {commodity_name} is approximately ${price:.2f} USD."
-    except Exception as e:
-        return f"An error occurred while fetching the price: {e}"
+        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}"
 
 @tool
 async def get_price_forecast(commodity_name: str, forecast_days: int) -> str:
-    """
-    Generates a price forecast for a commodity for 3 to 5 days.
-    Use this for questions about forecasts, predictions, or future outlooks for gold or silver.
-    """
-    commodity_name = commodity_name.lower()
-    ticker = COMMODITY_TICKERS.get(commodity_name)
-    if not ticker:
-        return f"Error: '{commodity_name}' is not supported. Please use 'gold' or 'silver'."
-    if not 3 <= forecast_days <= 5:
-        return "Error: Forecast can only be for 3, 4, or 5 days."
-
+    """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."
     try:
-        data = yf.download(ticker, period="6mo", interval="1d", progress=False)
-        if data.empty: return "Not enough historical data to generate a forecast."
-
-        model = ExponentialSmoothing(data['Close'], trend='add').fit()
-        forecast = model.forecast(steps=forecast_days)
-        
-        response = f"The {forecast_days}-day price forecast for {commodity_name} is:\n"
-        for i, val in enumerate(forecast, 1):
-            response += f"Day {i}: ${val:.2f} USD\n"
-        return response.strip()
-    except Exception as e:
-        return f"An error occurred during forecasting: {e}"
-
-# --- Main entry point for the server process ---
-async def main():
-    """Starts the tool server."""
-    tools = [get_current_price, get_price_forecast]
-    server = MCPMultiToolServer(tools)
-    await server.serve()
+        data = yf.download(ticker, period="6mo", progress=False)['Close']
+        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}"
 
+# --- Main Server Entrypoint ---
 if __name__ == "__main__":
-    asyncio.run(main())
+    server = MCPToolServer(tools=[get_current_price, get_price_forecast])
+    asyncio.run(server.serve())