Coverage for src/data_api/mcp/tools/buses.py: 100%

1"""Bus-related MCP tools.""" 

2 

3from datetime import datetime 

4from typing import Annotated, Literal 

5 

6from pydantic import Field 

7 

8from data_api.domain.buses.enums import BusDayWithCurrent, BusStopsName 

9from data_api.domain.buses.query import resolve_bus_time 

10from data_api.mcp.server import mcp 

11from data_api.mcp.tools.result_types import BusResult 

12from data_api.models.buses import BusCanonicalDetailedSchedule, BusCanonicalSchedule 

13from data_api.runtime import get_runtime 

14 

15 

16async def _get_bus_schedule( 

17 route: Literal["main", "nanda", "all"] = "all", 

18 direction: Literal["up", "down", "all"] = "all", 

19 limit: int = 5, 

20 stop: BusStopsName | None = None, 

21 day: BusDayWithCurrent = BusDayWithCurrent.current, 

22 time: str | None = None, 

23 details: bool = True, 

24) -> BusResult: 

25 """Query schedules using the same stop and time filters as REST.""" 

26 service = get_runtime().buses 

27 await service.update_data() 

28 

29 bus_time = resolve_bus_time(day, time, datetime.now()) 

30 

31 buses = service.query_schedule( 

32 route_type=route, 

33 day=bus_time.day, 

34 direction=direction, 

35 detailed=details, 

36 stop=stop, 

37 after_time=bus_time.after_time, 

38 limit=limit, 

39 ) 

40 schedule_model = BusCanonicalDetailedSchedule if details else BusCanonicalSchedule 

41 result = { 

42 "current_time": bus_time.current_time, 

43 "day_type": bus_time.day, 

44 "buses": [ 

45 schedule_model.model_validate(bus).model_dump(mode="json", exclude_none=True) 

46 for bus in buses 

47 ], 

48 "route_info": service.get_route_info( 

49 route if route != "all" else None, 

50 direction if direction != "all" else None, 

51 ), 

52 } 

53 if stop is not None: 

54 result["stop_name"] = stop.value 

55 result["stop_info"] = [ 

56 info for info in service.gen_bus_stops_info() if info["name"] == stop.value 

57 ] 

58 return result 

59 

60 

61@mcp.tool( 

62 name="get_bus_schedule", 

63 title="查詢公車時刻表", 

64 description="Get campus bus schedules, optionally filtered to buses serving a specific stop. " 

65 "Replaces get_next_buses and get_bus_stops. " 

66 "With stop, time filtering uses estimated arrival time at that stop (including buses " 

67 "that have departed but have not arrived yet) and each bus entry includes top-level " 

68 "arrive_time for that stop; without stop, arrive_time is omitted and filtering uses " 

69 "departure time. " 

70 "Use details=true for departure information and arrival times at all stops; " 

71 "details=false returns simple departure schedules. " 

72 "Line identifiers are main_red, main_green, nanda_route_1, and nanda_route_2. " 

73 "day=current uses the current day and time and ignores time. " 

74 "For a selected stop, stop_info includes its location. " 

75 "If going TO Nanda Campus, query direction=up; TO Main Campus, direction=down. " 

76 "Passengers TO Nanda cannot get off at other stops inside Main Campus; " 

77 "passengers TO Main Campus cannot board at stops inside Main Campus.", 

78) 

79async def get_bus_schedule( 

80 route: Literal["main", "nanda", "all"] = "all", 

81 direction: Literal["up", "down", "all"] = "all", 

82 limit: Annotated[int, Field(ge=1)] = 5, 

83 stop: BusStopsName | None = None, 

84 day: BusDayWithCurrent = BusDayWithCurrent.current, 

85 time: ( 

86 Annotated[ 

87 str, 

88 Field( 

89 pattern=r"^([01][0-9]|2[0-3]):[0-5][0-9]$", 

90 description="Time in HH:MM: arrival at the selected stop, otherwise departure", 

91 ), 

92 ] 

93 | None 

94 ) = None, 

95 details: bool = True, 

96) -> dict: 

97 """Get bus schedules and optional stop metadata.""" 

98 return await _get_bus_schedule(route, direction, limit, stop, day, time, details)