4c6a20e5fc
全局替换 EduBrain -> HuiBrain, edu-brain -> huibrain, edu_brain -> hui_brain, EDUBRAIN -> HUIBRAIN 涉及文件：README.md, pyproject.toml, docker-compose.yml, .env, .env.example, app/config.py, app/main.py, app/wework_bot.py, app/__init__.py, app/mcp/server.py, static/index.html, docs/IMAGE_API_GUIDE.md
298 lines
9.8 KiB
Python
298 lines
9.8 KiB
Python
"""
|
||
MCP Server 实现
|
||
基于 MCP SDK 提供知识库查询工具，供 AI Agent 调用
|
||
|
||
工具列表:
|
||
1. search_knowledge - 语义搜索知识库
|
||
2. get_page - 获取知识页面详情
|
||
3. list_courses - 列出所有课程
|
||
4. get_page_chunks - 获取页面的分块内容
|
||
"""
|
||
|
||
from __future__ import annotations
|
||
|
||
import json
|
||
import logging
|
||
|
||
from mcp.server.fastmcp import FastMCP
|
||
|
||
logger = logging.getLogger(__name__)
|
||
|
||
# ──────────────────────────── 创建 MCP Server ────────────────────────────
|
||
|
||
mcp = FastMCP(
|
||
name="HuiBrain-Knowledge",
|
||
instructions=(
|
||
"中文直播教育知识库助手。你可以使用以下工具来搜索和查询知识库中的内容：\n"
|
||
"- search_knowledge: 语义搜索知识库，支持按课程、讲师、日期过滤\n"
|
||
"- get_page: 获取指定知识页面的完整内容\n"
|
||
"- list_courses: 列出知识库中所有课程\n"
|
||
"- get_page_chunks: 获取页面的分块内容（用于精确定位）\n"
|
||
"- search_images: 根据关键字搜索已 OCR 识别的聊天截图，返回图片路径"
|
||
),
|
||
)
|
||
|
||
|
||
# ──────────────────────────── 工具定义 ────────────────────────────
|
||
|
||
@mcp.tool()
|
||
async def search_knowledge(
|
||
query: str,
|
||
top_k: int = 5,
|
||
course_name: str = "",
|
||
teacher_name: str = "",
|
||
threshold: float = 0.5,
|
||
) -> str:
|
||
"""
|
||
语义搜索知识库。
|
||
|
||
Args:
|
||
query: 搜索查询文本（自然语言描述你要查找的知识点）
|
||
top_k: 返回结果数量，默认 5
|
||
course_name: 按课程名称过滤（可选）
|
||
teacher_name: 按讲师名称过滤（可选）
|
||
threshold: 相似度阈值（0-1），默认 0.5
|
||
|
||
Returns:
|
||
搜索结果列表，包含标题、内容摘要和相似度分数
|
||
"""
|
||
from app.database import async_session_factory
|
||
from app.schemas.search import SearchRequest
|
||
from app.services.search_service import SearchService
|
||
|
||
async with async_session_factory() as session:
|
||
try:
|
||
request = SearchRequest(
|
||
query=query,
|
||
top_k=top_k,
|
||
course_name=course_name or None,
|
||
teacher_name=teacher_name or None,
|
||
threshold=threshold,
|
||
use_fulltext=True,
|
||
)
|
||
service = SearchService(session)
|
||
response = await service.search(request)
|
||
|
||
if not response.results:
|
||
return f"未找到与「{query}」相关的知识内容。"
|
||
|
||
# 格式化搜索结果
|
||
output_parts = [
|
||
f"找到 {response.total} 条相关结果（耗时 {response.elapsed_ms:.0f}ms）：\n"
|
||
]
|
||
for i, result in enumerate(response.results, 1):
|
||
output_parts.append(
|
||
f"【{i}】{result.page_title} "
|
||
f"(相似度: {result.score:.2%})"
|
||
)
|
||
if result.course_name:
|
||
output_parts.append(f" 课程: {result.course_name}")
|
||
if result.teacher_name:
|
||
output_parts.append(f" 讲师: {result.teacher_name}")
|
||
# 截断内容避免过长
|
||
content = result.content[:200]
|
||
if len(result.content) > 200:
|
||
content += "..."
|
||
output_parts.append(f" 内容: {content}")
|
||
output_parts.append("")
|
||
|
||
return "\n".join(output_parts)
|
||
|
||
except Exception as e:
|
||
logger.error("MCP 搜索失败: %s", e, exc_info=True)
|
||
return f"搜索出错: {str(e)}"
|
||
|
||
|
||
@mcp.tool()
|
||
async def get_page(page_id: int) -> str:
|
||
"""
|
||
获取指定知识页面的完整内容。
|
||
|
||
Args:
|
||
page_id: 知识页面 ID
|
||
|
||
Returns:
|
||
页面的完整内容，包括标题、正文、元数据
|
||
"""
|
||
from app.database import async_session_factory
|
||
from app.services.page_service import PageService
|
||
|
||
async with async_session_factory() as session:
|
||
try:
|
||
service = PageService(session)
|
||
page = await service.get_page(page_id)
|
||
|
||
if page is None:
|
||
return f"未找到 ID 为 {page_id} 的知识页面。"
|
||
|
||
parts = [
|
||
f"标题: {page.title}",
|
||
]
|
||
if page.course_name:
|
||
parts.append(f"课程: {page.course_name}")
|
||
if page.teacher_name:
|
||
parts.append(f"讲师: {page.teacher_name}")
|
||
if page.live_date:
|
||
parts.append(f"直播日期: {page.live_date}")
|
||
parts.append(f"\n{page.content}")
|
||
|
||
return "\n".join(parts)
|
||
|
||
except Exception as e:
|
||
logger.error("MCP 获取页面失败: %s", e, exc_info=True)
|
||
return f"获取页面出错: {str(e)}"
|
||
|
||
|
||
@mcp.tool()
|
||
async def list_courses() -> str:
|
||
"""
|
||
列出知识库中所有课程。
|
||
|
||
Returns:
|
||
课程列表，包含课程名称、讲师、页面数量等信息
|
||
"""
|
||
from sqlalchemy import text
|
||
from app.database import async_session_factory
|
||
|
||
async with async_session_factory() as session:
|
||
try:
|
||
result = await session.execute(text("""
|
||
SELECT
|
||
course_name,
|
||
teacher_name,
|
||
COUNT(*) as page_count,
|
||
MIN(live_date) as earliest_date,
|
||
MAX(live_date) as latest_date
|
||
FROM knowledge_pages
|
||
WHERE course_name IS NOT NULL
|
||
GROUP BY course_name, teacher_name
|
||
ORDER BY page_count DESC
|
||
"""))
|
||
rows = result.fetchall()
|
||
|
||
if not rows:
|
||
return "知识库中暂无课程数据。"
|
||
|
||
parts = [f"共 {len(rows)} 个课程：\n"]
|
||
for i, row in enumerate(rows, 1):
|
||
parts.append(
|
||
f"{i}. {row.course_name}"
|
||
f"（讲师: {row.teacher_name or '未知'}，"
|
||
f"页面: {row.page_count}，"
|
||
f"日期: {row.earliest_date or '未知'} ~ {row.latest_date or '未知'}）"
|
||
)
|
||
|
||
return "\n".join(parts)
|
||
|
||
except Exception as e:
|
||
logger.error("MCP 列出课程失败: %s", e, exc_info=True)
|
||
return f"列出课程出错: {str(e)}"
|
||
|
||
|
||
@mcp.tool()
|
||
async def search_images(
|
||
keyword: str,
|
||
top_k: int = 10,
|
||
) -> str:
|
||
"""
|
||
根据关键字搜索已 OCR 识别的聊天截图。
|
||
|
||
Args:
|
||
keyword: 搜索关键字（在 OCR 识别文本中查找）
|
||
top_k: 返回结果数量，默认 10
|
||
|
||
Returns:
|
||
匹配的图片列表，包含图片路径和 OCR 文本预览
|
||
"""
|
||
from sqlalchemy import select
|
||
from app.database import async_session_factory
|
||
from app.models.base import OCRImage
|
||
|
||
async with async_session_factory() as session:
|
||
try:
|
||
like_pattern = f"%{keyword}%"
|
||
result = await session.execute(
|
||
select(OCRImage)
|
||
.where(
|
||
OCRImage.status == "completed",
|
||
OCRImage.ocr_text.ilike(like_pattern),
|
||
)
|
||
.order_by(OCRImage.created_at.desc())
|
||
.limit(top_k)
|
||
)
|
||
images = result.scalars().all()
|
||
|
||
if not images:
|
||
return f"未找到包含「{keyword}」的图片。"
|
||
|
||
parts = [f"找到 {len(images)} 张包含「{keyword}」的图片：\n"]
|
||
for i, img in enumerate(images, 1):
|
||
preview = (img.ocr_text[:100] + "...") if img.ocr_text and len(img.ocr_text) > 100 else (img.ocr_text or "")
|
||
parts.append(
|
||
f"【{i}】图片路径: {img.file_path}\n"
|
||
f" OCR 预览: {preview}\n"
|
||
f" 置信度: {img.confidence:.0%} | 识别方式: {img.provider}\n"
|
||
)
|
||
|
||
return "\n".join(parts)
|
||
|
||
except Exception as e:
|
||
logger.error("MCP 图片搜索失败: %s", e, exc_info=True)
|
||
return f"图片搜索出错: {str(e)}"
|
||
|
||
|
||
@mcp.tool()
|
||
async def get_page_chunks(page_id: int) -> str:
|
||
"""
|
||
获取知识页面的分块内容（用于精确定位知识点）。
|
||
|
||
Args:
|
||
page_id: 知识页面 ID
|
||
|
||
Returns:
|
||
该页面的所有分块内容列表
|
||
"""
|
||
from app.database import async_session_factory
|
||
from app.services.page_service import PageService
|
||
|
||
async with async_session_factory() as session:
|
||
try:
|
||
service = PageService(session)
|
||
chunks = await service.get_page_chunks(page_id)
|
||
|
||
if not chunks:
|
||
return f"页面 {page_id} 没有分块数据。"
|
||
|
||
parts = [f"页面 {page_id} 共 {len(chunks)} 个分块：\n"]
|
||
for i, chunk in enumerate(chunks, 1):
|
||
content = chunk.content[:300]
|
||
if len(chunk.content) > 300:
|
||
content += "..."
|
||
parts.append(f"【分块 {i}】(ID: {chunk.id})\n{content}\n")
|
||
|
||
return "\n".join(parts)
|
||
|
||
except Exception as e:
|
||
logger.error("MCP 获取分块失败: %s", e, exc_info=True)
|
||
return f"获取分块出错: {str(e)}"
|
||
|
||
|
||
# ──────────────────────────── 启动入口 ────────────────────────────
|
||
|
||
def run_mcp_server(transport: str = "stdio", port: int = 8001):
|
||
"""
|
||
启动 MCP Server。
|
||
|
||
Args:
|
||
transport: 传输协议，支持 "stdio"（默认）或 "sse"
|
||
port: SSE 模式下的端口号
|
||
"""
|
||
if transport == "sse":
|
||
mcp.run(transport="sse", port=port)
|
||
else:
|
||
mcp.run(transport="stdio")
|
||
|
||
|
||
if __name__ == "__main__":
|
||
run_mcp_server()
|