Update app.py

app.py

@@ -6,7 +6,6 @@ from stable_baselines3 import PPO
 from gym import Env
 from gym.spaces import Box
 import matplotlib.pyplot as plt
-import io
 
 class PortfolioEnv(Env):
     def __init__(self, returns, initial_balance=10000):
@@ -46,15 +45,16 @@ class PortfolioEnv(Env):
         return self.state
 
 def fetch_data(tickers, start_date, end_date):
-    data = yf.download(tickers, start=start_date, end=end_date)['Adj Close']
-    return data
+    try:
+        data = yf.download(tickers, start=start_date, end=end_date)['Adj Close']
+        return data
+    except Exception as e:
+        raise ValueError(f"Failed to fetch data: {e}")
 
 def optimize_portfolio(tickers, start_date, end_date, initial_balance):
-    # Fetch real-time data
     data = fetch_data(tickers, start_date, end_date)
     returns = data.pct_change().dropna()
 
-    # Define the environment
     env = PortfolioEnv(returns, initial_balance=initial_balance)
     model = PPO("MlpPolicy", env, verbose=0)
     model.learn(total_timesteps=5000)
@@ -76,21 +76,23 @@ def optimize_portfolio(tickers, start_date, end_date, initial_balance):
     plt.ylabel("Portfolio Value")
     plt.legend()
     plt.grid()
-    plt.savefig("portfolio_chart.png")
+
+    buffer = io.BytesIO()
+    plt.savefig(buffer, format="png")
+    buffer.seek(0)
     plt.close()
 
-    # Prepare the output
     weights = {f"Asset_{i + 1} ({tickers.split(',')[i].strip()})": float(weight) for i, weight in enumerate(portfolio_weights)}
-    return weights, "portfolio_chart.png"
+    return weights, buffer
 
 def run_optimization(tickers, start_date, end_date, initial_balance):
     if not tickers or not start_date or not end_date:
-        return "Error: Please fill all the fields.", None
+        return {"error": "Please fill all the fields."}, None
     try:
-        weights, chart_path = optimize_portfolio(tickers, start_date, end_date, initial_balance)
-        return weights, chart_path
+        weights, chart_buffer = optimize_portfolio(tickers, start_date, end_date, initial_balance)
+        return weights, chart_buffer
     except Exception as e:
-        return f"Error: {e}", None
+        return {"error": str(e)}, None
 
 interface = gr.Interface(
     fn=run_optimization,
@@ -105,11 +107,9 @@ interface = gr.Interface(
         gr.Image(label="Portfolio Value Chart"),
     ],
     title="AI-Powered Portfolio Optimization",
-    description="""
-    Enter stock tickers (e.g., AAPL, MSFT, TSLA), a date range, and your initial investment amount. 
-    The app fetches real-time historical data, runs AI optimization, and returns the optimized portfolio weights 
-    along with a performance chart.
-    """
+    description="""Enter stock tickers (e.g., AAPL, MSFT, TSLA), a date range, and your initial investment amount.
+                   The app fetches historical data, runs AI optimization, and returns the optimized portfolio weights 
+                   along with a performance chart."""
 )
 
 if __name__ == "__main__":