agent_management/test_stock_agents.py

#!/usr/bin/env python3
"""
测试美股 Agents 功能
"""
import asyncio
import aiohttp
import os
import sys

# LLM 配置
LLM_BASE_URL = "https://litellm.graystone-fb459c5d.southeastasia.azurecontainerapps.io/v1"
LLM_API_KEY = "sk-mPV5MVVVVvfGSkXA-ASQXQ"
LLM_MODEL = "taiji/gpt-4o-mini"

async def fetch_stock_quote(symbol: str):
    """获取股票行情数据 (Stock Quote Agent)"""
    url = f"https://query1.finance.yahoo.com/v8/finance/chart/{symbol}"
    params = {"interval": "1d", "range": "1d"}
    headers = {"User-Agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36"}
    
    try:
        async with aiohttp.ClientSession() as session:
            async with session.get(url, params=params, headers=headers, timeout=aiohttp.ClientTimeout(total=15)) as response:
                if response.status == 200:
                    data = await response.json()
                    result = data.get("chart", {}).get("result", [])
                    
                    if not result:
                        return {"success": False, "error": f"未找到股票: {symbol}"}
                    
                    quote_data = result[0]
                    meta = quote_data.get("meta", {})
                    
                    current_price = meta.get("regularMarketPrice", 0)
                    previous_close = meta.get("previousClose", 0)
                    change = current_price - previous_close if previous_close else 0
                    change_percent = (change / previous_close * 100) if previous_close else 0
                    
                    return {
                        "success": True,
                        "symbol": symbol.upper(),
                        "name": meta.get("shortName", symbol),
                        "price": current_price,
                        "change": round(change, 2),
                        "change_percent": round(change_percent, 2),
                        "high_52week": meta.get("fiftyTwoWeekHigh"),
                        "low_52week": meta.get("fiftyTwoWeekLow"),
                        "market_cap": meta.get("marketCap"),
                    }
                else:
                    return {"success": False, "error": f"API 请求失败: HTTP {response.status}"}
    except Exception as e:
        return {"success": False, "error": str(e)}


async def fetch_stock_news(symbol: str):
    """获取股票新闻 (Stock News Agent)"""
    url = f"https://query1.finance.yahoo.com/v1/finance/search"
    params = {"q": symbol, "newsCount": 5}
    headers = {"User-Agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36"}
    
    try:
        async with aiohttp.ClientSession() as session:
            async with session.get(url, params=params, headers=headers, timeout=aiohttp.ClientTimeout(total=15)) as response:
                if response.status == 200:
                    data = await response.json()
                    news = data.get("news", [])
                    
                    return {
                        "success": True,
                        "symbol": symbol.upper(),
                        "news_count": len(news),
                        "news": [{"title": n.get("title", ""), "publisher": n.get("publisher", "")} for n in news[:5]]
                    }
                else:
                    return {"success": False, "error": f"API 请求失败: HTTP {response.status}"}
    except Exception as e:
        return {"success": False, "error": str(e)}


async def fetch_historical_data(symbol: str, period: str = "1mo"):
    """获取历史数据用于技术分析 (Stock Analysis Agent)"""
    url = f"https://query1.finance.yahoo.com/v8/finance/chart/{symbol}"
    params = {"interval": "1d", "range": period}
    headers = {"User-Agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36"}
    
    try:
        async with aiohttp.ClientSession() as session:
            async with session.get(url, params=params, headers=headers, timeout=aiohttp.ClientTimeout(total=15)) as response:
                if response.status == 200:
                    data = await response.json()
                    result = data.get("chart", {}).get("result", [])
                    
                    if not result:
                        return {"success": False, "error": f"未找到股票: {symbol}"}
                    
                    quote_data = result[0]
                    indicators = quote_data.get("indicators", {}).get("quote", [{}])[0]
                    timestamps = quote_data.get("timestamp", [])
                    
                    closes = indicators.get("close", [])
                    closes = [c for c in closes if c is not None]
                    
                    if len(closes) >= 5:
                        # 简单技术分析
                        sma_5 = sum(closes[-5:]) / 5
                        current = closes[-1]
                        trend = "看涨" if current > sma_5 else "看跌"
                        
                        return {
                            "success": True,
                            "symbol": symbol.upper(),
                            "data_points": len(closes),
                            "current_price": round(current, 2),
                            "sma_5": round(sma_5, 2),
                            "trend": trend,
                            "high": round(max(closes), 2),
                            "low": round(min(closes), 2),
                        }
                    else:
                        return {"success": False, "error": "数据点不足"}
                else:
                    return {"success": False, "error": f"API 请求失败: HTTP {response.status}"}
    except Exception as e:
        return {"success": False, "error": str(e)}


async def test_with_llm(symbol: str, context: str):
    """使用 LLM 生成分析报告"""
    try:
        async with aiohttp.ClientSession() as session:
            payload = {
                "model": LLM_MODEL,
                "messages": [
                    {"role": "system", "content": "你是一个专业的美股分析师，请根据提供的数据给出简洁的分析。"},
                    {"role": "user", "content": f"请分析以下 {symbol} 股票数据并给出简要建议（50字以内）:\n{context}"}
                ],
                "max_tokens": 200
            }
            headers = {
                "Authorization": f"Bearer {LLM_API_KEY}",
                "Content-Type": "application/json"
            }
            
            async with session.post(
                f"{LLM_BASE_URL}/chat/completions",
                json=payload,
                headers=headers,
                timeout=aiohttp.ClientTimeout(total=30)
            ) as response:
                if response.status == 200:
                    data = await response.json()
                    content = data.get("choices", [{}])[0].get("message", {}).get("content", "")
                    return {"success": True, "analysis": content}
                else:
                    error_text = await response.text()
                    return {"success": False, "error": f"LLM 请求失败: {response.status} - {error_text[:100]}"}
    except Exception as e:
        return {"success": False, "error": str(e)}


async def main():
    print("=" * 70)
    print("美股 AI Agents 本地测试")
    print("=" * 70)
    print(f"LLM: {LLM_MODEL}")
    print(f"API: {LLM_BASE_URL}")
    print("=" * 70)
    
    symbols = ["AAPL", "TSLA", "NVDA", "MSFT", "GOOGL"]
    
    # 测试 1: Stock Quote Agent
    print("\n📈 【测试 1: Stock Quote Agent - 实时行情】")
    print("-" * 70)
    for symbol in symbols:
        result = await fetch_stock_quote(symbol)
        if result.get("success"):
            print(f"✅ {result['symbol']:6} | {result['name'][:25]:25} | ${result['price']:>10.2f} | {result['change_percent']:+6.2f}%")
        else:
            print(f"❌ {symbol}: {result.get('error')}")
    
    # 测试 2: Stock News Agent
    print("\n📰 【测试 2: Stock News Agent - 新闻资讯】")
    print("-" * 70)
    for symbol in ["AAPL", "TSLA"]:
        result = await fetch_stock_news(symbol)
        if result.get("success"):
            print(f"✅ {result['symbol']} - 找到 {result['news_count']} 条新闻:")
            for news in result['news'][:2]:
                print(f"   • {news['title'][:60]}...")
        else:
            print(f"❌ {symbol}: {result.get('error')}")
    
    # 测试 3: Stock Analysis Agent
    print("\n📊 【测试 3: Stock Analysis Agent - 技术分析】")
    print("-" * 70)
    for symbol in ["AAPL", "NVDA", "TSLA"]:
        result = await fetch_historical_data(symbol)
        if result.get("success"):
            print(f"✅ {result['symbol']:6} | 当前: ${result['current_price']:>8.2f} | SMA5: ${result['sma_5']:>8.2f} | 趋势: {result['trend']} | 区间: ${result['low']:.2f}-${result['high']:.2f}")
        else:
            print(f"❌ {symbol}: {result.get('error')}")
    
    # 测试 4: LLM 综合分析
    print("\n🤖 【测试 4: LLM 综合分析】")
    print("-" * 70)
    
    # 获取一只股票的完整数据
    symbol = "AAPL"
    quote = await fetch_stock_quote(symbol)
    analysis = await fetch_historical_data(symbol)
    
    if quote.get("success") and analysis.get("success"):
        context = f"""
股票: {symbol} ({quote['name']})
当前价格: ${quote['price']:.2f}
涨跌幅: {quote['change_percent']:+.2f}%
52周范围: ${quote.get('low_52week', 0):.2f} - ${quote.get('high_52week', 0):.2f}
5日均线: ${analysis['sma_5']:.2f}
技术趋势: {analysis['trend']}
近期区间: ${analysis['low']:.2f} - ${analysis['high']:.2f}
"""
        print(f"📋 {symbol} 数据汇总:")
        print(context)
        
        print("🔄 调用 LLM 生成分析报告...")
        llm_result = await test_with_llm(symbol, context)
        if llm_result.get("success"):
            print(f"\n💡 AI 分析建议:")
            print(f"   {llm_result['analysis']}")
        else:
            print(f"❌ LLM 调用失败: {llm_result.get('error')}")
    else:
        print(f"❌ 获取数据失败")
    
    print("\n" + "=" * 70)
    print("✅ 测试完成!")
    print("=" * 70)


if __name__ == "__main__":
    asyncio.run(main())