Update app.py

app.py

@@ -1,15 +1,8 @@
 """
-AURA Chat — Gradio Space
+AURA Chat — Gradio Space + MCP Server Tab
 Single-file Gradio app that:
-- Accepts newline-separated prompts (data queries) from the user.
-- On "Analyze" scrapes those queries, sends the aggregated text to a locked LLM,
-  and returns a polished analysis with a ranked list of best stocks and an
-  "Investment Duration" (when to enter / when to exit) for each stock.
-- Seeds a chat component with the generated analysis; user can then chat about it.
-
-Notes:
-- Model, max tokens, and delay between scrapes are fixed and cannot be changed via UI.
-- Set OPENAI_API_KEY in environment (Space Secrets).
+- Chat / Analysis Tab: Enter prompts, analyze, and chat with the assistant.
+- MCP Server Tab: Call scraping and analysis functions directly with JSON output.
 """
 
 import os
@@ -19,12 +12,10 @@ import asyncio
 import requests
 import atexit
 import traceback
-from datetime import datetime
 from typing import List
 import gradio as gr
 
-
-# Defensive: ensure a fresh event loop early to avoid fd race on shutdown.
+# Defensive: fresh event loop early to avoid fd race on shutdown
 if sys.platform != "win32":
     try:
         loop = asyncio.new_event_loop()
@@ -32,7 +23,6 @@ if sys.platform != "win32":
     except Exception:
         traceback.print_exc()
 
-
 # =============================================================================
 # CONFIGURATION (fixed)
 # =============================================================================
@@ -42,7 +32,6 @@ SCRAPER_HEADERS = {
     "Content-Type": "application/json"
 }
 
-# FIXED model & tokens (cannot be changed from UI)
 LLM_MODEL = os.getenv("LLM_MODEL", "openai/gpt-oss-20b:free")
 MAX_TOKENS = int(os.getenv("LLM_MAX_TOKENS", "3000"))
 SCRAPE_DELAY = float(os.getenv("SCRAPE_DELAY", "1.0"))
@@ -50,12 +39,7 @@ SCRAPE_DELAY = float(os.getenv("SCRAPE_DELAY", "1.0"))
 OPENAI_API_KEY = os.getenv("OPENAI_API_KEY")
 OPENAI_BASE_URL = os.getenv("OPENAI_BASE_URL", "https://openrouter.ai/api/v1")
 
-
-# =============================================================================
-# PROMPT ENGINEERING (fixed)
-# =============================================================================
 PROMPT_TEMPLATE = f"""You are AURA, a concise, professional hedge-fund research assistant.
-
 Task:
 - Given scraped data below, produce a clear, readable analysis that:
   1) Lists the top 5 stock picks (or fewer if not enough data).
@@ -68,50 +52,33 @@ Task:
      highest conviction pick).
   5) Output in plain text, clean formatting, easy for humans to read. No JSON.
   6) After the list, include a concise "Assumptions & Risks" section (2-3 bullet points).
-
 Important: Be decisive. If data is insufficient, state that clearly and provide
 the best-available picks with lower confidence.
-
 Max tokens for the LLM response: {MAX_TOKENS}
 Model: {LLM_MODEL}"""
 
-
 # =============================================================================
 # SCRAPING HELPERS
 # =============================================================================
 def deep_scrape(query: str, retries: int = 3, timeout: int = 40) -> str:
-    """Post a query to SCRAPER_API_URL and return a readable aggregation (or an error string)."""
     payload = {"query": query}
     last_err = None
-    
     for attempt in range(1, retries + 1):
         try:
-            resp = requests.post(
-                SCRAPER_API_URL,
-                headers=SCRAPER_HEADERS,
-                json=payload,
-                timeout=timeout
-            )
+            resp = requests.post(SCRAPER_API_URL, headers=SCRAPER_HEADERS, json=payload, timeout=timeout)
             resp.raise_for_status()
             data = resp.json()
-            
-            # Format into readable text
             if isinstance(data, dict):
                 parts = [f"{k.upper()}:\n{v}\n" for k, v in data.items()]
                 return "\n".join(parts)
-            else:
-                return str(data)
-                
+            return str(data)
         except Exception as e:
             last_err = e
             if attempt < retries:
                 time.sleep(1.0)
-    
     return f"ERROR: Scraper failed: {last_err}"
 
-
 def multi_scrape(queries: List[str], delay: float = SCRAPE_DELAY) -> str:
-    """Scrape multiple queries and join results into one large string."""
     aggregated = []
     for q in queries:
         q = q.strip()
@@ -123,7 +90,6 @@ def multi_scrape(queries: List[str], delay: float = SCRAPE_DELAY) -> str:
         time.sleep(delay)
     return "\n".join(aggregated)
 
-
 # =============================================================================
 # LLM INTERACTION
 # =============================================================================
@@ -132,44 +98,28 @@ try:
 except Exception:
     OpenAI = None
 
-
-def run_llm_system_and_user(
-    system_prompt: str,
-    user_text: str,
-    model: str = LLM_MODEL,
-    max_tokens: int = MAX_TOKENS
-) -> str:
-    """Create the OpenAI client lazily, call the chat completions endpoint, then close."""
+def run_llm_system_and_user(system_prompt: str, user_text: str, model: str = LLM_MODEL, max_tokens: int = MAX_TOKENS) -> str:
     if OpenAI is None:
-        return "ERROR: openai package not installed or available. See requirements."
-    
+        return "ERROR: openai package not installed or available."
     if not OPENAI_API_KEY:
-        return "ERROR: OPENAI_API_KEY not set in environment. Please add it to Space Secrets."
-    
+        return "ERROR: OPENAI_API_KEY not set in environment."
     client = None
     try:
         client = OpenAI(base_url=OPENAI_BASE_URL, api_key=OPENAI_API_KEY)
         completion = client.chat.completions.create(
             model=model,
-            messages=[
-                {"role": "system", "content": system_prompt},
-                {"role": "user", "content": user_text},
-            ],
+            messages=[{"role": "system", "content": system_prompt}, {"role": "user", "content": user_text}],
             max_tokens=max_tokens,
         )
-        
-        # Extract content robustly
         if hasattr(completion, "choices") and len(completion.choices) > 0:
             try:
                 return completion.choices[0].message.content
             except Exception:
                 return str(completion.choices[0])
         return str(completion)
-        
     except Exception as e:
         return f"ERROR: LLM call failed: {e}"
     finally:
-        # Try to close client transport
         try:
             if client is not None:
                 try:
@@ -182,156 +132,137 @@ def run_llm_system_and_user(
         except Exception:
             pass
 
-
 # =============================================================================
 # MAIN PIPELINE
 # =============================================================================
 def analyze_and_seed_chat(prompts_text: str):
-    """Called when user clicks Analyze. Returns: (analysis_text, initial_chat_messages_list)"""
     if not prompts_text.strip():
-        return "Please enter at least one prompt (query) describing what data to gather.", []
-    
+        return "Please enter at least one prompt (query).", []
     queries = [line.strip() for line in prompts_text.splitlines() if line.strip()]
     scraped = multi_scrape(queries, delay=SCRAPE_DELAY)
-    
     if scraped.startswith("ERROR"):
         return scraped, []
-    
-    # Compose user payload for LLM
     user_payload = f"SCRAPED DATA:\n\n{scraped}\n\nPlease follow the system instructions and output the analysis."
     analysis = run_llm_system_and_user(PROMPT_TEMPLATE, user_payload)
-    
     if analysis.startswith("ERROR"):
         return analysis, []
-    
-    # Seed chat with user request and assistant analysis
     initial_chat = [
         {"role": "user", "content": f"Analyze the data I provided (prompts: {', '.join(queries)})"},
         {"role": "assistant", "content": analysis}
     ]
     return analysis, initial_chat
 
-
 def continue_chat(chat_messages, user_message: str, analysis_text: str):
-    """Handle chat follow-ups. Returns updated list of message dicts."""
     if chat_messages is None:
         chat_messages = []
-    if not user_message or not user_message.strip():
+    if not user_message.strip():
         return chat_messages
-    
-    # Append user's new message
     chat_messages.append({"role": "user", "content": user_message})
-    
-    # Build LLM input using analysis as reference context
     followup_system = (
-        "You are AURA, a helpful analyst. The conversation context includes a recently "
-        "generated analysis from scraped data. Use that analysis as ground truth context; "
-        "answer follow-up questions, explain rationale, and provide clarifications. "
-        "Be concise and actionable."
+        "You are AURA, a helpful analyst. Use the previous analysis as context; answer follow-ups concisely."
     )
-    user_payload = f"REFERENCE ANALYSIS:\n\n{analysis_text}\n\nUSER QUESTION: {user_message}\n\nRespond concisely and reference lines from the analysis where appropriate."
-    
+    user_payload = f"REFERENCE ANALYSIS:\n\n{analysis_text}\n\nUSER QUESTION: {user_message}"
     assistant_reply = run_llm_system_and_user(followup_system, user_payload)
-    if assistant_reply.startswith("ERROR"):
-        assistant_reply = assistant_reply
-    
-    # Append assistant reply
     chat_messages.append({"role": "assistant", "content": assistant_reply})
     return chat_messages
 
+def convert_to_gradio_chat_format(chat_messages):
+    return chat_messages or []
+
+# =============================================================================
+# MCP SERVER HELPERS
+# =============================================================================
+def mcp_scrape(query: str):
+    return {"query": query, "result": deep_scrape(query)}
+
+def mcp_multi_scrape(queries_text: str):
+    queries = [line.strip() for line in queries_text.splitlines() if line.strip()]
+    return {"queries": queries, "result": multi_scrape(queries)}
+
+def mcp_analyze(prompts_text: str):
+    analysis, seed_chat = analyze_and_seed_chat(prompts_text)
+    return {"prompts": prompts_text, "analysis": analysis, "seed_chat": seed_chat}
 
 # =============================================================================
 # GRADIO UI
 # =============================================================================
-def build_demo():
-    with gr.Blocks(title="AURA Chat — Hedge Fund Picks") as demo:
-        # Custom CSS
-        gr.HTML("""
-        <style>
-        .gradio-container { max-width: 1100px; margin: 18px auto; }
-        .header { text-align: left; margin-bottom: 6px; }
-        .muted { color: #7d8590; font-size: 14px; }
-        .analysis-box { background: #ffffff; border-radius: 8px; padding: 12px; box-shadow: 0 4px 14px rgba(0,0,0,0.06); }
-        </style>
-        """)
-        
-        gr.Markdown("# AURA Chat — Hedge Fund Picks")
-        gr.Markdown(
-            "**Enter one or more data prompts (one per line)** — e.g. SEC insider transactions october 2025 company XYZ.\n\n"
-            "Only input prompts; model, tokens and timing are fixed. Press **Analyze** to fetch & generate the picks. "
-            "After analysis you can chat with the assistant about the results."
-        )
-        
-        with gr.Row():
-            with gr.Column(scale=1):
-                prompts = gr.Textbox(
-                    lines=6,
-                    label="Data Prompts (one per line)",
-                    placeholder="SEC insider transactions october 2025\n13F filings Q3 2025\ncompany: ACME corp insider buys"
+def build_demo_with_mcp():
+    with gr.Blocks(title="AURA Chat — Hedge Fund Picks + MCP Server") as demo:
+        with gr.Tabs():
+            # Chat / Analysis Tab
+            with gr.TabItem("Chat / Analysis"):
+                with gr.Row():
+                    with gr.Column(scale=1):
+                        prompts = gr.Textbox(lines=6, label="Data Prompts")
+                        analyze_btn = gr.Button("Analyze", variant="primary")
+                        error_box = gr.Markdown("", visible=False)
+                    with gr.Column(scale=1):
+                        analysis_out = gr.Textbox(label="Generated Analysis", lines=18, interactive=False)
+                        gr.Markdown("**Chat with AURA**")
+                        chatbot = gr.Chatbot(height=420)
+                        user_input = gr.Textbox(placeholder="Ask follow-up...", label="Your question")
+                        send_btn = gr.Button("Send")
+                
+                analysis_state = gr.State("")
+                chat_state = gr.State([])
+
+                analyze_btn.click(
+                    fn=lambda txt: on_analyze(txt),
+                    inputs=[prompts],
+                    outputs=[analysis_out, error_box, analysis_state, chat_state]
                 )
-                analyze_btn = gr.Button("Analyze", variant="primary")
-                error_box = gr.Markdown("", visible=False)
-                gr.Markdown(f"**Fixed settings:** Model = {LLM_MODEL} • Max tokens = {MAX_TOKENS} • Scrape delay = {SCRAPE_DELAY}s")
-                gr.Markdown("**Important:** Add your OPENAI_API_KEY to Space Secrets before running.")
-            
-            with gr.Column(scale=1):
-                analysis_out = gr.Textbox(
-                    label="Generated Analysis (Top picks with Investment Duration)",
-                    lines=18,
-                    interactive=False
+                send_btn.click(
+                    fn=lambda chat_list, msg, analysis_txt: on_send(chat_list, msg, analysis_txt),
+                    inputs=[chat_state, user_input, analysis_state],
+                    outputs=[chat_state, user_input]
                 )
-                gr.Markdown("**Chat with AURA about this analysis**")
-                chatbot = gr.Chatbot(label="AURA Chat", height=420)
-                user_input = gr.Textbox(
-                    placeholder="Ask a follow-up question about the analysis...",
-                    label="Your question"
+                user_input.submit(
+                    fn=lambda chat_list, msg, analysis_txt: on_send(chat_list, msg, analysis_txt),
+                    inputs=[chat_state, user_input, analysis_state],
+                    outputs=[chat_state, user_input]
                 )
-                send_btn = gr.Button("Send")
-        
-        # States
-        analysis_state = gr.State("")
-        chat_state = gr.State([])
-        
-        # Handler functions
-        def on_analyze(prompts_text):
-            analysis_text, initial_chat = analyze_and_seed_chat(prompts_text)
-            if analysis_text.startswith("ERROR"):
-                return "", f"**Error:** {analysis_text}", "", []
-            return analysis_text, "", analysis_text, initial_chat
-        
-        def on_send(chat_state_list, user_msg, analysis_text):
-            if not user_msg or not user_msg.strip():
-                return chat_state_list or [], ""
-            updated_history = continue_chat(chat_state_list or [], user_msg, analysis_text)
-            return updated_history, ""
-        
-        def render_chat(chat_messages):
-            return chat_messages or []
-        
-        # Wire handlers
-        analyze_btn.click(
-            fn=on_analyze,
-            inputs=[prompts],
-            outputs=[analysis_out, error_box, analysis_state, chat_state]
-        )
-        send_btn.click(
-            fn=on_send,
-            inputs=[chat_state, user_input, analysis_state],
-            outputs=[chat_state, user_input]
-        )
-        user_input.submit(
-            fn=on_send,
-            inputs=[chat_state, user_input, analysis_state],
-            outputs=[chat_state, user_input]
-        )
-        chat_state.change(
-            fn=render_chat,
-            inputs=[chat_state],
-            outputs=[chatbot]
-        )
-    
+                chat_state.change(
+                    fn=convert_to_gradio_chat_format,
+                    inputs=[chat_state],
+                    outputs=[chatbot]
+                )
+
+            # MCP Server Tab
+            with gr.TabItem("MCP Server"):
+                gr.Markdown("**Call scraping and analysis functions directly:**")
+                with gr.Row():
+                    with gr.Column(scale=1):
+                        single_query = gr.Textbox(label="Single Scrape Query")
+                        scrape_btn = gr.Button("Scrape Query")
+                        scrape_out = gr.JSON()
+                        multi_queries = gr.Textbox(lines=6, label="Multi Scrape Queries")
+                        multi_scrape_btn = gr.Button("Multi Scrape")
+                        multi_scrape_out = gr.JSON()
+                    with gr.Column(scale=1):
+                        analysis_prompts = gr.Textbox(lines=6, label="Analysis Prompts")
+                        analyze_mcp_btn = gr.Button("Run Full Analysis")
+                        analyze_mcp_out = gr.JSON()
+                
+                scrape_btn.click(fn=mcp_scrape, inputs=[single_query], outputs=[scrape_out])
+                multi_scrape_btn.click(fn=mcp_multi_scrape, inputs=[multi_queries], outputs=[multi_scrape_out])
+                analyze_mcp_btn.click(fn=mcp_analyze, inputs=[analysis_prompts], outputs=[analyze_mcp_out])
+
     return demo
 
+# =============================================================================
+# Handlers for Chat Tab
+# =============================================================================
+def on_analyze(prompts_text):
+    analysis_text, initial_chat = analyze_and_seed_chat(prompts_text)
+    if analysis_text.startswith("ERROR"):
+        return "", f"**Error:** {analysis_text}", "", []
+    return analysis_text, "", analysis_text, initial_chat
+
+def on_send(chat_state_list, user_msg, analysis_text):
+    if not user_msg.strip():
+        return chat_state_list or [], ""
+    updated_history = continue_chat(chat_state_list or [], user_msg, analysis_text)
+    return updated_history, ""
 
 # =============================================================================
 # CLEAN SHUTDOWN
@@ -340,26 +271,17 @@ def _cleanup_on_exit():
     try:
         loop = asyncio.get_event_loop()
         if loop and not loop.is_closed():
-            try:
-                loop.stop()
-            except Exception:
-                pass
-            try:
-                loop.close()
-            except Exception:
-                pass
-    except Exception:
-        pass
+            try: loop.stop()
+            except Exception: pass
+            try: loop.close()
+            except Exception: pass
+    except Exception: pass
 
 atexit.register(_cleanup_on_exit)
 
-
 # =============================================================================
 # RUN
 # =============================================================================
 if __name__ == "__main__":
-    demo = build_demo()
-    demo.launch(
-        server_name="0.0.0.0",
-        server_port=int(os.environ.get("PORT", 7860))
-    )
+    demo = build_demo_with_mcp()
+    demo.launch(server_name="0.0.0.0", server_port=int(os.environ.get("PORT", 7860)))