MangoFanFanw
16bd8e9f9a
1. 将 jwxt() 移动到 njupt_api 下，实现根据设置选择教务系统登录方式。 2. 将 api_router.py 和 mcp_router.py 中的对 ZhengFang() 的调用全部替换为对 jwxt() 的调用。
135 lines
4.7 KiB
Python
135 lines
4.7 KiB
Python
from pathlib import Path
|
||
from typing import Annotated
|
||
|
||
from fastmcp import FastMCP
|
||
from fastmcp.utilities.types import Image
|
||
from mcp.types import ToolAnnotations
|
||
from pydantic import Field
|
||
from sqlmodel import Session, select
|
||
|
||
from njupt_api.baselib import LoggingMiddleware, logger
|
||
from njupt_api.zhengfang import (
|
||
course_dict_serializer,
|
||
course_list_serializer,
|
||
jwxt,
|
||
)
|
||
from njupt_api.zhengfang.exc import LoginError
|
||
from router.enhance.lib import ReturnDto, apply_enhance
|
||
from router.enhance.model import Course, engine
|
||
|
||
mcp = FastMCP("NJUPT API Suan")
|
||
|
||
mcp.add_middleware(LoggingMiddleware())
|
||
|
||
mcp_app = mcp.http_app("/")
|
||
|
||
|
||
# 统一参数文档
|
||
USERNAME_TYPE = Annotated[str, Field(description="用户名，也即学号，一般是一位字母接八位数字，字母需要大写。")]
|
||
PASSWORD_TYPE = Annotated[str, Field(description="密码，字符串。")]
|
||
WEEK_TYPE = Annotated[int, Field(description="获取第几周的课表，默认为 0 即获取全部。")]
|
||
IMG_TYPE = Annotated[
|
||
bool,
|
||
Field(
|
||
description="是否需要同时生成图片，默认为 False。如果为 True，图片链接将在 img_url 中提供，链接有效时间为两小时。", # noqa: E501
|
||
),
|
||
]
|
||
|
||
|
||
@mcp.tool(
|
||
name="tool_schedule_class",
|
||
title="获取默认班级课表",
|
||
description="获取存储在酸 API 中的默认班级课表，返回值包含 success result message 和 img_url 四个字段。",
|
||
annotations=ToolAnnotations(
|
||
title="获取默认课表",
|
||
readOnlyHint=True,
|
||
destructiveHint=False,
|
||
idempotentHint=True,
|
||
openWorldHint=False,
|
||
),
|
||
)
|
||
async def tool_schedule_class(
|
||
week: WEEK_TYPE = 0,
|
||
img: IMG_TYPE = False,
|
||
) -> ReturnDto:
|
||
with Session(engine) as session:
|
||
course_dtos = session.exec(select(Course)).all()
|
||
logger.success("从数据库中返回一次性存储的班级课表。")
|
||
course_list: list[dict] = [course_dict_serializer(course) for course in course_dtos]
|
||
return await apply_enhance(course_list, week, img)
|
||
|
||
|
||
@mcp.tool(
|
||
name="tool_schedule_class_special",
|
||
title="获取指定学生的班级课表",
|
||
description="获取指定学生的班级课表。需要提供学号和密码。返回值包含 success result message 和 img_url 四个字段。",
|
||
annotations=ToolAnnotations(
|
||
title="获取指定学生的班级课表",
|
||
readOnlyHint=True,
|
||
destructiveHint=False,
|
||
idempotentHint=True,
|
||
openWorldHint=False,
|
||
),
|
||
)
|
||
async def tool_schedule_class_special(
|
||
username: USERNAME_TYPE,
|
||
password: PASSWORD_TYPE,
|
||
week: WEEK_TYPE = 0,
|
||
img: IMG_TYPE = False,
|
||
) -> ReturnDto:
|
||
try:
|
||
async with jwxt(username, password) as zf:
|
||
final_course_list = course_list_serializer(await zf.get_class_schedule())
|
||
logger.success(f"{username} | 获取指定学生的班级课表成功。")
|
||
return await apply_enhance(final_course_list, week, img)
|
||
except LoginError as e:
|
||
return ReturnDto(success=False, message=str(e))
|
||
|
||
|
||
@mcp.tool(
|
||
name="tool_schedule_student_special",
|
||
title="获取指定学生的个人课表",
|
||
description="获取指定学生的个人课表。需要提供学号和密码。返回值包含 success result message 和 img_url 四个字段。",
|
||
annotations=ToolAnnotations(
|
||
title="获取指定学生的个人课表",
|
||
readOnlyHint=True,
|
||
destructiveHint=False,
|
||
idempotentHint=True,
|
||
openWorldHint=False,
|
||
),
|
||
)
|
||
async def tool_schedule_student_special(
|
||
username: USERNAME_TYPE,
|
||
password: PASSWORD_TYPE,
|
||
week: WEEK_TYPE = 0,
|
||
img: IMG_TYPE = False,
|
||
) -> ReturnDto:
|
||
try:
|
||
async with jwxt(username, password) as zf:
|
||
final_course_list = course_list_serializer(await zf.get_student_schedule())
|
||
logger.success(f"{username} | 获取指定学生的个人课表成功。")
|
||
return await apply_enhance(final_course_list, week, img)
|
||
except LoginError as e:
|
||
return ReturnDto(success=False, message=str(e))
|
||
|
||
|
||
@mcp.tool(
|
||
name="tool_schedule_image",
|
||
title="直接获取课表图片",
|
||
description="接收使用其他课表工具得到的图片的文件名，返回 ImageContent。",
|
||
annotations=ToolAnnotations(
|
||
title="直接获取课表图片",
|
||
readOnlyHint=True,
|
||
destructiveHint=False,
|
||
idempotentHint=True,
|
||
openWorldHint=False,
|
||
),
|
||
)
|
||
async def tool_schedule_image(
|
||
img_name: Annotated[str, Field(description="课表图片的文件名，形如 schedule-{uuid4}.png")],
|
||
) -> Image | ReturnDto:
|
||
img_path = Path.cwd() / "temp" / img_name
|
||
if img_path.exists():
|
||
return Image(path=img_path)
|
||
return ReturnDto(success=False, message=f"未找到指定的课表图片 {img_name}")
|