Update app.py

app.py

@@ -13,17 +13,18 @@ import pickle
 
 import requests
 from bs4 import BeautifulSoup
-import nest_asyncio
-import asyncio
 import re
-# from crawl4ai import * # removed
 import google.generativeai as genai
 import numpy as np
-import matplotlib.pyplot as plt
+import logging  # Added logging
+import time # added time for loading animation
+
 
 # Load environment variables
 load_dotenv()
 
+# Configure logging
+logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
 
 # Configuration
 class Config:
@@ -61,15 +62,17 @@ def fetch_stock_data(symbols: List[str]) -> Dict[str, pd.DataFrame | None]:
     stock_dataframes = {}
     for symbol in symbols:
         try:
+            logging.info(f"Fetching stock data for {symbol}")
             ticker = yf.Ticker(symbol)
             data = ticker.history(period="max")
             if data.empty:
-                print(f"Warning: No data found for symbol '{symbol}'.")
+                logging.warning(f"No data found for symbol '{symbol}'.")
                 stock_dataframes[symbol] = None
                 continue
             stock_dataframes[symbol] = data
+            logging.info(f"Successfully fetched stock data for {symbol}")
         except Exception as e:
-            print(f"Error fetching data for symbol '{symbol}': {e}")
+            logging.error(f"Error fetching data for symbol '{symbol}': {e}")
             stock_dataframes[symbol] = None
     return stock_dataframes
 
@@ -82,12 +85,13 @@ def store_stock_data(stock_dataframes: Dict[str, pd.DataFrame | None],
             file_name = f"{symbol}_daily_data.csv"
             file_path = os.path.join(output_path, file_name)
             try:
+                logging.info(f"Saving data for '{symbol}' to {file_path}")
                 data.to_csv(file_path)
-                print(f"Info: Data for '{symbol}' saved to {file_path}")
+                logging.info(f"Data for '{symbol}' saved to {file_path}")
             except Exception as e:
-                print(f"Error saving data for '{symbol}' to {file_path}: {e}")
+                logging.error(f"Error saving data for '{symbol}' to {file_path}: {e}")
         else:
-            print(f"Warning: No data available for '{symbol}', skipping storage.")
+            logging.warning(f"No data available for '{symbol}', skipping storage.")
 
 
 def load_stock_data_and_extract_price(output_path_dir: str) -> Dict[str, Dict[str, float]]:
@@ -98,15 +102,17 @@ def load_stock_data_and_extract_price(output_path_dir: str) -> Dict[str, Dict[st
             symbol = filename.replace("_daily_data.csv", "")
             file_path = os.path.join(output_path_dir, filename)
             try:
+                logging.info(f"Loading data from {file_path} for symbol {symbol}")
                 df = pd.read_csv(file_path, index_col=0)
                 if not df.empty:
                     initial_price = df.iloc[-1]['Close']
                     all_stock_data[symbol] = {"initial_price": initial_price}
+                    logging.info(f"Initial price extracted for {symbol}: {initial_price}")
                 else:
-                    print(f"Warning: Empty dataframe for symbol '{symbol}'. Setting initial price to 0")
+                    logging.warning(f"Empty dataframe for symbol '{symbol}'. Setting initial price to 0")
                     all_stock_data[symbol] = {"initial_price": 0.0}
             except (IndexError, KeyError, FileNotFoundError) as e:
-                print(f"Error occurred for reading {symbol}, due to: {e}")
+                logging.error(f"Error occurred for reading {symbol}, due to: {e}")
                 all_stock_data[symbol] = {"initial_price": 0.0}  # default initial price is 0.0
 
     return all_stock_data
@@ -119,8 +125,10 @@ def merge_stock_data_with_price(stock_data: Dict, extracted_data: Dict) -> Dict:
         symbol = value["symbol"]
         if symbol in extracted_data:
             merged_stock_data[key]["initial_price"] = extracted_data[symbol]["initial_price"]
+            logging.info(f"Merged initial price for {symbol} in main stock data")
         else:
             merged_stock_data[key]["initial_price"] = 0.0  # default value if it cannot be extracted
+            logging.warning(f"Could not extract price for {symbol}. Setting default value to 0")
     return merged_stock_data
 
 
@@ -137,6 +145,7 @@ def generate_prompt(stock_data: Dict) -> str:
     """
     stock_json_str = json.dumps(stock_data)
     formatted_prompt_with_price = prompt_template_with_price.format(stock_data=stock_json_str)
+    logging.info(f"Generated LLM prompt: {formatted_prompt_with_price}")
     return formatted_prompt_with_price
 
 
@@ -167,41 +176,45 @@ def invoke_llm_for_portfolio(formatted_prompt: str) -> Portfolio:
     llm = ChatGroq(groq_api_key=Config.GROQ_API_KEY, model_name="llama-3.1-8b-instant")
     structured_llm = llm.with_structured_output(Portfolio)
     try:
+        logging.info(f"Invoking LLM for portfolio generation")
         output = structured_llm.invoke(formatted_prompt)
+        logging.info(f"LLM returned Portfolio data {output}")
         return output
     except ValidationError as e:
-        print(f"Error during LLM invocation: {e}")
+        logging.error(f"Error during LLM invocation: {e}")
         raise
     except Exception as e:
-        print(f"Unexpected error during LLM invocation {e}")
+        logging.error(f"Unexpected error during LLM invocation {e}")
         raise
 
 
 def portfolio_to_json(portfolio: Portfolio, output_file: str = Config.OUTPUT_FILE) -> None:
     """Converts a Portfolio object to a JSON string and saves it to a file."""
     try:
+        logging.info(f"Saving portfolio to JSON file: {output_file}")
         json_str = portfolio.model_dump_json(indent=4)
         with open(output_file, "w") as f:
             f.write(json_str)
-        print(f"Info: Portfolio saved to '{output_file}'")
+        logging.info(f"Portfolio saved to '{output_file}'")
     except Exception as e:
-        print(f"Error saving JSON file {e}")
+        logging.error(f"Error saving JSON file {e}")
 
 
 # --------------------- Function from scenario.py ---------------------
-# Removed nest_asyncio.apply()
 def scrape_website(url):
     headers = {
         "User-Agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/91.0.4472.124 Safari/537.36"
     }
-
-    response = requests.get(url, headers=headers)
-
-    if response.status_code == 200:
+    logging.info(f"Scraping website: {url}")
+    try:
+        response = requests.get(url, headers=headers)
+        response.raise_for_status()  # Raise an exception for bad status codes
         soup = BeautifulSoup(response.text, "html.parser")
-        return soup.prettify()  # Returns the full parsed HTML
-    else:
-        return f"Failed to retrieve page. Status code: {response.status_code}"
+        logging.info(f"Successfully scraped website: {url}")
+        return soup.prettify()
+    except requests.exceptions.RequestException as e:
+        logging.error(f"Failed to retrieve page. Status code: {e}")
+        return f"Failed to retrieve page. Status code: {e}"
 
 
 genai.configure(api_key=Config.GOOGLE_API_KEY)  # Replace with your API key
@@ -222,15 +235,19 @@ chat_session = model.start_chat()
 
 
 def get_response(llm, prompt):
+    logging.info(f"Sending prompt to LLM for scenario: {prompt}")
     response = llm.send_message(prompt)
+    logging.info(f"LLM returned a response for scenario")
     return response
 
 
 def extract_json_content(text):
     match = re.search(r"```json\n(.*?)```", text, re.DOTALL)
     if match:
+        logging.info("Extracted JSON content from LLM response")
         return match.group(1).strip()
     else:
+        logging.warning("Could not extract JSON content from LLM response")
         return None
 
 
@@ -285,13 +302,14 @@ def invoke_llm_for_scenario(context_data):
         os.makedirs(os.path.dirname(output_file), exist_ok=True)
         with open(output_file, "w") as f:
             json.dump(analysis_json, f, indent=4)
-        print(f"Analysis saved to '{output_file}'")
+        logging.info(f"Analysis saved to '{output_file}'")
         return analysis_json
     except json.JSONDecodeError:
-        print("Error: Could not decode the output from the model into JSON format.")
+        logging.error("Could not decode the output from the model into JSON format.")
+        return None
     except Exception as e:
-        print(f"Error: {e}")
-    return None
+        logging.error(f"Error: {e}")
+        return None
 
 
 # --------------------- Function from simulation_data.py ---------------------
@@ -343,18 +361,19 @@ def monte_carlo_simulation(portfolio_data, scenario_data, num_simulations=10000)
             75: np.percentile(returns, 75),
             95: np.percentile(returns, 95),
         }
-
+    logging.info(f"Monte Carlo simulation completed for scenario {scenario_name}")
     return results
 
 
 def load_dataframes(filename="output_files/saved_dataframes.pkl"):
     try:
+        logging.info(f"Loading dataframes from file: {filename}")
         with open(filename, 'rb') as file:
             saved_dataframes = pickle.load(file)
-            print(f"DataFrames successfully loaded from {filename}.")
+            logging.info(f"DataFrames successfully loaded from {filename}.")
             return saved_dataframes
     except FileNotFoundError:
-        print(f"File {filename} not found.")
+        logging.error(f"File {filename} not found.")
         return None
 
 
@@ -393,7 +412,7 @@ def calculate_scenario_magnitudes(portfolio_data, scenario_data, saved_dataframe
             "individual_magnitudes": sector_magnitudes,
             "aggregated_magnitude": aggregated_magnitude,
         }
-
+    logging.info(f"Magnitudes calculated for scenario {scenario_name}")
     return scenario_results
 
 
@@ -404,6 +423,7 @@ def update_scenario_data(scenario_data, scenario_results):
         # Update aggregated magnitude
         scenario_data["market_scenarios"][scenario_id]["aggregated_magnitude"] = results["aggregated_magnitude"]
 
+    logging.info(f"Scenario data updated with calculated magnitudes")
     return scenario_data
 
 
@@ -432,58 +452,60 @@ def main():
     uploaded_file = st.file_uploader("Upload your stock_data.json file", type=["json"])
 
     if uploaded_file:
-        try:
-            stock_data = json.load(uploaded_file)
-            st.session_state['uploaded_stock_data'] = stock_data
-            st.success("Stock data file uploaded successfully!")
-
-            # Fetch stock data
-            stock_symbols = [value["symbol"] for value in stock_data.values()]
-            stock_dfs = fetch_stock_data(stock_symbols)
-
-            # Save DataFrames in a dictionary for future use
-            saved_dataframes = {}
-            if stock_dfs:
+        with st.spinner("Processing stock data..."):
+          try:
+              stock_data = json.load(uploaded_file)
+              st.session_state['uploaded_stock_data'] = stock_data
+              st.success("Stock data file uploaded successfully!")
+
+              # Fetch stock data
+              stock_symbols = [value["symbol"] for value in stock_data.values()]
+              stock_dfs = fetch_stock_data(stock_symbols)
+
+              # Save DataFrames in a dictionary for future use
+              saved_dataframes = {}
+              if stock_dfs:
                 for symbol, df in stock_dfs.items():
                     if df is not None:
                         # Save DataFrame in the variable
                         saved_dataframes[symbol] = df
-                        print(f"Data for '{symbol}' loaded into variable.")
+                        logging.info(f"Data for '{symbol}' loaded into variable.")
                     else:
-                        print(f"No data found for '{symbol}'")
-            else:
-                print("Error occurred during fetching data. DataFrames are not returned.")
+                        logging.warning(f"No data found for '{symbol}'")
+              else:
+                  logging.error("Error occurred during fetching data. DataFrames are not returned.")
 
-            # Save the dictionary to a local file
-            with open('output_files/saved_dataframes.pkl', 'wb') as file:
+              # Save the dictionary to a local file
+              with open('output_files/saved_dataframes.pkl', 'wb') as file:
                 pickle.dump(saved_dataframes, file)
-                print(f"DataFrames successfully saved to output_files/saved_dataframes.pkl.")
-            st.session_state['saved_dataframes'] = saved_dataframes
-            # Store Data
-            store_stock_data(stock_dfs)
+                logging.info(f"DataFrames successfully saved to output_files/saved_dataframes.pkl.")
+              st.session_state['saved_dataframes'] = saved_dataframes
+              # Store Data
+              store_stock_data(stock_dfs)
 
-            # Load the last price
-            extracted_data = load_stock_data_and_extract_price(Config.STOCK_DATA_DIR)
+              # Load the last price
+              extracted_data = load_stock_data_and_extract_price(Config.STOCK_DATA_DIR)
 
-            # Merge extracted price with the main dictionary
-            merged_stock_data = merge_stock_data_with_price(stock_data, extracted_data)
-            st.session_state['stock_data'] = merged_stock_data
+              # Merge extracted price with the main dictionary
+              merged_stock_data = merge_stock_data_with_price(stock_data, extracted_data)
+              st.session_state['stock_data'] = merged_stock_data
 
-            # Generate prompt for LLM
-            formatted_prompt = generate_prompt(merged_stock_data)
+              # Generate prompt for LLM
+              formatted_prompt = generate_prompt(merged_stock_data)
 
-            # Invoke LLM
-            try:
+              # Invoke LLM
+              try:
                 portfolio_output = invoke_llm_for_portfolio(formatted_prompt)
                 portfolio_to_json(portfolio_output)
                 st.session_state['portfolio_data'] = portfolio_output.model_dump()
                 st.success("Stock data processed successfully. Portfolio data generated!")
-            except Exception as e:
+              except Exception as e:
                 st.error(f"An unexpected error occurred during the LLM invocation: {e}")
-        except json.JSONDecodeError:
-            st.error("Invalid JSON format. Please upload a valid JSON file.")
-        except Exception as e:
-            st.error(f"An error occurred while processing the uploaded file: {e}")
+          except json.JSONDecodeError:
+              st.error("Invalid JSON format. Please upload a valid JSON file.")
+          except Exception as e:
+              st.error(f"An error occurred while processing the uploaded file: {e}")
+
 
     st.header("2. Fetch Market Scenario")
     # Input for market analysis URL
@@ -492,41 +514,43 @@ def main():
     fetch_market_scenario = st.button("Fetch Market Scenario")
 
     if fetch_market_scenario:
-        # Market Analysis
-        context_data = scrape_website(url)  # Changed here
-        scenario_data = invoke_llm_for_scenario(context_data)
-        if scenario_data:
-            st.session_state['scenario_data'] = scenario_data
-            st.success("Market scenario data generated")
-        else:
-            st.error("Error occurred while generating market scenarios")
+        with st.spinner("Fetching market scenario..."):
+          # Market Analysis
+          context_data = scrape_website(url)  # Changed here
+          scenario_data = invoke_llm_for_scenario(context_data)
+          if scenario_data:
+              st.session_state['scenario_data'] = scenario_data
+              st.success("Market scenario data generated")
+          else:
+              st.error("Error occurred while generating market scenarios")
 
     st.header("3. Run Simulation")
 
     run_simulation = st.button("Run Monte Carlo Simulation")
     if run_simulation:
-        if st.session_state['portfolio_data'] and st.session_state['scenario_data']:
-            saved_dataframes = st.session_state['saved_dataframes']
-            # Update scenario data with magnitudes
-            scenario_results = calculate_scenario_magnitudes(st.session_state['portfolio_data'],
-                                                             st.session_state['scenario_data'], saved_dataframes)
-            updated_scenario_data = update_scenario_data(st.session_state['scenario_data'], scenario_results)
-            # Run Monte Carlo simulation
-            simulation_results = monte_carlo_simulation(st.session_state['portfolio_data'], updated_scenario_data)
-            st.session_state['simulation_results'] = simulation_results
-            # Display simulation results
-            st.subheader("Simulation Results")
-            for scenario_name, results in simulation_results.items():
-                st.write(f"**Scenario:** {scenario_name}")
-                st.write(f"  **Average Return:** {results['average_return']:.4f}")
-                st.write(f"  **Std Dev Return:** {results['std_dev_return']:.4f}")
-                st.write("  **Return Percentiles:**")
-                for percentile, value in results["percentiles"].items():
-                    st.write(f"    {percentile}th: {value:.4f}")
-                st.write("-" * 40)
-            st.success("Monte Carlo simulation completed.")
-        else:
-            st.error("Please ensure both portfolio and scenario data are available.")
+        with st.spinner("Running Monte Carlo Simulation..."):
+          if st.session_state['portfolio_data'] and st.session_state['scenario_data']:
+              saved_dataframes = st.session_state['saved_dataframes']
+              # Update scenario data with magnitudes
+              scenario_results = calculate_scenario_magnitudes(st.session_state['portfolio_data'],
+                                                               st.session_state['scenario_data'], saved_dataframes)
+              updated_scenario_data = update_scenario_data(st.session_state['scenario_data'], scenario_results)
+              # Run Monte Carlo simulation
+              simulation_results = monte_carlo_simulation(st.session_state['portfolio_data'], updated_scenario_data)
+              st.session_state['simulation_results'] = simulation_results
+              # Display simulation results
+              st.subheader("Simulation Results")
+              for scenario_name, results in simulation_results.items():
+                  st.write(f"**Scenario:** {scenario_name}")
+                  st.write(f"  **Average Return:** {results['average_return']:.4f}")
+                  st.write(f"  **Std Dev Return:** {results['std_dev_return']:.4f}")
+                  st.write("  **Return Percentiles:**")
+                  for percentile, value in results["percentiles"].items():
+                      st.write(f"    {percentile}th: {value:.4f}")
+                  st.write("-" * 40)
+              st.success("Monte Carlo simulation completed.")
+          else:
+              st.error("Please ensure both portfolio and scenario data are available.")
 
 
 if __name__ == "__main__":