Update app.py

app.py

@@ -1,34 +1,50 @@
+import os
+import sys
 import gradio as gr
 import pandas as pd
 import numpy as np
 import plotly.graph_objects as go
 from datetime import datetime
-import os
 import json
 from typing import Dict, List, Tuple, Optional
 
-# LangChain imports for conversational AI
-from langchain.schema import HumanMessage, AIMessage, SystemMessage
-from langchain_openai import ChatOpenAI
+# ========== Optional memory monitoring ==========
+try:
+    import psutil
+    print(f"🧠 Memory at startup: {psutil.virtual_memory().percent}% used, "
+          f"{psutil.virtual_memory().available / 1e9:.2f} GB free")
+except ImportError:
+    pass
+
+# ========== Lazy imports for heavy modules ==========
+def get_llm():
+    from langchain_openai import ChatOpenAI
+    from langchain_core.messages import HumanMessage, AIMessage, SystemMessage
+    api_key = os.getenv("OPENROUTER_API_KEY")
+    if not api_key:
+        raise ValueError("OPENROUTER_API_KEY environment variable not set.")
+    return ChatOpenAI(
+        base_url="https://openrouter.ai/api/v1",
+        api_key=api_key,
+        model="openai/gpt-3.5-turbo:free",
+        temperature=0.3,
+    )
 
-# Import our modules
-from index_calculator import IndexCalculator
-from data_sources import RealTimeIndexCalculator, IndianStockAPI, GlobalStockAPI
+def get_realtime_calculator():
+    from data_sources import RealTimeIndexCalculator, IndianStockAPI, GlobalStockAPI
+    return RealTimeIndexCalculator(), IndianStockAPI(), GlobalStockAPI()
 
-# Initialize memory for conversation context using a custom memory handler
+# ========== Simple memory (optional, not heavily used) ==========
 class SimpleMemory:
     def __init__(self):
         self.messages = []
-    
     def add_message(self, message: str, role: str):
         self.messages.append({'role': role, 'content': message})
-    
     def get_history(self):
         return self.messages
 
 memory = SimpleMemory()
 
-# System prompt for the AI assistant
 SYSTEM_PROMPT = """You are an expert financial analyst specializing in index numbers and real-time market data. 
 You help users understand, calculate, and analyze various types of indices including:
 - Laspeyres, Paasche, and Fisher price indices
@@ -41,174 +57,199 @@ You have access to:
 3. A powerful index calculator for custom analysis
 When users ask about stocks or indices, fetch live data. Provide clear explanations with numbers."""
 
-# Initialize LLM
-def get_llm():
-    """Initialize LLM with OpenRouter for free access"""
-    return ChatOpenAI(
-        base_url="https://openrouter.ai/api/v1",
-        api_key=os.getenv("OPENROUTER_API_KEY", "sk-or-v1-64e1068c3d61c7c1d7f5d9b8a2e3f4c5d6e7f8g9h0i1j2k3l4m5n6o7p8q9r0"),
-        model="openai/gpt-3.5-turbo:free",
-        temperature=0.3,
-    )
-
-# Initialize real-time calculator
-rt_calc = RealTimeIndexCalculator()
-indian_api = IndianStockAPI()
-global_api = GlobalStockAPI()
-
+# ========== Core conversational function ==========
 def process_query(query: str, history: List, exchange: str) -> Tuple[str, List, pd.DataFrame, str]:
     """
     Process user query with real-time data awareness
     """
-    llm = get_llm()
+    rt_calc, indian_api, global_api = get_realtime_calculator()
     
-    # Check if query is about real-time data
     response_text = ""
     data_preview = pd.DataFrame()
     
-    # Detect Indian stocks (common patterns)
+    # Common Indian stocks
     indian_stocks = ['RELIANCE', 'TCS', 'INFY', 'HDFCBANK', 'ICICIBANK', 'SBIN', 
                      'ITC', 'BHARTIARTL', 'KOTAKBANK', 'LT', 'WIPRO', 'AXISBANK']
     
     query_lower = query.lower()
     
-    # Handle real-time queries
+    # Real-time data handling
     if any(stock.lower() in query_lower for stock in indian_stocks) or 'nifty' in query_lower or 'sensex' in query_lower:
-        # Extract stock symbol from query
         for stock in indian_stocks:
             if stock.lower() in query_lower:
-                data = indian_api.get_stock_price(stock, exchange)
-                if data and 'data' in data:
-                    formatted = indian_api.format_for_index(data)
-                    response_text = f"**{formatted['company_name']}**\n\n"
-                    response_text += f"💰 Current Price: ₹{formatted['last_price']:,.2f}\n"
-                    response_text += f"📊 Change: {formatted['change']:+.2f} ({formatted['percent_change']:+.2f}%)\n"
-                    response_text += f"📈 Day Range: ₹{formatted['day_low']} - ₹{formatted['day_high']}\n"
-                    response_text += f"📉 52-Week Range: ₹{formatted['year_low']} - ₹{formatted['year_high']}\n"
-                    response_text += f"💹 Market Cap: ₹{formatted['market_cap']:,.0f} Cr\n"
-                    response_text += f"📊 P/E Ratio: {formatted['pe_ratio']:.2f}\n"
-                    
-                    # Create preview DataFrame
-                    data_preview = pd.DataFrame([formatted])
+                try:
+                    data = indian_api.get_stock_price(stock, exchange)
+                    if data and 'data' in data:
+                        formatted = indian_api.format_for_index(data)
+                        response_text = f"**{formatted['company_name']}**\n\n"
+                        response_text += f"💰 Current Price: ₹{formatted['last_price']:,.2f}\n"
+                        response_text += f"📊 Change: {formatted['change']:+.2f} ({formatted['percent_change']:+.2f}%)\n"
+                        response_text += f"📈 Day Range: ₹{formatted['day_low']} - ₹{formatted['day_high']}\n"
+                        response_text += f"📉 52-Week Range: ₹{formatted['year_low']} - ₹{formatted['year_high']}\n"
+                        response_text += f"💹 Market Cap: ₹{formatted['market_cap']:,.0f} Cr\n"
+                        response_text += f"📊 P/E Ratio: {formatted['pe_ratio']:.2f}\n"
+                        data_preview = pd.DataFrame([formatted])
+                except Exception as e:
+                    response_text = f"Error fetching data: {str(e)}"
                 break
         
-        if 'nifty' in query_lower:
-            nifty_data = global_api.get_stock_info('^NSEI')
-            if nifty_data:
-                response_text += f"\n\n**Nifty 50 Index**\n"
-                response_text += f"Current: {nifty_data['last_price']:,.2f}\n"
-                response_text += f"Change: {nifty_data['change']:+.2f} ({nifty_data['percent_change']:+.2f}%)"
+        if 'nifty' in query_lower and not response_text:
+            try:
+                nifty_data = global_api.get_stock_info('^NSEI')
+                if nifty_data:
+                    response_text = f"**Nifty 50 Index**\nCurrent: {nifty_data['last_price']:,.2f}\n"
+                    response_text += f"Change: {nifty_data['change']:+.2f} ({nifty_data['percent_change']:+.2f}%)"
+            except Exception as e:
+                response_text = f"Error fetching Nifty: {str(e)}"
     
-    # Handle global stocks
     elif any(ticker in query.upper() for ticker in ['AAPL', 'GOOGL', 'MSFT', 'AMZN', 'TSLA']):
         tickers = ['AAPL', 'GOOGL', 'MSFT', 'AMZN', 'TSLA']
         for ticker in tickers:
             if ticker in query.upper():
-                data = global_api.get_stock_info(ticker)
-                if data:
-                    response_text = f"**{data['company_name']} ({ticker})**\n\n"
-                    response_text += f"💰 Current Price: ${data['last_price']:,.2f}\n"
-                    response_text += f"📊 Change: ${data['change']:+.2f} ({data['percent_change']:+.2f}%)\n"
-                    response_text += f"📈 Day Range: ${data['day_low']} - ${data['day_high']}\n"
-                    response_text += f"💹 Market Cap: ${data['market_cap']/1e9:,.2f}B\n"
-                    data_preview = pd.DataFrame([data])
+                try:
+                    data = global_api.get_stock_info(ticker)
+                    if data:
+                        response_text = f"**{data['company_name']} ({ticker})**\n\n"
+                        response_text += f"💰 Current Price: ${data['last_price']:,.2f}\n"
+                        response_text += f"📊 Change: ${data['change']:+.2f} ({data['percent_change']:+.2f}%)\n"
+                        response_text += f"📈 Day Range: ${data['day_low']} - ${data['day_high']}\n"
+                        response_text += f"💹 Market Cap: ${data['market_cap']/1e9:,.2f}B\n"
+                        data_preview = pd.DataFrame([data])
+                except Exception as e:
+                    response_text = f"Error fetching data: {str(e)}"
                 break
     
     # If no real-time match, use LLM
     if not response_text:
-        messages = [
-            SystemMessage(content=SYSTEM_PROMPT),
-        ] + [
-            HumanMessage(content=msg) if i % 2 == 0 else AIMessage(content=msg) 
-            for i, msg in enumerate(history)
-        ] if history else [] + [HumanMessage(content=query)]
-        
         try:
+            llm = get_llm()
+            from langchain_core.messages import HumanMessage, AIMessage, SystemMessage
+            
+            messages = [SystemMessage(content=SYSTEM_PROMPT)]
+            if history:
+                # history is a flat list: [user1, assistant1, user2, assistant2, ...]
+                for i in range(0, len(history), 2):
+                    if i < len(history):
+                        messages.append(HumanMessage(content=history[i]))
+                    if i+1 < len(history):
+                        messages.append(AIMessage(content=history[i+1]))
+            messages.append(HumanMessage(content=query))
+            
             response = llm.invoke(messages)
             response_text = response.content
         except Exception as e:
-            response_text = f"I encountered an error: {str(e)}"
+            response_text = f"LLM error: {str(e)}"
     
-    # Update history
+    # Update history (keep last 10 exchanges to avoid memory bloat)
     updated_history = history + [query, response_text]
+    if len(updated_history) > 20:
+        updated_history = updated_history[-20:]
     
     return response_text, updated_history, data_preview, ""
 
-def create_realtime_ui():
-    """Create the real-time market data tab UI"""
-    with gr.Row():
-        with gr.Column(scale=1):
-            gr.Markdown("""
-            ### 📡 Real-Time Market Data
-            
-            Fetch live stock prices from:
-            - **NSE** (National Stock Exchange of India) [citation:1]
-            - **BSE** (Bombay Stock Exchange) [citation:1]
-            - **Global Markets** via Yahoo Finance [citation:2]
-            """)
-            
-            exchange = gr.Radio(
-                choices=["NSE", "BSE", "Global"],
-                value="NSE",
-                label="Select Exchange"
-            )
-            
-            symbol_input = gr.Textbox(
-                label="Stock Symbol",
-                placeholder="RELIANCE, TCS, INFY, or AAPL, GOOGL for global"
-            )
-            
-            fetch_btn = gr.Button("Fetch Live Data", variant="primary")
-            
-            gr.Markdown("""
-            ### 📊 Build Custom Index
-            Enter up to 5 symbols to create your own index
-            """)
-            
-            symbols_list = gr.Dataframe(
-                headers=["Symbol", "Weight (%)"],
-                datatype=["str", "number"],
-                row_count=5,
-                col_count=2,
-                value=[["RELIANCE", 20], ["TCS", 20], ["INFY", 20], ["HDFCBANK", 20], ["ITC", 20]]
-            )
-            
-            index_type = gr.Dropdown(
-                choices=["Price-Weighted", "Market Cap-Weighted", "Equal-Weighted"],
-                value="Market Cap-Weighted",
-                label="Index Type"
-            )
-            
-            build_index_btn = gr.Button("Build Custom Index", variant="secondary")
-        
-        with gr.Column(scale=1):
-            stock_output = gr.JSON(label="Stock Details")
-            
-            with gr.Row():
-                price = gr.Number(label="Current Price")
-                change = gr.Number(label="Change")
-                pct_change = gr.Number(label="% Change")
-            
-            chart_output = gr.Plot(label="Price Chart")
-            
-            index_result = gr.JSON(label="Custom Index Result")
+# ========== Real-time UI functions ==========
+def fetch_live_stock(symbol: str, exchange: str):
+    """Fetch live stock data and return JSON, price, change, pct, chart"""
+    _, indian_api, global_api = get_realtime_calculator()
+    if exchange in ["NSE", "BSE"]:
+        data = indian_api.get_stock_price(symbol, exchange)
+        if data and 'data' in data:
+            formatted = indian_api.format_for_index(data)
+            # Simple chart
+            fig = go.Figure()
+            fig.add_trace(go.Indicator(
+                mode="number+delta",
+                value=formatted['last_price'],
+                number={'prefix': '₹'},
+                delta={'reference': formatted['last_price'] - formatted['change']}
+            ))
+            fig.update_layout(height=300)
+            return formatted, formatted['last_price'], formatted['change'], formatted['percent_change'], fig
+    else:
+        data = global_api.get_stock_info(symbol)
+        if data:
+            fig = go.Figure()
+            fig.add_trace(go.Indicator(
+                mode="number+delta",
+                value=data['last_price'],
+                number={'prefix': '$'},
+                delta={'reference': data['last_price'] - data['change']}
+            ))
+            fig.update_layout(height=300)
+            return data, data['last_price'], data['change'], data['percent_change'], fig
+    return {}, 0, 0, 0, go.Figure()
+
+def build_custom_index(symbols_df: pd.DataFrame, index_type: str, exchange: str):
+    """Build custom index from symbols and weights"""
+    _, indian_api, global_api = get_realtime_calculator()
+    weights = []
+    prices = []
+    total_weight = 0
+    for _, row in symbols_df.iterrows():
+        symbol = row['Symbol']
+        weight = float(row['Weight (%)']) / 100.0
+        if exchange in ["NSE", "BSE"]:
+            data = indian_api.get_stock_price(symbol, exchange)
+            if data and 'data' in data:
+                # Adjust based on actual API structure
+                price = data['data'][0]['last_price'] if 'data' in data and data['data'] else 0
+            else:
+                price = 0
+        else:
+            data = global_api.get_stock_info(symbol)
+            price = data.get('last_price', 0) if data else 0
+        weights.append(weight)
+        prices.append(price)
+        total_weight += weight
+    
+    # Normalize weights if total != 1
+    if total_weight != 1:
+        weights = [w/total_weight for w in weights]
+    
+    if index_type == "Price-Weighted":
+        index_value = sum(prices) / len(prices) if prices else 0
+    elif index_type == "Equal-Weighted":
+        index_value = sum(prices) / len(prices) if prices else 0
+    elif index_type == "Market Cap-Weighted":
+        index_value = sum(w * p for w, p in zip(weights, prices)) if prices else 0
+    else:
+        index_value = 0
+    
+    return {"index_value": round(index_value, 2), "components": prices, "weights": weights}
+
+# ========== Comparison UI functions ==========
+def create_comparison_ui():
+    """Return UI components for index comparison tab"""
+    with gr.Blocks() as comp_ui:
+        with gr.Row():
+            symbols_input = gr.Textbox(label="Symbols (comma-separated)", placeholder="RELIANCE, TCS, INFY")
+            exchange = gr.Radio(["NSE", "BSE", "Global"], value="NSE", label="Exchange")
+        compare_btn = gr.Button("Compare Indices")
+        results_output = gr.JSON(label="Comparison Results")
+        chart_output = gr.Plot(label="Comparison Chart")
+    return {"symbols": symbols_input, "exchange": exchange, "btn": compare_btn,
+            "results": results_output, "chart": chart_output}
+
+def compare_live_indices(symbols_str: str, exchange: str):
+    """Compare multiple indices"""
+    symbols = [s.strip() for s in symbols_str.split(",") if s.strip()]
+    _, indian_api, global_api = get_realtime_calculator()
+    data = {}
+    for sym in symbols:
+        if exchange in ["NSE", "BSE"]:
+            d = indian_api.get_stock_price(sym, exchange)
+            if d and 'data' in d and d['data']:
+                data[sym] = d['data'][0]['last_price']
+        else:
+            d = global_api.get_stock_info(sym)
+            if d:
+                data[sym] = d['last_price']
     
-    return {
-        "exchange": exchange,
-        "symbol": symbol_input,
-        "fetch_btn": fetch_btn,
-        "symbols_list": symbols_list,
-        "index_type": index_type,
-        "build_btn": build_index_btn,
-        "stock_output": stock_output,
-        "price": price,
-        "change": change,
-        "pct_change": pct_change,
-        "chart": chart_output,
-        "index_result": index_result
-    }
+    fig = go.Figure(data=[go.Bar(x=list(data.keys()), y=list(data.values()))])
+    fig.update_layout(title="Index Comparison", yaxis_title="Price")
+    return data, fig
 
-# Build the main Gradio interface
+# ========== Main Gradio interface ==========
 with gr.Blocks(title="AI-Powered Index Analyst with Real-Time Data", theme=gr.themes.Soft()) as demo:
     gr.Markdown("""
     # 📈 AI-Powered Index Analyst with Real-Time Market Data
@@ -243,6 +284,7 @@ with gr.Blocks(title="AI-Powered Index Analyst with Real-Time Data", theme=gr.th
         data_output = gr.DataFrame(label="Data Preview")
         calc_output = gr.Textbox(label="Calculation Results", lines=3)
         
+        # Output mapping: (response, updated_history, data_preview, status)
         submit.click(
             process_query,
             inputs=[msg, chat_history, exchange_selector],
@@ -261,42 +303,108 @@ with gr.Blocks(title="AI-Powered Index Analyst with Real-Time Data", theme=gr.th
         )
     
     with gr.Tab("📡 Real-Time Market Data"):
-        rt_ui = create_realtime_ui()
+        with gr.Row():
+            with gr.Column(scale=1):
+                gr.Markdown("""
+                ### 📡 Real-Time Market Data
+                
+                Fetch live stock prices from:
+                - **NSE** (National Stock Exchange of India) [citation:1]
+                - **BSE** (Bombay Stock Exchange) [citation:1]
+                - **Global Markets** via Yahoo Finance [citation:2]
+                """)
+                
+                exchange = gr.Radio(
+                    choices=["NSE", "BSE", "Global"],
+                    value="NSE",
+                    label="Select Exchange"
+                )
+                
+                symbol_input = gr.Textbox(
+                    label="Stock Symbol",
+                    placeholder="RELIANCE, TCS, INFY, or AAPL, GOOGL for global"
+                )
+                
+                fetch_btn = gr.Button("Fetch Live Data", variant="primary")
+                
+                gr.Markdown("""
+                ### 📊 Build Custom Index
+                Enter up to 5 symbols to create your own index
+                """)
+                
+                symbols_list = gr.Dataframe(
+                    headers=["Symbol", "Weight (%)"],
+                    datatype=["str", "number"],
+                    row_count=5,
+                    col_count=2,
+                    value=[["RELIANCE", 20], ["TCS", 20], ["INFY", 20], ["HDFCBANK", 20], ["ITC", 20]]
+                )
+                
+                index_type = gr.Dropdown(
+                    choices=["Price-Weighted", "Market Cap-Weighted", "Equal-Weighted"],
+                    value="Market Cap-Weighted",
+                    label="Index Type"
+                )
+                
+                build_index_btn = gr.Button("Build Custom Index", variant="secondary")
+            
+            with gr.Column(scale=1):
+                stock_output = gr.JSON(label="Stock Details")
+                
+                with gr.Row():
+                    price = gr.Number(label="Current Price")
+                    change = gr.Number(label="Change")
+                    pct_change = gr.Number(label="% Change")
+                
+                chart_output = gr.Plot(label="Price Chart")
+                
+                index_result = gr.JSON(label="Custom Index Result")
         
-        rt_ui["fetch_btn"].click(
+        fetch_btn.click(
             fetch_live_stock,
-            inputs=[rt_ui["symbol"], rt_ui["exchange"]],
-            outputs=[rt_ui["stock_output"], rt_ui["price"], rt_ui["change"], rt_ui["pct_change"], rt_ui["chart"]]
+            inputs=[symbol_input, exchange],
+            outputs=[stock_output, price, change, pct_change, chart_output]
         )
         
-        rt_ui["build_btn"].click(
+        build_index_btn.click(
             build_custom_index,
-            inputs=[rt_ui["symbols_list"], rt_ui["index_type"], rt_ui["exchange"]],
-            outputs=[rt_ui["index_result"]]
+            inputs=[symbols_list, index_type, exchange],
+            outputs=[index_result]
         )
     
     with gr.Tab("📊 Index Calculator"):
-        # Import from original app
-        from app import create_index_ui, calculate_indices
+        try:
+            from app import create_index_ui, calculate_indices
+        except ImportError:
+            # Stub functions if module not available
+            def create_index_ui():
+                with gr.Blocks():
+                    gr.Markdown("Index Calculator placeholder – implement if needed.")
+                return {}
+            def calculate_indices(file, items, p0, q0, p1, q1):
+                return "Calculation placeholder", 0, 0, 0, go.Figure()
+        
         calc_ui = create_index_ui()
-        calc_ui["btn"].click(
-            calculate_indices,
-            inputs=[
-                calc_ui["file"],
-                calc_ui["items"],
-                calc_ui["p0"],
-                calc_ui["q0"],
-                calc_ui["p1"],
-                calc_ui["q1"]
-            ],
-            outputs=[
-                calc_ui["results"],
-                calc_ui["laspeyres"],
-                calc_ui["paasche"],
-                calc_ui["fisher"],
-                calc_ui["plot"]
-            ]
-        )
+        if calc_ui and 'btn' in calc_ui:
+            # Connect the calculator button if UI defined
+            calc_ui["btn"].click(
+                calculate_indices,
+                inputs=[
+                    calc_ui.get("file"),
+                    calc_ui.get("items"),
+                    calc_ui.get("p0"),
+                    calc_ui.get("q0"),
+                    calc_ui.get("p1"),
+                    calc_ui.get("q1")
+                ],
+                outputs=[
+                    calc_ui.get("results"),
+                    calc_ui.get("laspeyres"),
+                    calc_ui.get("paasche"),
+                    calc_ui.get("fisher"),
+                    calc_ui.get("plot")
+                ]
+            )
     
     with gr.Tab("🔄 Index Comparison"):
         comp_ui = create_comparison_ui()