app.py

import gradio as gr
import requests
import json
import os
import time
from requests.adapters import HTTPAdapter
from urllib3.util.retry import Retry

MCP_SPACE = "JC321/EasyReportsMCPServer"
MCP_URL = "https://jc321-easyreportsmcpserver.hf.space"

# 设置请求头
HEADERS = {
    "Content-Type": "application/json",
    "User-Agent": "SEC-Query-Assistant/1.0 (jtyxabc@gmail.com)"
}

# 创建带重试的 requests session
def create_session_with_retry():
    """创建带重试机制的 requests session"""
    session = requests.Session()
    retry = Retry(
        total=3,  # 最多重试3次
        backoff_factor=1,  # 重试间隔：1秒, 2秒, 4秒
        status_forcelist=[500, 502, 503, 504],  # 这些状态码会触发重试
    )
    adapter = HTTPAdapter(max_retries=retry)
    session.mount('http://', adapter)
    session.mount('https://', adapter)
    return session

# 创建全局 session
session = create_session_with_retry()

# 格式化数值显示
def format_value(value, value_type="money"):
    """
    格式化数值：0显示为N/A，其他显示为带单位的格式
    value_type: "money" (金额), "eps" (每股收益), "number" (普通数字)
    """
    if value is None or value == 0:
        return "N/A"
    
    if value_type == "money":
        return f"${value:.2f}B"
    elif value_type == "eps":
        return f"${value:.2f}"
    else:  # number
        return f"{value:.2f}"

def normalize_cik(cik):
    """
    格式化 CIK 为标准的 10 位格式
    """
    if not cik:
        return None
    # 转换为字符串并移除非数字字符
    cik_str = str(cik).replace('-', '').replace(' ', '')
    # 仅保留数字
    cik_str = ''.join(c for c in cik_str if c.isdigit())
    # 填充前导 0 至 10 位
    return cik_str.zfill(10) if cik_str else None

def parse_mcp_response(response_data):
    """
    解析 MCP 协议响应数据
    支持格式:
    1. {"result": {"content": [{"type": "text", "text": "{...}"}]}}
    2. {"content": [{"type": "text", "text": "{...}"}]}
    3. 直接的 JSON 数据
    """
    if not isinstance(response_data, dict):
        return response_data
    
    # 格式 1: {"result": {"content": [...]}}
    if "result" in response_data and "content" in response_data["result"]:
        content = response_data["result"]["content"]
        if content and len(content) > 0:
            text_content = content[0].get("text", "{}")
            # 直接解析 JSON（MCP Server 已移除 emoji 前缀）
            try:
                return json.loads(text_content)
            except json.JSONDecodeError:
                return text_content
        return {}
    
    # 格式 2: {"content": [...]}
    elif "content" in response_data:
        content = response_data.get("content", [])
        if content and len(content) > 0:
            text_content = content[0].get("text", "{}")
            # 直接解析 JSON
            try:
                return json.loads(text_content)
            except json.JSONDecodeError:
                return text_content
        return {}
    
    # 格式 3: 直接返回
    return response_data

# MCP 工具定义
def create_mcp_tools():
    """创建 MCP 工具列表"""
    return [
        {
            "name": "query_financial_data",
            "description": "Query SEC financial data for US listed companies",
            "parameters": {
                "type": "object",
                "properties": {
                    "company_name": {
                        "type": "string",
                        "description": "Company name or stock symbol (e.g., Apple, NVIDIA, AAPL)"
                    },
                    "query_type": {
                        "type": "string",
                        "enum": ["Latest Financial Data", "3-Year Trends", "5-Year Trends"],
                        "description": "Type of financial query"
                    }
                },
                "required": ["company_name", "query_type"]
            }
        }
    ]

# 工具执行函数
def execute_tool(tool_name, **kwargs):
    """执行 MCP 工具"""
    if tool_name == "query_financial_data":
        return query_financial_data(kwargs.get("company_name"), kwargs.get("query_type"))
    return f"Unknown tool: {tool_name}"
# 创建超链接
def create_source_link(source_form, source_url=None):
    """为Source Form创建超链接，使用MCP后端返回的URL"""
    if not source_form or source_form == 'N/A':
        return source_form
    
    # 如果后端提供了URL，使用后端的URL
    if source_url and source_url != 'N/A':
        return f"[{source_form}]({source_url})"
    
    # 如果没有URL，只显示文本
    return source_form

def query_financial_data(company_name, query_type):
    """查询财务数据的主函数"""
    
    if not company_name:
        return "Please enter a company name or stock symbol"
    
    # 翻译英文查询类型为中文（用于后端处理）
    query_type_mapping = {
        "Latest Financial Data": "最新财务数据",
        "3-Year Trends": "3年趋势",
        "5-Year Trends": "5年趋势",
        "Company Filings": "公司报表列表"
    }
    internal_query_type = query_type_mapping.get(query_type, query_type)
    
    try:
        # 使用 MCP 协议调用工具
        # 先搜索公司（使用 advanced_search_company）
        try:
            search_resp = session.post(
                f"{MCP_URL}/message",
                json={
                    "method": "tools/call",
                    "params": {
                        "name": "advanced_search_company",
                        "arguments": {"company_input": company_name}
                    }
                },
                headers=HEADERS,
                timeout=60  # 增加到60秒
            )
        except requests.exceptions.Timeout:
            return f"❌ MCP Server Timeout: The server took too long to respond (>60s).\n\n**Possible reasons**:\n1. MCP Server is cold starting (first request after idle)\n2. Server is overloaded\n3. Network issues\n\n**Suggestion**: Please try again in a few moments. If the problem persists, the MCP Server at {MCP_URL} may be down."
        
        if search_resp.status_code != 200:
            return f"❌ Server Error: HTTP {search_resp.status_code}\n\nResponse: {search_resp.text[:500]}"
        
        try:
            result = search_resp.json()
            # 使用统一的 MCP 响应解析函数
            company = parse_mcp_response(result)
        except (ValueError, KeyError, json.JSONDecodeError) as e:
            return f"❌ JSON Parse Error: {str(e)}\n\nResponse: {search_resp.text[:500]}"
        
        if isinstance(company, dict) and company.get("error"):
            return f"❌ Error: {company['error']}"
        
        # advanced_search 返回的字段: cik, name, ticker
        # 注意: 不是 tickers 和 sic_description
        company_name = company.get('name', 'Unknown')
        ticker = company.get('ticker', 'N/A')
        
        result = f"# {company_name}\n\n"
        result += f"**Stock Symbol**: {ticker}\n"
        # sic_description 需要后续通过 get_company_info 获取，这里暂时不显示
        result += "\n---\n\n"
        
        # 获取并格式化 CIK 为 10 位标准格式
        cik = normalize_cik(company.get('cik'))
        if not cik:
            return result + f"❌ Error: Invalid CIK from company search\n\nDebug: company data = {json.dumps(company, indent=2)}"
        
        # 根据查询类型获取数据
        if internal_query_type == "最新财务数据":
            data_resp = session.post(
                f"{MCP_URL}/message",
                json={
                    "method": "tools/call",
                    "params": {
                        "name": "get_latest_financial_data",
                        "arguments": {"cik": cik}
                    }
                },
                headers=HEADERS,
                timeout=60  # 增加到60秒
            )
            
            if data_resp.status_code != 200:
                return result + f"❌ Server Error: HTTP {data_resp.status_code}\n\n{data_resp.text[:500]}"
            
            try:
                data_result = data_resp.json()
                # 使用统一的 MCP 响应解析函数
                data = parse_mcp_response(data_result)
            except (ValueError, KeyError, json.JSONDecodeError) as e:
                return result + f"❌ JSON Parse Error: {str(e)}\n\n{data_resp.text[:500]}"
            
            if isinstance(data, dict) and data.get("error"):
                return result + f"❌ {data['error']}"
            
            cik = data.get('cik')
            result += f"## Fiscal Year {data.get('period', 'N/A')}\n\n"
            
            total_revenue = data.get('total_revenue', 0) / 1e9 if data.get('total_revenue') else 0
            net_income = data.get('net_income', 0) / 1e9 if data.get('net_income') else 0
            eps = data.get('earnings_per_share', 0) if data.get('earnings_per_share') else 0
            opex = data.get('operating_expenses', 0) / 1e9 if data.get('operating_expenses') else 0
            ocf = data.get('operating_cash_flow', 0) / 1e9 if data.get('operating_cash_flow') else 0
            
            result += f"- **Total Revenue**: {format_value(total_revenue)}\n"
            result += f"- **Net Income**: {format_value(net_income)}\n"
            result += f"- **Earnings Per Share**: {format_value(eps, 'eps')}\n"
            result += f"- **Operating Expenses**: {format_value(opex)}\n"
            result += f"- **Operating Cash Flow**: {format_value(ocf)}\n"
            # 使用后端返回的 source_url
            source_form = data.get('source_form', 'N/A')
            source_url = data.get('source_url', None)  # 从后端获取URL
            result += f"- **Source Form**: {create_source_link(source_form, source_url)}\n"
        
        elif internal_query_type == "3年趋势":
            metrics_resp = session.post(
                f"{MCP_URL}/message",
                json={
                    "method": "tools/call",
                    "params": {
                        "name": "extract_financial_metrics",
                        "arguments": {"cik": cik, "years": 3}
                    }
                },
                headers=HEADERS,
                timeout=120  # 3年趋势需要更长时间，增加到120秒
            )
            
            if metrics_resp.status_code != 200:
                return result + f"❌ Server Error: HTTP {metrics_resp.status_code}\n\n{metrics_resp.text[:500]}"
            
            try:
                metrics_result = metrics_resp.json()
                # 使用统一的 MCP 响应解析函数
                metrics = parse_mcp_response(metrics_result)
            except (ValueError, KeyError, json.JSONDecodeError) as e:
                return result + f"❌ JSON Parse Error: {str(e)}\n\nResponse: {metrics_resp.text[:500]}"
            
            if isinstance(metrics, dict) and metrics.get("error"):
                return result + f"❌ {metrics['error']}"
            
            result += f"## 3-Year Financial Trends ({metrics.get('periods', 0)} periods)\n\n"
            
            # 显示所有数据（包括年度和季度）
            all_data = metrics.get('data', [])  # MCP Server 返回的字段是 'data'
            
            # 去重：根据period和source_form去重
            seen = set()
            unique_data = []
            for m in all_data:
                key = (m.get('period', 'N/A'), m.get('source_form', 'N/A'))
                if key not in seen:
                    seen.add(key)
                    unique_data.append(m)
            
            # 按期间降序排序，确保显示最近的3年数据
            # 使用更智能的排序：先按年份，再按是否是季度
            # 正确顺序：FY2024 → 2024Q3 → 2024Q2 → 2024Q1 → FY2023
            def sort_key(x):
                period = x.get('period', '0000')
                # 提取年份（前4位）
                year = period[:4] if len(period) >= 4 else '0000'
                # 如果有Q，提取季度号
                if 'Q' in period:
                    quarter = period[period.index('Q')+1] if period.index('Q')+1 < len(period) else '0'
                    return (year, 1, 4 - int(quarter))  # Q在FY后面：Q3, Q2, Q1 (4-3=1, 4-2=2, 4-1=3)
                else:
                    return (year, 0, 0)  # FY 排在同年的所有Q之前
            
            unique_data = sorted(unique_data, key=sort_key, reverse=True)
            
            result += "| Period | Revenue (B) | Net Income (B) | EPS | Operating Expenses (B) | Operating Cash Flow (B) | Source Form |\n"
            result += "|--------|-------------|----------------|-----|------------------------|-------------------------|-------------|\n"
            
            for m in unique_data:
                period = m.get('period', 'N/A')
                rev = (m.get('total_revenue') or 0) / 1e9
                inc = (m.get('net_income') or 0) / 1e9
                eps_val = m.get('earnings_per_share') or 0
                opex = (m.get('operating_expenses') or 0) / 1e9
                ocf = (m.get('operating_cash_flow') or 0) / 1e9
                source_form = m.get('source_form', 'N/A')
                source_url = m.get('source_url', None)  # 从后端获取URL
                
                # 区分年度和季度，修复双重FY前缀问题
                if 'Q' in period:
                    # 季度数据，不添加前缀
                    display_period = period
                else:
                    # 年度数据，只在没有FY的情况下添加
                    display_period = period if period.startswith('FY') else f"FY{period}"
                
                source_link = create_source_link(source_form, source_url)
                
                result += f"| {display_period} | {format_value(rev)} | {format_value(inc)} | {format_value(eps_val, 'eps')} | {format_value(opex)} | {format_value(ocf)} | {source_link} |\n"
        
        elif internal_query_type == "5年趋势":
            metrics_resp = session.post(
                f"{MCP_URL}/message",
                json={
                    "method": "tools/call",
                    "params": {
                        "name": "extract_financial_metrics",
                        "arguments": {"cik": cik, "years": 5}
                    }
                },
                headers=HEADERS,
                timeout=180  # 5年趋势需要更长时间，增加到180秒
            )
            
            if metrics_resp.status_code != 200:
                return result + f"❌ Server Error: HTTP {metrics_resp.status_code}\n\n{metrics_resp.text[:500]}"
            
            try:
                metrics_result = metrics_resp.json()
                # 使用统一的 MCP 响应解析函数
                metrics = parse_mcp_response(metrics_result)
            except (ValueError, KeyError, json.JSONDecodeError) as e:
                return result + f"❌ JSON Parse Error: {str(e)}\n\nResponse: {metrics_resp.text[:500]}"
            
            if isinstance(metrics, dict) and metrics.get("error"):
                return result + f"❌ {metrics['error']}"
            
            # 显示所有数据（包括年度和季度）
            all_data = metrics.get('data', [])  # MCP Server 返回的字段是 'data'
            
            # 去重：根据period和source_form去重
            seen = set()
            unique_data = []
            for m in all_data:
                key = (m.get('period', 'N/A'), m.get('source_form', 'N/A'))
                if key not in seen:
                    seen.add(key)
                    unique_data.append(m)
            
            # 按期间降序排序，确保显示最近的5年数据
            # 使用更智能的排序：先按年份，再按是否是季度
            # 正确顺序：FY2024 → 2024Q3 → 2024Q2 → 2024Q1 → FY2023
            def sort_key(x):
                period = x.get('period', '0000')
                # 提取年份（前4位）
                year = period[:4] if len(period) >= 4 else '0000'
                # 如果有Q，提取季度号
                if 'Q' in period:
                    quarter = period[period.index('Q')+1] if period.index('Q')+1 < len(period) else '0'
                    return (year, 1, 4 - int(quarter))  # Q在FY后面：Q3, Q2, Q1 (4-3=1, 4-2=2, 4-1=3)
                else:
                    return (year, 0, 0)  # FY 排在同年的所有Q之前
            
            unique_data = sorted(unique_data, key=sort_key, reverse=True)
            
            result += f"## 5-Year Financial Trends ({metrics.get('periods', 0)} periods)\n\n"
            result += "| Period | Revenue (B) | Net Income (B) | EPS | Operating Expenses (B) | Operating Cash Flow (B) | Source Form |\n"
            result += "|--------|-------------|----------------|-----|------------------------|-------------------------|-------------|\n"
            
            for m in unique_data:
                period = m.get('period', 'N/A')
                rev = (m.get('total_revenue') or 0) / 1e9
                inc = (m.get('net_income') or 0) / 1e9
                eps_val = m.get('earnings_per_share') or 0
                opex = (m.get('operating_expenses') or 0) / 1e9
                ocf = (m.get('operating_cash_flow') or 0) / 1e9
                source_form = m.get('source_form', 'N/A')
                source_url = m.get('source_url', None)  # 从后端获取URL
                
                # 区分年度和季度，修复双重FY前缀问题
                if 'Q' in period:
                    # 季度数据，不添加前缀
                    display_period = period
                else:
                    # 年度数据，只在没有FY的情况下添加
                    display_period = period if period.startswith('FY') else f"FY{period}"
                
                source_link = create_source_link(source_form, source_url)
                
                result += f"| {display_period} | {format_value(rev)} | {format_value(inc)} | {format_value(eps_val, 'eps')} | {format_value(opex)} | {format_value(ocf)} | {source_link} |\n"
        
        elif internal_query_type == "公司报表列表":
            # 查询公司所有报表
            filings_resp = session.post(
                f"{MCP_URL}/message",
                json={
                    "method": "tools/call",
                    "params": {
                        "name": "get_company_filings",
                        "arguments": {"cik": cik, "limit": 50}
                    }
                },
                headers=HEADERS,
                timeout=90  # 增加到90秒
            )
            
            if filings_resp.status_code != 200:
                return result + f"❌ Server Error: HTTP {filings_resp.status_code}\n\n{filings_resp.text[:500]}"
            
            try:
                filings_result = filings_resp.json()
                # 使用统一的 MCP 响应解析函数
                filings_data = parse_mcp_response(filings_result)
            except (ValueError, KeyError, json.JSONDecodeError) as e:
                return result + f"❌ JSON Parse Error: {str(e)}\n\n{filings_resp.text[:500]}"
            
            if isinstance(filings_data, dict) and filings_data.get("error"):
                return result + f"❌ {filings_data['error']}"
            
            filings = filings_data.get('filings', []) if isinstance(filings_data, dict) else filings_data
            
            result += f"## Company Filings ({len(filings)} records)\n\n"
            result += "| Form Type | Filing Date | Accession Number | Primary Document |\n"
            result += "|-----------|-------------|------------------|------------------|\n"
            
            for filing in filings:
                form_type = filing.get('form_type', 'N/A')
                filing_date = filing.get('filing_date', 'N/A')
                accession_num = filing.get('accession_number', 'N/A')
                primary_doc = filing.get('primary_document', 'N/A')
                filing_url = filing.get('filing_url', None)  # 从后端获取URL
                
                # 使用后端返回的URL创建链接
                if filing_url and filing_url != 'N/A':
                    form_link = f"[{form_type}]({filing_url})"
                    primary_doc_link = f"[{primary_doc}]({filing_url})"
                else:
                    form_link = form_type
                    primary_doc_link = primary_doc
                
                result += f"| {form_link} | {filing_date} | {accession_num} | {primary_doc_link} |\n"
        
        return result
        
    except requests.exceptions.RequestException as e:
        return f"❌ Network Error: {str(e)}\n\nMCP Server: {MCP_URL}"
    except Exception as e:
        import traceback
        return f"❌ Unexpected Error: {str(e)}\n\nTraceback:\n{traceback.format_exc()}"

# Chatbot 功能：使用MCP工具
def chatbot_response(message, history):
    """聊天机器人响应函数，集成MCP工具"""
    try:
        # 检查是否是财务查询相关问题
        financial_keywords = ['financial', 'revenue', 'income', 'earnings', 'cash flow', 'expenses', '财务', '收入', '利润', 'data', 'trend', 'performance']
        
        if any(keyword in message.lower() for keyword in financial_keywords):
            # 提取公司名称和查询类型
            company_keywords = ['apple', 'microsoft', 'nvidia', 'tesla', 'alibaba', 'google', 'amazon', 'meta', 'tsla', 'aapl', 'msft', 'nvda', 'googl', 'amzn']
            detected_company = None
            
            for company in company_keywords:
                if company in message.lower():
                    if company in ['aapl']: detected_company = 'Apple'
                    elif company in ['msft']: detected_company = 'Microsoft' 
                    elif company in ['nvda']: detected_company = 'NVIDIA'
                    elif company in ['tsla']: detected_company = 'Tesla'
                    elif company in ['googl']: detected_company = 'Google'
                    elif company in ['amzn']: detected_company = 'Amazon'
                    else: detected_company = company.capitalize()
                    break
            
            if detected_company:
                # 根据问题内容选择查询类型
                if any(word in message.lower() for word in ['trend', '趋势', 'history', 'historical', 'over time']):
                    if any(word in message for word in ['5', 'five', '五年']):
                        query_type = '5-Year Trends'
                    else:
                        query_type = '3-Year Trends'
                else:
                    query_type = 'Latest Financial Data'
                
                # 调用财务查询函数
                result = query_financial_data(detected_company, query_type)
                return f"Here's the financial information for {detected_company}:\n\n{result}"
            else:
                return "I can help you query financial data! Please specify a company name. For example: 'Show me Apple's latest financial data' or 'What's NVIDIA's 3-year trend?' \n\nSupported companies include: Apple, Microsoft, NVIDIA, Tesla, Alibaba, Google, Amazon, and more."
        
        # 如果不是财务查询，返回通用回复
        return "Hello! I'm a financial data assistant powered by SEC EDGAR data. I can help you query financial information for US listed companies.\n\n**What I can do:**\n- Get latest financial data (revenue, income, EPS, etc.)\n- Show 3-year or 5-year financial trends\n- Provide detailed financial metrics\n\n**Try asking:**\n- 'Show me Apple's latest financial data'\n- 'What's NVIDIA's 3-year financial trend?'\n- 'How is Microsoft performing financially?'"
    
    except Exception as e:
        return f"Sorry, I encountered an error: {str(e)}. Please try asking about financial data for specific companies like Apple, Microsoft, NVIDIA, Tesla, etc."

# 包装函数，显示加载状态
def query_with_status(company, query_type):
    """Query with loading status indicator"""
    try:
        # 返回加载状态和结果
        yield "<div style='padding: 10px; background: #e3f2fd; border-left: 4px solid #2196f3; margin: 10px 0;'>🔄 <strong>Loading...</strong> Querying SEC EDGAR data for <strong>{}</strong>...</div>".format(company), ""
        
        # 执行实际查询
        result = query_financial_data(company, query_type)
        
        # 返回成功状态和结果
        yield "<div style='padding: 10px; background: #e8f5e9; border-left: 4px solid #4caf50; margin: 10px 0;'>✅ <strong>Query completed successfully!</strong></div>", result
    
    except Exception as e:
        # 返回错误状态
        yield "<div style='padding: 10px; background: #ffebee; border-left: 4px solid #f44336; margin: 10px 0;'>❌ <strong>Error:</strong> {}</div>".format(str(e)), ""

# 创建 Gradio 界面
with gr.Blocks(title="SEC Financial Data Query Assistant") as demo:
    gr.Markdown("# 🤖 SEC Financial Data Query Assistant")
    gr.Markdown("Query SEC financial data for US listed companies through MCP Server")
    
    # 添加 MCP Server 状态提示
    with gr.Row():
        gr.Markdown(f"🔗 **MCP Server**: [{MCP_URL}]({MCP_URL})")
    with gr.Row():
        gr.Markdown("⚠️ **Note**: First query after idle may take 1-2 minutes (server cold start). Please be patient.")
    
    with gr.Tab("AI Assistant"):
        # 使用 Gradio ChatInterface（兼容 4.44.1）
        chat = gr.ChatInterface(
            fn=chatbot_response,
            examples=[
                "Show me Apple's latest financial data",
                "What's NVIDIA's 3-year financial trend?",
                "Get Microsoft's 5-year financial trends",
                "How is Tesla performing financially?"
            ],
            cache_examples=False
        )
    
    with gr.Tab("Direct Query"):
        gr.Markdown("## 🔍 Direct Financial Data Query")
        gr.Markdown("Select a company and query type to retrieve financial information.")
        
        with gr.Row():
            company_input = gr.Textbox(
                label="Company Name or Stock Symbol",
                placeholder="e.g., NVIDIA, Apple, Alibaba, AAPL",
                scale=2
            )
            query_type = gr.Radio(
                ["Latest Financial Data", "3-Year Trends", "5-Year Trends", "Company Filings"],
                label="Query Type",
                value="Latest Financial Data",
                scale=1
            )
        
        submit_btn = gr.Button("🔍 Query Financial Data", variant="primary", size="lg")
        
        # 添加加载状态指示器
        with gr.Row():
            status_text = gr.Markdown("")
        
        output = gr.Markdown(label="Query Results")
        
        # 示例
        gr.Examples(
            examples=[
                ["NVIDIA", "Latest Financial Data"],
                ["Apple", "3-Year Trends"],
                ["Microsoft", "5-Year Trends"],
                ["Alibaba", "Company Filings"],
                ["Tesla", "3-Year Trends"]
            ],
            inputs=[company_input, query_type],
            outputs=output,
            fn=query_financial_data,
            cache_examples=False
        )
        
        submit_btn.click(
            fn=query_with_status,
            inputs=[company_input, query_type],
            outputs=[status_text, output],
            show_progress="full"  # 显示完整的进度条
        )
    
    gr.Markdown("---")
    gr.Markdown(f"**Data Source**: SEC EDGAR | **MCP Server**: `{MCP_URL}`")

# Launch the app for Hugging Face Space
if __name__ == "__main__":
    demo.launch(
        server_name="0.0.0.0",
        server_port=7860,
        show_error=True
    )