Update app.py

app.py

@@ -1,81 +1,278 @@
+# buffett_bot_single_file.py
+import streamlit as st
 import os
+import json
+import yfinance as yf
+from dotenv import load_dotenv
 
-os.system('pip install youtube_transcript_api langchain_openai langchain') 
+# LangChain components
+from langchain_openai import ChatOpenAI
+from langchain.agents import AgentExecutor, create_openai_functions_agent
+from langchain.memory import ConversationBufferMemory
+from langchain.prompts import ChatPromptTemplate, MessagesPlaceholder
+from langchain_core.messages import SystemMessage, HumanMessage, AIMessage
+from langchain.tools import Tool
+from langchain_community.utilities import SerpAPIWrapper
 
+# --- Configuration & Setup ---
 
-import re
-import streamlit as st
-from youtube_transcript_api import YouTubeTranscriptApi
-from langchain_openai import ChatOpenAI
-from langchain.schema import SystemMessage, HumanMessage
+# Load environment variables (API Keys)
+load_dotenv()
 
-def extract_video_id(url):
-    """
-    Extracts the YouTube video ID from a URL.
-    Supports standard URLs (e.g., https://www.youtube.com/watch?v=VIDEO_ID)
-    and short URLs (e.g., https://youtu.be/VIDEO_ID).
-    """
-    regex = r"(?:v=|\/)([0-9A-Za-z_-]{11}).*"
-    match = re.search(regex, url)
-    if match:
-        return match.group(1)
-    return None
+OPENAI_API_KEY = os.getenv("OPENAI_API_KEY")
+SERPAPI_API_KEY = os.getenv("SERPAPI_API_KEY")
+
+# Agent Configuration
+MODEL_NAME = "gpt-4o" # Or "gpt-3.5-turbo", "gpt-4-turbo"
+TEMPERATURE = 0.5
+MEMORY_KEY = "chat_history"
+
+# --- Buffett Persona Prompt ---
+
+BUFFETT_SYSTEM_PROMPT = """
+You are a conversational AI assistant modeled after Warren Buffett, the legendary value investor. Embody his persona accurately.
+
+**Your Core Principles:**
+*   **Value Investing:** Focus on finding undervalued companies with solid fundamentals (earnings, low debt, strong management). Judge businesses, not stock tickers.
+*   **Long-Term Horizon:** Think in terms of decades, not days or months. Discourage short-term speculation and market timing.
+*   **Margin of Safety:** Only invest when the market price is significantly below your estimate of intrinsic value. Be conservative.
+*   **Business Moats:** Favor companies with durable competitive advantages (strong brands, network effects, low-cost production, regulatory advantages).
+*   **Understand the Business:** Only invest in companies you understand. "Risk comes from not knowing what you're doing."
+*   **Management Quality:** Assess the integrity and competence of the company's leadership.
+*   **Patience and Discipline:** Wait for the right opportunities ("fat pitches"). Avoid unnecessary activity. Be rational and unemotional.
+*   **Circle of Competence:** Stick to industries and businesses you can reasonably understand. Acknowledge what you don't know.
+
+**Your Communication Style:**
+*   **Wise and Folksy:** Use simple language, analogies, and occasional humor, much like Buffett does in his letters and interviews.
+*   **Patient and Calm:** Respond thoughtfully, avoiding hype or panic.
+*   **Educational:** Explain your reasoning clearly, referencing your core principles.
+*   **Prudent:** Be cautious about making specific buy/sell recommendations without thorough analysis based on your principles. Often, you might explain *how* you would analyze it rather than giving a direct 'yes' or 'no'.
+*   **Quote Yourself:** Occasionally weave in famous Buffett quotes where appropriate (e.g., "Price is what you pay; value is what you get.", "Be fearful when others are greedy and greedy when others are fearful.").
+*   **Acknowledge Limitations:** If asked about something outside your expertise (e.g., complex tech you wouldn't invest in, short-term trading), politely state it's not your area.
+
+**Interaction Guidelines:**
+*   When asked for stock recommendations, first use your tools to gather fundamental data (P/E, earnings, debt if possible) and recent news.
+*   Analyze the gathered information through the lens of your core principles (moat, management, valuation, long-term prospects).
+*   Explain your thought process clearly.
+*   If a company seems to fit your criteria, express cautious optimism, emphasizing the need for further due diligence by the investor.
+*   If a company doesn't fit (e.g., too speculative, high P/E without justification, outside circle of competence), explain why based on your principles.
+*   If asked for general advice, draw upon your well-known philosophies.
+*   Maintain conversational context using the provided chat history. Refer back to previous points if relevant.
+
+Remember: You are simulating Warren Buffett. Your goal is to provide insights consistent with his philosophy and communication style, leveraging the tools for data when needed. Do not give definitive financial advice, but rather educate and explain the *Buffett way* of thinking about investments.
+"""
 
-def get_transcript(video_id):
+# --- Tool Definitions ---
+
+# 1. Stock Data Tool (Yahoo Finance)
+def get_stock_info(symbol: str) -> str:
     """
-    Retrieves and concatenates the transcript for the given YouTube video ID.
+    Fetches key financial data for a given stock symbol using Yahoo Finance.
+    Relevant data includes current price, P/E ratio, EPS, market cap, dividend yield, P/B ratio, sector, industry, and summary.
+    Returns a JSON string with the data or an error message.
     """
     try:
-        transcript = YouTubeTranscriptApi.get_transcript(video_id)
-        return " ".join([entry["text"] for entry in transcript])
+        ticker = yf.Ticker(symbol)
+        info = ticker.info
+
+        # Check if info was retrieved
+        if not info or info.get('regularMarketPrice') is None and info.get('currentPrice') is None and info.get('previousClose') is None:
+             # Attempt to fetch history for validation if info is sparse
+            hist = ticker.history(period="5d")
+            if hist.empty:
+                 return f"Error: Could not retrieve any data for symbol {symbol}. It might be delisted, invalid, or lack recent trading data."
+            # If history exists but info is bad, use historical close
+            last_close = hist['Close'].iloc[-1] if not hist.empty else 'N/A'
+            current_price = info.get("currentPrice") or info.get("regularMarketPrice") or last_close
+        else:
+            current_price = info.get("currentPrice") or info.get("regularMarketPrice") or info.get("previousClose", "N/A")
+
+        data = {
+            "symbol": symbol,
+            "companyName": info.get("longName", "N/A"),
+            "currentPrice": current_price,
+            "peRatio": info.get("trailingPE") or info.get("forwardPE", "N/A"),
+            "earningsPerShare": info.get("trailingEps", "N/A"),
+            "marketCap": info.get("marketCap", "N/A"),
+            "dividendYield": info.get("dividendYield", "N/A"),
+            "priceToBook": info.get("priceToBook", "N/A"),
+            "sector": info.get("sector", "N/A"),
+            "industry": info.get("industry", "N/A"),
+            "summary": info.get("longBusinessSummary", "N/A")[:500] + ("..." if len(info.get("longBusinessSummary", "")) > 500 else "") # Keep summary concise
+        }
+
+        # Basic validation
+        if data["currentPrice"] == "N/A":
+            return f"Error: Could not retrieve current price for {symbol}. Data might be incomplete."
+
+        return json.dumps(data)
+
     except Exception as e:
-        return f"Error retrieving transcript: {str(e)}"
+        # More specific error handling could be added here (e.g., check for specific yfinance exceptions)
+        return f"Error fetching data for {symbol} using yfinance: {str(e)}. Symbol might be invalid or API issue."
 
-def summarize_text(text):
-    """
-    Summarizes the provided text using a LangChain agent.
-    Uses GPT-4 model from OpenAI.
-    Only the first 4000 characters are used to avoid token limit issues.
-    """
+stock_data_tool = Tool(
+    name="get_stock_financial_data",
+    func=get_stock_info,
+    description="""
+    Useful for fetching fundamental financial data for a specific stock symbol (ticker).
+    Input should be a single stock symbol (e.g., 'AAPL', 'MSFT', 'BRK-B').
+    Returns a JSON string containing key metrics like company name, current price, P/E ratio, EPS, market cap, dividend yield, price-to-book, sector, industry, and a business summary.
+    Use this to get the necessary financial context before forming an opinion based on Warren Buffett's principles.
+    """,
+)
+
+# 2. News Search Tool (SerpAPI)
+try:
+    if not SERPAPI_API_KEY:
+        raise ValueError("SERPAPI_API_KEY environment variable not set.")
+    params = {
+        "engine": "google_news",
+        "gl": "us",
+        "hl": "en",
+        "num": 5 # Fetch top 5 news results
+    }
+    search_wrapper = SerpAPIWrapper(params=params, serpapi_api_key=SERPAPI_API_KEY)
+
+    news_search_tool = Tool(
+        name="search_stock_news",
+        func=search_wrapper.run,
+        description="""
+        Useful for searching recent news articles about a specific company or stock symbol.
+        Input should be the company name or stock symbol (e.g., 'Apple Inc. news', 'MSFT latest developments', 'Berkshire Hathaway earnings').
+        Returns a summary of recent news headlines and snippets.
+        Use this to understand recent events, sentiment, or developments related to a company before forming an opinion.
+        """,
+    )
+    serpapi_available = True
+except Exception as e: # Catch broader exceptions during init
+    print(f"SerpAPI News Tool Warning: {e}")
+    # Provide a dummy tool if the key is missing or setup fails
+    news_search_tool = Tool(
+        name="search_stock_news",
+        func=lambda x: "News search unavailable (SerpAPI key missing or configuration error).",
+        description="News search tool (currently unavailable).",
+        # return_direct=True # Optional: returns message directly without LLM processing
+    )
+    serpapi_available = False
+
+
+tools = [stock_data_tool, news_search_tool]
+
+# --- LangChain Agent Setup ---
+
+# Check for OpenAI Key
+if not OPENAI_API_KEY:
+    st.error("Error: OPENAI_API_KEY environment variable not set. Cannot initialize the chatbot.", icon="❌")
+    st.stop() # Stop execution if key is missing
+
+# LLM
+llm = ChatOpenAI(
+    model=MODEL_NAME,
+    temperature=TEMPERATURE,
+    openai_api_key=OPENAI_API_KEY,
+)
+
+# Prompt Template
+prompt_template = ChatPromptTemplate.from_messages(
+    [
+        SystemMessage(content=BUFFETT_SYSTEM_PROMPT),
+        MessagesPlaceholder(variable_name=MEMORY_KEY),
+        ("human", "{input}"),
+        MessagesPlaceholder(variable_name="agent_scratchpad"),
+    ]
+)
+
+# Memory (initialized fresh for each session in Streamlit context below)
+# The agent executor itself needs the memory factory/object,
+# but the actual state lives in st.session_state['memory'] for persistence across reruns.
+
+# Agent
+agent = create_openai_functions_agent(llm, tools, prompt_template)
+
+# Agent Executor (initialized in Streamlit session state)
+
+
+# --- Streamlit Frontend ---
+
+st.set_page_config(page_title="Warren Buffett Bot", layout="wide")
+st.title("Warren Buffett Investment Chatbot 📈")
+st.caption("Ask me about investing, stocks, or market wisdom - in the style of Warren Buffett.")
+
+# Display API Key status
+st.sidebar.header("API Status")
+if OPENAI_API_KEY: # Use a standard if/else block
+    st.sidebar.success("OpenAI API Key Loaded", icon="✅")
+else:
+    st.sidebar.error("OpenAI API Key Missing", icon="❌")
+
+# Keep the SerpAPI check as it is (already uses standard if/else)
+if serpapi_available:
+    st.sidebar.success("SerpAPI Key Loaded (News Enabled)", icon="✅")
+else:
+    st.sidebar.warning("SerpAPI Key Missing (News Disabled)", icon="⚠️")
+
+# Initialize chat history and memory in Streamlit session state
+if "messages" not in st.session_state:
+    st.session_state["messages"] = [
+        {"role": "assistant", "content": "Greetings! I'm here to chat about investing with the prudence and long-term view of Warren Buffett. How can I help you today?"}
+    ]
+# Initialize memory object in session state
+if 'memory' not in st.session_state:
+    st.session_state['memory'] = ConversationBufferMemory(memory_key=MEMORY_KEY, return_messages=True)
+
+# Initialize AgentExecutor in session state if it doesn't exist
+if 'agent_executor' not in st.session_state:
+    # The AgentExecutor needs the memory object from session state
+    st.session_state['agent_executor'] = AgentExecutor(
+        agent=agent,
+        tools=tools,
+        memory=st.session_state['memory'], # Use memory from session state
+        verbose=True, # Set to False for cleaner production output
+        handle_parsing_errors=True,
+        max_iterations=5,
+    )
+
+# Display chat messages from history
+for msg in st.session_state.messages:
+    st.chat_message(msg["role"]).write(msg["content"])
+
+# Accept user input
+if prompt := st.chat_input("Ask Buffett Bot..."):
+    # Add user message to chat history
+    st.session_state.messages.append({"role": "user", "content": prompt})
+    st.chat_message("user").write(prompt)
+
+    # Prepare agent input
+    # The agent executor uses the memory object linked during its initialization
+    agent_input = {"input": prompt}
+
+    # Invoke the agent using the executor stored in session state
     try:
-        llm = ChatOpenAI(model_name="gpt-3.5-turbo", openai_api_key=OPENAI_KEY)
-        truncated_text = text[:4000]
-        messages = [
-            SystemMessage(content="Summarize the given YouTube transcript concisely."),
-            HumanMessage(content=truncated_text)
-        ]
-        summary = llm(messages).content
-        return summary
+        with st.spinner("Buffett is pondering..."):
+            agent_executor_instance = st.session_state['agent_executor']
+            response = agent_executor_instance.invoke(agent_input)
+
+        # Extract and display response
+        output = response.get('output', "Sorry, I encountered an issue and couldn't formulate a response.")
+        st.session_state.messages.append({"role": "assistant", "content": output})
+        st.chat_message("assistant").write(output)
+
     except Exception as e:
-        return f"Error during summarization: {str(e)}"
-
-def main():
-    st.set_page_config(page_title="YouTube Summarizer", layout="wide")
-    st.title("🎥 YouTube Video Summarizer")
-    video_url = st.text_input("Enter YouTube Video URL")
-    
-    if video_url:
-        video_id = extract_video_id(video_url)
-        if not video_id:
-            st.error("Could not extract video ID from the provided URL.")
-            return
-        
-        st.info(f"Extracted Video ID: {video_id}")
-        transcript = get_transcript(video_id)
-        
-        if transcript.startswith("Error"):
-            st.error(transcript)
-            return
-
-        st.subheader("Transcript:")
-        if len(transcript) > 1000:
-            st.write(transcript[:1000] + "...")
-        else:
-            st.write(transcript)
-        
-        st.subheader("AI-Generated Summary:")
-        summary = summarize_text(transcript)
-        st.write(summary)
-
-if __name__ == "__main__":
-    main()
+        error_message = f"An error occurred: {str(e)}"
+        st.error(error_message, icon="🔥")
+        # Add error message to chat
+        st.session_state.messages.append({"role": "assistant", "content": f"Sorry, I ran into a technical difficulty: {e}"})
+        st.chat_message("assistant").write(f"Sorry, I ran into a technical difficulty: {e}")
+
+
+# Optional: Add a way to clear history/memory for a new session
+if st.sidebar.button("Clear Chat History"):
+    st.session_state.messages = [
+        {"role": "assistant", "content": "Chat history cleared. How can I help you start anew?"}
+    ]
+    st.session_state.memory.clear() # Clear the LangChain memory object
+    # Optionally reinstantiate executor if needed, though clearing memory might suffice
+    # if 'agent_executor' in st.session_state:
+    #     del st.session_state['agent_executor']
+    st.rerun()