251 lines
8.9 KiB
Python
251 lines
8.9 KiB
Python
"""MCP 工具实现。
|
|
|
|
每个工具:接收自然语言查询 → DashScope 嵌入 → Qdrant 检索 → 格式化返回。
|
|
嵌入 provider 复用 M5,检索复用 M6。
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
from dataclasses import dataclass
|
|
from typing import Any
|
|
|
|
from dotenv import load_dotenv
|
|
from loguru import logger
|
|
from mcp.server.fastmcp import FastMCP
|
|
|
|
from embedding import make_sync_provider
|
|
from vectorstore import SearchFilter, VectorStore, make_qdrant_client
|
|
|
|
# 加载 .env(API key 等)
|
|
load_dotenv()
|
|
|
|
# --------------------------------------------------------------------------- #
|
|
# 单例(模块加载时初始化,所有工具共用)
|
|
# --------------------------------------------------------------------------- #
|
|
|
|
@dataclass
|
|
class _Backend:
|
|
embedder: Any # EmbeddingProvider
|
|
vector_store: VectorStore
|
|
|
|
_backend: _Backend | None = None
|
|
|
|
|
|
def _get_backend() -> _Backend:
|
|
global _backend
|
|
if _backend is None:
|
|
emb = make_sync_provider() # 读取 EMBEDDING_PROVIDER 环境变量
|
|
logger.info("MCP embedder 就绪: dim={}", emb.dim)
|
|
client = make_qdrant_client()
|
|
store = VectorStore(client)
|
|
logger.info("MCP vector_store 就绪: count={}", store.count())
|
|
_backend = _Backend(embedder=emb, vector_store=store)
|
|
return _backend
|
|
|
|
|
|
# --------------------------------------------------------------------------- #
|
|
# 嵌入 + 检索
|
|
# --------------------------------------------------------------------------- #
|
|
|
|
def _search(
|
|
query: str,
|
|
top_k: int = 10,
|
|
filter: SearchFilter | None = None,
|
|
) -> list[dict[str, Any]]:
|
|
"""嵌入查询文本 → Qdrant 检索 → 返回 dict 列表。"""
|
|
be = _get_backend()
|
|
vec = be.embedder.embed_one(query)
|
|
results = be.vector_store.query(
|
|
query_vector=vec, top_k=top_k, filter=filter, score_threshold=0.3,
|
|
)
|
|
return [
|
|
{
|
|
"title": r.title,
|
|
"url": r.url,
|
|
"source": r.source_id,
|
|
"score": round(r.score, 4),
|
|
"publish_time": r.publish_time.isoformat() if r.publish_time else None,
|
|
"event": {
|
|
"sentiment": (r.event or {}).get("sentiment"),
|
|
"importance": (r.event or {}).get("importance"),
|
|
"event_type": (r.event or {}).get("event_type"),
|
|
"stock_codes": (r.event or {}).get("stock_codes", []),
|
|
"company_names": (r.event or {}).get("company_names", []),
|
|
"industries": (r.event or {}).get("industries", []),
|
|
"summary": (r.event or {}).get("summary"),
|
|
},
|
|
}
|
|
for r in results
|
|
]
|
|
|
|
|
|
# --------------------------------------------------------------------------- #
|
|
# MCP 服务器 & 工具
|
|
# --------------------------------------------------------------------------- #
|
|
|
|
mcp = FastMCP(
|
|
name="A股DeepResearch",
|
|
instructions="A 股 Deep Research 知识库——语义检索财经新闻与投资事件。",
|
|
)
|
|
|
|
|
|
def _fmt_results(hits: list[dict[str, Any]], query: str) -> str:
|
|
"""把检索结果格式化为 Markdown 文本。"""
|
|
if not hits:
|
|
return f"未找到与「{query}」相关的结果。"
|
|
lines = [f"# 检索结果: {query}", "", f"共 {len(hits)} 条:", ""]
|
|
for i, h in enumerate(hits, 1):
|
|
ev = h["event"]
|
|
sentiment = {"positive": "🟢利好", "neutral": "⚪中性", "negative": "🔴利空"}.get(
|
|
ev.get("sentiment"), ""
|
|
)
|
|
lines.append(f"### {i}. {h['title']}")
|
|
lines.append(f"- 来源: {h['source']} | 相似度: {h['score']} | {sentiment}")
|
|
lines.append(f"- 时间: {h['publish_time'] or '未知'}")
|
|
if ev.get("company_names"):
|
|
lines.append(f"- 公司: {', '.join(ev['company_names'][:5])}")
|
|
if ev.get("stock_codes"):
|
|
lines.append(f"- 代码: {', '.join(ev['stock_codes'][:5])}")
|
|
if ev.get("industries"):
|
|
lines.append(f"- 行业: {', '.join(ev['industries'][:3])}")
|
|
if ev.get("summary"):
|
|
lines.append(f"- 摘要: {ev['summary']}")
|
|
lines.append(f"- URL: {h['url']}")
|
|
lines.append("")
|
|
return "\n".join(lines)
|
|
|
|
|
|
# --------------------------------------------------------------------------- #
|
|
# Tool 1: search_news
|
|
# --------------------------------------------------------------------------- #
|
|
|
|
@mcp.tool()
|
|
def search_news(query: str, top_k: int = 10) -> str:
|
|
"""语义检索 A 股财经新闻知识库。
|
|
|
|
参数:
|
|
query: 自然语言查询(如 "宁德时代最新动态" "AI 行业政策")
|
|
top_k: 返回条数(默认 10)
|
|
|
|
返回: Markdown 格式的检索结果,含标题、来源、相似度、URL。
|
|
"""
|
|
logger.info("search_news query={!r} top_k={}", query, top_k)
|
|
hits = _search(query, top_k=top_k)
|
|
return _fmt_results(hits, query)
|
|
|
|
|
|
# --------------------------------------------------------------------------- #
|
|
# Tool 2: search_company_news
|
|
# --------------------------------------------------------------------------- #
|
|
|
|
@mcp.tool()
|
|
def search_company_news(query: str, company: str, top_k: int = 10) -> str:
|
|
"""检索指定公司的相关新闻。
|
|
|
|
参数:
|
|
query: 自然语言查询
|
|
company: 公司名称(如 "宁德时代" "贵州茅台")
|
|
top_k: 返回条数(默认 10)
|
|
|
|
返回: Markdown 格式的检索结果。
|
|
"""
|
|
logger.info("search_company_news query={!r} company={!r}", query, company)
|
|
hits = _search(
|
|
query, top_k=top_k,
|
|
filter=SearchFilter(company_names=[company]),
|
|
)
|
|
if not hits:
|
|
# 降级:无精确命中时做纯语义搜索
|
|
logger.info("company 精确命中 0 条,降级为纯语义搜索")
|
|
hits = _search(query, top_k=top_k)
|
|
return _fmt_results(hits, f"{query} (公司:{company})")
|
|
|
|
|
|
# --------------------------------------------------------------------------- #
|
|
# Tool 3: search_industry_news
|
|
# --------------------------------------------------------------------------- #
|
|
|
|
@mcp.tool()
|
|
def search_industry_news(query: str, industry: str, top_k: int = 10) -> str:
|
|
"""检索指定行业的相关新闻。
|
|
|
|
参数:
|
|
query: 自然语言查询
|
|
industry: 行业名(如 "动力电池" "白酒" "半导体")
|
|
top_k: 返回条数(默认 10)
|
|
|
|
返回: Markdown 格式的检索结果。
|
|
"""
|
|
logger.info("search_industry_news query={!r} industry={!r}", query, industry)
|
|
hits = _search(
|
|
query, top_k=top_k,
|
|
filter=SearchFilter(industries=[industry]),
|
|
)
|
|
if not hits:
|
|
logger.info("industry 精确命中 0 条,降级为纯语义搜索")
|
|
hits = _search(query, top_k=top_k)
|
|
return _fmt_results(hits, f"{query} (行业:{industry})")
|
|
|
|
|
|
# --------------------------------------------------------------------------- #
|
|
# Tool 4: search_stock_events
|
|
# --------------------------------------------------------------------------- #
|
|
|
|
@mcp.tool()
|
|
def search_stock_events(query: str, stock_code: str, top_k: int = 10) -> str:
|
|
"""检索指定股票代码相关的投资事件。
|
|
|
|
参数:
|
|
query: 自然语言查询
|
|
stock_code: 6 位 A 股代码(如 "300750" "000001",可带 .SH/.SZ 后缀)
|
|
top_k: 返回条数(默认 10)
|
|
|
|
返回: Markdown 格式的检索结果。
|
|
"""
|
|
# 标准化:去掉后缀,补一致性
|
|
code = stock_code.strip().split(".")[0].upper()
|
|
logger.info("search_stock_events query={!r} code={!r}", query, code)
|
|
hits = _search(
|
|
query, top_k=top_k,
|
|
filter=SearchFilter(stock_codes=[code]),
|
|
)
|
|
if not hits:
|
|
logger.info("stock_code 精确命中 0 条,降级为纯语义搜索")
|
|
hits = _search(query, top_k=top_k)
|
|
return _fmt_results(hits, f"{query} (代码:{code})")
|
|
|
|
|
|
# --------------------------------------------------------------------------- #
|
|
# Tool 5: search_sentiment_trend
|
|
# --------------------------------------------------------------------------- #
|
|
|
|
@mcp.tool()
|
|
def search_sentiment_trend(query: str, sentiment: str = "all", top_k: int = 20) -> str:
|
|
"""检索并统计特定情绪倾向的新闻。
|
|
|
|
参数:
|
|
query: 自然语言查询
|
|
sentiment: positive(利好) / negative(利空) / neutral(中性) / all(全部,默认)
|
|
top_k: 返回条数(默认 20)
|
|
|
|
返回: Markdown 格式的检索结果 + 情绪分布统计。
|
|
"""
|
|
logger.info("search_sentiment_trend query={!r} sentiment={!r}", query, sentiment)
|
|
filt = None
|
|
if sentiment in ("positive", "negative", "neutral"):
|
|
filt = SearchFilter(sentiment=sentiment)
|
|
|
|
hits = _search(query, top_k=top_k, filter=filt)
|
|
|
|
# 统计情绪分布
|
|
pos = sum(1 for h in hits if h["event"].get("sentiment") == "positive")
|
|
neg = sum(1 for h in hits if h["event"].get("sentiment") == "negative")
|
|
neu = sum(1 for h in hits if h["event"].get("sentiment") == "neutral")
|
|
|
|
header = (
|
|
f"# 情绪趋势: {query}\n\n"
|
|
f"共 {len(hits)} 条 | "
|
|
f"🟢利好 {pos} | 🔴利空 {neg} | ⚪中性 {neu}\n"
|
|
)
|
|
return header + "\n" + _fmt_results(hits, query)
|