train

app.py

@@ -1,25 +1,94 @@
 import gradio as gr
+from nyct_gtfs import NYCTFeed
+from datetime import datetime
+from typing import Any
 
 
-def letter_counter(word, letter):
-    """Count the occurrences of a specific letter in a word.
+# Helper function to convert datetime objects to string
+def convert_datetime_to_string(datetime_obj: Any) -> str | None:
+    if datetime_obj is None:
+        return None
+    return datetime_obj.strftime("%H:%M:%S")
+
+
+# Helper function to process train data from NYCTFeed
+def get_train_info_list(trains_data: Any) -> list[dict[str, Any]]:
+    train_info_list_data = []
+    for train in trains_data:
+        train_info = {
+            "name": str(train),
+            "line": train.route_id,
+            "direction": train.direction,
+            "stop_time_updates": [],
+        }
+        train_info["stop_time_updates"] = [
+            {"stop_name": x.stop_name, "arrival": convert_datetime_to_string(x.arrival)}
+            for x in train.stop_time_updates
+        ]
+        train_info_list_data.append(train_info)
+    return train_info_list_data
+
+
+def get_next_mta_train(
+    target_station: str, target_direction: str, feed_id: str = "1"
+) -> str:
+    """Get the next train arrival information for a given station and direction.
 
     Args:
-        word: The word or phrase to analyze
-        letter: The letter to count occurrences of
+        target_station: The name of the target station (e.g., "Times Sq-42 St", "14 St-Union Sq").
+        target_direction: The direction of the train ("N" for Northbound / Uptown, "S" for Southbound / Downtown).
+        feed_id: The GTFS feed ID for the subway lines (e.g., "1" for 1,2,3,4,5,6,S lines).
+                 Common Feed IDs:
+                 "1": 1, 2, 3, 4, 5, 6, 7, S (42 St Shuttle)
+                 "A": A, C, E, S (Rockaway Shuttle)
+                 "N": N, Q, R, W
+                 "B": B, D, F, M, S (Franklin Ave)
+                 "L": L
+                 "G": G
+                 "J": J, Z
+                 "7": 7
+                 "SIR": Staten Island Railway
 
-    Returns:
-        The number of times the letter appears in the word
+    This function returns a string with the next train arrival information for the given station and direction.
+    You can use this tool to get the next train arrival information for a given station and direction.
     """
-    return word.lower().count(letter.lower())
+    try:
+        feed = NYCTFeed(feed_id)
+        trains_data = feed.trips
+    except Exception as e:
+        return f"Failed to load MTA feed data for feed ID {feed_id}: {e}"
+
+    if not trains_data:
+        return f"No train data found for feed ID {feed_id}."
+
+    train_info_processed = get_train_info_list(trains_data)
+    current_time = datetime.now().strftime("%H:%M:%S")
+    train_info_string = f"Current time: {current_time}\n"
+    for train in train_info_processed:
+        for stop in train["stop_time_updates"]:
+            if (
+                stop["stop_name"] == target_station
+                and train["direction"] == target_direction
+            ):
+                # train_name = train["name"] # Original notebook had this, but it's often complex like "14:50 S 1 to South Ferry"
+                train_line = train["line"]
+                # train_direction = train["direction"] # Already have target_direction
+                train_arrival = stop["arrival"]
+                if train_arrival:
+                    train_info_string += f"The next {target_direction} bound {train_line} train arriving at {target_station} will arrive at {train_arrival}.\n"
+                else:
+                    train_info_string += f"The next {target_direction} bound {train_line} train is scheduled at {target_station}, but arrival time is not currently available.\n"
+    if train_info_string == "":
+        return f"No {target_direction} bound trains found for {target_station} on feed {feed_id} at this time."
+    return train_info_string
 
 
 demo = gr.Interface(
-    fn=letter_counter,
-    inputs=["text", "text"],
-    outputs="number",
-    title="Letter Counter",
-    description="Count how many times a letter appears in a word",
+    fn=get_next_mta_train,
+    inputs=["text", "text", "text"],
+    outputs="text",
+    title="MTA Subway Tracker",
+    description="Get the next train arrival information for a given station and direction",
 )
 
 demo.launch(mcp_server=True)