Capstone: AI Stock Analyst

# -*- coding: utf-8 -*-

"""Shared state schema for the stock-analyst LangGraph pipeline."""

from typing import Annotated
from typing import Dict
from typing import List
from typing import Optional
from typing import TypedDict

from langchain_core.messages import BaseMessage
from langgraph.graph.message import add_messages


class AnalystState(TypedDict):
    """Shared state passed between all agents in the stock analysis graph."""

    ticker: str
    messages: Annotated[List[BaseMessage], add_messages]
    company_info: Optional[Dict]
    earnings_data: Optional[Dict]
    sentiment_summary: Optional[str]
    final_analysis: Optional[str]
    error: Optional[str]

# -*- coding: utf-8 -*-

"""LangChain tools for fetching stock data and web-search sentiment."""

from __future__ import annotations

from typing import Dict
from typing import TypedDict

import yfinance as yf  # type: ignore[import-untyped]
from langchain_community.tools import DuckDuckGoSearchRun
from langchain_core.tools import tool


class CompanyInfo(TypedDict):
    """Company overview fields returned by get_company_info."""

    name: str
    description: str
    sector: str
    industry: str
    market_cap: str


@tool
def get_company_info(ticker: str) -> CompanyInfo:
    """
    Fetch company overview for a given stock ticker
    using yfinance.

    :returns: The `CompanyInfo` values.
    :raises: ValueError if the ticker is not found or the data is empty.
    """

    stock = yf.Ticker(ticker)
    info = stock.info

    if not info or info.get("quoteType") is None:
        raise ValueError(f"Ticker '{ticker}' not found or returned no data.")

    return {
        "name": info.get("longName", "N/A"),
        "sector": info.get("sector", "N/A"),
        "industry": info.get("industry", "N/A"),
        "market_cap": info.get("marketCap", "N/A"),
        "description": info.get("longBusinessSummary", "N/A")[:800],
    }


@tool
def get_earnings_data(ticker: str) -> Dict:
    """
    Fetch the most recent quarterly earnings for a stock ticker
    using yfinance. And returns the last four quarters of EPS and
    revenue figures.
    """

    stock = yf.Ticker(ticker)
    earnings = stock.quarterly_income_stmt

    if earnings is None or earnings.empty:
        return {"error": f"No earnings data available for '{ticker}'."}

    quarters = []
    for col in list(earnings.columns)[:4]:
        eps_row = earnings.loc["Diluted EPS"] if "Diluted EPS" in earnings.index else None
        rev_row = earnings.loc["Total Revenue"] if "Total Revenue" in earnings.index else None
        quarters.append({
            "quarter": str(col.date()),
            "eps": float(eps_row[col]) if eps_row is not None else None,
            "revenue": float(rev_row[col]) if rev_row is not None else None,
        })

    return {"quarters": quarters}


@tool
def search_sentiment(query: str) -> str:
    """
    Search the web for recent news and sentiment about
    a given query and returns a text summary of the top search
    results (up to ~1 200 characters).
    """

    search = DuckDuckGoSearchRun()
    results = search.run(query)
    return results[:1200]

# -*- coding: utf-8 -*-

"""LangGraph agent nodes for the stock-analyst pipeline."""

from __future__ import annotations

from langchain_core.messages import HumanMessage
from langchain_google_genai import ChatGoogleGenerativeAI

from state import AnalystState
from tools import get_company_info
from tools import get_earnings_data
from tools import search_sentiment

# Single shared LLM instance, only the analyst
# node uses it...
_llm = ChatGoogleGenerativeAI(
    model="gemini-2.5-flash",
    temperature=0.3,
)


# Validator node (entry point in the graph)...
def validate_ticker(state: AnalystState) -> dict:
    """Validate ticker format before any API calls are made."""

    ticker = state.get("ticker", "").strip().upper()
    if not ticker or not ticker.isalpha() or len(ticker) > 5:
        return {"error": f"'{ticker}' is not a valid stock ticker."}
    return {"ticker": ticker}


# The below three functions are the `Research Agents` and
# run in parallel via fan-out...

def company_info_node(state: AnalystState) -> dict:
    """Fetch company overview and write the result to state."""

    try:
        info = get_company_info.invoke({"ticker": state["ticker"]})
        return {"company_info": info}

    except Exception as exc:  # pylint: disable=broad-exception-caught
        return {"company_info": {"error": str(exc)}}


def earnings_node(state: AnalystState) -> dict:
    """Fetch quarterly earnings and write the result to state."""

    try:
        data = get_earnings_data.invoke({"ticker": state["ticker"]})
        return {"earnings_data": data}

    except Exception as exc:  # pylint: disable=broad-exception-caught
        return {"earnings_data": {"error": str(exc)}}


def sentiment_node(state: AnalystState) -> dict:
    """Search for recent news sentiment and write the result to state."""

    try:
        query = f"{state['ticker']} stock news sentiment 2025"
        summary = search_sentiment.invoke({"query": query})
        return {"sentiment_summary": summary}

    except Exception as exc:  # pylint: disable=broad-exception-caught
        return {"sentiment_summary": f"Sentiment search failed: {exc}"}


# The `Analyst Agent` (fan-in synthesis)...

_ANALYST_PROMPT = """You are a senior equity analyst. Three specialised research agents
have gathered data for you independently and in parallel:

- Agent 1 (Company Info): {company_info}
- Agent 2 (Earnings):     {earnings_data}
- Agent 3 (Sentiment):    {sentiment_summary}

Based on these multi-agent findings, provide a concise investment insight for: {ticker}

Structure your response with:
1. Company snapshot (2-3 sentences)
2. Financial health (earnings trend interpretation)
3. Market sentiment (based on recent news)
4. Investment outlook (bullish / neutral / bearish with rationale)

Be direct and analytical."""


def analyst_node(state: AnalystState) -> dict:
    """Synthesize all research findings into an investment analysis using Gemini Flash."""

    prompt = _ANALYST_PROMPT.format(
        ticker=state["ticker"],
        company_info=state.get("company_info", "Not available"),
        earnings_data=state.get("earnings_data", "Not available"),
        sentiment_summary=state.get("sentiment_summary", "Not available"),
    )

    try:
        response = _llm.invoke([HumanMessage(content=prompt)])
        return {"final_analysis": response.content}

    except Exception as exc:  # pylint: disable=broad-exception-caught
        return {"error": f"Analyst agent failed: {exc}"}

# -*- coding: utf-8 -*-

"""Graph definition and compilation for the stock-analyst multi-agent pipeline."""

from __future__ import annotations

from langgraph.graph import StateGraph, START, END

from agents import (
    validate_ticker,
    company_info_node,
    earnings_node,
    sentiment_node,
    analyst_node,
)
from state import AnalystState


def _route_after_validation(state: AnalystState) -> list[str] | str:
    """
    Short-circuit to END on error; fan-out to 3 parallel research
    agents otherwise. Returning a list triggers LangGraph's parallel
    fan-out, all named nodes start concurrently without any
    additional configuration.
    """

    if state.get("error"):
        return END
    return ["company_info", "earnings", "sentiment"]


def build_graph():
    """
    Graph topology:
      START -> validate
      validate -> END                                   (invalid ticker - short circuit)
      validate -> company_info, earnings, sentiment     (parallel fan-out)
      company_info, earnings, sentiment -> analyst      (implicit fan-in)
      analyst -> END
    """

    graph = StateGraph(AnalystState)  # ty: ignore[invalid-argument-type]

    graph.add_node("validate", validate_ticker)
    graph.add_node("company_info", company_info_node)
    graph.add_node("earnings", earnings_node)
    graph.add_node("sentiment", sentiment_node)
    graph.add_node("analyst", analyst_node)

    graph.add_edge(START, "validate")
    graph.add_conditional_edges("validate", _route_after_validation)

    # Fan-in: all three research nodes must complete
    # before analyst runs...
    graph.add_edge("company_info", "analyst")
    graph.add_edge("earnings", "analyst")
    graph.add_edge("sentiment", "analyst")

    # After the analysis, just end process...
    graph.add_edge("analyst", END)
    return graph.compile()


# Module-level compiled graph - Streamlit won't recompile on every interaction
stock_analyst_graph = build_graph()

# -*- coding: utf-8 -*-

"""Streamlit front-end for the AI Stock Analyst multi-agent application."""

from __future__ import annotations

import streamlit as st
from dotenv import load_dotenv

# `load_dotenv()` MUST run before importing `graph.py` because
# the `ChatGoogleGenerativeAI` module reads `GOOGLE_API_KEY` at
# module init, not at call time. Which is probably a BAD
# design...
load_dotenv()

from graph import stock_analyst_graph  # noqa: E402  # pylint: disable=wrong-import-position


st.set_page_config(
    page_title="AI Stock Analyst",
    page_icon="📈",
    layout="centered",
)

st.title("AI Stock Analyst")

st.markdown(
    "Powered by a **LangGraph multi-agent system**: three research agents run in "
    "parallel, then a **Gemini Flash** analyst syntheses their findings."
)

st.divider()

ticker_input = st.text_input(
    "Enter a stock ticker symbol",
    placeholder="e.g. AAPL, MSFT, NVDA, TSLA",
    max_chars=10,
)

analyze_btn = st.button("Analyse", type="primary", use_container_width=True)

if analyze_btn and ticker_input.strip():
    ticker = ticker_input.strip().upper()

    with st.spinner(f"Running multi-agent analysis for {ticker}..."):
        result = stock_analyst_graph.invoke({
            "ticker": ticker,
            "messages": [],
            "company_info": None,
            "earnings_data": None,
            "sentiment_summary": None,
            "final_analysis": None,
            "error": None,
        })

    if result.get("error"):
        st.error(f"Error: {result['error']}")
        st.stop()

    # Company overview...
    st.subheader("Company Overview")
    info = result.get("company_info") or {}

    if "error" in info:
        st.warning(info["error"])
    else:
        col1, col2 = st.columns(2)
        with col1:
            st.metric("Company", info.get("name", "N/A"))
            st.metric("Sector", info.get("sector", "N/A"))
        with col2:
            st.metric("Industry", info.get("industry", "N/A"))
            market_cap = info.get("market_cap")

            st.metric(
                "Market Cap",
                ("$" + f"{market_cap:,.0f}") if isinstance(market_cap, (int, float)) else "N/A",
            )

        if info.get("description"):
            st.caption(info["description"])

    st.divider()

    # Earnings summary...
    st.subheader("Earnings Summary (Last 4 Quarters)")
    earnings = result.get("earnings_data") or {}

    if "error" in earnings:
        st.warning(earnings["error"])
    elif "quarters" in earnings:
        for q in earnings["quarters"]:
            eps = q.get("eps")
            rev = q.get("revenue")
            eps_str = f"EPS: {eps:.2f}" if eps is not None else "EPS: N/A"
            rev_str = ("Revenue: $" + f"{rev:,.0f}") if rev is not None else "Revenue: N/A"
            st.markdown(f"**{q['quarter']}** - {eps_str} | {rev_str}")
    else:
        st.info("No earnings data available.")

    st.divider()

    # Market sentiment...
    st.subheader("Market Sentiment (Recent News)")
    sentiment = result.get("sentiment_summary")

    if sentiment:
        st.markdown(sentiment)
    else:
        st.info("No sentiment data available.")

    st.divider()

    # AI analysis...
    st.subheader("AI Investment Analysis")
    analysis = result.get("final_analysis")
    if analysis:
        st.markdown(analysis)

elif analyze_btn:
    st.warning("Please enter a ticker symbol before clicking Analyse.")