mirror of
https://github.com/ZhuLinsen/daily_stock_analysis
synced 2026-09-20 10:53:33 +08:00
feat: deep research agent & event monitor — /research command, alert rules, scheduler integration (#649)
* feat: deep research agent & event monitor — /research command, alert rules, scheduler integration - ResearchAgent with 3-phase approach (decompose → research → synthesise) - Token budget tracking for deep research queries - /research bot command with aliases (/深研, /deepsearch) - REST endpoint POST /api/v1/agent/research - EventMonitor with PriceAlert/VolumeAlert rules - Scheduler background task system for periodic alert polling - EventMonitor config validation in system_config_service - main.py: EventMonitor wiring into schedule mode - Unsupported rule type rejection at config and runtime level * fix: address PR #649 review comments - events.py: fix to_dict_list docstring, validate created_at as float, raise ValueError on non-dict rules - research.py: remove unused import uuid - agent.py: replace nested ThreadPoolExecutor with asyncio.wait_for, use Field(default_factory=list), pass stock_code context - scheduler.py: clamp minimum interval_seconds to 30s - test_multi_agent.py: reduce sleep 3s→0.05s and timeout 1s→0.01s to speed up tests * fix: normalize price alert direction and tighten research ticker regex - events.py from_dict_list: .lower() on deserialized direction so _check_price matches correctly (fixes silent missed alerts for mixed-case config like {"direction":"Above"}) - research.py: restrict US ticker auto-detection to 1-4 chars (was 1-5); avoids misclassifying common words (MACRO, TREND) as stock codes in /research topic queries
This commit is contained in:
@@ -234,6 +234,83 @@ def _build_executor(config, strategies: Optional[List[str]] = None):
|
||||
return build_agent_executor(config, skills=strategies)
|
||||
|
||||
|
||||
# ============================================================
|
||||
# Deep research endpoint
|
||||
# ============================================================
|
||||
|
||||
class ResearchRequest(BaseModel):
|
||||
question: str
|
||||
stock_code: Optional[str] = None
|
||||
|
||||
class ResearchResponse(BaseModel):
|
||||
success: bool
|
||||
content: str
|
||||
sources: List[str] = Field(default_factory=list)
|
||||
token_usage: int = 0
|
||||
error: Optional[str] = None
|
||||
|
||||
|
||||
@router.post("/research", response_model=ResearchResponse)
|
||||
async def agent_research(request: ResearchRequest):
|
||||
"""Run a deep-research query via the ResearchAgent.
|
||||
|
||||
Similar to the ``/research`` bot command but exposed as a REST endpoint.
|
||||
"""
|
||||
config = get_config()
|
||||
if not config.is_agent_available():
|
||||
raise HTTPException(status_code=400, detail="Agent mode is not enabled")
|
||||
|
||||
question = request.question
|
||||
context: Optional[Dict[str, Any]] = None
|
||||
if request.stock_code:
|
||||
question = f"[Stock: {request.stock_code}] {question}"
|
||||
context = {"stock_code": request.stock_code}
|
||||
|
||||
try:
|
||||
from src.agent.research import ResearchAgent
|
||||
from src.agent.factory import get_tool_registry
|
||||
from src.agent.llm_adapter import LLMToolAdapter
|
||||
|
||||
registry = get_tool_registry()
|
||||
llm_adapter = LLMToolAdapter(config)
|
||||
budget = getattr(config, "agent_deep_research_budget", 30000)
|
||||
|
||||
agent = ResearchAgent(
|
||||
tool_registry=registry,
|
||||
llm_adapter=llm_adapter,
|
||||
token_budget=budget,
|
||||
)
|
||||
|
||||
research_timeout = getattr(config, "agent_deep_research_timeout", 180)
|
||||
loop = asyncio.get_running_loop()
|
||||
|
||||
task = loop.run_in_executor(None, agent.research, question, context)
|
||||
|
||||
try:
|
||||
result = await asyncio.wait_for(task, timeout=research_timeout)
|
||||
except asyncio.TimeoutError:
|
||||
logger.warning("Agent research API timed out after %ss", research_timeout)
|
||||
return ResearchResponse(
|
||||
success=False,
|
||||
content="",
|
||||
sources=[],
|
||||
token_usage=0,
|
||||
error=f"Deep research timed out after {research_timeout}s",
|
||||
)
|
||||
|
||||
return ResearchResponse(
|
||||
success=result.success,
|
||||
content=result.report,
|
||||
sources=[f"Sub-question {i+1}: {q}" for i, q in enumerate(result.sub_questions)],
|
||||
token_usage=result.total_tokens,
|
||||
error=result.error if not result.success else None,
|
||||
)
|
||||
except Exception as e:
|
||||
logger.error("Agent research API failed: %s", e)
|
||||
logger.exception("Agent research error details:")
|
||||
raise HTTPException(status_code=500, detail=str(e))
|
||||
|
||||
|
||||
@router.post("/chat/stream")
|
||||
async def agent_chat_stream(request: ChatRequest):
|
||||
"""
|
||||
|
||||
@@ -15,6 +15,7 @@ from bot.commands.market import MarketCommand
|
||||
from bot.commands.batch import BatchCommand
|
||||
from bot.commands.ask import AskCommand
|
||||
from bot.commands.chat import ChatCommand
|
||||
from bot.commands.research import ResearchCommand
|
||||
from bot.commands.strategies import StrategiesCommand
|
||||
from bot.commands.history import HistoryCommand
|
||||
|
||||
@@ -27,6 +28,7 @@ ALL_COMMANDS = [
|
||||
BatchCommand,
|
||||
AskCommand,
|
||||
ChatCommand,
|
||||
ResearchCommand,
|
||||
StrategiesCommand,
|
||||
HistoryCommand,
|
||||
]
|
||||
@@ -40,6 +42,7 @@ __all__ = [
|
||||
'BatchCommand',
|
||||
'AskCommand',
|
||||
'ChatCommand',
|
||||
'ResearchCommand',
|
||||
'StrategiesCommand',
|
||||
'HistoryCommand',
|
||||
'ALL_COMMANDS',
|
||||
|
||||
160
bot/commands/research.py
Normal file
160
bot/commands/research.py
Normal file
@@ -0,0 +1,160 @@
|
||||
# -*- coding: utf-8 -*-
|
||||
"""
|
||||
Research command — deep research on a stock or market topic.
|
||||
|
||||
Usage:
|
||||
/research 600519 -> Deep research on Kweichow Moutai
|
||||
/research 600519 近期业绩风险 -> Focused research with specific question
|
||||
/research 新能源板块前景分析 -> Topic-based research
|
||||
"""
|
||||
|
||||
import logging
|
||||
import time
|
||||
from concurrent.futures import Future, ThreadPoolExecutor, TimeoutError as FuturesTimeoutError
|
||||
from typing import List, Optional
|
||||
|
||||
from bot.commands.base import BotCommand
|
||||
from bot.models import BotMessage, BotResponse
|
||||
from src.config import get_config
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class ResearchCommand(BotCommand):
|
||||
"""
|
||||
Research command handler — invoke the deep research agent.
|
||||
|
||||
Usage:
|
||||
/research 600519 -> Deep research on a stock
|
||||
/research 600519 业绩风险分析 -> Focused question
|
||||
/research 新能源板块 发展前景 -> Sector research
|
||||
"""
|
||||
|
||||
@property
|
||||
def name(self) -> str:
|
||||
return "research"
|
||||
|
||||
@property
|
||||
def aliases(self) -> List[str]:
|
||||
return ["深研", "deepsearch"]
|
||||
|
||||
@property
|
||||
def description(self) -> str:
|
||||
return "Deep research on a stock or market topic"
|
||||
|
||||
@property
|
||||
def usage(self) -> str:
|
||||
return "/research <stock_code|topic> [specific question]"
|
||||
|
||||
def execute(self, message: BotMessage, args: List[str]) -> BotResponse:
|
||||
if not args:
|
||||
return BotResponse.text_response(
|
||||
f"Usage: {self.usage}\n"
|
||||
"Example: /research 600519 近期有哪些风险\n"
|
||||
"Example: /research 新能源板块前景分析"
|
||||
)
|
||||
|
||||
config = get_config()
|
||||
|
||||
# Check agent availability (consistent with /chat and API)
|
||||
if not config.is_agent_available():
|
||||
return BotResponse.text_response(
|
||||
"⚠️ Agent mode is not available. Configure LITELLM_MODEL or set AGENT_MODE=true to use /research."
|
||||
)
|
||||
|
||||
# Parse arguments — first arg may be stock code, rest is the question
|
||||
query_parts = list(args)
|
||||
stock_code: Optional[str] = None
|
||||
|
||||
# Try to detect a stock code in the first argument
|
||||
first = query_parts[0].upper().replace(",", ",")
|
||||
import re
|
||||
if re.match(r"^\d{6}$", first) or re.match(r"^HK\d{5}$", first) or re.match(r"^[A-Z]{1,4}(\.[A-Z]{1,2})?$", first):
|
||||
stock_code = first
|
||||
query_parts = query_parts[1:]
|
||||
|
||||
# Build the research query
|
||||
if query_parts:
|
||||
question = " ".join(query_parts)
|
||||
elif stock_code:
|
||||
question = f"Comprehensive deep research on stock {stock_code}: fundamentals, technicals, news sentiment, and risk factors"
|
||||
else:
|
||||
question = " ".join(args)
|
||||
|
||||
if stock_code:
|
||||
question = f"[Stock: {stock_code}] {question}"
|
||||
|
||||
# Run the research agent
|
||||
try:
|
||||
from src.agent.research import ResearchAgent
|
||||
from src.agent.factory import get_tool_registry
|
||||
from src.agent.llm_adapter import LLMToolAdapter
|
||||
|
||||
registry = get_tool_registry()
|
||||
llm_adapter = LLMToolAdapter(config)
|
||||
budget = getattr(config, "agent_deep_research_budget", 30000)
|
||||
|
||||
agent = ResearchAgent(
|
||||
tool_registry=registry,
|
||||
llm_adapter=llm_adapter,
|
||||
token_budget=budget,
|
||||
)
|
||||
|
||||
# Deep research can take minutes; cap with a timeout to prevent
|
||||
# indefinite blocking on Bot platforms with response-time limits.
|
||||
# IMPORTANT: we must NOT use `with ThreadPoolExecutor(...)` because
|
||||
# __exit__ calls shutdown(wait=True), which blocks until the thread
|
||||
# finishes — defeating the timeout. Instead we create the pool
|
||||
# manually and call shutdown(wait=False) on the timeout path so the
|
||||
# caller returns immediately (the orphan thread finishes in the
|
||||
# background).
|
||||
research_timeout = getattr(config, "agent_deep_research_timeout", 180)
|
||||
logger.info("[ResearchCommand] Starting deep research (timeout=%ds): %s", research_timeout, question[:100])
|
||||
t0 = time.time()
|
||||
|
||||
pool = ThreadPoolExecutor(max_workers=1)
|
||||
future: Future = pool.submit(
|
||||
agent.research,
|
||||
question,
|
||||
{"stock_code": stock_code, "stock_name": ""} if stock_code else None,
|
||||
)
|
||||
try:
|
||||
result = future.result(timeout=research_timeout)
|
||||
except FuturesTimeoutError:
|
||||
duration = round(time.time() - t0, 1)
|
||||
logger.warning("[ResearchCommand] Deep research timed out after %ss", duration)
|
||||
return BotResponse.text_response(
|
||||
f"⏳ 深度研究超时({duration}s / {research_timeout}s),请稍后重试或缩小研究范围。"
|
||||
)
|
||||
finally:
|
||||
pool.shutdown(wait=False, cancel_futures=True)
|
||||
|
||||
duration = round(time.time() - t0, 1)
|
||||
|
||||
if result.success:
|
||||
# Build rich response
|
||||
header = f"🔬 **Deep Research Report**\n"
|
||||
if stock_code:
|
||||
header += f"Stock: {stock_code}\n"
|
||||
header += f"Sub-questions: {len(result.sub_questions)} | Sources: {result.findings_count}\n"
|
||||
header += f"Time: {duration}s | Tokens: {result.total_tokens:,}\n"
|
||||
header += "─" * 40 + "\n\n"
|
||||
|
||||
report = header + result.report
|
||||
|
||||
# Truncate if too long for bot message
|
||||
max_len = 4000
|
||||
if len(report) > max_len:
|
||||
report = report[:max_len] + "\n\n... (report truncated, full report available via API)"
|
||||
|
||||
return BotResponse.markdown_response(report)
|
||||
else:
|
||||
return BotResponse.text_response(
|
||||
f"⚠️ Research did not complete successfully.\n"
|
||||
f"Partial results: {result.findings_count} findings collected.\n"
|
||||
f"Time: {duration}s"
|
||||
)
|
||||
|
||||
except Exception as exc:
|
||||
logger.error("[ResearchCommand] Error: %s", exc, exc_info=True)
|
||||
return BotResponse.text_response(f"❌ Research failed: {exc}")
|
||||
25
main.py
25
main.py
@@ -680,10 +680,33 @@ def main() -> int:
|
||||
def scheduled_task():
|
||||
run_full_analysis(config, args, stock_codes)
|
||||
|
||||
background_tasks = []
|
||||
if getattr(config, 'agent_event_monitor_enabled', False):
|
||||
from src.agent.events import build_event_monitor_from_config, run_event_monitor_once
|
||||
|
||||
monitor = build_event_monitor_from_config(config)
|
||||
if monitor is not None:
|
||||
interval_minutes = max(1, getattr(config, 'agent_event_monitor_interval_minutes', 5))
|
||||
|
||||
def event_monitor_task():
|
||||
triggered = run_event_monitor_once(monitor)
|
||||
if triggered:
|
||||
logger.info("[EventMonitor] 本轮触发 %d 条提醒", len(triggered))
|
||||
|
||||
background_tasks.append({
|
||||
"task": event_monitor_task,
|
||||
"interval_seconds": interval_minutes * 60,
|
||||
"run_immediately": True,
|
||||
"name": "agent_event_monitor",
|
||||
})
|
||||
else:
|
||||
logger.info("EventMonitor 已启用,但未加载到有效规则,跳过后台提醒任务")
|
||||
|
||||
run_with_schedule(
|
||||
task=scheduled_task,
|
||||
schedule_time=config.schedule_time,
|
||||
run_immediately=should_run_immediately
|
||||
run_immediately=should_run_immediately,
|
||||
background_tasks=background_tasks,
|
||||
)
|
||||
return 0
|
||||
|
||||
|
||||
457
src/agent/events.py
Normal file
457
src/agent/events.py
Normal file
@@ -0,0 +1,457 @@
|
||||
# -*- coding: utf-8 -*-
|
||||
"""
|
||||
EventMonitor — lightweight event-driven alert system.
|
||||
|
||||
Monitors a set of stocks for threshold events and triggers
|
||||
notifications when conditions are met. Designed to run as a
|
||||
background task (e.g. via ``--schedule`` or a dedicated loop).
|
||||
|
||||
Currently supported runtime events:
|
||||
- Price crossing threshold (above / below)
|
||||
- Volume spike (> N× average)
|
||||
|
||||
Other alert types remain defined as enum placeholders for future
|
||||
extension, but config validation rejects them until the monitor can
|
||||
actually evaluate them.
|
||||
|
||||
Usage::
|
||||
|
||||
from src.agent.events import EventMonitor, PriceAlert
|
||||
monitor = EventMonitor()
|
||||
monitor.add_alert(PriceAlert(stock_code="600519", direction="above", price=1800.0))
|
||||
triggered = await monitor.check_all()
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import json
|
||||
import logging
|
||||
import time
|
||||
from dataclasses import dataclass, field
|
||||
from enum import Enum
|
||||
from typing import Any, Callable, Dict, List, Optional
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class AlertType(str, Enum):
|
||||
PRICE_CROSS = "price_cross"
|
||||
VOLUME_SPIKE = "volume_spike"
|
||||
SENTIMENT_SHIFT = "sentiment_shift"
|
||||
RISK_FLAG = "risk_flag"
|
||||
CUSTOM = "custom"
|
||||
|
||||
|
||||
class AlertStatus(str, Enum):
|
||||
ACTIVE = "active"
|
||||
TRIGGERED = "triggered"
|
||||
EXPIRED = "expired"
|
||||
DISMISSED = "dismissed"
|
||||
|
||||
|
||||
_RUNTIME_SUPPORTED_ALERT_TYPES = frozenset({
|
||||
AlertType.PRICE_CROSS,
|
||||
AlertType.VOLUME_SPIKE,
|
||||
})
|
||||
|
||||
|
||||
def _supported_alert_type_names() -> str:
|
||||
return ", ".join(sorted(alert_type.value for alert_type in _RUNTIME_SUPPORTED_ALERT_TYPES))
|
||||
|
||||
|
||||
def _ensure_runtime_supported_alert_type(alert_type: AlertType) -> None:
|
||||
if alert_type not in _RUNTIME_SUPPORTED_ALERT_TYPES:
|
||||
raise ValueError(
|
||||
f"unsupported alert_type for current EventMonitor runtime: {alert_type.value} "
|
||||
f"(supported: {_supported_alert_type_names()})"
|
||||
)
|
||||
|
||||
|
||||
@dataclass
|
||||
class AlertRule:
|
||||
"""Base alert rule definition."""
|
||||
stock_code: str
|
||||
alert_type: AlertType
|
||||
description: str = ""
|
||||
status: AlertStatus = AlertStatus.ACTIVE
|
||||
created_at: float = field(default_factory=time.time)
|
||||
triggered_at: Optional[float] = None
|
||||
ttl_hours: float = 24.0 # auto-expire after this many hours
|
||||
metadata: Dict[str, Any] = field(default_factory=dict)
|
||||
|
||||
|
||||
@dataclass
|
||||
class PriceAlert(AlertRule):
|
||||
"""Alert when price crosses a threshold."""
|
||||
alert_type: AlertType = AlertType.PRICE_CROSS
|
||||
direction: str = "above" # "above" or "below"
|
||||
price: float = 0.0
|
||||
|
||||
def __post_init__(self):
|
||||
if not self.description:
|
||||
self.description = f"{self.stock_code} price {self.direction} {self.price}"
|
||||
|
||||
|
||||
@dataclass
|
||||
class VolumeAlert(AlertRule):
|
||||
"""Alert when volume exceeds N× average."""
|
||||
alert_type: AlertType = AlertType.VOLUME_SPIKE
|
||||
multiplier: float = 2.0 # trigger when volume > multiplier × avg
|
||||
|
||||
def __post_init__(self):
|
||||
if not self.description:
|
||||
self.description = f"{self.stock_code} volume > {self.multiplier}× average"
|
||||
|
||||
|
||||
@dataclass
|
||||
class SentimentAlert(AlertRule):
|
||||
"""Alert on sentiment direction change."""
|
||||
alert_type: AlertType = AlertType.SENTIMENT_SHIFT
|
||||
from_sentiment: str = "positive" # "positive", "negative", "neutral"
|
||||
to_sentiment: str = "negative"
|
||||
|
||||
def __post_init__(self):
|
||||
if not self.description:
|
||||
self.description = f"{self.stock_code} sentiment shift: {self.from_sentiment} → {self.to_sentiment}"
|
||||
|
||||
|
||||
@dataclass
|
||||
class TriggeredAlert:
|
||||
"""An alert that was triggered, ready for notification."""
|
||||
rule: AlertRule
|
||||
triggered_at: float = field(default_factory=time.time)
|
||||
current_value: Any = None
|
||||
message: str = ""
|
||||
|
||||
|
||||
class EventMonitor:
|
||||
"""Monitor stocks for event-driven alerts.
|
||||
|
||||
This class manages a list of :class:`AlertRule` objects and checks
|
||||
them against current market data. Triggered alerts are collected
|
||||
and can be forwarded to the notification system.
|
||||
"""
|
||||
|
||||
def __init__(self):
|
||||
self.rules: List[AlertRule] = []
|
||||
self._callbacks: List[Callable[[TriggeredAlert], None]] = []
|
||||
|
||||
def add_alert(self, rule: AlertRule) -> None:
|
||||
"""Register a new alert rule."""
|
||||
_ensure_runtime_supported_alert_type(rule.alert_type)
|
||||
self.rules.append(rule)
|
||||
logger.info("[EventMonitor] Added alert: %s", rule.description)
|
||||
|
||||
def remove_expired(self) -> int:
|
||||
"""Remove alerts that have expired based on TTL.
|
||||
|
||||
Returns:
|
||||
Number of expired alerts removed.
|
||||
"""
|
||||
now = time.time()
|
||||
before = len(self.rules)
|
||||
self.rules = [
|
||||
r for r in self.rules
|
||||
if r.status != AlertStatus.EXPIRED
|
||||
and (now - r.created_at) < r.ttl_hours * 3600
|
||||
]
|
||||
removed = before - len(self.rules)
|
||||
if removed:
|
||||
logger.info("[EventMonitor] Removed %d expired alerts", removed)
|
||||
return removed
|
||||
|
||||
def on_trigger(self, callback: Callable[[TriggeredAlert], None]) -> None:
|
||||
"""Register a callback for when an alert triggers."""
|
||||
self._callbacks.append(callback)
|
||||
|
||||
async def check_all(self) -> List[TriggeredAlert]:
|
||||
"""Check all active rules against current market data.
|
||||
|
||||
Returns:
|
||||
List of triggered alerts.
|
||||
"""
|
||||
self.remove_expired()
|
||||
triggered: List[TriggeredAlert] = []
|
||||
|
||||
for rule in self.rules:
|
||||
if rule.status != AlertStatus.ACTIVE:
|
||||
continue
|
||||
|
||||
try:
|
||||
result = await self._check_rule(rule)
|
||||
if result:
|
||||
triggered.append(result)
|
||||
rule.status = AlertStatus.TRIGGERED
|
||||
rule.triggered_at = time.time()
|
||||
# Notify callbacks (offload slow/sync ones to thread)
|
||||
for cb in self._callbacks:
|
||||
try:
|
||||
if asyncio.iscoroutinefunction(cb):
|
||||
await cb(result)
|
||||
else:
|
||||
await asyncio.to_thread(cb, result)
|
||||
except Exception as exc:
|
||||
logger.warning("[EventMonitor] Callback error: %s", exc)
|
||||
except Exception as exc:
|
||||
logger.debug("[EventMonitor] Check failed for %s: %s", rule.description, exc)
|
||||
|
||||
return triggered
|
||||
|
||||
async def _check_rule(self, rule: AlertRule) -> Optional[TriggeredAlert]:
|
||||
"""Check a single rule. Returns TriggeredAlert if condition met."""
|
||||
if isinstance(rule, PriceAlert):
|
||||
return await self._check_price(rule)
|
||||
elif isinstance(rule, VolumeAlert):
|
||||
return await self._check_volume(rule)
|
||||
# SentimentAlert and custom alerts require more context —
|
||||
# implemented as hooks for future extension
|
||||
return None
|
||||
|
||||
async def _check_price(self, rule: PriceAlert) -> Optional[TriggeredAlert]:
|
||||
"""Check price alert against realtime quote."""
|
||||
try:
|
||||
def _fetch_quote():
|
||||
from data_provider import DataFetcherManager
|
||||
|
||||
fm = DataFetcherManager()
|
||||
return fm.get_realtime_quote(rule.stock_code)
|
||||
|
||||
quote = await asyncio.to_thread(_fetch_quote)
|
||||
if quote is None:
|
||||
return None
|
||||
|
||||
current_price = float(getattr(quote, "price", 0) or 0)
|
||||
if current_price <= 0:
|
||||
return None
|
||||
|
||||
triggered = False
|
||||
if rule.direction == "above" and current_price >= rule.price:
|
||||
triggered = True
|
||||
elif rule.direction == "below" and current_price <= rule.price:
|
||||
triggered = True
|
||||
|
||||
if triggered:
|
||||
return TriggeredAlert(
|
||||
rule=rule,
|
||||
current_value=current_price,
|
||||
message=f"🔔 {rule.stock_code} price {rule.direction} {rule.price}: "
|
||||
f"current = {current_price}",
|
||||
)
|
||||
except Exception as exc:
|
||||
logger.debug("[EventMonitor] _check_price error: %s", exc)
|
||||
return None
|
||||
|
||||
async def _check_volume(self, rule: VolumeAlert) -> Optional[TriggeredAlert]:
|
||||
"""Check volume spike against recent average."""
|
||||
try:
|
||||
def _fetch_daily_data():
|
||||
from data_provider import DataFetcherManager
|
||||
|
||||
fm = DataFetcherManager()
|
||||
return fm.get_daily_data(rule.stock_code, days=20)
|
||||
|
||||
result = await asyncio.to_thread(_fetch_daily_data)
|
||||
# get_daily_data returns (df, source) tuple or None
|
||||
if result is None:
|
||||
return None
|
||||
df, _source = result
|
||||
if df is None or df.empty:
|
||||
return None
|
||||
|
||||
avg_vol = df["volume"].mean()
|
||||
latest_vol = df["volume"].iloc[-1]
|
||||
|
||||
if avg_vol > 0 and latest_vol > avg_vol * rule.multiplier:
|
||||
return TriggeredAlert(
|
||||
rule=rule,
|
||||
current_value=latest_vol,
|
||||
message=f"📊 {rule.stock_code} volume spike: "
|
||||
f"{latest_vol:,.0f} ({latest_vol / avg_vol:.1f}× avg)",
|
||||
)
|
||||
except Exception as exc:
|
||||
logger.debug("[EventMonitor] _check_volume error: %s", exc)
|
||||
return None
|
||||
|
||||
# -----------------------------------------------------------------
|
||||
# Persistence helpers
|
||||
# -----------------------------------------------------------------
|
||||
|
||||
def to_dict_list(self) -> List[Dict[str, Any]]:
|
||||
"""Serialize all rules for persistence."""
|
||||
results = []
|
||||
for rule in self.rules:
|
||||
entry: Dict[str, Any] = {
|
||||
"stock_code": rule.stock_code,
|
||||
"alert_type": rule.alert_type.value,
|
||||
"description": rule.description,
|
||||
"status": rule.status.value,
|
||||
"created_at": rule.created_at,
|
||||
"ttl_hours": rule.ttl_hours,
|
||||
}
|
||||
if isinstance(rule, PriceAlert):
|
||||
entry["direction"] = rule.direction
|
||||
entry["price"] = rule.price
|
||||
elif isinstance(rule, VolumeAlert):
|
||||
entry["multiplier"] = rule.multiplier
|
||||
results.append(entry)
|
||||
return results
|
||||
|
||||
@classmethod
|
||||
def from_dict_list(cls, data: List[Dict[str, Any]]) -> "EventMonitor":
|
||||
"""Restore an EventMonitor from serialized data."""
|
||||
monitor = cls()
|
||||
for index, entry in enumerate(data, start=1):
|
||||
try:
|
||||
validate_event_alert_rule(entry)
|
||||
|
||||
alert_type = entry.get("alert_type", "custom")
|
||||
stock_code = entry.get("stock_code", "")
|
||||
if alert_type == AlertType.PRICE_CROSS.value:
|
||||
rule = PriceAlert(
|
||||
stock_code=stock_code,
|
||||
direction=entry.get("direction", "above").lower(),
|
||||
price=float(entry.get("price", 0.0)),
|
||||
)
|
||||
elif alert_type == AlertType.VOLUME_SPIKE.value:
|
||||
rule = VolumeAlert(
|
||||
stock_code=stock_code,
|
||||
multiplier=float(entry.get("multiplier", 2.0)),
|
||||
)
|
||||
else:
|
||||
raise ValueError(f"unsupported alert_type: {alert_type}")
|
||||
rule.status = AlertStatus(entry.get("status", "active"))
|
||||
raw_created = entry.get("created_at")
|
||||
try:
|
||||
rule.created_at = float(raw_created) if raw_created is not None else time.time()
|
||||
except (TypeError, ValueError):
|
||||
rule.created_at = time.time()
|
||||
rule.ttl_hours = float(entry.get("ttl_hours", 24.0))
|
||||
monitor.add_alert(rule)
|
||||
except Exception as exc:
|
||||
logger.warning("[EventMonitor] Skip invalid rule #%d: %s", index, exc)
|
||||
return monitor
|
||||
|
||||
|
||||
def parse_event_alert_rules(raw_rules: Any) -> List[Dict[str, Any]]:
|
||||
"""Parse event alert rules from config JSON or already-loaded objects."""
|
||||
if raw_rules is None:
|
||||
return []
|
||||
|
||||
parsed = raw_rules
|
||||
if isinstance(raw_rules, str):
|
||||
cleaned = raw_rules.strip()
|
||||
if not cleaned:
|
||||
return []
|
||||
parsed = json.loads(cleaned)
|
||||
|
||||
if isinstance(parsed, dict):
|
||||
parsed = parsed.get("rules", [])
|
||||
|
||||
if not isinstance(parsed, list):
|
||||
raise ValueError("Event alert rules must be a JSON array")
|
||||
|
||||
invalid_indices = [idx for idx, entry in enumerate(parsed) if not isinstance(entry, dict)]
|
||||
if invalid_indices:
|
||||
raise ValueError(
|
||||
"Event alert rules list must contain only objects; "
|
||||
f"invalid entries at positions: {invalid_indices}"
|
||||
)
|
||||
|
||||
return parsed
|
||||
|
||||
|
||||
def validate_event_alert_rule(rule: Dict[str, Any]) -> None:
|
||||
"""Validate one serialized EventMonitor rule."""
|
||||
if not isinstance(rule, dict):
|
||||
raise ValueError("Event alert rule must be an object")
|
||||
|
||||
stock_code = str(rule.get("stock_code") or "").strip()
|
||||
if not stock_code:
|
||||
raise ValueError("stock_code is required")
|
||||
|
||||
try:
|
||||
alert_type = AlertType(rule.get("alert_type", ""))
|
||||
except ValueError as exc:
|
||||
raise ValueError(f"invalid alert_type: {rule.get('alert_type')}") from exc
|
||||
_ensure_runtime_supported_alert_type(alert_type)
|
||||
|
||||
status = rule.get("status")
|
||||
if status is not None:
|
||||
try:
|
||||
AlertStatus(status)
|
||||
except ValueError as exc:
|
||||
raise ValueError(f"invalid status: {status}") from exc
|
||||
|
||||
ttl_hours = rule.get("ttl_hours")
|
||||
if ttl_hours is not None:
|
||||
try:
|
||||
ttl_value = float(ttl_hours)
|
||||
except (TypeError, ValueError) as exc:
|
||||
raise ValueError(f"invalid ttl_hours: {ttl_hours}") from exc
|
||||
if ttl_value <= 0:
|
||||
raise ValueError("ttl_hours must be > 0")
|
||||
|
||||
if alert_type == AlertType.PRICE_CROSS:
|
||||
direction = str(rule.get("direction", "above")).lower()
|
||||
if direction not in {"above", "below"}:
|
||||
raise ValueError(f"invalid direction: {direction}")
|
||||
try:
|
||||
price = float(rule.get("price"))
|
||||
except (TypeError, ValueError) as exc:
|
||||
raise ValueError(f"invalid price: {rule.get('price')}") from exc
|
||||
if price <= 0:
|
||||
raise ValueError("price must be > 0")
|
||||
elif alert_type == AlertType.VOLUME_SPIKE:
|
||||
try:
|
||||
multiplier = float(rule.get("multiplier", 2.0))
|
||||
except (TypeError, ValueError) as exc:
|
||||
raise ValueError(f"invalid multiplier: {rule.get('multiplier')}") from exc
|
||||
if multiplier <= 0:
|
||||
raise ValueError("multiplier must be > 0")
|
||||
|
||||
|
||||
def build_event_monitor_from_config(config=None, notifier=None) -> Optional[EventMonitor]:
|
||||
"""Build an EventMonitor from runtime config and attach notification callbacks."""
|
||||
if config is None:
|
||||
from src.config import get_config
|
||||
config = get_config()
|
||||
|
||||
if not getattr(config, "agent_event_monitor_enabled", False):
|
||||
return None
|
||||
|
||||
raw_rules = getattr(config, "agent_event_alert_rules_json", "")
|
||||
try:
|
||||
rules = parse_event_alert_rules(raw_rules)
|
||||
except Exception as exc:
|
||||
logger.warning("[EventMonitor] Failed to parse configured alert rules: %s", exc)
|
||||
return None
|
||||
|
||||
if not rules:
|
||||
logger.info("[EventMonitor] Enabled but no alert rules configured")
|
||||
return None
|
||||
|
||||
monitor = EventMonitor.from_dict_list(rules)
|
||||
if not monitor.rules:
|
||||
return None
|
||||
|
||||
from src.notification import NotificationBuilder, NotificationService
|
||||
|
||||
notification_service = notifier or NotificationService()
|
||||
|
||||
def _notify(triggered: TriggeredAlert) -> None:
|
||||
title = f"Event Alert | {triggered.rule.stock_code}"
|
||||
content = triggered.message or triggered.rule.description or "Alert triggered"
|
||||
alert_text = NotificationBuilder.build_simple_alert(title=title, content=content, alert_type="warning")
|
||||
sent = notification_service.send(alert_text)
|
||||
if not sent:
|
||||
logger.info("[EventMonitor] No notification channel available for alert: %s", title)
|
||||
|
||||
monitor.on_trigger(_notify)
|
||||
logger.info("[EventMonitor] Loaded %d configured alert rule(s)", len(monitor.rules))
|
||||
return monitor
|
||||
|
||||
|
||||
def run_event_monitor_once(monitor: EventMonitor) -> List[TriggeredAlert]:
|
||||
"""Run one synchronous monitor cycle."""
|
||||
return asyncio.run(monitor.check_all())
|
||||
302
src/agent/research.py
Normal file
302
src/agent/research.py
Normal file
@@ -0,0 +1,302 @@
|
||||
# -*- coding: utf-8 -*-
|
||||
"""
|
||||
ResearchAgent — deep research specialist for in-depth analysis.
|
||||
|
||||
Responsible for:
|
||||
- Decomposing a complex research query into sub-questions
|
||||
- Iterative search and information gathering
|
||||
- Cross-verification of findings
|
||||
- Producing a structured research report
|
||||
|
||||
Triggered by ``/research`` command or API async task interface.
|
||||
Designed for long-running analysis (up to ``AGENT_DEEP_RESEARCH_BUDGET``
|
||||
tokens).
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import logging
|
||||
import re
|
||||
import time
|
||||
from dataclasses import dataclass, field
|
||||
from typing import Any, Callable, Dict, List, Optional
|
||||
|
||||
from src.agent.llm_adapter import LLMToolAdapter
|
||||
from src.agent.runner import RunLoopResult, run_agent_loop
|
||||
from src.agent.tools.registry import ToolRegistry
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# Default token budget for deep research
|
||||
_DEFAULT_TOKEN_BUDGET = 30000
|
||||
|
||||
|
||||
class ResearchAgent:
|
||||
"""Multi-turn deep research agent.
|
||||
|
||||
Unlike the standard agent loop which runs a fixed number of steps,
|
||||
the ResearchAgent:
|
||||
1. Decomposes the query into sub-questions (planning phase)
|
||||
2. Researches each sub-question with dedicated searches
|
||||
3. Synthesises findings into a comprehensive report
|
||||
4. Tracks total token usage against a configurable budget
|
||||
"""
|
||||
|
||||
agent_name = "research"
|
||||
tool_names = [
|
||||
"search_stock_news",
|
||||
"search_comprehensive_intel",
|
||||
"get_stock_info",
|
||||
"get_realtime_quote",
|
||||
"get_daily_history",
|
||||
"get_sector_rankings",
|
||||
"get_market_indices",
|
||||
]
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
tool_registry: ToolRegistry,
|
||||
llm_adapter: LLMToolAdapter,
|
||||
token_budget: int = _DEFAULT_TOKEN_BUDGET,
|
||||
max_sub_questions: int = 5,
|
||||
):
|
||||
self.tool_registry = tool_registry
|
||||
self.llm_adapter = llm_adapter
|
||||
self.token_budget = token_budget
|
||||
self.max_sub_questions = max_sub_questions
|
||||
|
||||
def research(
|
||||
self,
|
||||
query: str,
|
||||
context: Optional[Dict[str, Any]] = None,
|
||||
progress_callback: Optional[Callable[[Dict[str, Any]], None]] = None,
|
||||
) -> ResearchResult:
|
||||
"""Execute a deep research task.
|
||||
|
||||
Args:
|
||||
query: The research question or topic.
|
||||
context: Optional context (stock_code, stock_name, etc.).
|
||||
progress_callback: Optional progress updates.
|
||||
|
||||
Returns:
|
||||
A :class:`ResearchResult` containing the report and metadata.
|
||||
"""
|
||||
t0 = time.time()
|
||||
tokens_used = 0
|
||||
all_findings: List[Dict[str, Any]] = []
|
||||
|
||||
# Phase 1: Decompose
|
||||
if progress_callback:
|
||||
progress_callback({"type": "research_phase", "phase": "decompose", "message": "Decomposing research query..."})
|
||||
|
||||
sub_questions = self._decompose_query(query, context)
|
||||
tokens_used += sub_questions.get("tokens", 0)
|
||||
|
||||
questions = sub_questions.get("questions", [query])[:self.max_sub_questions]
|
||||
logger.info("[ResearchAgent] decomposed into %d sub-questions", len(questions))
|
||||
|
||||
# Phase 2: Research each sub-question
|
||||
for i, question in enumerate(questions):
|
||||
if tokens_used >= self.token_budget:
|
||||
logger.warning("[ResearchAgent] token budget exceeded (%d/%d), stopping", tokens_used, self.token_budget)
|
||||
break
|
||||
|
||||
if progress_callback:
|
||||
progress_callback({
|
||||
"type": "research_phase",
|
||||
"phase": "search",
|
||||
"message": f"Researching ({i + 1}/{len(questions)}): {question[:60]}...",
|
||||
"progress": (i + 1) / len(questions),
|
||||
})
|
||||
|
||||
finding = self._research_sub_question(question, context, tokens_used)
|
||||
tokens_used += finding.get("tokens", 0)
|
||||
all_findings.append(finding)
|
||||
|
||||
# Phase 3: Synthesise
|
||||
if progress_callback:
|
||||
progress_callback({"type": "research_phase", "phase": "synthesize", "message": "Synthesising research report..."})
|
||||
|
||||
report = self._synthesise_report(query, all_findings, context) if all_findings else {"content": "No findings gathered.", "tokens": 0}
|
||||
tokens_used += report.get("tokens", 0)
|
||||
|
||||
duration = round(time.time() - t0, 2)
|
||||
|
||||
return ResearchResult(
|
||||
success=bool(report.get("content")),
|
||||
report=report.get("content", ""),
|
||||
sub_questions=questions,
|
||||
findings_count=len(all_findings),
|
||||
total_tokens=tokens_used,
|
||||
duration_s=duration,
|
||||
error=report.get("error"),
|
||||
)
|
||||
|
||||
def _call_text_completion(
|
||||
self,
|
||||
messages: List[Dict[str, Any]],
|
||||
*,
|
||||
temperature: float,
|
||||
max_tokens: int,
|
||||
timeout: int,
|
||||
) -> Dict[str, Any]:
|
||||
"""Run a text-only LLM completion via the shared adapter."""
|
||||
response = self.llm_adapter.call_text(
|
||||
messages,
|
||||
temperature=temperature,
|
||||
max_tokens=max_tokens,
|
||||
timeout=timeout,
|
||||
)
|
||||
if response.provider == "error":
|
||||
raise RuntimeError(response.content or "LLM completion failed")
|
||||
return {
|
||||
"content": (response.content or "").strip(),
|
||||
"tokens": response.usage.get("total_tokens", 0),
|
||||
}
|
||||
|
||||
def _decompose_query(self, query: str, context: Optional[Dict[str, Any]]) -> Dict[str, Any]:
|
||||
"""Use LLM to decompose a research query into sub-questions."""
|
||||
stock_hint = ""
|
||||
if context and context.get("stock_code"):
|
||||
stock_hint = f"\nStock context: {context['stock_code']} ({context.get('stock_name', '')})"
|
||||
|
||||
system = """\
|
||||
You are a research planning assistant. Given a research query, decompose it \
|
||||
into 3-5 specific, searchable sub-questions.
|
||||
|
||||
Return a JSON object:
|
||||
{"questions": ["question 1", "question 2", ...]}
|
||||
"""
|
||||
messages = [
|
||||
{"role": "system", "content": system},
|
||||
{"role": "user", "content": f"Research query: {query}{stock_hint}"},
|
||||
]
|
||||
|
||||
try:
|
||||
completion = self._call_text_completion(
|
||||
messages,
|
||||
temperature=0.3,
|
||||
max_tokens=400,
|
||||
timeout=15,
|
||||
)
|
||||
raw = completion["content"]
|
||||
tokens = completion["tokens"]
|
||||
|
||||
# Parse JSON
|
||||
if raw.startswith("```"):
|
||||
raw = re.sub(r'^```(?:json)?\s*', '', raw)
|
||||
raw = re.sub(r'\s*```$', '', raw)
|
||||
parsed = json.loads(raw)
|
||||
return {"questions": parsed.get("questions", [query]), "tokens": tokens}
|
||||
except Exception as exc:
|
||||
logger.warning("[ResearchAgent] decompose failed: %s", exc)
|
||||
return {"questions": [query], "tokens": 0}
|
||||
|
||||
def _research_sub_question(
|
||||
self,
|
||||
question: str,
|
||||
context: Optional[Dict[str, Any]],
|
||||
current_tokens: int,
|
||||
) -> Dict[str, Any]:
|
||||
"""Research a single sub-question using the agent loop."""
|
||||
remaining_budget = self.token_budget - current_tokens
|
||||
|
||||
system = f"""\
|
||||
You are a research agent investigating a specific question.
|
||||
Use your tools to search for relevant information, then summarise \
|
||||
your findings in 2-4 paragraphs. Be factual and cite sources.
|
||||
Token budget remaining: ~{remaining_budget}
|
||||
"""
|
||||
stock_context = ""
|
||||
if context and context.get("stock_code"):
|
||||
stock_context = f" (related to stock {context['stock_code']})"
|
||||
|
||||
messages = [
|
||||
{"role": "system", "content": system},
|
||||
{"role": "user", "content": f"Research question: {question}{stock_context}"},
|
||||
]
|
||||
|
||||
try:
|
||||
registry = self._filtered_registry()
|
||||
result: RunLoopResult = run_agent_loop(
|
||||
messages=messages,
|
||||
tool_registry=registry,
|
||||
llm_adapter=self.llm_adapter,
|
||||
max_steps=4,
|
||||
)
|
||||
return {
|
||||
"question": question,
|
||||
"content": result.content,
|
||||
"tokens": result.total_tokens,
|
||||
"success": result.success,
|
||||
}
|
||||
except Exception as exc:
|
||||
logger.warning("[ResearchAgent] sub-question failed: %s", exc)
|
||||
return {"question": question, "content": "", "tokens": 0, "success": False, "error": str(exc)}
|
||||
|
||||
def _synthesise_report(
|
||||
self,
|
||||
original_query: str,
|
||||
findings: List[Dict[str, Any]],
|
||||
context: Optional[Dict[str, Any]],
|
||||
) -> Dict[str, Any]:
|
||||
"""Synthesise all findings into a coherent research report."""
|
||||
findings_text = "\n\n".join(
|
||||
f"### Sub-question: {f['question']}\n{f.get('content', 'No data')}"
|
||||
for f in findings if f.get("content")
|
||||
)
|
||||
|
||||
system = """\
|
||||
You are a senior research analyst. Synthesise the following research \
|
||||
findings into a comprehensive, well-structured report.
|
||||
|
||||
## Report Structure
|
||||
1. **Executive Summary** (2-3 sentences)
|
||||
2. **Key Findings** (bullet points)
|
||||
3. **Detailed Analysis** (sections per topic)
|
||||
4. **Risk Factors** (if applicable)
|
||||
5. **Conclusion & Recommendations**
|
||||
|
||||
Use Markdown formatting. Be concise but thorough.
|
||||
"""
|
||||
messages = [
|
||||
{"role": "system", "content": system},
|
||||
{"role": "user", "content": f"Original query: {original_query}\n\n## Research Findings\n\n{findings_text}"},
|
||||
]
|
||||
|
||||
try:
|
||||
completion = self._call_text_completion(
|
||||
messages,
|
||||
temperature=0.3,
|
||||
max_tokens=2000,
|
||||
timeout=30,
|
||||
)
|
||||
content = completion["content"]
|
||||
tokens = completion["tokens"]
|
||||
return {"content": content, "tokens": tokens}
|
||||
except Exception as exc:
|
||||
logger.warning("[ResearchAgent] synthesis failed: %s", exc)
|
||||
return {"content": findings_text, "tokens": 0, "error": str(exc)}
|
||||
|
||||
def _filtered_registry(self) -> ToolRegistry:
|
||||
"""Return a registry restricted to research-related tools.
|
||||
|
||||
Reuses the same filtering logic as :meth:`BaseAgent._filtered_registry`.
|
||||
"""
|
||||
from src.agent.agents.base_agent import BaseAgent
|
||||
# Borrow the shared implementation; it respects self.tool_names / self.tool_registry.
|
||||
return BaseAgent._filtered_registry(self)
|
||||
|
||||
|
||||
@dataclass
|
||||
class ResearchResult:
|
||||
"""Output from a deep research task."""
|
||||
|
||||
success: bool = False
|
||||
report: str = ""
|
||||
sub_questions: List[str] = field(default_factory=list)
|
||||
findings_count: int = 0
|
||||
total_tokens: int = 0
|
||||
duration_s: float = 0.0
|
||||
error: Optional[str] = None
|
||||
@@ -19,7 +19,7 @@ import sys
|
||||
import time
|
||||
import threading
|
||||
from datetime import datetime
|
||||
from typing import Callable, Optional
|
||||
from typing import Any, Callable, Dict, List, Optional
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
@@ -80,6 +80,7 @@ class Scheduler:
|
||||
self.schedule_time = schedule_time
|
||||
self.shutdown_handler = GracefulShutdown()
|
||||
self._task_callback: Optional[Callable] = None
|
||||
self._background_tasks: List[Dict[str, Any]] = []
|
||||
self._running = False
|
||||
|
||||
def set_daily_task(self, task: Callable, run_immediately: bool = True):
|
||||
@@ -116,6 +117,89 @@ class Scheduler:
|
||||
|
||||
except Exception as e:
|
||||
logger.exception(f"定时任务执行失败: {e}")
|
||||
|
||||
def add_background_task(
|
||||
self,
|
||||
task: Callable,
|
||||
interval_seconds: int,
|
||||
run_immediately: bool = False,
|
||||
name: Optional[str] = None,
|
||||
) -> None:
|
||||
"""Register a periodic background task executed inside the scheduler loop.
|
||||
|
||||
Note: The scheduler loop polls every 30 seconds, so *interval_seconds*
|
||||
below 30 will be clamped to 30 to avoid promising unreachable precision.
|
||||
"""
|
||||
clamped_interval = max(30, int(interval_seconds))
|
||||
if int(interval_seconds) < 30:
|
||||
logger.warning(
|
||||
"后台任务 %s 请求间隔 %ds,但调度循环每 30s 轮询一次,已自动调整为 30s",
|
||||
name or getattr(task, "__name__", "background_task"),
|
||||
interval_seconds,
|
||||
)
|
||||
entry = {
|
||||
"task": task,
|
||||
"interval_seconds": clamped_interval,
|
||||
"last_run": 0.0,
|
||||
"name": name or getattr(task, "__name__", "background_task"),
|
||||
"thread": None,
|
||||
"running": False,
|
||||
}
|
||||
if not run_immediately:
|
||||
entry["last_run"] = time.time()
|
||||
self._background_tasks.append(entry)
|
||||
logger.info(
|
||||
"已注册后台任务: %s(间隔 %s 秒,立即执行=%s)",
|
||||
entry["name"],
|
||||
entry["interval_seconds"],
|
||||
run_immediately,
|
||||
)
|
||||
if run_immediately:
|
||||
self._start_background_task(entry)
|
||||
|
||||
def _start_background_task(self, entry: Dict[str, Any]) -> bool:
|
||||
"""Start one background task in a dedicated daemon thread."""
|
||||
worker = entry.get("thread")
|
||||
if worker is not None and worker.is_alive():
|
||||
return False
|
||||
|
||||
def _runner() -> None:
|
||||
try:
|
||||
logger.info("后台任务开始执行: %s", entry["name"])
|
||||
entry["task"]()
|
||||
except Exception as exc:
|
||||
logger.exception("后台任务执行失败 [%s]: %s", entry["name"], exc)
|
||||
finally:
|
||||
entry["running"] = False
|
||||
entry["thread"] = None
|
||||
|
||||
entry["last_run"] = time.time()
|
||||
entry["running"] = True
|
||||
worker = threading.Thread(
|
||||
target=_runner,
|
||||
daemon=True,
|
||||
name=f"scheduler-bg-{entry['name']}",
|
||||
)
|
||||
entry["thread"] = worker
|
||||
worker.start()
|
||||
return True
|
||||
|
||||
def _run_background_tasks(self) -> None:
|
||||
"""Execute any background tasks whose interval has elapsed."""
|
||||
if not self._background_tasks:
|
||||
return
|
||||
|
||||
now = time.time()
|
||||
for entry in self._background_tasks:
|
||||
worker = entry.get("thread")
|
||||
if worker is not None and worker.is_alive():
|
||||
continue
|
||||
if entry.get("running"):
|
||||
entry["running"] = False
|
||||
entry["thread"] = None
|
||||
if now - entry["last_run"] < entry["interval_seconds"]:
|
||||
continue
|
||||
self._start_background_task(entry)
|
||||
|
||||
def run(self):
|
||||
"""
|
||||
@@ -129,6 +213,7 @@ class Scheduler:
|
||||
|
||||
while self._running and not self.shutdown_handler.should_shutdown:
|
||||
self.schedule.run_pending()
|
||||
self._run_background_tasks()
|
||||
time.sleep(30) # 每30秒检查一次
|
||||
|
||||
# 每小时打印一次心跳
|
||||
@@ -153,7 +238,8 @@ class Scheduler:
|
||||
def run_with_schedule(
|
||||
task: Callable,
|
||||
schedule_time: str = "18:00",
|
||||
run_immediately: bool = True
|
||||
run_immediately: bool = True,
|
||||
background_tasks: Optional[List[Dict[str, Any]]] = None,
|
||||
):
|
||||
"""
|
||||
便捷函数:使用定时调度运行任务
|
||||
@@ -164,6 +250,13 @@ def run_with_schedule(
|
||||
run_immediately: 是否立即执行一次
|
||||
"""
|
||||
scheduler = Scheduler(schedule_time=schedule_time)
|
||||
for entry in background_tasks or []:
|
||||
scheduler.add_background_task(
|
||||
task=entry["task"],
|
||||
interval_seconds=entry["interval_seconds"],
|
||||
run_immediately=entry.get("run_immediately", False),
|
||||
name=entry.get("name"),
|
||||
)
|
||||
scheduler.set_daily_task(task, run_immediately=run_immediately)
|
||||
scheduler.run()
|
||||
|
||||
|
||||
@@ -4,6 +4,7 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
import json
|
||||
import re
|
||||
import time
|
||||
from typing import Any, Dict, List, Optional, Sequence, Set, Tuple
|
||||
@@ -383,6 +384,40 @@ class SystemConfigService:
|
||||
}
|
||||
)
|
||||
|
||||
elif data_type == "json":
|
||||
try:
|
||||
parsed = json.loads(value)
|
||||
except json.JSONDecodeError:
|
||||
issues.append(
|
||||
{
|
||||
"key": key,
|
||||
"code": "invalid_json",
|
||||
"message": "Value must be valid JSON",
|
||||
"severity": "error",
|
||||
"expected": "valid JSON",
|
||||
"actual": value[:120],
|
||||
}
|
||||
)
|
||||
else:
|
||||
if key == "AGENT_EVENT_ALERT_RULES_JSON":
|
||||
try:
|
||||
from src.agent.events import parse_event_alert_rules, validate_event_alert_rule
|
||||
|
||||
rule_index = 0
|
||||
for rule_index, rule in enumerate(parse_event_alert_rules(parsed), start=1):
|
||||
validate_event_alert_rule(rule)
|
||||
except ValueError as exc:
|
||||
issues.append(
|
||||
{
|
||||
"key": key,
|
||||
"code": "invalid_event_rule",
|
||||
"message": f"Rule validation failed: {exc}",
|
||||
"severity": "error",
|
||||
"expected": "supported EventMonitor rule fields and enum values",
|
||||
"actual": f"rule #{rule_index or 1}",
|
||||
}
|
||||
)
|
||||
|
||||
if "enum" in validation and value and value not in validation["enum"]:
|
||||
issues.append(
|
||||
{
|
||||
|
||||
@@ -870,6 +870,162 @@ class TestBaseAgentMessageAssembly(unittest.TestCase):
|
||||
self.assertEqual(messages[-1], {"role": "user", "content": "current turn"})
|
||||
|
||||
|
||||
# ============================================================
|
||||
# EventMonitor serialization
|
||||
# ============================================================
|
||||
|
||||
class TestEventMonitor(unittest.TestCase):
|
||||
"""Test EventMonitor serialize/deserialize round-trip."""
|
||||
|
||||
def test_round_trip(self):
|
||||
from src.agent.events import EventMonitor, PriceAlert, VolumeAlert
|
||||
monitor = EventMonitor()
|
||||
monitor.add_alert(PriceAlert(stock_code="600519", direction="above", price=1800.0))
|
||||
monitor.add_alert(VolumeAlert(stock_code="000858", multiplier=3.0))
|
||||
|
||||
data = monitor.to_dict_list()
|
||||
self.assertEqual(len(data), 2)
|
||||
|
||||
restored = EventMonitor.from_dict_list(data)
|
||||
self.assertEqual(len(restored.rules), 2)
|
||||
self.assertEqual(restored.rules[0].stock_code, "600519")
|
||||
self.assertEqual(restored.rules[1].stock_code, "000858")
|
||||
|
||||
def test_remove_expired(self):
|
||||
import time
|
||||
from src.agent.events import EventMonitor, PriceAlert
|
||||
monitor = EventMonitor()
|
||||
alert = PriceAlert(stock_code="600519", direction="above", price=1800.0, ttl_hours=0.0)
|
||||
alert.created_at = time.time() - 3600 # 1 hour ago
|
||||
monitor.rules.append(alert)
|
||||
removed = monitor.remove_expired()
|
||||
self.assertEqual(removed, 1)
|
||||
self.assertEqual(len(monitor.rules), 0)
|
||||
|
||||
def test_add_alert_rejects_unsupported_rule_type(self):
|
||||
from src.agent.events import EventMonitor, SentimentAlert
|
||||
|
||||
monitor = EventMonitor()
|
||||
|
||||
with self.assertRaises(ValueError):
|
||||
monitor.add_alert(SentimentAlert(stock_code="600519"))
|
||||
|
||||
|
||||
class TestEventMonitorAsync(unittest.IsolatedAsyncioTestCase):
|
||||
"""Test async EventMonitor checks offload blocking fetches."""
|
||||
|
||||
async def test_check_price_uses_to_thread_and_triggers(self):
|
||||
from src.agent.events import EventMonitor, PriceAlert
|
||||
|
||||
monitor = EventMonitor()
|
||||
rule = PriceAlert(stock_code="600519", direction="above", price=1800.0)
|
||||
quote = SimpleNamespace(price=1810.0)
|
||||
|
||||
with patch("src.agent.events.asyncio.to_thread", new=AsyncMock(return_value=quote)) as to_thread:
|
||||
triggered = await monitor._check_price(rule)
|
||||
|
||||
self.assertIsNotNone(triggered)
|
||||
self.assertEqual(triggered.rule.stock_code, "600519")
|
||||
to_thread.assert_awaited_once()
|
||||
|
||||
async def test_check_volume_safe_when_fetch_returns_none(self):
|
||||
"""_check_volume must not crash when get_daily_data returns None."""
|
||||
from src.agent.events import EventMonitor, VolumeAlert
|
||||
|
||||
monitor = EventMonitor()
|
||||
rule = VolumeAlert(stock_code="600519", multiplier=2.0)
|
||||
|
||||
with patch("src.agent.events.asyncio.to_thread", new=AsyncMock(return_value=None)):
|
||||
result = await monitor._check_volume(rule)
|
||||
|
||||
self.assertIsNone(result)
|
||||
|
||||
async def test_check_all_async_callback(self):
|
||||
"""on_trigger callbacks should be properly awaited if coroutine."""
|
||||
from src.agent.events import EventMonitor, PriceAlert
|
||||
|
||||
monitor = EventMonitor()
|
||||
rule = PriceAlert(stock_code="600519", direction="above", price=1800.0)
|
||||
monitor.add_alert(rule)
|
||||
|
||||
callback_values = []
|
||||
async_cb = AsyncMock(side_effect=lambda alert: callback_values.append(alert.rule.stock_code))
|
||||
monitor.on_trigger(async_cb)
|
||||
|
||||
quote = SimpleNamespace(price=1810.0)
|
||||
with patch("src.agent.events.asyncio.to_thread", new=AsyncMock(return_value=quote)):
|
||||
triggered = await monitor.check_all()
|
||||
|
||||
self.assertEqual(len(triggered), 1)
|
||||
async_cb.assert_awaited_once()
|
||||
|
||||
|
||||
class TestEventMonitorConfigIntegration(unittest.TestCase):
|
||||
"""Test config-driven EventMonitor construction."""
|
||||
|
||||
def test_build_event_monitor_from_config(self):
|
||||
from src.agent.events import build_event_monitor_from_config
|
||||
|
||||
config = SimpleNamespace(
|
||||
agent_event_monitor_enabled=True,
|
||||
agent_event_alert_rules_json='[{"stock_code":"600519","alert_type":"price_cross","direction":"above","price":1800}]',
|
||||
)
|
||||
|
||||
with patch("src.notification.NotificationService", return_value=MagicMock()):
|
||||
monitor = build_event_monitor_from_config(config=config)
|
||||
|
||||
self.assertIsNotNone(monitor)
|
||||
self.assertEqual(len(monitor.rules), 1)
|
||||
self.assertEqual(monitor.rules[0].stock_code, "600519")
|
||||
|
||||
def test_build_event_monitor_returns_none_on_invalid_json(self):
|
||||
from src.agent.events import build_event_monitor_from_config
|
||||
|
||||
config = SimpleNamespace(
|
||||
agent_event_monitor_enabled=True,
|
||||
agent_event_alert_rules_json='[invalid',
|
||||
)
|
||||
|
||||
monitor = build_event_monitor_from_config(config=config)
|
||||
self.assertIsNone(monitor)
|
||||
|
||||
def test_build_event_monitor_skips_invalid_rule_entries(self):
|
||||
from src.agent.events import build_event_monitor_from_config
|
||||
|
||||
config = SimpleNamespace(
|
||||
agent_event_monitor_enabled=True,
|
||||
agent_event_alert_rules_json=(
|
||||
'[{"stock_code":"600519","alert_type":"price_cross","direction":"above","price":1800},'
|
||||
'{"stock_code":"000858","alert_type":"price_cross","status":"bad","direction":"above","price":120}]'
|
||||
),
|
||||
)
|
||||
|
||||
with patch("src.notification.NotificationService", return_value=MagicMock()):
|
||||
monitor = build_event_monitor_from_config(config=config)
|
||||
|
||||
self.assertIsNotNone(monitor)
|
||||
self.assertEqual(len(monitor.rules), 1)
|
||||
self.assertEqual(monitor.rules[0].stock_code, "600519")
|
||||
|
||||
def test_build_event_monitor_skips_unsupported_rule_types(self):
|
||||
from src.agent.events import build_event_monitor_from_config
|
||||
|
||||
config = SimpleNamespace(
|
||||
agent_event_monitor_enabled=True,
|
||||
agent_event_alert_rules_json=(
|
||||
'[{"stock_code":"600519","alert_type":"sentiment_shift"},'
|
||||
'{"stock_code":"000858","alert_type":"price_cross","direction":"above","price":120}]'
|
||||
),
|
||||
)
|
||||
|
||||
with patch("src.notification.NotificationService", return_value=MagicMock()):
|
||||
monitor = build_event_monitor_from_config(config=config)
|
||||
|
||||
self.assertIsNotNone(monitor)
|
||||
self.assertEqual(len(monitor.rules), 1)
|
||||
self.assertEqual(monitor.rules[0].stock_code, "000858")
|
||||
|
||||
|
||||
# ============================================================
|
||||
# AgentMemory
|
||||
# ============================================================
|
||||
@@ -1132,5 +1288,135 @@ class TestRiskOverride(unittest.TestCase):
|
||||
self.assertEqual(dashboard["decision_type"], "buy")
|
||||
|
||||
|
||||
# ============================================================
|
||||
# ResearchCommand timeout guard
|
||||
# ============================================================
|
||||
|
||||
class TestResearchCommandTimeout(unittest.TestCase):
|
||||
"""Verify that ResearchCommand respects the configured timeout."""
|
||||
|
||||
def test_research_timeout_returns_timeout_response(self):
|
||||
"""When research takes longer than the timeout, a timeout message is returned."""
|
||||
import time as _time
|
||||
from bot.commands.research import ResearchCommand
|
||||
from bot.models import BotMessage
|
||||
|
||||
cmd = ResearchCommand()
|
||||
|
||||
slow_result = SimpleNamespace(
|
||||
success=True, report="ok", sub_questions=["q"], findings_count=1,
|
||||
total_tokens=100, duration_s=1.0, error=None,
|
||||
)
|
||||
|
||||
def _slow_research(query, context=None):
|
||||
_time.sleep(0.05)
|
||||
return slow_result
|
||||
|
||||
msg = MagicMock(spec=BotMessage)
|
||||
msg.platform = "test"
|
||||
msg.user_id = "u1"
|
||||
|
||||
config = SimpleNamespace(
|
||||
agent_deep_research_budget=30000,
|
||||
agent_deep_research_timeout=0.01, # 10ms — will trigger timeout
|
||||
litellm_model="test-model",
|
||||
agent_mode=True,
|
||||
)
|
||||
config.is_agent_available = lambda: True
|
||||
|
||||
with patch("bot.commands.research.get_config", return_value=config), \
|
||||
patch("src.agent.factory.get_tool_registry", return_value=MagicMock()), \
|
||||
patch("src.agent.llm_adapter.LLMToolAdapter", return_value=MagicMock()), \
|
||||
patch("src.agent.research.ResearchAgent.research", side_effect=_slow_research):
|
||||
response = cmd.execute(msg, ["600519"])
|
||||
|
||||
self.assertIn("超时", response.text)
|
||||
|
||||
|
||||
# ============================================================
|
||||
# ResearchAgent filtered registry & API endpoint
|
||||
# ============================================================
|
||||
|
||||
class TestResearchAgentFilteredRegistry(unittest.TestCase):
|
||||
"""Test that ResearchAgent._filtered_registry delegates to BaseAgent's implementation."""
|
||||
|
||||
def test_filtered_registry_delegates_to_base(self):
|
||||
from src.agent.research import ResearchAgent
|
||||
from src.agent.tools.registry import ToolRegistry
|
||||
|
||||
registry = ToolRegistry()
|
||||
fake_tool = MagicMock()
|
||||
fake_tool.name = "search_stock_news"
|
||||
registry.register(fake_tool)
|
||||
|
||||
llm_adapter = MagicMock()
|
||||
agent = ResearchAgent(tool_registry=registry, llm_adapter=llm_adapter)
|
||||
|
||||
filtered = agent._filtered_registry()
|
||||
self.assertIsInstance(filtered, ToolRegistry)
|
||||
self.assertIsNotNone(filtered.get("search_stock_news"))
|
||||
|
||||
def test_decompose_query_uses_shared_adapter(self):
|
||||
from src.agent.research import ResearchAgent
|
||||
|
||||
llm_adapter = MagicMock()
|
||||
llm_adapter.call_text.return_value = SimpleNamespace(
|
||||
provider="gemini",
|
||||
content='{"questions":["Q1","Q2"]}',
|
||||
usage={"total_tokens": 42},
|
||||
)
|
||||
agent = ResearchAgent(tool_registry=MagicMock(), llm_adapter=llm_adapter)
|
||||
|
||||
result = agent._decompose_query("分析 600519", {"stock_code": "600519"})
|
||||
|
||||
self.assertEqual(result["questions"], ["Q1", "Q2"])
|
||||
llm_adapter.call_text.assert_called_once()
|
||||
|
||||
def test_synthesise_report_uses_shared_adapter(self):
|
||||
from src.agent.research import ResearchAgent
|
||||
|
||||
llm_adapter = MagicMock()
|
||||
llm_adapter.call_text.return_value = SimpleNamespace(
|
||||
provider="gemini",
|
||||
content="Final research report",
|
||||
usage={"total_tokens": 88},
|
||||
)
|
||||
agent = ResearchAgent(tool_registry=MagicMock(), llm_adapter=llm_adapter)
|
||||
|
||||
result = agent._synthesise_report(
|
||||
"分析 600519",
|
||||
[{"question": "Q1", "content": "A1"}],
|
||||
{"stock_code": "600519"},
|
||||
)
|
||||
|
||||
self.assertEqual(result["content"], "Final research report")
|
||||
llm_adapter.call_text.assert_called_once()
|
||||
|
||||
|
||||
class TestAgentResearchEndpoint(unittest.IsolatedAsyncioTestCase):
|
||||
async def test_agent_research_returns_timeout_response(self):
|
||||
from api.v1.endpoints.agent import ResearchRequest, agent_research
|
||||
from concurrent.futures import TimeoutError as FuturesTimeoutError
|
||||
|
||||
config = SimpleNamespace(
|
||||
litellm_model="gemini/test-model",
|
||||
agent_deep_research_budget=30000,
|
||||
agent_deep_research_timeout=1,
|
||||
is_agent_available=lambda: True,
|
||||
)
|
||||
fake_loop = SimpleNamespace(
|
||||
run_in_executor=AsyncMock(side_effect=FuturesTimeoutError),
|
||||
)
|
||||
|
||||
with patch("api.v1.endpoints.agent.get_config", return_value=config), \
|
||||
patch("api.v1.endpoints.agent.asyncio.get_running_loop", return_value=fake_loop), \
|
||||
patch("src.agent.factory.get_tool_registry", return_value=MagicMock()), \
|
||||
patch("src.agent.llm_adapter.LLMToolAdapter", return_value=MagicMock()):
|
||||
response = await agent_research(ResearchRequest(question="600519 风险"))
|
||||
|
||||
self.assertFalse(response.success)
|
||||
self.assertIn("timed out", response.error)
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
unittest.main()
|
||||
|
||||
105
tests/test_scheduler_background.py
Normal file
105
tests/test_scheduler_background.py
Normal file
@@ -0,0 +1,105 @@
|
||||
# -*- coding: utf-8 -*-
|
||||
"""Tests for Scheduler background task support."""
|
||||
|
||||
import sys
|
||||
import unittest
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
|
||||
class _FakeJob:
|
||||
def __init__(self):
|
||||
self.next_run = None
|
||||
|
||||
@property
|
||||
def day(self):
|
||||
return self
|
||||
|
||||
def at(self, _value):
|
||||
return self
|
||||
|
||||
def do(self, _fn):
|
||||
return self
|
||||
|
||||
|
||||
class _FakeScheduleModule:
|
||||
def every(self):
|
||||
return _FakeJob()
|
||||
|
||||
def get_jobs(self):
|
||||
return []
|
||||
|
||||
def run_pending(self):
|
||||
return None
|
||||
|
||||
|
||||
class SchedulerBackgroundTaskTestCase(unittest.TestCase):
|
||||
def test_background_task_runs_when_interval_elapsed(self):
|
||||
fake_schedule = _FakeScheduleModule()
|
||||
with patch.dict(sys.modules, {"schedule": fake_schedule}):
|
||||
from src.scheduler import Scheduler
|
||||
|
||||
scheduler = Scheduler(schedule_time="18:00")
|
||||
calls = []
|
||||
fake_thread = MagicMock()
|
||||
fake_thread.is_alive.return_value = False
|
||||
|
||||
def _make_thread(target=None, **kwargs):
|
||||
fake_thread.start.side_effect = target
|
||||
return fake_thread
|
||||
|
||||
with patch("src.scheduler.threading.Thread", side_effect=_make_thread):
|
||||
scheduler.add_background_task(lambda: calls.append("ran"), interval_seconds=1, run_immediately=True, name="test")
|
||||
|
||||
self.assertEqual(calls, ["ran"])
|
||||
|
||||
def test_background_task_waits_for_interval(self):
|
||||
fake_schedule = _FakeScheduleModule()
|
||||
with patch.dict(sys.modules, {"schedule": fake_schedule}):
|
||||
from src.scheduler import Scheduler
|
||||
|
||||
scheduler = Scheduler(schedule_time="18:00")
|
||||
calls = []
|
||||
scheduler.add_background_task(lambda: calls.append("ran"), interval_seconds=60, run_immediately=False, name="test")
|
||||
|
||||
with patch("src.scheduler.time.time", return_value=scheduler._background_tasks[0]["last_run"] + 10):
|
||||
scheduler._run_background_tasks()
|
||||
|
||||
self.assertEqual(calls, [])
|
||||
|
||||
def test_run_with_schedule_registers_background_tasks_before_immediate_daily_task(self):
|
||||
fake_schedule = _FakeScheduleModule()
|
||||
with patch.dict(sys.modules, {"schedule": fake_schedule}):
|
||||
from src import scheduler as scheduler_module
|
||||
|
||||
order = []
|
||||
|
||||
class FakeScheduler:
|
||||
def __init__(self, schedule_time="18:00"):
|
||||
order.append(("init", schedule_time))
|
||||
|
||||
def add_background_task(self, **kwargs):
|
||||
order.append(("background", kwargs["name"]))
|
||||
|
||||
def set_daily_task(self, task, run_immediately=True):
|
||||
order.append(("daily", run_immediately))
|
||||
|
||||
def run(self):
|
||||
order.append(("run", None))
|
||||
|
||||
with patch.object(scheduler_module, "Scheduler", FakeScheduler):
|
||||
scheduler_module.run_with_schedule(
|
||||
task=lambda: None,
|
||||
run_immediately=True,
|
||||
background_tasks=[{
|
||||
"task": lambda: None,
|
||||
"interval_seconds": 60,
|
||||
"run_immediately": True,
|
||||
"name": "event_monitor",
|
||||
}],
|
||||
)
|
||||
|
||||
self.assertEqual(order[1:3], [("background", "event_monitor"), ("daily", True)])
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -144,6 +144,18 @@ class SystemConfigServiceTestCase(unittest.TestCase):
|
||||
self.assertFalse(validation["valid"])
|
||||
self.assertTrue(any(issue["code"] == "invalid_enum" for issue in validation["issues"]))
|
||||
|
||||
def test_validate_reports_invalid_json(self) -> None:
|
||||
validation = self.service.validate(items=[{"key": "AGENT_EVENT_ALERT_RULES_JSON", "value": "[invalid"}])
|
||||
|
||||
self.assertFalse(validation["valid"])
|
||||
self.assertTrue(any(issue["code"] == "invalid_json" for issue in validation["issues"]))
|
||||
|
||||
def test_validate_accepts_blank_optional_json(self) -> None:
|
||||
validation = self.service.validate(items=[{"key": "AGENT_EVENT_ALERT_RULES_JSON", "value": ""}])
|
||||
|
||||
self.assertTrue(validation["valid"])
|
||||
self.assertEqual(validation["issues"], [])
|
||||
|
||||
@patch.object(
|
||||
Config,
|
||||
"_parse_litellm_yaml",
|
||||
@@ -228,6 +240,24 @@ class SystemConfigServiceTestCase(unittest.TestCase):
|
||||
self.assertEqual(payload["resolved_protocol"], "openai")
|
||||
self.assertEqual(payload["resolved_model"], "openai/deepseek-chat")
|
||||
|
||||
def test_validate_reports_invalid_event_rule_semantics(self) -> None:
|
||||
validation = self.service.validate(items=[{
|
||||
"key": "AGENT_EVENT_ALERT_RULES_JSON",
|
||||
"value": '[{"stock_code":"600519","alert_type":"price_cross","status":"bad","direction":"above","price":1800}]',
|
||||
}])
|
||||
|
||||
self.assertFalse(validation["valid"])
|
||||
self.assertTrue(any(issue["code"] == "invalid_event_rule" for issue in validation["issues"]))
|
||||
|
||||
def test_validate_rejects_unsupported_event_rule_type(self) -> None:
|
||||
validation = self.service.validate(items=[{
|
||||
"key": "AGENT_EVENT_ALERT_RULES_JSON",
|
||||
"value": '[{"stock_code":"600519","alert_type":"sentiment_shift"}]',
|
||||
}])
|
||||
|
||||
self.assertFalse(validation["valid"])
|
||||
self.assertTrue(any(issue["code"] == "invalid_event_rule" for issue in validation["issues"]))
|
||||
|
||||
@patch("src.search_service.reset_search_service")
|
||||
def test_update_with_reload_resets_search_service_singleton(self, mock_reset_search_service) -> None:
|
||||
response = self.service.update(
|
||||
|
||||
Reference in New Issue
Block a user