Update app.py

app.py

@@ -1,69 +1,171 @@
-import gradio as gr
-import groq
+# Standard library imports
 import os
 import tempfile
 import uuid
-import yfinance as yf
-import pandas as pd
-import plotly.graph_objects as go
-from dotenv import load_dotenv
-from langchain.text_splitter import RecursiveCharacterTextSplitter
-from langchain.vectorstores import FAISS
-from langchain.embeddings import HuggingFaceEmbeddings
-import fitz  # PyMuPDF
 import base64
-from PIL import Image
 import io
-import requests
 import json
+import re
+from datetime import datetime, timedelta
+
+# Third-party imports
+import gradio as gr
+import groq
+import numpy as np
+import pandas as pd
+import requests
+import fitz  # PyMuPDF
+from PIL import Image
+from dotenv import load_dotenv
+import yfinance as yf
+import plotly.graph_objects as go
+from pydantic import BaseModel, validator
+
+# LangChain imports
+from langchain_community.embeddings import HuggingFaceEmbeddings
+from langchain_community.vectorstores import FAISS
+from langchain.text_splitter import RecursiveCharacterTextSplitter
 
 # Load environment variables
 load_dotenv()
 client = groq.Client(api_key=os.getenv("GROQ_LEGAL_API_KEY"))
 embeddings = HuggingFaceEmbeddings(model_name="sentence-transformers/all-MiniLM-L6-v2")
+SERPER_API_KEY = os.getenv("SERPER_API_KEY")
+BRAVE_API_KEY = os.getenv("BRAVE_API_KEY")
 
 # Directory to store FAISS indexes
 FAISS_INDEX_DIR = "faiss_indexes_finance"
 if not os.path.exists(FAISS_INDEX_DIR):
     os.makedirs(FAISS_INDEX_DIR)
 
-# Dictionary to store user-specific vectorstores
+# Dictionaries for state management
 user_vectorstores = {}
+chart_data_store = {}
+
+# Pydantic Models
+class PDFState(BaseModel):
+    page_images: list[str]
+    total_pages: int
+    total_words: int
+
+class StockAnalysisInput(BaseModel):
+    ticker: str
+    period: str
+
+    @validator('ticker')
+    def validate_ticker(cls, v):
+        v = v.strip().lstrip("$")
+        if not v.isalpha() or len(v) > 5:
+            raise ValueError('Invalid ticker symbol')
+        return v.upper()
+
+    @validator('period')
+    def validate_period(cls, v):
+        valid_periods = ["1mo", "3mo", "6mo", "1y", "2y", "5y", "max"]
+        if v not in valid_periods:
+            raise ValueError('Invalid period')
+        return v
 
-# Custom CSS for Finance theme
+class ModelName(str):
+    @classmethod
+    def __get_validators__(cls):
+        yield cls.validate
+
+    @classmethod
+    def validate(cls, v):
+        allowed_models = ["llama3-70b-8192", "llama3-8b-8192", "mixtral-8x7b-32768", "gemma-7b-it"]
+        if v not in allowed_models:
+            raise ValueError(f"Invalid model name: {v}")
+        return v
+
+# Custom CSS with dark mode support
 custom_css = """
 :root {
-    --primary-color: #FFD700;  /* Gold */
-    --secondary-color: #008000;  /* Dark Green */
-    --light-background: #F0FFF0;  /* Honeydew */
-    --dark-text: #333333;
-    --white: #FFFFFF;
+    --bg-color: #FFFFFF;
+    --text-color: #333333;
+    --primary-color: #0C4160;
+    --secondary-color: #0D6980;
+    --accent-color: #16A6DB;
+    --light-color: #EBF5FA;
     --border-color: #E5E7EB;
 }
-body { background-color: var(--light-background); font-family: 'Inter', sans-serif; }
+
+[data-theme="dark"] {
+    --bg-color: #1E1E1E;
+    --text-color: #F5F5F5;
+    --primary-color: #16A6DB;
+    --secondary-color: #0D6980;
+    --accent-color: #0C4160;
+    --light-color: #333333;
+    --border-color: #444444;
+}
+
+body { background-color: var(--bg-color); color: var(--text-color); font-family: 'IBM Plex Sans', sans-serif; }
 .container { max-width: 1200px !important; margin: 0 auto !important; padding: 10px; }
-.header { background-color: var(--white); border-bottom: 2px solid var(--border-color); padding: 15px 0; margin-bottom: 20px; border-radius: 12px 12px 0 0; box-shadow: 0 2px 4px rgba(0,0,0,0.05); }
-.header-title { color: var(--secondary-color); font-size: 1.8rem; font-weight: 700; text-align: center; }
-.header-subtitle { color: var(--dark-text); font-size: 1rem; text-align: center; margin-top: 5px; }
-.chat-container { border-radius: 12px !important; box-shadow: 0 4px 6px rgba(0,0,0,0.1) !important; background-color: var(--white) !important; border: 1px solid var(--border-color) !important; min-height: 500px; }
-.message-user { background-color: var(--primary-color) !important; color: var(--dark-text) !important; border-radius: 18px 18px 4px 18px !important; padding: 12px 16px !important; margin-left: auto !important; max-width: 80% !important; }
-.message-bot { background-color: #F0F0F0 !important; color: var(--dark-text) !important; border-radius: 18px 18px 18px 4px !important; padding: 12px 16px !important; margin-right: auto !important; max-width: 80% !important; }
-.input-area { background-color: var(--white) !important; border-top: 1px solid var(--border-color) !important; padding: 12px !important; border-radius: 0 0 12px 12px !important; }
+.header { background-color: var(--primary-color); padding: 20px 0; margin-bottom: 20px; border-radius: 12px 12px 0 0; box-shadow: 0 2px 4px rgba(0,0,0,0.1); }
+.header-title { color: var(--light-color); font-size: 1.8rem; font-weight: 700; text-align: center; }
+.header-subtitle { color: var(--light-color); opacity: 0.8; font-size: 1rem; text-align: center; margin-top: 5px; }
+.chat-container { border-radius: 12px !important; box-shadow: 0 4px 6px rgba(0,0,0,0.1) !important; background-color: var(--bg-color) !important; border: 1px solid var(--border-color) !important; min-height: 500px; }
+.message-user { background-color: var(--accent-color) !important; color: var(--light-color) !important; border-radius: 18px 18px 4px 18px !important; padding: 12px 16px !important; margin-left: auto !important; max-width: 80% !important; }
+.message-bot { background-color: var(--light-color) !important; color: var(--text-color) !important; border-radius: 18px 18px 18px 4px !important; padding: 12px 16px !important; margin-right: auto !important; max-width: 80% !important; }
+.input-area { background-color: var(--bg-color) !important; border-top: 1px solid var(--border-color) !important; padding: 12px !important; border-radius: 0 0 12px 12px !important; }
 .input-box { border: 1px solid var(--border-color) !important; border-radius: 24px !important; padding: 12px 16px !important; box-shadow: 0 2px 4px rgba(0,0,0,0.05) !important; }
-.send-btn { background-color: var(--secondary-color) !important; border-radius: 24px !important; color: var(--white) !important; padding: 10px 20px !important; font-weight: 500 !important; }
-.clear-btn { background-color: #F0F0F0 !important; border: 1px solid var(--border-color) !important; border-radius: 24px !important; color: var(--dark-text) !important; padding: 8px 16px !important; font-weight: 500 !important; }
-.pdf-viewer-container { border-radius: 12px !important; box-shadow: 0 4px 6px rgba(0,0,0,0.1) !important; background-color: var(--white) !important; border: 1px solid var(--border-color) !important; padding: 20px; }
+.send-btn { background-color: var(--accent-color) !important; border-radius: 24px !important; color: var(--light-color) !important; padding: 10px 20px !important; font-weight: 500 !important; }
+.clear-btn { background-color: var(--light-color) !important; border: 1px solid var(--border-color) !important; border-radius: 24px !important; color: var(--text-color) !important; padding: 8px 16px !important; font-weight: 500 !important; }
+.pdf-viewer-container { border-radius: 12px !important; box-shadow: 0 4px 6px rgba(0,0,0,0.1) !important; background-color: var(--bg-color) !important; border: 1px solid var(--border-color) !important; padding: 20px; }
 .pdf-viewer-image { max-width: 100%; height: auto; border: 1px solid var(--border-color); border-radius: 12px; box-shadow: 0 2px 4px rgba(0,0,0,0.05); }
-.stats-box { background-color: #E6F2E6; padding: 10px; border-radius: 8px; margin-top: 10px; }
-.tool-container { background-color: var(--white); border-radius: 12px; box-shadow: 0 4px 6px rgba(0,0,0,0.1); padding: 15px; margin-bottom: 20px; }
-.tool-title { color: var(--secondary-color); font-size: 1.2rem; font-weight: 600; margin-bottom: 10px; }
+.stats-box { background-color: var(--light-color); padding: 10px; border-radius: 8px; margin-top: 10px; }
+.tool-container { background-color: var(--bg-color); border-radius: 12px; box-shadow: 0 2px 4px rgba(0,0,0,0.05); padding: 15px; margin-bottom: 20px; }
+.tool-title { font-weight: bold; color: var(--primary-color); margin-bottom: 10px; font-size: 1.1rem; }
 .chart-container { height: 400px; width: 100%; border-radius: 8px; overflow: hidden; }
+.toggle-container { display: flex; align-items: center; margin-bottom: 15px; }
+.toggle-label { margin-right: 10px; font-weight: 500; }
+.search-toggle { margin-left: 5px; }
+.spinner { border: 4px solid #f3f3f3; border-top: 4px solid var(--primary-color); border-radius: 50%; width: 40px; height: 40px; animation: spin 1s linear infinite; }
+@keyframes spin { 0% { transform: rotate(0deg); } 100% { transform: rotate(360deg); } }
+"""
+
+# Custom JavaScript for UI enhancements
+custom_js = """
+function toggleTheme() {
+    const currentTheme = document.body.getAttribute('data-theme');
+    const newTheme = currentTheme === 'dark' ? 'light' : 'dark';
+    document.body.setAttribute('data-theme', newTheme);
+    localStorage.setItem('theme', newTheme);
+}
+
+function showSpinner() {
+    document.getElementById('spinner').style.display = 'block';
+}
+
+function hideSpinner() {
+    document.getElementById('spinner').style.display = 'none';
+}
+
+document.addEventListener('DOMContentLoaded', () => {
+    const savedTheme = localStorage.getItem('theme');
+    if (savedTheme) {
+        document.body.setAttribute('data-theme', savedTheme);
+    }
+    tippy('#pdf_file', { content: 'Upload a PDF document for analysis', placement: 'top' });
+    tippy('#ticker_input', { content: 'Enter a stock ticker symbol (e.g., AAPL)', placement: 'top' });
+});
+"""
+
+# Spinner HTML and Tippy.js for tooltips
+custom_html = """
+<div id="spinner" style="display: none; position: fixed; top: 50%; left: 50%; transform: translate(-50%, -50%);">
+    <div class="spinner"></div>
+</div>
+<script src="https://unpkg.com/@popperjs/core@2"></script>
+<script src="https://unpkg.com/tippy.js@6"></script>
 """
 
-# Function to process PDF files (unchanged)
+# Helper Functions
 def process_pdf(pdf_file):
     if pdf_file is None:
-        return None, "No file uploaded", {"page_images": [], "total_pages": 0, "total_words": 0}
+        return None, "No file uploaded", PDFState(page_images=[], total_pages=0, total_words=0)
+    
     try:
         session_id = str(uuid.uuid4())
         with tempfile.NamedTemporaryFile(suffix=".pdf", delete=False) as temp_file:
@@ -72,12 +174,7 @@ def process_pdf(pdf_file):
         
         doc = fitz.open(pdf_path)
         texts = [page.get_text() for page in doc]
-        page_images = []
-        for page in doc:
-            pix = page.get_pixmap()
-            img_bytes = pix.tobytes("png")
-            img_base64 = base64.b64encode(img_bytes).decode("utf-8")
-            page_images.append(img_base64)
+        page_images = [base64.b64encode(page.get_pixmap().tobytes("png")).decode("utf-8") for page in doc]
         total_pages = len(doc)
         total_words = sum(len(text.split()) for text in texts)
         doc.close()
@@ -90,89 +187,85 @@ def process_pdf(pdf_file):
         user_vectorstores[session_id] = vectorstore
 
         os.unlink(pdf_path)
-        pdf_state = {"page_images": page_images, "total_pages": total_pages, "total_words": total_words}
+        pdf_state = PDFState(page_images=page_images, total_pages=total_pages, total_words=total_words)
         return session_id, f"✅ Successfully processed {len(chunks)} text chunks from your PDF", pdf_state
     except Exception as e:
         if "pdf_path" in locals() and os.path.exists(pdf_path):
             os.unlink(pdf_path)
-        return None, f"Error processing PDF: {str(e)}", {"page_images": [], "total_pages": 0, "total_words": 0}
+        return None, f"Error processing PDF: {str(e)}", PDFState(page_images=[], total_pages=0, total_words=0)
 
-# Function to generate chatbot responses with Finance theme
-def generate_response(message, session_id, model_name, history):
-    if not message:
-        return history
+def serper_search(query):
+    if not SERPER_API_KEY:
+        return {"error": "Serper API key not configured."}
+    url = "https://google.serper.dev/search"
+    payload = json.dumps({"q": query, "gl": "us", "hl": "en", "autocorrect": True})
+    headers = {'X-API-KEY': SERPER_API_KEY, 'Content-Type': 'application/json'}
     try:
-        context = ""
-        if session_id and session_id in user_vectorstores:
-            vectorstore = user_vectorstores[session_id]
-            docs = vectorstore.similarity_search(message, k=3)
-            if docs:
-                context = "\n\nRelevant information from uploaded PDF:\n" + "\n".join(f"- {doc.page_content}" for doc in docs)
-        
-        # Check if it's a stock ticker query
-        if message.startswith("$") and len(message) > 1 and len(message) <= 6:
-            ticker = message[1:].upper()
-            try:
-                stock_data = get_stock_data(ticker)
-                response = f"**Stock Information for {ticker}**\n\n"
-                response += f"Current Price: ${stock_data['current_price']}\n"
-                response += f"52-Week High: ${stock_data['52wk_high']}\n"
-                response += f"Market Cap: ${stock_data['market_cap']:,}\n"
-                response += f"P/E Ratio: {stock_data['pe_ratio']}\n"
-                response += f"More data available in the Stock Analysis tab."
-                history.append((message, response))
-                return history
-            except Exception as e:
-                history.append((message, f"Error retrieving stock data for {ticker}: {str(e)}"))
-                return history
-                
-        system_prompt = "You are a financial assistant specializing in analyzing financial reports, statements, and market trends."
-        system_prompt += " You can help with stock market information, financial terminology, ratio analysis, and investment concepts."
-        if context:
-            system_prompt += " Use the following context to answer the question if relevant: " + context
-        
-        completion = client.chat.completions.create(
-            model=model_name,
-            messages=[
-                {"role": "system", "content": system_prompt},
-                {"role": "user", "content": message}
-            ],
-            temperature=0.7,
-            max_tokens=1024
-        )
-        response = completion.choices[0].message.content
-        history.append((message, response))
-        return history
+        response = requests.post(url, headers=headers, data=payload)
+        return response.json()
     except Exception as e:
-        history.append((message, f"Error generating response: {str(e)}"))
-        return history
+        return {"error": str(e)}
 
-# Functions to update PDF viewer (unchanged)
-def update_pdf_viewer(pdf_state):
-    if not pdf_state["total_pages"]:
-        return 0, None, "No PDF uploaded yet"
+def brave_search(query):
+    if not BRAVE_API_KEY:
+        return {"error": "Brave Search API key not configured."}
+    url = "https://api.search.brave.com/res/v1/web/search"
+    params = {"q": query, "count": 10, "search_lang": "en", "country": "us"}
+    headers = {'Accept': 'application/json', 'X-Subscription-Token': BRAVE_API_KEY}
     try:
-        img_data = base64.b64decode(pdf_state["page_images"][0])
-        img = Image.open(io.BytesIO(img_data))
-        return pdf_state["total_pages"], img, f"**Total Pages:** {pdf_state['total_pages']}\n**Total Words:** {pdf_state['total_words']}"
+        response = requests.get(url, params=params, headers=headers)
+        return response.json()
     except Exception as e:
-        print(f"Error decoding image: {e}")
-        return 0, None, "Error displaying PDF"
+        return {"error": str(e)}
 
-def update_image(page_num, pdf_state):
-    if not pdf_state["total_pages"] or page_num < 1 or page_num > pdf_state["total_pages"]:
-        return None
-    try:
-        img_data = base64.b64decode(pdf_state["page_images"][page_num - 1])
-        img = Image.open(io.BytesIO(img_data))
-        return img
-    except Exception as e:
-        print(f"Error decoding image: {e}")
-        return None
+def get_financial_news(ticker, use_brave_search, enable_search, model_name="llama3-8b-8192"):
+    if not enable_search:
+        return [{"title": "Real-time search disabled", "snippet": "Enable real-time search to fetch news"}]
+    query = f"{ticker} stock news financial analysis latest"
+    if use_brave_search and BRAVE_API_KEY:
+        results = brave_search(query)
+        if "web" in results and "results" in results["web"]:
+            return [{"title": item.get("title", ""), "link": item.get("url", ""), "snippet": item.get("description", ""), "source": item.get("source", "")} 
+                    for item in results["web"]["results"][:5]]
+    if SERPER_API_KEY:
+        results = serper_search(query)
+        if "organic" in results:
+            return [{"title": item.get("title", ""), "link": item.get("link", ""), "snippet": item.get("snippet", ""), "source": item.get("source", "")} 
+                    for item in results["organic"][:5]]
+    return [{"title": "No news available", "snippet": "Search APIs not configured or failed."}]
+
+def get_market_sentiment(ticker, use_brave_search, enable_search, model_name="llama3-8b-8192"):
+    if not enable_search:
+        return "Real-time search is disabled"
+    query = f"{ticker} stock market sentiment analysis"
+    snippets = []
+    if use_brave_search and BRAVE_API_KEY:
+        results = brave_search(query)
+        if "web" in results and "results" in results["web"]:
+            snippets = [item["description"] for item in results["web"]["results"][:3] if "description" in item]
+    elif SERPER_API_KEY:
+        results = serper_search(query)
+        if "organic" in results:
+            snippets = [item["snippet"] for item in results["organic"][:3] if "snippet" in item]
+    if snippets:
+        try:
+            completion = client.chat.completions.create(
+                model=ModelName.validate(model_name),
+                messages=[
+                    {"role": "system", "content": "Analyze the sentiment based on the provided text."},
+                    {"role": "user", "content": "\n".join(snippets)}
+                ],
+                temperature=0.2,
+                max_tokens=150
+            )
+            return completion.choices[0].message.content
+        except Exception as e:
+            return "Unable to determine sentiment"
+    return "No sentiment data available"
 
-# New Finance-specific tools
-def get_stock_data(ticker):
-    """Tool to fetch latest stock data for a given ticker"""
+def get_stock_data(ticker, enable_stock_data):
+    if not enable_stock_data:
+        return {"message": "Real-time stock data is disabled"}
     try:
         stock = yf.Ticker(ticker)
         info = stock.info
@@ -180,244 +273,241 @@ def get_stock_data(ticker):
             "current_price": info.get("currentPrice", info.get("regularMarketPrice", "N/A")),
             "52wk_high": info.get("fiftyTwoWeekHigh", "N/A"),
             "market_cap": info.get("marketCap", "N/A"),
-            "pe_ratio": info.get("trailingPE", "N/A"),
-            "dividend_yield": info.get("dividendYield", "N/A"),
-            "beta": info.get("beta", "N/A"),
-            "average_volume": info.get("averageVolume", "N/A")
+            "pe_ratio": info.get("trailingPE", "N/A")
         }
     except Exception as e:
-        print(f"Error fetching stock data: {e}")
         raise e
 
-def get_stock_history(ticker, period="1y"):
-    """Get historical data for charting"""
+def get_stock_history(ticker, period, enable_stock_data):
+    if not enable_stock_data:
+        return pd.DataFrame()
     try:
         stock = yf.Ticker(ticker)
-        hist = stock.history(period=period)
-        return hist
+        return stock.history(period=period)
     except Exception as e:
-        print(f"Error fetching stock history: {e}")
         return pd.DataFrame()
 
-def get_fred_data(indicator):
-    """Get economic data from FRED API"""
-    api_key = os.getenv("FRED_API_KEY", "")
-    if not api_key:
-        return "FRED API key not configured"
-    
-    base_url = "https://api.stlouisfed.org/fred/series/observations"
-    params = {
-        "series_id": indicator,
-        "api_key": api_key,
-        "file_type": "json",
-        "sort_order": "desc",
-        "limit": 100
-    }
-    
-    try:
-        response = requests.get(base_url, params=params)
-        data = response.json()
-        return data.get("observations", [])
-    except Exception as e:
-        print(f"Error fetching FRED data: {e}")
-        return []
-
-def create_stock_chart(ticker, period="1y"):
-    """Create an interactive stock chart using Plotly"""
-    try:
-        df = get_stock_history(ticker, period)
-        if df.empty:
-            return None
-            
-        fig = go.Figure()
-        
-        # Add candlestick chart
-        fig.add_trace(
-            go.Candlestick(
-                x=df.index,
-                open=df['Open'],
-                high=df['High'],
-                low=df['Low'],
-                close=df['Close'],
-                name=ticker
-            )
-        )
-        
-        # Add volume as bar chart on secondary y-axis
-        fig.add_trace(
-            go.Bar(
-                x=df.index,
-                y=df['Volume'],
-                name='Volume',
-                marker_color='rgba(0, 128, 0, 0.3)',
-                yaxis='y2'
-            )
-        )
-        
-        # Update layout for dual y-axis
-        fig.update_layout(
-            title=f'{ticker} Stock Price',
-            yaxis_title='Price (USD)',
-            xaxis_title='Date',
-            template='plotly_white',
-            yaxis=dict(
-                domain=[0.3, 1.0]
-            ),
-            yaxis2=dict(
-                domain=[0, 0.2],
-                title='Volume'
-            ),
-            legend=dict(
-                orientation="h",
-                yanchor="bottom",
-                y=1.02,
-                xanchor="right",
-                x=1
-            ),
-            height=500
-        )
-        
-        return fig
-    except Exception as e:
-        print(f"Error creating stock chart: {e}")
+def create_stock_chart(ticker, period, enable_stock_data):
+    if not enable_stock_data:
+        return None
+    df = get_stock_history(ticker, period, enable_stock_data)
+    if df.empty:
         return None
+    fig = go.Figure()
+    fig.add_trace(go.Candlestick(x=df.index, open=df['Open'], high=df['High'], low=df['Low'], close=df['Close'], name=ticker))
+    fig.add_trace(go.Bar(x=df.index, y=df['Volume'], name='Volume', marker_color='rgba(0, 128, 0, 0.3)', yaxis='y2'))
+    fig.update_layout(
+        title=f'{ticker} Stock Price',
+        yaxis_title='Price (USD)',
+        xaxis_title='Date',
+        template='plotly_white',
+        yaxis=dict(domain=[0.3, 1.0]),
+        yaxis2=dict(domain=[0, 0.2], title='Volume'),
+        height=500
+    )
+    return fig
 
-def analyze_ticker(ticker_input, period):
-    """Process the ticker input and return analysis"""
-    if not ticker_input:
-        return None, "Please enter a valid ticker symbol", None
+def analyze_ticker(ticker_input, period, use_brave_search, enable_stock_data, enable_search):
+    try:
+        input_data = StockAnalysisInput(ticker=ticker_input, period=period)
+    except ValueError as e:
+        return None, str(e), None
+    
+    ticker = input_data.ticker
+    period = input_data.period
     
-    ticker = ticker_input.strip().upper()
-    if ticker.startswith("$"):
-        ticker = ticker[1:]
+    if not enable_stock_data:
+        return None, "Real-time stock data is disabled.", None
     
     try:
-        stock_data = get_stock_data(ticker)
-        chart = create_stock_chart(ticker, period)
-        
-        # Create a formatted summary
+        stock_data = get_stock_data(ticker, enable_stock_data)
+        stock_history = get_stock_history(ticker, period, enable_stock_data)
+        chart = create_stock_chart(ticker, period, enable_stock_data)
+        chart_data_store[ticker] = {"history": stock_history, "stats": stock_data, "period": period}
+        sentiment = get_market_sentiment(ticker, use_brave_search, enable_search, model_name="llama3-8b-8192")
         summary = f"""
 ### {ticker} Analysis
 **Current Price:** ${stock_data['current_price']}  
 **52-Week High:** ${stock_data['52wk_high']}  
 **Market Cap:** ${stock_data['market_cap']:,}  
 **P/E Ratio:** {stock_data['pe_ratio']}  
-**Dividend Yield:** {stock_data['dividend_yield'] * 100 if stock_data['dividend_yield'] != 'N/A' else 'N/A'}%  
-**Beta:** {stock_data['beta']}  
-**Avg Volume:** {stock_data['average_volume']:,}  
-        """
-        
+**Market Sentiment:** {sentiment}
+"""
         return chart, summary, ticker
     except Exception as e:
         return None, f"Error analyzing ticker {ticker}: {str(e)}", None
 
-# Gradio interface
-with gr.Blocks(css=custom_css, theme=gr.themes.Soft()) as demo:
-    current_session_id = gr.State(None)
-    pdf_state = gr.State({"page_images": [], "total_pages": 0, "total_words": 0})
-    current_ticker = gr.State(None)
+def generate_response(message, session_id, model_name, history, current_ticker, use_brave_search, enable_search, enable_stock_data):
+    if not message:
+        return history
     
-    gr.HTML("""
-    <div class="header">
-        <div class="header-title">Fin-Vision</div>
-        <div class="header-subtitle">Analyze financial documents with Groq's LLM API.</div>
-    </div>
-    """)
+    try:
+        model_name = ModelName.validate(model_name)
+    except ValueError as e:
+        return history + [(message, str(e))]
     
-    with gr.Row(elem_classes="container"):
-        with gr.Column(scale=1, min_width=300):
-            pdf_file = gr.File(label="Upload PDF Document", file_types=[".pdf"], type="binary")
-            upload_button = gr.Button("Process PDF", variant="primary")
-            pdf_status = gr.Markdown("No PDF uploaded yet")
-            model_dropdown = gr.Dropdown(
-                choices=["llama3-70b-8192", "llama3-8b-8192", "mixtral-8x7b-32768", "gemma-7b-it"],
-                value="llama3-70b-8192",
-                label="Select Groq Model"
-            )
+    try:
+        context = ""
+        if session_id in user_vectorstores:
+            vectorstore = user_vectorstores[session_id]
+            docs = vectorstore.similarity_search(message, k=3)
+            if docs:
+                context = "\n\nRelevant information from uploaded PDF:\n" + "\n".join(f"- {doc.page_content}" for doc in docs)
+        
+        if message.startswith("$") and len(message) > 1 and len(message) <= 6:
+            ticker = message[1:].upper()
+            stock_data = get_stock_data(ticker, enable_stock_data)
+            news = get_financial_news(ticker, use_brave_search, enable_search, model_name)
+            sentiment = get_market_sentiment(ticker, use_brave_search, enable_search, model_name)
+            response = f"**Stock Information for {ticker}**\n\n"
+            if "message" in stock_data:
+                response += stock_data["message"] + "\n"
+            else:
+                response += f"Current Price: ${stock_data['current_price']}\n52-Week High: ${stock_data['52wk_high']}\nMarket Cap: ${stock_data['market_cap']:,}\nP/E Ratio: {stock_data['pe_ratio']}\n"
+            response += f"**Market Sentiment:**\n{sentiment}\n\n**Recent News:**\n"
+            for i, item in enumerate(news[:3]):
+                response += f"{i+1}. [{item['title']}]({item['link']})\n   {item['snippet'][:100]}...\n"
+            history.append((message, response))
+            return history
+        
+        if message.lower().startswith("/news "):
+            topic = message[6:].strip()
+            news = get_financial_news(topic, use_brave_search, enable_search, model_name)
+            response = f"**Latest Financial News on {topic}:**\n\n"
+            for i, item in enumerate(news[:5]):
+                response += f"{i+1}. **{item['title']}**\n   {item['snippet']}\n   [Read more]({item['link']})\n\n"
+            history.append((message, response))
+            return history
+        
+        system_prompt = "You are a financial assistant specializing in analyzing financial reports, statements, and market trends."
+        if context:
+            system_prompt += " Use the following context if relevant: " + context
+        completion = client.chat.completions.create(
+            model=model_name,
+            messages=[{"role": "system", "content": system_prompt}, {"role": "user", "content": message}],
+            temperature=0.7,
+            max_tokens=1024
+        )
+        response = completion.choices[0].message.content
+        history.append((message, response))
+        return history
+    except Exception as e:
+        history.append((message, f"Error generating response: {str(e)}"))
+        return history
+
+def update_pdf_viewer(pdf_state: PDFState):
+    if not pdf_state.total_pages:
+        return 0, None, "No PDF uploaded yet"
+    img_data = base64.b64decode(pdf_state.page_images[0])
+    img = Image.open(io.BytesIO(img_data))
+    return pdf_state.total_pages, img, f"**Total Pages:** {pdf_state.total_pages}\n**Total Words:** {pdf_state.total_words}"
+
+def update_image(page_num, pdf_state: PDFState):
+    if not pdf_state.total_pages or page_num < 1 or page_num > pdf_state.total_pages:
+        return None
+    img_data = base64.b64decode(pdf_state.page_images[page_num - 1])
+    return Image.open(io.BytesIO(img_data))
+
+# Gradio Interface
+def create_interface():
+    with gr.Blocks(css=custom_css, js=custom_js) as demo:
+        gr.HTML(custom_html)
+        current_session_id = gr.State(None)
+        pdf_state = gr.State(PDFState(page_images=[], total_pages=0, total_words=0))
+        current_ticker = gr.State(None)
+        
+        gr.HTML("""
+        <div class="header">
+            <div class="header-title">Fin-Vision</div>
+            <div class="header-subtitle">Analyze financial documents with Groq's LLM API.</div>
+        </div>
+        """)
+        
+        with gr.Row(elem_classes="container"):
+            with gr.Column(scale=1, min_width=300):
+                pdf_file = gr.File(label="Upload PDF Document", file_types=[".pdf"], type="binary", elem_id="pdf_file")
+                upload_button = gr.Button("Process PDF", variant="primary")
+                pdf_status = gr.Markdown("No PDF uploaded yet")
+                
+                with gr.Group():
+                    gr.Markdown("### Real-Time Data Settings")
+                    enable_stock_data = gr.Checkbox(label="Enable Real-Time Stock Data", value=True)
+                    enable_search = gr.Checkbox(label="Enable Real-Time Search", value=True)
+                    use_brave_search = gr.Checkbox(label="Use Brave Search (unchecked = Serper)", value=False)
+                
+                model_dropdown = gr.Dropdown(
+                    choices=["llama3-70b-8192", "llama3-8b-8192", "mixtral-8x7b-32768", "gemma-7b-it"],
+                    value="llama3-70b-8192",
+                    label="Select Groq Model"
+                )
+                theme_button = gr.Button("Toggle Theme")
+                
+                with gr.Group(elem_classes="tool-container"):
+                    ticker_input = gr.Textbox(label="Enter Ticker Symbol (e.g., AAPL)", placeholder="AAPL", elem_id="ticker_input")
+                    period_dropdown = gr.Dropdown(choices=["1mo", "3mo", "6mo", "1y", "2y", "5y", "max"], value="1y", label="Time Period")
+                    analyze_button = gr.Button("Analyze Stock")
             
-            # Finance Tools Section
-            gr.Markdown("### Financial Tools", elem_classes="tool-title")
-            with gr.Group(elem_classes="tool-container"):
+            with gr.Column(scale=2, min_width=600):
                 with gr.Tabs():
-                    with gr.TabItem("Stock Analysis"):
-                        ticker_input = gr.Textbox(label="Enter Ticker Symbol (e.g., AAPL)", placeholder="AAPL")
-                        period_dropdown = gr.Dropdown(
-                            choices=["1mo", "3mo", "6mo", "1y", "2y", "5y", "max"],
-                            value="1y",
-                            label="Time Period"
-                        )
-                        analyze_button = gr.Button("Analyze Stock")
-            
-        with gr.Column(scale=2, min_width=600):
-            with gr.Tabs():
-                with gr.TabItem("PDF Viewer"):
-                    with gr.Column(elem_classes="pdf-viewer-container"):
+                    with gr.TabItem("PDF Viewer"):
                         page_slider = gr.Slider(minimum=1, maximum=1, step=1, label="Page Number", value=1)
                         pdf_image = gr.Image(label="PDF Page", type="pil", elem_classes="pdf-viewer-image")
                         stats_display = gr.Markdown("No PDF uploaded yet", elem_classes="stats-box")
-                
-                with gr.TabItem("Stock Analysis"):
-                    with gr.Column(elem_classes="pdf-viewer-container"):
+                    
+                    with gr.TabItem("Stock Analysis"):
                         stock_chart = gr.Plot(label="Stock Price Chart", elem_classes="chart-container")
                         stock_summary = gr.Markdown("Enter a ticker symbol to see analysis")
-    
-    with gr.Row(elem_classes="container"):
-        with gr.Column(scale=2, min_width=600):
+        
+        with gr.Row(elem_classes="container"):
             chatbot = gr.Chatbot(height=500, bubble_full_width=False, show_copy_button=True, elem_classes="chat-container")
             with gr.Row():
-                msg = gr.Textbox(show_label=False, placeholder="Ask about your financial document or type $TICKER for stock info...", scale=5)
+                msg = gr.Textbox(show_label=False, placeholder="Ask about your financial document...", scale=5)
                 send_btn = gr.Button("Send", scale=1)
-            clear_btn = gr.Button("Clear Conversation")
-    
-    # Event Handlers
-    upload_button.click(
-        process_pdf,
-        inputs=[pdf_file],
-        outputs=[current_session_id, pdf_status, pdf_state]
-    ).then(
-        update_pdf_viewer,
-        inputs=[pdf_state],
-        outputs=[page_slider, pdf_image, stats_display]
-    )
-    
-    msg.submit(
-        generate_response,
-        inputs=[msg, current_session_id, model_dropdown, chatbot],
-        outputs=[chatbot]
-    ).then(lambda: "", None, [msg])
-    
-    send_btn.click(
-        generate_response,
-        inputs=[msg, current_session_id, model_dropdown, chatbot],
-        outputs=[chatbot]
-    ).then(lambda: "", None, [msg])
-    
-    clear_btn.click(
-        lambda: ([], None, "No PDF uploaded yet", {"page_images": [], "total_pages": 0, "total_words": 0}, 0, None, "No PDF uploaded yet", None),
-        None,
-        [chatbot, current_session_id, pdf_status, pdf_state, page_slider, pdf_image, stats_display, current_ticker]
-    )
-    
-    page_slider.change(
-        update_image,
-        inputs=[page_slider, pdf_state],
-        outputs=[pdf_image]
-    )
+        
+        # Event Handlers
+        upload_button.click(
+            js="showSpinner",
+            fn=process_pdf,
+            inputs=[pdf_file],
+            outputs=[current_session_id, pdf_status, pdf_state]
+        ).then(
+            update_pdf_viewer,
+            inputs=[pdf_state],
+            outputs=[page_slider, pdf_image, stats_display]
+        ).then(
+            js="hideSpinner",
+            inputs=[],
+            outputs=[]
+        )
+        
+        analyze_button.click(
+            js="showSpinner",
+            fn=analyze_ticker,
+            inputs=[ticker_input, period_dropdown, use_brave_search, enable_stock_data, enable_search],
+            outputs=[stock_chart, stock_summary, current_ticker]
+        ).then(
+            js="hideSpinner",
+            inputs=[],
+            outputs=[]
+        )
+        
+        msg.submit(
+            generate_response,
+            inputs=[msg, current_session_id, model_dropdown, chatbot, current_ticker, use_brave_search, enable_search, enable_stock_data],
+            outputs=[chatbot]
+        ).then(lambda: "", None, [msg])
+        
+        send_btn.click(
+            generate_response,
+            inputs=[msg, current_session_id, model_dropdown, chatbot, current_ticker, use_brave_search, enable_search, enable_stock_data],
+            outputs=[chatbot]
+        ).then(lambda: "", None, [msg])
+        
+        page_slider.change(update_image, inputs=[page_slider, pdf_state], outputs=[pdf_image])
+        theme_button.click(js="toggleTheme")
     
-    # Stock analysis handler
-    analyze_button.click(
-        analyze_ticker,
-        inputs=[ticker_input, period_dropdown],
-        outputs=[stock_chart, stock_summary, current_ticker]
-    )
-
-# Add footer with attribution
-gr.HTML("""
-<div style="text-align: center; margin-top: 20px; padding: 10px; color: #666; font-size: 0.8rem; border-top: 1px solid #eee;">
-    Created by Calvin Allen Crawford
-</div>
-""")
+    return demo
 
-# Launch the app
 if __name__ == "__main__":
+    demo = create_interface()
     demo.launch()