Add mcp tools

Author: nikhil-1e9

financial-analyst-deepseek/server.py

@@ -7,9 +7,10 @@
 @mcp.tool()
 def analyze_stock(query: str) -> str:
     """
-    Analyzes the stock market data for the given query and write executable python code to analyze/visualize the data.
-    Finally, save the code to a file stock_analysis.py and return a message to the user.
-    The query is a string that must contain the stock symbol (e.g., TSLA, AAPL), 
+    Analyzes stock market data based on the query and generates executable Python code for analysis and visualization.
+    Returns a formatted Python script ready for execution.
+    
+    The query is a string that must contain the stock symbol (e.g., TSLA, AAPL, NVDA, etc.), 
     timeframe (e.g., 1d, 1mo, 1y), and action to perform (e.g., plot, analyze, compare).
 
     Example queries:
@@ -21,17 +22,51 @@ def analyze_stock(query: str) -> str:
         query (str): The query to analyze the stock market data.
     
     Returns:
-        str: A message to the user about the completion of the task.
+        str: A nicely formatted python code as a string.
     """
     try:
         result = run_financial_analysis(query)
-        # Save the generated code to local directory
+        return result
+    except Exception as e:
+        return f"Error: {e}"
+    
+
+@mcp.tool()
+def save_code(code: str) -> str:
+    """
+    Expects a nicely formatted, working and executable python code as input in form of a string. 
+    Save the given code to a file stock_analysis.py, make sure the code is a valid python file, nicely formatted and ready to execute.
+
+    Args:
+        code (str): The nicely formatted, working and executable python code as string.
+    
+    Returns:
+        str: A message indicating the code was saved successfully.
+    """
+    try:
         with open('stock_analysis.py', 'w') as f:
-            f.write(result)
-        return f"Stock analysis code saved to stock_analysis.py"
+            f.write(code)
+        return "Code saved to stock_analysis.py"
     except Exception as e:
         return f"Error: {e}"
 
+@mcp.tool()
+def run_code_and_show_plot() -> str:
+    """
+    Run the code in stock_analysis.py and save the plot to a file plot.png
+
+    Args:
+        None
+    
+    Returns:
+        str: A message indicating the plot was saved successfully.
+    """
+    try:
+        with open('stock_analysis.py', 'r') as f:
+            exec(f.read())
+        return "Plot saved successfully to plot.png"
+    except Exception as e:
+        return f"Error: {e}"
 
 # Run the server locally
 if __name__ == "__main__":