app.py

import streamlit as st
import os
import pandas as pd
import yfinance as yf
from pydantic import BaseModel, Field
from typing import List, Literal, Optional
from llama_index.core import VectorStoreIndex, Settings
from llama_index.vector_stores.pinecone import PineconeVectorStore
from pinecone import Pinecone
from llama_index.embeddings.openai import OpenAIEmbedding
from llama_index.program.openai import OpenAIPydanticProgram
from llama_index.llms.openai import OpenAI
from llama_index.core.vector_stores import MetadataFilters, ExactMatchFilter

# --- 1. CONFIGURATION ---
st.set_page_config(page_title="Financial Agent (Strict Logic)", page_icon="📈", layout="wide")

# Ensure keys exist
if "OPENAI_API_KEY" not in os.environ:
    st.error("❌ OPENAI_API_KEY missing.")
    st.stop()

# --- 2. DATA MODELS (From your snippet) ---
class AgentResponse(BaseModel):
    answer: str
    sources: List[str]
    context_used: List[str]

class TickerExtraction(BaseModel):
    symbols: List[str] = Field(description="List of stock tickers.")

class RoutePrediction(BaseModel):
    tools: List[Literal["financial_rag", "market_data", "general_chat"]] = Field(description="Tools list")

# --- 3. CACHED INITIALIZATION ---
@st.cache_resource(show_spinner=False)
def initialize_resources():
    print("🔌 Initializing Strict-Boundary Agent...")
    
    # Setup LlamaIndex Settings
    Settings.llm = OpenAI(model="gpt-4o-mini", temperature=0)
    Settings.embed_model = OpenAIEmbedding(model="text-embedding-3-small")

    # Load CSV
    try:
        nasdaq_df = pd.read_csv('nasdaq-listed.csv')
        nasdaq_df.columns = [c.strip() for c in nasdaq_df.columns]
    except:
        nasdaq_df = pd.DataFrame()

    # Connect to Pinecone
    api_key = os.environ.get("PINECONE_API_KEY")
    if not api_key: raise ValueError("Pinecone Key Missing")
    
    pc = Pinecone(api_key=api_key)
    index = VectorStoreIndex.from_vector_store(
        vector_store=PineconeVectorStore(pinecone_index=pc.Index("financial-rag-agent"))
    )
    
    return nasdaq_df, index

# --- 4. HELPER FUNCTIONS (From your snippet) ---
def get_symbol_from_csv(query_str: str, df) -> Optional[str]:
    if df.empty: return None
    query_str = query_str.strip().upper()
    if query_str in df['Symbol'].values: return query_str
    matches = df[df['Security Name'].str.upper().str.contains(query_str, na=False)]
    if not matches.empty: return matches.loc[matches['Symbol'].str.len().idxmin()]['Symbol']
    return None

def get_tickers_from_query(query: str, index, df) -> List[str]:
    program = OpenAIPydanticProgram.from_defaults(
        output_cls=TickerExtraction,
        prompt_template_str="Identify all companies in query: {query_str}. Return list.",
        llm=Settings.llm
    )
    raw_entities = program(query_str=query).symbols
    valid_tickers = []
    for entity in raw_entities:
        ticker = get_symbol_from_csv(entity, df)
        if not ticker and len(entity) <= 5: ticker = entity.upper()
        if ticker: valid_tickers.append(ticker)
    
    if not valid_tickers:
        try:
            nodes = index.as_retriever(similarity_top_k=1).retrieve(query)
            if nodes and nodes[0].metadata.get("ticker"):
                valid_tickers.append(nodes[0].metadata.get("ticker"))
        except: pass
    return list(set(valid_tickers))

# --- 5. TOOLS (From your snippet) ---
def get_market_data(query: str, index, df):
    tickers = get_tickers_from_query(query, index, df)
    if not tickers: return "No companies found."
    results = []
    for ticker in tickers:
        try:
            stock = yf.Ticker(ticker)
            info = stock.info
            data = {
                "Ticker": ticker,
                "Price": info.get('currentPrice', 'N/A'),
                "Market Cap": info.get('marketCap', 'N/A'),
                "PE Ratio": info.get('trailingPE', 'N/A'),
                "52w High": info.get('fiftyTwoWeekHigh', 'N/A'),
                "52w Low": info.get('fiftyTwoWeekLow', 'N/A'),
                "Volume": info.get('volume', 'N/A'),
                "Currency": info.get('currency', 'USD')
            }
            results.append(str(data))
        except Exception as e:
            results.append(f"{ticker}: Data Error ({e})")
    return "\n".join(results)

def get_financial_rag(query: str, index, df):
    target_tickers = get_tickers_from_query(query, index, df)
    SUPPORTED = ["AAPL", "TSLA", "NVDA"]
    payload = {"content": "", "sources": [], "raw_nodes": []}
    
    for ticker in target_tickers:
        if ticker not in SUPPORTED:
            payload["content"] += f"\n[NOTE: No 10-K report available for {ticker}.]\n"
            continue
            
        filters = MetadataFilters(filters=[ExactMatchFilter(key="ticker", value=ticker)])
        # Using logic from your snippet (similarity_top_k=3)
        engine = index.as_query_engine(similarity_top_k=3, filters=filters)
        resp = engine.query(query)
        
        payload["content"] += f"\n--- {ticker} 10-K Data ---\n{resp.response}\n"
        for n in resp.source_nodes:
            payload["sources"].append(f"{n.metadata.get('company')} 10-K")
            payload["raw_nodes"].append(n.node.get_text())
            
    return payload

# --- 6. AGENT LOGIC (From your snippet) ---
def run_agent(user_query: str, index, df) -> AgentResponse:
    # THE STRICT PROMPT YOU PROVIDED
    router_prompt = """
    Route the user query to the correct tool based on these strict definitions:
    
    1. "financial_rag":
       - Use for ANY question about a specific company's internal details.
       - INCLUDES: Revenue, Profit, Income, CEO, Board Members, Risks, Strategy, Competitors, Legal Issues, History.
       - Key Trigger: If the answer would be found in a PDF report or Wikipedia page, use this.
       
    2. "market_data":
       - Use ONLY for Real-Time Trading Metrics.
       - INCLUDES: Current Price, Market Cap, PE Ratio, Trading Volume, 52-Week High/Low.
       - EXCLUDES: Historical revenue or annual profit (Use financial_rag for those).
       
    3. "general_chat":
       - Use ONLY for non-business questions (e.g. "Hi", "Help").
       - NEVER use this if a specific company (Tesla, Apple, Nvidia) is mentioned.
    
    Query: {query_str}
    """
    
    router = OpenAIPydanticProgram.from_defaults(
        output_cls=RoutePrediction,
        prompt_template_str=router_prompt,
        llm=Settings.llm
    )
    tools = router(query_str=user_query).tools
    
    results = {}
    sources = []
    context_used = []
    
    if "market_data" in tools:
        res = get_market_data(user_query, index, df)
        results["market_data"] = res
        context_used.append(res)
        sources.append("Real-time Market Data")
        
    if "financial_rag" in tools:
        res = get_financial_rag(user_query, index, df)
        results["financial_rag"] = res["content"]
        sources.extend(res["sources"])
        context_used.extend(res["raw_nodes"])
        
    final_prompt = f"""
    You are a Wall Street Financial Analyst. Answer the user request using the provided context.
    
    Context Data:
    {results}
    
    Instructions:
    1. Compare Metrics if multiple companies are listed.
    2. Synthesize qualitative (Risks) and quantitative (Price) data.
    3. Explicitly state if a report is missing.
    4. Cite sources.
    
    User Query: {user_query}
    """
    
    response_text = Settings.llm.complete(final_prompt).text
    
    return AgentResponse(
        answer=response_text,
        sources=list(set(sources)),
        context_used=context_used
    )

# --- 7. STREAMLIT UI ---
# Initialize Logic
with st.sidebar:
    st.title("🔧 System Status")
    with st.spinner("Initializing Strict-Boundary Agent..."):
        try:
            nasdaq_df, pinecone_index = initialize_resources()
            st.success("✅ Brain Loaded")
            st.success(f"✅ {len(nasdaq_df)} Tickers Indexed")
        except Exception as e:
            st.error(f"Initialization Failed: {e}")
            st.stop()
            
    st.markdown("---")
    st.markdown("### 🎯 RAG Coverage")
    st.code("AAPL\nTSLA\nNVDA")

st.title("📈 Financial Agent (Strict Logic)")

if "messages" not in st.session_state:
    st.session_state.messages = []

# Display History
for message in st.session_state.messages:
    with st.chat_message(message["role"]):
        st.markdown(message["content"])
        if "sources" in message:
             with st.expander("📚 Sources & Context"):
                 st.write(message["sources"])
                 for i, c in enumerate(message["context"][:3]): # Limit preview
                     st.text(f"Snippet {i+1}: {str(c)[:300]}...")

# Input Handler
if prompt := st.chat_input("Enter query..."):
    st.session_state.messages.append({"role": "user", "content": prompt})
    with st.chat_message("user"):
        st.markdown(prompt)

    with st.chat_message("assistant"):
        with st.status("🧠 Analyst is thinking...", expanded=True) as status:
            try:
                # RUN THE SAVED LOGIC
                response = run_agent(prompt, pinecone_index, nasdaq_df)
                
                status.update(label="✅ Complete", state="complete", expanded=False)
                st.markdown(response.answer)
                
                # Audit Trail
                with st.expander("🔍 Audit Trail (Full Context)"):
                    st.write("**Sources:**", response.sources)
                    st.write("**Raw Retrieval:**")
                    for ctx in response.context_used:
                        st.text(str(ctx))
                
                st.session_state.messages.append({
                    "role": "assistant", 
                    "content": response.answer,
                    "sources": response.sources,
                    "context": response.context_used
                })
                
            except Exception as e:
                st.error(f"Error: {e}")
                status.update(label="❌ Error", state="error")