# buffett_bot_single_file_with_key_input.py import streamlit as st import os import json import yfinance as yf from dotenv import load_dotenv # 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 # No need for HumanMessage/AIMessage here anymore from langchain.tools import Tool from langchain_community.utilities import SerpAPIWrapper # --- Page Config --- 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.") # --- Load .env file (as a fallback) --- load_dotenv() # --- API Key Input in Sidebar --- st.sidebar.header("API Configuration") # Use session state to store keys entered by the user if 'openai_api_key' not in st.session_state: st.session_state.openai_api_key = "" if 'serpapi_api_key' not in st.session_state: st.session_state.serpapi_api_key = "" # Get keys from user input, using session state for persistence input_openai_key = st.sidebar.text_input( "OpenAI API Key", type="password", value=st.session_state.openai_api_key, # Pre-fill if already entered help="Get your key from https://platform.openai.com/account/api-keys", key="openai_input" # Assign a key to the widget itself ) input_serpapi_key = st.sidebar.text_input( "SerpAPI API Key", type="password", value=st.session_state.serpapi_api_key, # Pre-fill if already entered help="Get your key from https://serpapi.com/manage-api-key", key="serpapi_input" # Assign a key to the widget itself ) # Update session state when input changes (Streamlit handles this implicitly with widget keys) st.session_state.openai_api_key = input_openai_key st.session_state.serpapi_api_key = input_serpapi_key # Determine active keys: prioritize user input, fallback to env variables active_openai_key = st.session_state.openai_api_key or os.getenv("OPENAI_API_KEY") active_serpapi_key = st.session_state.serpapi_api_key or os.getenv("SERPAPI_API_KEY") # --- Display API Status --- st.sidebar.header("API Status") if active_openai_key: st.sidebar.success("OpenAI API Key Loaded", icon="✅") else: st.sidebar.error("OpenAI API Key Missing", icon="❌") serpapi_available = False if active_serpapi_key: try: # Test SerpAPI key validity briefly if needed (optional) # test_search = SerpAPIWrapper(serpapi_api_key=active_serpapi_key) # test_search.run("test") serpapi_available = True st.sidebar.success("SerpAPI Key Loaded (News Enabled)", icon="✅") except Exception as e: st.sidebar.warning(f"SerpAPI Key Provided but Error Occurred: {e}", icon="⚠️") else: st.sidebar.warning("SerpAPI Key Missing (News Disabled)", icon="⚠️") # --- Constants & Prompt --- MODEL_NAME = "gpt-4o" # Or "gpt-3.5-turbo", "gpt-4-turbo" TEMPERATURE = 0.5 MEMORY_KEY = "chat_history" 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. """ # --- Tool Definitions --- # 1. Stock Data Tool (Yahoo Finance) - No changes needed here @st.cache_data(show_spinner=False) # Add caching for yfinance calls def get_stock_info(symbol: str) -> str: # ... (keep the existing get_stock_info function code) ... """ Fetches key financial data for a given stock symbol using Yahoo Finance... """ try: ticker = yf.Ticker(symbol) info = ticker.info if not info or info.get('regularMarketPrice') is None and info.get('currentPrice') is None and info.get('previousClose') is None: hist = ticker.history(period="5d") if hist.empty: return f"Error: Could not retrieve any data for symbol {symbol}." 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 "") } if data["currentPrice"] == "N/A": return f"Error: Could not retrieve current price for {symbol}." return json.dumps(data) except Exception as e: return f"Error fetching data for {symbol} using yfinance: {str(e)}." 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)..." # Keep description ) # 2. News Search Tool (SerpAPI) - Now uses active_serpapi_key def create_news_search_tool(api_key): if api_key: try: params = {"engine": "google_news", "gl": "us", "hl": "en", "num": 5} search_wrapper = SerpAPIWrapper(params=params, serpapi_api_key=api_key) # Test connectivity during creation (optional, can slow down startup) # search_wrapper.run("test query") return Tool( name="search_stock_news", func=search_wrapper.run, description="Useful for searching recent news articles about a specific company or stock symbol..." # Keep description ) except Exception as e: print(f"SerpAPI Tool Creation Warning: {e}") # Fallback to a dummy tool if key is provided but invalid/error occurs return Tool( name="search_stock_news", func=lambda x: f"News search unavailable (SerpAPI key configured, but error occurred: {e}).", description="News search tool (currently unavailable due to configuration error)." ) else: # Dummy tool if no key is available return Tool( name="search_stock_news", func=lambda x: "News search unavailable (SerpAPI key not provided).", description="News search tool (unavailable - API key needed)." ) news_search_tool = create_news_search_tool(active_serpapi_key) tools = [stock_data_tool, news_search_tool] # --- Main App Logic --- # Only proceed if OpenAI key is available if not active_openai_key: st.warning("Please enter your OpenAI API Key in the sidebar to activate the bot.", icon="🔑") st.stop() # Stop execution if no OpenAI key # --- LangChain Agent Setup (conditional on key) --- try: # LLM llm = ChatOpenAI( model=MODEL_NAME, temperature=TEMPERATURE, openai_api_key=active_openai_key, # Use the active 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"), ] ) # Initialize Memory and AgentExecutor in session state if they don't exist or if keys changed # We need to re-initialize if the keys change to ensure the agent uses the new ones. # A simple way is to check if the current active key matches what might be stored from a previous init. reinitialize_agent = False if 'agent_executor' not in st.session_state: reinitialize_agent = True elif 'agent_openai_key' in st.session_state and st.session_state.agent_openai_key != active_openai_key: reinitialize_agent = True elif 'agent_serpapi_key' in st.session_state and st.session_state.agent_serpapi_key != active_serpapi_key: reinitialize_agent = True # Re-init if SerpAPI key changed to update the tool if reinitialize_agent: st.session_state['memory'] = ConversationBufferMemory(memory_key=MEMORY_KEY, return_messages=True) agent = create_openai_functions_agent(llm, tools, prompt_template) st.session_state['agent_executor'] = AgentExecutor( agent=agent, tools=tools, # Pass the potentially updated news tool memory=st.session_state['memory'], verbose=True, # Set to False for cleaner production output handle_parsing_errors=True, max_iterations=5, ) # Store the keys used for this initialization st.session_state.agent_openai_key = active_openai_key st.session_state.agent_serpapi_key = active_serpapi_key st.experimental_rerun() # Rerun ensures UI uses the newly created agent state correctly # Initialize chat history 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?"} ] # 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 agent_input = {"input": prompt} # Invoke the agent try: 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: error_message = f"An error occurred: {str(e)}" st.error(error_message, icon="🔥") 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}") except Exception as e: st.error(f"Failed to initialize the LangChain agent: {e}", icon="🚨") st.warning("Please ensure your OpenAI API Key is correct and valid.") # 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?"} ] if 'memory' in st.session_state: st.session_state.memory.clear() # Clear the LangChain memory object # We might need to re-initialize the agent executor if memory state is critical beyond buffer st.rerun()