402 lines
13 KiB
Python
402 lines
13 KiB
Python
#!/usr/bin/env python3
|
||
"""
|
||
Funstat BOT MCP Server
|
||
|
||
基于 Telethon 的 MCP 服务器，用于与 @openaiw_bot 交互
|
||
提供搜索、查询、统计等功能
|
||
"""
|
||
|
||
import asyncio
|
||
import json
|
||
import logging
|
||
import os
|
||
import time
|
||
from typing import Any, Dict, List, Optional
|
||
from datetime import datetime, timedelta
|
||
from collections import deque
|
||
|
||
from mcp.server import Server
|
||
from mcp.types import (
|
||
Resource,
|
||
Tool,
|
||
TextContent,
|
||
ImageContent,
|
||
EmbeddedResource,
|
||
)
|
||
from pydantic import AnyUrl
|
||
from telethon import TelegramClient
|
||
from telethon.tl.types import Message
|
||
|
||
# 配置日志
|
||
logging.basicConfig(
|
||
level=logging.INFO,
|
||
format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
|
||
)
|
||
logger = logging.getLogger("funstat_mcp")
|
||
|
||
# 配置
|
||
API_ID = 24660516
|
||
API_HASH = "eae564578880a59c9963916ff1bbbd3a"
|
||
# Session 文件路径 - 使用独立的安全目录，防止被意外删除
|
||
SESSION_PATH = os.path.expanduser("~/telegram_sessions/funstat_bot")
|
||
BOT_USERNAME = "@openaiw_bot"
|
||
|
||
# 速率限制配置
|
||
RATE_LIMIT_PER_SECOND = 18 # 每秒最多18个请求
|
||
RATE_LIMIT_WINDOW = 1.0 # 1秒时间窗口
|
||
|
||
# 缓存配置
|
||
CACHE_TTL = 3600 # 缓存1小时
|
||
|
||
|
||
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(f"速率限制: 等待 {sleep_time:.2f} 秒")
|
||
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(f"缓存命中: {key}")
|
||
return value
|
||
else:
|
||
# 过期，删除
|
||
del self.cache[key]
|
||
return None
|
||
|
||
def set(self, key: str, value: Any):
|
||
"""设置缓存"""
|
||
self.cache[key] = (value, time.time())
|
||
logger.info(f"缓存保存: {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(f"清理了 {len(expired_keys)} 个过期缓存")
|
||
|
||
|
||
class FunstatMCPServer:
|
||
"""Funstat MCP 服务器"""
|
||
|
||
def __init__(self):
|
||
self.server = Server("funstat-mcp")
|
||
self.client: Optional[TelegramClient] = None
|
||
self.bot_entity = None
|
||
self.rate_limiter = RateLimiter(RATE_LIMIT_PER_SECOND, RATE_LIMIT_WINDOW)
|
||
self.cache = ResponseCache()
|
||
|
||
# 注册处理器
|
||
self.server.list_tools()(self.list_tools)
|
||
self.server.call_tool()(self.call_tool)
|
||
|
||
async def initialize(self):
|
||
"""初始化 Telegram 客户端"""
|
||
logger.info("初始化 Telegram 客户端...")
|
||
|
||
# 检查 session 文件
|
||
session_file = f"{SESSION_PATH}.session"
|
||
if not os.path.exists(session_file):
|
||
raise FileNotFoundError(
|
||
f"Session 文件不存在: {session_file}\n"
|
||
f"请先运行 create_session.py 创建 session 文件\n"
|
||
f"或者将现有 session 文件复制到: ~/telegram_sessions/"
|
||
)
|
||
|
||
logger.info(f"使用 Session 文件: {session_file}")
|
||
|
||
# 创建客户端
|
||
self.client = TelegramClient(SESSION_PATH, API_ID, API_HASH)
|
||
await self.client.start()
|
||
|
||
# 获取 bot 实体
|
||
logger.info(f"连接到 {BOT_USERNAME}...")
|
||
self.bot_entity = await self.client.get_entity(BOT_USERNAME)
|
||
logger.info(f"✅ 已连接到: {self.bot_entity.first_name}")
|
||
|
||
# 获取当前用户信息
|
||
me = await self.client.get_me()
|
||
logger.info(f"✅ 当前账号: @{me.username} (ID: {me.id})")
|
||
|
||
async def send_command_and_wait(
|
||
self,
|
||
command: str,
|
||
timeout: int = 10,
|
||
use_cache: bool = True
|
||
) -> str:
|
||
"""发送命令到 BOT 并等待响应"""
|
||
|
||
# 检查缓存
|
||
cache_key = f"cmd:{command}"
|
||
if use_cache:
|
||
cached = self.cache.get(cache_key)
|
||
if cached:
|
||
return cached
|
||
|
||
# 速率限制
|
||
await self.rate_limiter.acquire()
|
||
|
||
logger.info(f"📤 发送命令: {command}")
|
||
|
||
# 记录发送前的最新消息 ID
|
||
last_message_id = 0
|
||
async for message in self.client.iter_messages(self.bot_entity, limit=1):
|
||
last_message_id = message.id
|
||
break
|
||
|
||
# 发送消息
|
||
send_time = datetime.now()
|
||
await self.client.send_message(self.bot_entity, command)
|
||
|
||
# 等待响应（稍等一下让 BOT 有时间响应）
|
||
await asyncio.sleep(1.5)
|
||
|
||
# 获取新消息
|
||
start_time = time.time()
|
||
while time.time() - start_time < timeout:
|
||
# 获取最新消息
|
||
async for message in self.client.iter_messages(self.bot_entity, limit=5):
|
||
# 检查是否是新消息
|
||
if message.id > last_message_id:
|
||
# 检查是否是 BOT 的消息
|
||
if not message.out and message.text:
|
||
response_text = message.text
|
||
logger.info(f"✅ 收到响应 ({len(response_text)} 字符)")
|
||
|
||
# 保存到缓存
|
||
if use_cache:
|
||
self.cache.set(cache_key, response_text)
|
||
|
||
return response_text
|
||
|
||
# 继续等待
|
||
await asyncio.sleep(0.5)
|
||
|
||
raise TimeoutError(f"等待 BOT 响应超时 ({timeout}秒)")
|
||
|
||
async def list_tools(self) -> List[Tool]:
|
||
"""列出所有可用工具"""
|
||
return [
|
||
Tool(
|
||
name="funstat_search",
|
||
description="搜索 Telegram 群组、频道。支持关键词搜索，返回相关的群组列表",
|
||
inputSchema={
|
||
"type": "object",
|
||
"properties": {
|
||
"query": {
|
||
"type": "string",
|
||
"description": "搜索关键词，例如: 'python', '区块链', 'AI'"
|
||
}
|
||
},
|
||
"required": ["query"]
|
||
}
|
||
),
|
||
Tool(
|
||
name="funstat_topchat",
|
||
description="获取热门群组/频道列表，按成员数或活跃度排序",
|
||
inputSchema={
|
||
"type": "object",
|
||
"properties": {
|
||
"category": {
|
||
"type": "string",
|
||
"description": "分类筛选（可选），例如: 'tech', 'crypto', 'news'"
|
||
}
|
||
}
|
||
}
|
||
),
|
||
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": "用户标识: 用户名(@username)、用户ID、或手机号"
|
||
}
|
||
},
|
||
"required": ["identifier"]
|
||
}
|
||
),
|
||
Tool(
|
||
name="funstat_balance",
|
||
description="查询当前账号的积分余额和使用统计",
|
||
inputSchema={
|
||
"type": "object",
|
||
"properties": {}
|
||
}
|
||
),
|
||
Tool(
|
||
name="funstat_menu",
|
||
description="显示 funstat BOT 的主菜单和所有可用功能",
|
||
inputSchema={
|
||
"type": "object",
|
||
"properties": {}
|
||
}
|
||
),
|
||
Tool(
|
||
name="funstat_start",
|
||
description="获取 funstat BOT 的欢迎信息和使用说明",
|
||
inputSchema={
|
||
"type": "object",
|
||
"properties": {}
|
||
}
|
||
)
|
||
]
|
||
|
||
async def call_tool(self, name: str, arguments: Dict[str, Any]) -> List[TextContent]:
|
||
"""调用工具"""
|
||
logger.info(f"🔧 调用工具: {name} with {arguments}")
|
||
|
||
try:
|
||
if name == "funstat_search":
|
||
query = arguments["query"]
|
||
response = await self.send_command_and_wait(f"/search {query}")
|
||
return [TextContent(type="text", text=response)]
|
||
|
||
elif name == "funstat_topchat":
|
||
category = arguments.get("category", "")
|
||
if category:
|
||
response = await self.send_command_and_wait(f"/topchat {category}")
|
||
else:
|
||
response = await self.send_command_and_wait("/topchat")
|
||
return [TextContent(type="text", text=response)]
|
||
|
||
elif name == "funstat_text":
|
||
text = arguments["text"]
|
||
response = await self.send_command_and_wait(f"/text {text}")
|
||
return [TextContent(type="text", text=response)]
|
||
|
||
elif name == "funstat_human":
|
||
name_query = arguments["name"]
|
||
response = await self.send_command_and_wait(f"/human {name_query}")
|
||
return [TextContent(type="text", text=response)]
|
||
|
||
elif name == "funstat_user_info":
|
||
identifier = arguments["identifier"]
|
||
# 直接发送用户标识（用户名、ID等）
|
||
response = await self.send_command_and_wait(identifier)
|
||
return [TextContent(type="text", text=response)]
|
||
|
||
elif name == "funstat_balance":
|
||
response = await self.send_command_and_wait("/balance")
|
||
return [TextContent(type="text", text=response)]
|
||
|
||
elif name == "funstat_menu":
|
||
response = await self.send_command_and_wait("/menu")
|
||
return [TextContent(type="text", text=response)]
|
||
|
||
elif name == "funstat_start":
|
||
response = await self.send_command_and_wait("/start")
|
||
return [TextContent(type="text", text=response)]
|
||
|
||
else:
|
||
raise ValueError(f"未知工具: {name}")
|
||
|
||
except Exception as e:
|
||
logger.error(f"❌ 工具调用失败: {e}")
|
||
return [TextContent(
|
||
type="text",
|
||
text=f"❌ 错误: {str(e)}"
|
||
)]
|
||
|
||
async def run(self):
|
||
"""运行服务器"""
|
||
await self.initialize()
|
||
|
||
# 启动定期清理过期缓存的任务
|
||
async def cache_cleanup_task():
|
||
while True:
|
||
await asyncio.sleep(300) # 每5分钟清理一次
|
||
self.cache.clear_expired()
|
||
|
||
asyncio.create_task(cache_cleanup_task())
|
||
|
||
logger.info("🚀 Funstat MCP Server 已启动")
|
||
|
||
# 运行 MCP 服务器
|
||
from mcp.server.stdio import stdio_server
|
||
|
||
async with stdio_server() as (read_stream, write_stream):
|
||
await self.server.run(
|
||
read_stream,
|
||
write_stream,
|
||
self.server.create_initialization_options()
|
||
)
|
||
|
||
|
||
async def main():
|
||
"""主函数"""
|
||
server = FunstatMCPServer()
|
||
await server.run()
|
||
|
||
|
||
if __name__ == "__main__":
|
||
asyncio.run(main())
|