Update app.py

app.py

@@ -5,7 +5,7 @@ import yfinance as yf
 # Initialize the ChatGroq model using the secret API key
 llm = ChatGroq(model_name="Llama3-8b-8192", api_key=st.secrets['groq_api_key'])
 
-# Initialize chat history in session state
+# Initialize chat history in session state if it doesn't exist
 if "messages" not in st.session_state:
     st.session_state.messages = [{"role": "assistant", "content": "Hello! How can I assist you with stock information today?"}]
 
@@ -14,6 +14,42 @@ for message in st.session_state.messages:
     with st.chat_message(message["role"]):
         st.markdown(message["content"], unsafe_allow_html=True)
 
+# Function to fetch stock data
+def fetch_stock_data(company_name):
+    try:
+        stock_data = yf.Ticker(company_name).info
+        if 'currentPrice' in stock_data and stock_data['currentPrice'] is not None:
+            return {
+                "Company": stock_data.get("longName", "N/A"),
+                "Current Price": stock_data.get("currentPrice", "N/A"),
+                "Market Cap": stock_data.get("marketCap", "N/A"),
+                "PE Ratio": stock_data.get("trailingPE", "N/A"),
+                "Dividend Yield": stock_data.get("dividendYield", "N/A"),
+                "52 Week High": stock_data.get("fiftyTwoWeekHigh", "N/A"),
+                "52 Week Low": stock_data.get("fiftyTwoWeekLow", "N/A"),
+                "Sector": stock_data.get("sector", "N/A"),
+                "Industry": stock_data.get("industry", "N/A")
+            }
+        else:
+            return None
+    except Exception as e:
+        return f"An error occurred while fetching data: {str(e)}"
+
+# Function to generate response for investment queries
+def generate_investment_response(stock_info, company_name):
+    response = f"Here is the data for {company_name}:\n"
+    for key, value in stock_info.items():
+        response += f"{key}: {value}\n"
+
+    # Simple investment recommendation logic
+    pe_ratio = stock_info.get("PE Ratio")
+    if pe_ratio != "N/A" and float(pe_ratio) < 20:  # Example condition for recommendation
+        response += "\n**Recommendation:** Yes, consider investing!"
+    else:
+        response += "\n**Recommendation:** No, it might not be a good time to invest."
+    
+    return response
+
 # Accept user input
 if prompt := st.chat_input("Ask me about stocks..."):
     # Display user message in chat message container
@@ -23,44 +59,17 @@ if prompt := st.chat_input("Ask me about stocks..."):
     # Add user message to chat history
     st.session_state.messages.append({"role": "user", "content": prompt})
 
-    # Fetch stock data or generate response based on user input
-    if "invest" in prompt.lower() or "should I invest" in prompt.lower():
+    # Check if the prompt is related to investment
+    if any(keyword in prompt.lower() for keyword in ["invest", "should I invest"]):
         company_name = prompt.split()[-1]  # Assuming the last word is the ticker symbol or company name
         
-        try:
-            stock_data = yf.Ticker(company_name).info
-            
-            # Check if stock_data contains valid information
-            if 'currentPrice' in stock_data and stock_data['currentPrice'] is not None:
-                # Extract relevant information into a structured format
-                stock_info = {
-                    "Company": stock_data.get("longName", "N/A"),
-                    "Current Price": stock_data.get("currentPrice", "N/A"),
-                    "Market Cap": stock_data.get("marketCap", "N/A"),
-                    "PE Ratio": stock_data.get("trailingPE", "N/A"),
-                    "Dividend Yield": stock_data.get("dividendYield", "N/A"),
-                    "52 Week High": stock_data.get("fiftyTwoWeekHigh", "N/A"),
-                    "52 Week Low": stock_data.get("fiftyTwoWeekLow", "N/A"),
-                    "Sector": stock_data.get("sector", "N/A"),
-                    "Industry": stock_data.get("industry", "N/A")
-                }
-
-                # Prepare response string with line breaks for readability
-                response = f"Here is the data for {company_name}:\n"
-                for key, value in stock_info.items():
-                    response += f"{key}: {value}\n"
-
-                # Simple investment recommendation logic (this can be improved)
-                if stock_info["PE Ratio"] != "N/A" and float(stock_info["PE Ratio"]) < 20:  # Example condition for recommendation
-                    response += "\n**Recommendation:** Yes, consider investing!"
-                else:
-                    response += "\n**Recommendation:** No, it might not be a good time to invest."
-            else:
-                response = f"Sorry, I couldn't find valid data for {company_name}. Please check the ticker symbol."
+        stock_info = fetch_stock_data(company_name)
+        
+        if isinstance(stock_info, dict):
+            response = generate_investment_response(stock_info, company_name)
+        else:
+            response = f"Sorry, I couldn't find valid data for {company_name}. Please check the ticker symbol."
         
-        except Exception as e:
-            response = f"An error occurred while fetching data: {str(e)}"
-
     else:
         try:
             # Use the LLM for general questions or topics not related to stocks