chore: initial commit
This commit is contained in:
你的用户名
481
core/server.py
Normal file
481
core/server.py
Normal file
@@ -0,0 +1,481 @@
|
||||
#!/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 = int(os.getenv("TELEGRAM_API_ID", "24660516"))
|
||||
API_HASH = os.getenv("TELEGRAM_API_HASH", "eae564578880a59c9963916ff1bbbd3a")
|
||||
# Session 文件路径 - 使用独立的安全目录，防止被意外删除
|
||||
SESSION_PATH = os.path.expanduser(
|
||||
os.getenv("TELEGRAM_SESSION_PATH", "~/telegram_sessions/funstat_bot")
|
||||
)
|
||||
BOT_USERNAME = os.getenv("FUNSTAT_BOT_USERNAME", "@openaiw_bot")
|
||||
|
||||
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")
|
||||
|
||||
# 速率限制配置
|
||||
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}")
|
||||
|
||||
proxy = None
|
||||
if PROXY_HOST and PROXY_PORT:
|
||||
try:
|
||||
proxy_port = int(PROXY_PORT)
|
||||
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,
|
||||
)
|
||||
except ValueError:
|
||||
logger.warning(
|
||||
"代理端口无效，忽略代理配置: %s",
|
||||
PROXY_PORT,
|
||||
)
|
||||
|
||||
# 创建客户端
|
||||
self.client = TelegramClient(SESSION_PATH, API_ID, API_HASH, proxy=proxy)
|
||||
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"].strip()
|
||||
if not identifier:
|
||||
raise ValueError("用户标识不能为空")
|
||||
|
||||
# funstat BOT 需要显式的 /user_info 命令
|
||||
response = await self.send_command_and_wait(f"/user_info {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 服务器 - Streamable HTTP 模式
|
||||
from mcp.server.streamable_http import StreamableHTTPServerTransport
|
||||
from starlette.applications import Starlette
|
||||
from starlette.routing import Mount
|
||||
import uvicorn
|
||||
import uuid
|
||||
|
||||
# 是否启用会话校验
|
||||
require_session = os.getenv("FUNSTAT_REQUIRE_SESSION", "false").lower() in ("1", "true", "yes")
|
||||
|
||||
# 创建 Streamable HTTP 传输（生成唯一 session ID，默认关闭强校验以兼容旧客户端）
|
||||
session_id = str(uuid.uuid4()) if require_session else None
|
||||
transport = StreamableHTTPServerTransport(
|
||||
mcp_session_id=session_id,
|
||||
is_json_response_enabled=True, # 启用 JSON 响应
|
||||
)
|
||||
|
||||
# 在后台运行 MCP 服务器
|
||||
async def run_mcp_server():
|
||||
async with transport.connect() as streams:
|
||||
await self.server.run(
|
||||
streams[0],
|
||||
streams[1],
|
||||
self.server.create_initialization_options(),
|
||||
)
|
||||
|
||||
# 启动 MCP 服务器任务
|
||||
asyncio.create_task(run_mcp_server())
|
||||
|
||||
# 创建 Starlette 应用（transport.handle_request 是 ASGI 应用）
|
||||
app = Starlette()
|
||||
app.mount("/", transport.handle_request)
|
||||
|
||||
# 获取端口配置
|
||||
port = int(os.getenv("FUNSTAT_PORT", "8091"))
|
||||
host = os.getenv("FUNSTAT_HOST", "127.0.0.1")
|
||||
|
||||
logger.info(f"🌐 启动 SSE 服务器: http://{host}:{port}")
|
||||
logger.info(f"📡 SSE 端点: http://{host}:{port}/sse")
|
||||
logger.info(f"📨 消息端点: http://{host}:{port}/messages")
|
||||
if session_id:
|
||||
logger.info(f"🔒 Session ID: {session_id}")
|
||||
|
||||
# 启动服务器
|
||||
config = uvicorn.Config(
|
||||
app,
|
||||
host=host,
|
||||
port=port,
|
||||
log_level="info"
|
||||
)
|
||||
server_instance = uvicorn.Server(config)
|
||||
await server_instance.serve()
|
||||
|
||||
|
||||
async def main():
|
||||
"""主函数"""
|
||||
server = FunstatMCPServer()
|
||||
await server.run()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
asyncio.run(main())
|
||||
Reference in New Issue
Block a user