app.py

# 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()