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 if it doesn't exist
+# Initialize chat history in session state
 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,42 +14,6 @@ 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
@@ -59,17 +23,44 @@ if prompt := st.chat_input("Ask me about stocks..."):
     # Add user message to chat history
     st.session_state.messages.append({"role": "user", "content": prompt})
 
-    # Check if the prompt is related to investment
-    if any(keyword in prompt.lower() for keyword in ["invest", "should I invest"]):
+    # Fetch stock data or generate response based on user input
+    if "invest" in prompt.lower() or "should I invest" in prompt.lower():
         company_name = prompt.split()[-1]  # Assuming the last word is the ticker symbol or company name
         
-        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."
+        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."
         
+        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