549 lines
20 KiB
Python
549 lines
20 KiB
Python
#!/usr/bin/env python3
|
||
"""
|
||
Funstat BOT MCP Server（PostgreSQL + 远程同步版）
|
||
"""
|
||
|
||
import asyncio
|
||
import logging
|
||
import os
|
||
import time
|
||
from collections import deque
|
||
from typing import Any, Dict, List, Optional
|
||
|
||
from mcp.server import Server
|
||
from mcp.types import Tool, TextContent
|
||
from telethon import TelegramClient
|
||
from telethon.tl.types import Message
|
||
|
||
from config import get_settings
|
||
from parsers import extract_entities
|
||
from storage import StorageManager
|
||
from models import BotResponse, PageRecord
|
||
from uploader import RemoteUploader
|
||
|
||
|
||
logging.basicConfig(
|
||
level=logging.INFO,
|
||
format="%(asctime)s - %(name)s - %(levelname)s - %(message)s",
|
||
)
|
||
logger = logging.getLogger("funstat_mcp")
|
||
|
||
settings = get_settings()
|
||
|
||
# 速率限制配置
|
||
RATE_LIMIT_PER_SECOND = 18
|
||
RATE_LIMIT_WINDOW = 1.0
|
||
CACHE_TTL = 3600
|
||
|
||
# 代理配置（可选）
|
||
PROXY_TYPE = os.getenv("FUNSTAT_PROXY_TYPE", "socks5")
|
||
PROXY_HOST = os.getenv("FUNSTAT_PROXY_HOST")
|
||
PROXY_PORT = os.getenv("FUNSTAT_PROXY_PORT")
|
||
PROXY_USERNAME = os.getenv("FUNSTAT_PROXY_USERNAME")
|
||
PROXY_PASSWORD = os.getenv("FUNSTAT_PROXY_PASSWORD")
|
||
|
||
|
||
class RateLimiter:
|
||
def __init__(self, max_requests: int, time_window: float):
|
||
self.max_requests = max_requests
|
||
self.time_window = time_window
|
||
self.requests = deque()
|
||
|
||
async def acquire(self):
|
||
now = time.time()
|
||
while self.requests and self.requests[0] < now - self.time_window:
|
||
self.requests.popleft()
|
||
|
||
if len(self.requests) >= self.max_requests:
|
||
sleep_time = self.requests[0] + self.time_window - now
|
||
if sleep_time > 0:
|
||
logger.info("速率限制: 等待 %.2f 秒", sleep_time)
|
||
await asyncio.sleep(sleep_time)
|
||
return await self.acquire()
|
||
|
||
self.requests.append(now)
|
||
|
||
|
||
class ResponseCache:
|
||
def __init__(self, ttl: int = CACHE_TTL):
|
||
self.cache: Dict[str, tuple[Any, float]] = {}
|
||
self.ttl = ttl
|
||
|
||
def get(self, key: str) -> Optional[Any]:
|
||
if key in self.cache:
|
||
value, timestamp = self.cache[key]
|
||
if time.time() - timestamp < self.ttl:
|
||
logger.info("缓存命中: %s", key)
|
||
return value
|
||
del self.cache[key]
|
||
return None
|
||
|
||
def set(self, key: str, value: Any):
|
||
self.cache[key] = (value, time.time())
|
||
logger.info("缓存保存: %s", key)
|
||
|
||
def clear_expired(self):
|
||
now = time.time()
|
||
expired_keys = [
|
||
key for key, (_, timestamp) in self.cache.items()
|
||
if now - timestamp >= self.ttl
|
||
]
|
||
for key in expired_keys:
|
||
del self.cache[key]
|
||
if expired_keys:
|
||
logger.info("清理了 %s 个过期缓存", len(expired_keys))
|
||
|
||
|
||
class FunstatMCPServer:
|
||
def __init__(self):
|
||
self.server = Server("funstat-mcp")
|
||
self.client: Optional[TelegramClient] = None
|
||
self.bot_entity = None
|
||
self.account_display: Optional[str] = None
|
||
self.rate_limiter = RateLimiter(RATE_LIMIT_PER_SECOND, RATE_LIMIT_WINDOW)
|
||
self.cache = ResponseCache()
|
||
self.storage = StorageManager(settings.database_url)
|
||
self.uploader = RemoteUploader(self.storage, settings) if settings.remote_upload_enabled else None
|
||
|
||
self.server.list_tools()(self.list_tools)
|
||
self.server.call_tool()(self.call_tool)
|
||
|
||
async def initialize(self):
|
||
logger.info("初始化 Telegram 客户端...")
|
||
session_base = os.path.expanduser(settings.telegram_session_path)
|
||
session_file = f"{session_base}.session"
|
||
if not os.path.exists(session_file):
|
||
raise FileNotFoundError(
|
||
f"Session 文件不存在: {session_file}\n"
|
||
"请先运行 create_session_safe.py 使用自己的 Telegram 账号创建 session 文件"
|
||
)
|
||
|
||
proxy = None
|
||
if PROXY_HOST and PROXY_PORT:
|
||
try:
|
||
proxy_port = int(PROXY_PORT)
|
||
except ValueError:
|
||
logger.warning("代理端口无效，忽略代理配置: %s", PROXY_PORT)
|
||
else:
|
||
if PROXY_USERNAME:
|
||
proxy = (
|
||
PROXY_TYPE,
|
||
PROXY_HOST,
|
||
proxy_port,
|
||
PROXY_USERNAME,
|
||
PROXY_PASSWORD or ""
|
||
)
|
||
else:
|
||
proxy = (PROXY_TYPE, PROXY_HOST, proxy_port)
|
||
logger.info("使用代理连接: %s://%s:%s", PROXY_TYPE, PROXY_HOST, proxy_port)
|
||
|
||
self.client = TelegramClient(
|
||
session_base,
|
||
settings.telegram_api_id,
|
||
settings.telegram_api_hash,
|
||
proxy=proxy
|
||
)
|
||
await self.client.start()
|
||
|
||
logger.info("连接到 %s ...", settings.telegram_bot_username)
|
||
self.bot_entity = await self.client.get_entity(settings.telegram_bot_username)
|
||
logger.info("✅ 已连接到: %s", getattr(self.bot_entity, "first_name", settings.telegram_bot_username))
|
||
|
||
me = await self.client.get_me()
|
||
self.account_display = f"@{me.username}" if me.username else f"ID:{me.id}"
|
||
logger.info("✅ 当前账号: %s (ID: %s)", self.account_display, me.id)
|
||
|
||
async def send_command_and_wait(
|
||
self,
|
||
command: str,
|
||
timeout: int = 12,
|
||
use_cache: bool = True,
|
||
paginate: bool = False,
|
||
max_pages: Optional[int] = None
|
||
) -> BotResponse:
|
||
paginate = paginate and settings.enable_pagination
|
||
effective_max_pages = max(1, settings.pagination_max_pages)
|
||
if max_pages:
|
||
effective_max_pages = min(effective_max_pages, max_pages)
|
||
|
||
cache_key = f"cmd:{command}:paginate={int(paginate)}:max={effective_max_pages}"
|
||
if use_cache:
|
||
cached = self.cache.get(cache_key)
|
||
if cached:
|
||
return cached
|
||
|
||
if not self.client:
|
||
raise RuntimeError("Telegram 客户端尚未初始化")
|
||
|
||
await self.rate_limiter.acquire()
|
||
logger.info("📤 发送命令: %s", command)
|
||
|
||
last_message_id = await self._get_latest_message_id()
|
||
await self.client.send_message(self.bot_entity, command)
|
||
await asyncio.sleep(1.5)
|
||
|
||
start_time = time.time()
|
||
last_seen_text = None
|
||
while time.time() - start_time < timeout:
|
||
async for message in self.client.iter_messages(self.bot_entity, limit=5):
|
||
if message.out or not message.text:
|
||
continue
|
||
|
||
is_new = message.id > last_message_id
|
||
is_updated = message.id == last_message_id and message.text != last_seen_text
|
||
|
||
if is_new or is_updated:
|
||
last_seen_text = message.text
|
||
logger.info("✅ 收到响应 (%s 字符)", len(message.text))
|
||
|
||
if paginate and message.reply_markup:
|
||
pages = await self._collect_paginated_pages(message, effective_max_pages)
|
||
response_text = self._format_pages(pages)
|
||
else:
|
||
pages = [
|
||
PageRecord(
|
||
page_number=1,
|
||
text=message.text,
|
||
entities=extract_entities(message.text)
|
||
)
|
||
]
|
||
response_text = message.text
|
||
|
||
result = BotResponse(text=response_text, pages=pages)
|
||
if use_cache:
|
||
self.cache.set(cache_key, result)
|
||
return result
|
||
|
||
await asyncio.sleep(0.5)
|
||
|
||
raise TimeoutError(f"等待 BOT 响应超时 ({timeout}秒)")
|
||
|
||
async def _get_latest_message_id(self) -> int:
|
||
assert self.client is not None
|
||
async for message in self.client.iter_messages(self.bot_entity, limit=1):
|
||
return message.id
|
||
return 0
|
||
|
||
async def _collect_paginated_pages(self, initial_message: Message, max_pages: int) -> List[PageRecord]:
|
||
pages: List[PageRecord] = [
|
||
PageRecord(
|
||
page_number=1,
|
||
text=initial_message.text or "",
|
||
entities=extract_entities(initial_message.text or "")
|
||
)
|
||
]
|
||
current_message = initial_message
|
||
current_text = initial_message.text or ""
|
||
|
||
for next_page in range(2, max_pages + 1):
|
||
button_info = self._find_next_page_button(current_message)
|
||
if not button_info:
|
||
break
|
||
|
||
logger.info("➡️ 点击翻页按钮: %s", button_info["text"])
|
||
await self.rate_limiter.acquire()
|
||
try:
|
||
await current_message.click(button_info["index"])
|
||
except Exception as exc:
|
||
logger.warning("翻页按钮点击失败: %s", exc)
|
||
break
|
||
|
||
await asyncio.sleep(settings.pagination_delay)
|
||
new_message = await self._wait_for_updated_message(current_text)
|
||
if not new_message or not new_message.text:
|
||
break
|
||
|
||
current_message = new_message
|
||
current_text = new_message.text
|
||
pages.append(
|
||
PageRecord(
|
||
page_number=next_page,
|
||
text=current_text,
|
||
entities=extract_entities(current_text)
|
||
)
|
||
)
|
||
|
||
logger.info("📚 自动翻页完成，共获取 %s 页", len(pages))
|
||
return pages
|
||
|
||
async def _wait_for_updated_message(self, previous_text: str, timeout: float = None) -> Optional[Message]:
|
||
assert self.client is not None
|
||
timeout = timeout or settings.pagination_timeout
|
||
start_time = time.time()
|
||
|
||
while time.time() - start_time < timeout:
|
||
messages = await self.client.get_messages(self.bot_entity, limit=1)
|
||
if not messages:
|
||
await asyncio.sleep(0.4)
|
||
continue
|
||
|
||
candidate = messages[0]
|
||
if candidate.out or not candidate.text:
|
||
await asyncio.sleep(0.4)
|
||
continue
|
||
|
||
if candidate.text.strip() != (previous_text or "").strip():
|
||
return candidate
|
||
|
||
await asyncio.sleep(0.4)
|
||
|
||
return None
|
||
|
||
def _find_next_page_button(self, message: Message) -> Optional[Dict[str, Any]]:
|
||
if not message.reply_markup or not hasattr(message.reply_markup, "rows"):
|
||
return None
|
||
|
||
button_index = 0
|
||
for row in message.reply_markup.rows:
|
||
for button in row.buttons:
|
||
text = getattr(button, "text", "")
|
||
if text and any(keyword in text for keyword in settings.pagination_keywords):
|
||
return {"index": button_index, "text": text}
|
||
button_index += 1
|
||
return None
|
||
|
||
@staticmethod
|
||
def _format_pages(pages: List[PageRecord]) -> str:
|
||
if not pages:
|
||
return ""
|
||
formatted = [
|
||
f"【第{page.page_number}页】\n{page.text}"
|
||
for page in pages
|
||
]
|
||
return "\n\n".join(formatted)
|
||
|
||
async def list_tools(self) -> List[Tool]:
|
||
return [
|
||
Tool(
|
||
name="funstat_search",
|
||
description="搜索 Telegram 群组、频道，支持自动翻页",
|
||
inputSchema={
|
||
"type": "object",
|
||
"properties": {
|
||
"query": {"type": "string", "description": "搜索关键词"},
|
||
"paginate": {
|
||
"type": "boolean",
|
||
"description": "是否自动翻页（默认开启）"
|
||
},
|
||
"max_pages": {
|
||
"type": "integer",
|
||
"minimum": 1,
|
||
"maximum": settings.pagination_max_pages,
|
||
"description": "最大翻页数量"
|
||
}
|
||
},
|
||
"required": ["query"]
|
||
}
|
||
),
|
||
Tool(
|
||
name="funstat_topchat",
|
||
description="获取热门群组/频道列表，按成员数或活跃度排序",
|
||
inputSchema={
|
||
"type": "object",
|
||
"properties": {
|
||
"category": {"type": "string", "description": "分类筛选（可选）"}
|
||
}
|
||
}
|
||
),
|
||
Tool(
|
||
name="funstat_text",
|
||
description="通过消息文本搜索，查找包含特定文本的消息和来源群组",
|
||
inputSchema={
|
||
"type": "object",
|
||
"properties": {
|
||
"text": {"type": "string", "description": "搜索文本"}
|
||
},
|
||
"required": ["text"]
|
||
}
|
||
),
|
||
Tool(
|
||
name="funstat_human",
|
||
description="通过姓名搜索用户，查找 Telegram 用户信息",
|
||
inputSchema={
|
||
"type": "object",
|
||
"properties": {
|
||
"name": {"type": "string", "description": "姓名或关键词"}
|
||
},
|
||
"required": ["name"]
|
||
}
|
||
),
|
||
Tool(
|
||
name="funstat_user_info",
|
||
description="查询用户详细信息，支持用户名、用户ID或手机号",
|
||
inputSchema={
|
||
"type": "object",
|
||
"properties": {
|
||
"identifier": {"type": "string", "description": "用户标识"}
|
||
},
|
||
"required": ["identifier"]
|
||
}
|
||
),
|
||
Tool(name="funstat_balance", description="查询当前账号的积分余额和使用统计"),
|
||
Tool(name="funstat_menu", description="显示 funstat BOT 的主菜单"),
|
||
Tool(name="funstat_start", description="获取 funstat BOT 的欢迎信息"),
|
||
]
|
||
|
||
async def call_tool(self, name: str, arguments: Dict[str, Any]) -> List[TextContent]:
|
||
logger.info("🔧 调用工具: %s 参数=%s", name, arguments)
|
||
try:
|
||
if name == "funstat_search":
|
||
query = arguments["query"]
|
||
paginate_flag = arguments.get("paginate", True)
|
||
max_pages = arguments.get("max_pages")
|
||
bot_command = f"/search {query}"
|
||
response = await self.send_command_and_wait(
|
||
bot_command,
|
||
paginate=paginate_flag,
|
||
max_pages=max_pages
|
||
)
|
||
await self._persist_response(
|
||
name,
|
||
bot_command,
|
||
{"query": query, "paginate": paginate_flag, "max_pages": max_pages},
|
||
response
|
||
)
|
||
return [TextContent(type="text", text=response.text)]
|
||
|
||
if name == "funstat_topchat":
|
||
category = arguments.get("category", "").strip()
|
||
bot_command = f"/topchat {category}".strip()
|
||
response = await self.send_command_and_wait(bot_command)
|
||
await self._persist_response(
|
||
name,
|
||
bot_command,
|
||
{"category": category} if category else {},
|
||
response
|
||
)
|
||
return [TextContent(type="text", text=response.text)]
|
||
|
||
if name == "funstat_text":
|
||
text = arguments["text"]
|
||
bot_command = f"/text {text}"
|
||
response = await self.send_command_and_wait(bot_command)
|
||
await self._persist_response(name, bot_command, {"text": text}, response)
|
||
return [TextContent(type="text", text=response.text)]
|
||
|
||
if name == "funstat_human":
|
||
name_query = arguments["name"]
|
||
bot_command = f"/human {name_query}"
|
||
response = await self.send_command_and_wait(bot_command)
|
||
await self._persist_response(name, bot_command, {"name": name_query}, response)
|
||
return [TextContent(type="text", text=response.text)]
|
||
|
||
if name == "funstat_user_info":
|
||
identifier = arguments["identifier"].strip()
|
||
if not identifier:
|
||
raise ValueError("用户标识不能为空")
|
||
bot_command = f"/user_info {identifier}"
|
||
response = await self.send_command_and_wait(bot_command, use_cache=False)
|
||
await self._persist_response(name, bot_command, {"identifier": identifier}, response)
|
||
return [TextContent(type="text", text=response.text)]
|
||
|
||
if name == "funstat_balance":
|
||
bot_command = "/balance"
|
||
response = await self.send_command_and_wait(bot_command, use_cache=False)
|
||
await self._persist_response(name, bot_command, {}, response)
|
||
return [TextContent(type="text", text=response.text)]
|
||
|
||
if name == "funstat_menu":
|
||
bot_command = "/menu"
|
||
response = await self.send_command_and_wait(bot_command)
|
||
await self._persist_response(name, bot_command, {}, response)
|
||
return [TextContent(type="text", text=response.text)]
|
||
|
||
if name == "funstat_start":
|
||
bot_command = "/start"
|
||
response = await self.send_command_and_wait(bot_command)
|
||
await self._persist_response(name, bot_command, {}, response)
|
||
return [TextContent(type="text", text=response.text)]
|
||
|
||
raise ValueError(f"未知工具: {name}")
|
||
|
||
except Exception as exc:
|
||
logger.error("❌ 工具调用失败: %s", exc, exc_info=exc)
|
||
return [TextContent(type="text", text=f"❌ 错误: {exc}")]
|
||
|
||
async def _persist_response(
|
||
self,
|
||
mcp_command: str,
|
||
bot_command: str,
|
||
arguments: Dict[str, Any],
|
||
response: BotResponse
|
||
):
|
||
if not response.pages:
|
||
return
|
||
inserted = await self.storage.save_response(
|
||
command=mcp_command,
|
||
bot_command=bot_command,
|
||
arguments=arguments,
|
||
pages=response.pages,
|
||
source_account=self.account_display
|
||
)
|
||
if inserted:
|
||
logger.info("💾 已写入 %s 条记录到本地数据库", len(inserted))
|
||
|
||
async def run(self):
|
||
await self.storage.initialize()
|
||
if self.uploader:
|
||
await self.uploader.start()
|
||
|
||
await self.initialize()
|
||
|
||
async def cache_cleanup_task():
|
||
while True:
|
||
await asyncio.sleep(300)
|
||
self.cache.clear_expired()
|
||
|
||
asyncio.create_task(cache_cleanup_task())
|
||
logger.info("🚀 Funstat MCP Server 已启动")
|
||
|
||
from mcp.server.streamable_http import StreamableHTTPServerTransport
|
||
from starlette.applications import Starlette
|
||
import uvicorn
|
||
import uuid
|
||
|
||
require_session = os.getenv("FUNSTAT_REQUIRE_SESSION", "false").lower() in ("1", "true", "yes")
|
||
session_id = str(uuid.uuid4()) if require_session else None
|
||
transport = StreamableHTTPServerTransport(
|
||
mcp_session_id=session_id,
|
||
is_json_response_enabled=True,
|
||
)
|
||
|
||
async def run_mcp_server():
|
||
async with transport.connect() as streams:
|
||
await self.server.run(
|
||
streams[0],
|
||
streams[1],
|
||
self.server.create_initialization_options(),
|
||
)
|
||
|
||
asyncio.create_task(run_mcp_server())
|
||
|
||
app = Starlette()
|
||
app.mount("/", transport.handle_request)
|
||
|
||
logger.info("🌐 启动 SSE 服务器: http://%s:%s", settings.host, settings.port)
|
||
logger.info("📡 SSE 端点: http://%s:%s/sse", settings.host, settings.port)
|
||
logger.info("📨 消息端点: http://%s:%s/messages", settings.host, settings.port)
|
||
|
||
config = uvicorn.Config(
|
||
app,
|
||
host=settings.host,
|
||
port=settings.port,
|
||
log_level="info"
|
||
)
|
||
server_instance = uvicorn.Server(config)
|
||
await server_instance.serve()
|
||
|
||
|
||
async def main():
|
||
server = FunstatMCPServer()
|
||
try:
|
||
await server.run()
|
||
finally:
|
||
await server.storage.close()
|
||
|
||
|
||
if __name__ == "__main__":
|
||
asyncio.run(main())
|