Merge pull request #6 from robholland/rh-trains

WIP: Add some train searches.

pyproject.toml

@@ -16,7 +16,7 @@ packages = [
 "Bug Tracker" = "https://github.com/temporalio/samples-python/issues"
 
 [tool.poetry.dependencies]
-python = ">=3.9,<4.0"
+python = ">=3.10,<4.0"
 temporalio = "^1.8.0"
 
 # Standard library modules (e.g. asyncio, collections) don't need to be added
@@ -32,6 +32,8 @@ google-generativeai = "^0.8.4"
 anthropic = "^0.45.0"
 deepseek = "^1.0.0"
 requests = "^2.32.3"
+pandas = "^2.2.3"
+gtfs-kit = "^10.1.1"
 
 [tool.poetry.group.dev.dependencies]
 pytest = "^7.3"

scripts/find_events_test.py

@@ -1,4 +1,4 @@
-from tools.find_events import find_events
+from tools.search_events import find_events
 import json
 
 # Example usage

thirdparty/train_api.py

@@ -0,0 +1,207 @@
+# Version-compatible imports
+try:
+    # Modern Python
+    from http.server import HTTPServer, BaseHTTPRequestHandler
+    from urllib.parse import parse_qs, urlparse
+    import json
+except ImportError:
+    # Python 1.5.2
+    from BaseHTTPServer import HTTPServer, BaseHTTPRequestHandler
+    from urlparse import parse_qs, urlparse
+
+import time
+import random
+import string
+
+
+def parse_datetime(datetime_str):
+    # Parse YYYY-MM-DDTHH:MM format
+    try:
+        date_part, time_part = datetime_str.split('T')
+        year, month, day = map(int, date_part.split('-'))
+        hour, minute = map(int, time_part.split(':'))
+        return year, month, day, hour, minute
+    except:
+        return None, None, None, None, None
+
+
+class TrainServer(BaseHTTPRequestHandler):
+    def format_json(self, obj):
+        # Simple JSON-like string formatter for 1.5.2 compatibility
+        try:
+            return json.dumps(obj)
+        except NameError:
+            if isinstance(obj, dict):
+                pairs = []
+                for k, v in obj.items():
+                    if isinstance(v, str):
+                        pairs.append('"%s": "%s"' % (k, v))
+                    else:
+                        pairs.append('"%s": %s' % (k, str(v)))
+                return "{" + ", ".join(pairs) + "}"
+            elif isinstance(obj, list):
+                return "[" + ", ".join([self.format_json(x) for x in obj]) + "]"
+            else:
+                return str(obj)
+
+    def write_response(self, response):
+        try:
+            # Python 3
+            self.wfile.write(response.encode('utf-8'))
+        except AttributeError:
+            # Python 1.5.2
+            self.wfile.write(response)
+
+    def generate_journeys(self, origin, destination, out_datetime, ret_datetime):
+        journeys = []
+
+        # Helper to format datetime
+        def format_datetime(year, month, day, hour, minute):
+            return "%04d-%02d-%02dT%02d:%02d" % (year, month, day, hour, minute)
+
+        # Generate outbound journeys
+        year, month, day, hour, minute = out_datetime
+        for offset in [-60, -30, 0, 30, 60]:
+            # Calculate journey times
+            adj_minutes = minute + offset
+            adj_hour = hour + (adj_minutes // 60)
+            adj_minute = adj_minutes % 60
+
+            # Simple handling of day rollover
+            adj_day = day + (adj_hour // 24)
+            adj_hour = adj_hour % 24
+
+            # Journey takes 1-2 hours
+            duration = 60 + random.randint(0, 60)
+            arr_hour = (adj_hour + (duration // 60))
+            arr_minute = (adj_minute + (duration % 60)) % 60
+            arr_day = adj_day + (arr_hour // 24)
+            arr_hour = arr_hour % 24
+
+            journey = {
+                "id": "T%d" % random.randint(1000, 9999),
+                "type": "outbound",
+                "departure": origin,
+                "arrival": destination,
+                "departure_time": format_datetime(year, month, adj_day, adj_hour, adj_minute),
+                "arrival_time": format_datetime(year, month, arr_day, arr_hour, arr_minute),
+                "platform": str(random.randint(1, 8)),
+                "price": round(30 + random.random() * 50, 2)
+            }
+            journeys.append(journey)
+
+        # Generate return journeys if return datetime provided
+        if ret_datetime[0] is not None:
+            year, month, day, hour, minute = ret_datetime
+            for offset in [-60, -30, 0, 30, 60]:
+                adj_minutes = minute + offset
+                adj_hour = hour + (adj_minutes // 60)
+                adj_minute = adj_minutes % 60
+
+                adj_day = day + (adj_hour // 24)
+                adj_hour = adj_hour % 24
+
+                duration = 60 + random.randint(0, 60)
+                arr_hour = (adj_hour + (duration // 60))
+                arr_minute = (adj_minute + (duration % 60)) % 60
+                arr_day = adj_day + (arr_hour // 24)
+                arr_hour = arr_hour % 24
+
+                journey = {
+                    "id": "T%d" % random.randint(1000, 9999),
+                    "type": "return",
+                    "departure": destination,
+                    "arrival": origin,
+                    "departure_time": format_datetime(year, month, adj_day, adj_hour, adj_minute),
+                    "arrival_time": format_datetime(year, month, arr_day, arr_hour, arr_minute),
+                    "platform": str(random.randint(1, 8)),
+                    "price": round(30 + random.random() * 50, 2)
+                }
+                journeys.append(journey)
+
+        return journeys
+
+    def do_GET(self):
+        parsed_url = urlparse(self.path)
+
+        if parsed_url.path == '/api/journeys':
+            try:
+                params = parse_qs(parsed_url.query)
+                origin = params.get('from', [''])[0]
+                destination = params.get('to', [''])[0]
+                outbound_datetime = params.get('outbound_time', [''])[0]
+                return_datetime = params.get('return_time', [''])[0]
+
+                if not origin or not destination or not outbound_datetime:
+                    self.send_response(400)
+                    self.send_header('Content-Type', 'application/json')
+                    self.end_headers()
+                    self.write_response(self.format_json({
+                        "error": "Required parameters: 'from', 'to', and 'outbound_time'"
+                    }))
+                    return
+
+                # Parse datetimes
+                out_dt = parse_datetime(outbound_datetime)
+                ret_dt = parse_datetime(return_datetime) if return_datetime else (
+                    None, None, None, None, None)
+
+                if out_dt[0] is None:
+                    self.send_response(400)
+                    self.send_header('Content-Type', 'application/json')
+                    self.end_headers()
+                    self.write_response(self.format_json({
+                        "error": "Invalid datetime format. Use YYYY-MM-DDTHH:MM"
+                    }))
+                    return
+
+                self.send_response(200)
+                self.send_header('Content-Type', 'application/json')
+                self.end_headers()
+
+                journeys = self.generate_journeys(
+                    origin, destination, out_dt, ret_dt)
+                response = self.format_json({"journeys": journeys})
+                self.write_response(response)
+
+            except Exception as e:
+                self.send_response(500)
+                self.send_header('Content-Type', 'application/json')
+                self.end_headers()
+                self.write_response(self.format_json({"error": str(e)}))
+        else:
+            self.send_response(404)
+            self.end_headers()
+
+    def do_POST(self):
+        parsed_url = urlparse(self.path)
+
+        if parsed_url.path.startswith('/api/book/'):
+            journey_id = parsed_url.path.split('/')[-1]
+            self.send_response(200)
+            self.send_header('Content-Type', 'application/json')
+            self.end_headers()
+
+            booking_ref = "BR" + \
+                "".join([random.choice(string.digits) for _ in range(5)])
+
+            response = self.format_json({
+                "booking_reference": booking_ref,
+                "journey_id": journey_id,
+                "status": "confirmed"
+            })
+            self.write_response(response)
+
+        else:
+            self.send_response(404)
+            self.end_headers()
+
+
+def run_server():
+    server = HTTPServer(('', 8080), TrainServer)
+    print("Train booking server starting on port 8080...")
+    server.serve_forever()
+
+
+if __name__ == '__main__':
+    run_server()

tools/__init__.py

@@ -1,16 +1,19 @@
-from .find_events import find_events
+from .search_fixtures import search_fixtures
-from .find_fixtures import find_fixtures
 from .search_flights import search_flights
+from .search_trains import search_trains
+from .search_trains import book_train
 from .create_invoice import create_invoice
 
 
 def get_handler(tool_name: str):
-    if tool_name == "FindEvents":
+    if tool_name == "SearchFixtures":
-        return find_events
+        return search_fixtures
-    if tool_name == "FindFixtures":
-        return find_fixtures
     if tool_name == "SearchFlights":
         return search_flights
+    if tool_name == "SearchTrains":
+        return search_trains
+    if tool_name == "BookTrain":
+        return book_train
     if tool_name == "CreateInvoice":
         return create_invoice
 

tools/data/find_events_data.json

@@ -1,252 +0,0 @@
-{
-  "Melbourne": [
-    {
-      "eventName": "Australian Open",
-      "dateFrom": "2025-01-13",
-      "dateTo": "2025-01-26",
-      "description": "A two-week Grand Slam tennis tournament featuring the world's top players, accompanied by various entertainment options including live music and family-friendly activities."
-    },
-    {
-      "eventName": "Melbourne International Comedy Festival",
-      "dateFrom": "2025-03-26",
-      "dateTo": "2025-04-20",
-      "description": "One of the world's largest comedy festivals, showcasing stand-up, cabaret, theatre, and street performances across numerous city venues."
-    },
-    {
-      "eventName": "Melbourne International Film Festival (MIFF)",
-      "dateFrom": "2025-08-07",
-      "dateTo": "2025-08-23",
-      "description": "Established in 1952, MIFF presents a diverse selection of Australian and international films, including features, documentaries, and shorts."
-    },
-    {
-      "eventName": "Melbourne Fringe Festival",
-      "dateFrom": "2025-09-17",
-      "dateTo": "2025-10-04",
-      "description": "An open-access arts festival featuring a wide array of art forms such as theatre, comedy, music, and digital art across various venues."
-    },
-    {
-      "eventName": "Moomba Festival",
-      "dateFrom": "2025-03-07",
-      "dateTo": "2025-03-10",
-      "description": "Australia's largest free community festival, celebrated over four days during the Labour Day long weekend, including a parade, live music, fireworks, and the famous Birdman Rally along the Yarra River."
-    },
-    {
-      "eventName": "White Night Melbourne",
-      "dateFrom": "2025-08-22",
-      "dateTo": "2025-08-24",
-      "description": "A dusk-to-dawn arts and cultural festival transforming the city with light installations, projections, music, and performances."
-    },
-    {
-      "eventName": "Melbourne Food and Wine Festival",
-      "dateFrom": "2025-03-19",
-      "dateTo": "2025-03-29",
-      "description": "A celebration of Victoria's culinary scene, featuring food and wine events, masterclasses, and dining experiences."
-    }
-  ],
-  "Sydney": [
-    {
-      "eventName": "Sydney Gay and Lesbian Mardi Gras",
-      "dateFrom": "2025-02-14",
-      "dateTo": "2025-03-01",
-      "description": "One of the largest LGBTQ+ festivals globally, featuring a vibrant parade, parties, and cultural events celebrating diversity and inclusion."
-    },
-    {
-      "eventName": "Vivid Sydney",
-      "dateFrom": "2025-05-22",
-      "dateTo": "2025-06-13",
-      "description": "An annual festival of light, music, and ideas, transforming the city with mesmerizing light installations and projections."
-    },
-    {
-      "eventName": "Sydney Festival",
-      "dateFrom": "2025-01-08",
-      "dateTo": "2025-01-26",
-      "description": "A major arts festival presenting a diverse program of theatre, dance, music, and visual arts across the city."
-    },
-    {
-      "eventName": "Sculpture by the Sea, Bondi",
-      "dateFrom": "2025-10-23",
-      "dateTo": "2025-11-09",
-      "description": "An outdoor sculpture exhibition along the Bondi to Tamarama coastal walk, showcasing works by Australian and international artists."
-    },
-    {
-      "eventName": "Sydney Writers' Festival",
-      "dateFrom": "2025-04-27",
-      "dateTo": "2025-05-03",
-      "description": "An annual literary festival featuring talks, panel discussions, and workshops with acclaimed authors and thinkers."
-    },
-    {
-      "eventName": "Sydney Film Festival",
-      "dateFrom": "2025-06-04",
-      "dateTo": "2025-06-15",
-      "description": "One of the longest-running film festivals in the world, showcasing a diverse selection of local and international films."
-    }
-  ],
-  "Auckland": [
-    {
-      "eventName": "Pasifika Festival",
-      "dateFrom": "2025-03-08",
-      "dateTo": "2025-03-09",
-      "description": "The largest Pacific Islands-themed festival globally, celebrating the diverse cultures of the Pacific with traditional cuisine, performances, and arts."
-    },
-    {
-      "eventName": "Auckland Arts Festival",
-      "dateFrom": "2025-03-11",
-      "dateTo": "2025-03-29",
-      "description": "A biennial multi-arts festival showcasing local and international artists in theatre, dance, music, and visual arts."
-    },
-    {
-      "eventName": "Auckland Writers Festival",
-      "dateFrom": "2025-05-13",
-      "dateTo": "2025-05-18",
-      "description": "An annual event bringing together international and local writers for discussions, readings, and workshops."
-    },
-    {
-      "eventName": "Auckland Diwali Festival",
-      "dateFrom": "2025-10-26",
-      "dateTo": "2025-10-27",
-      "description": "A vibrant celebration of Indian culture and the Hindu festival of Diwali, featuring performances, food stalls, and traditional activities."
-    }
-  ],
-  "Brisbane": [
-    {
-      "eventName": "Brisbane Festival",
-      "dateFrom": "2025-09-05",
-      "dateTo": "2025-09-26",
-      "description": "A major international arts festival featuring theatre, music, dance, and visual arts, culminating in the Riverfire fireworks display."
-    },
-    {
-      "eventName": "NRL Magic Round",
-      "dateFrom": "2025-05-02",
-      "dateTo": "2025-05-04",
-      "description": "A rugby league extravaganza where all NRL matches for the round are played at Suncorp Stadium, attracting fans nationwide."
-    },
-    {
-      "eventName": "Brisbane International Film Festival",
-      "dateFrom": "2025-10-01",
-      "dateTo": "2025-10-11",
-      "description": "Showcasing a curated selection of films from around the world, including premieres and special events."
-    },
-    {
-      "eventName": "Brisbane Comedy Festival",
-      "dateFrom": "2025-02-22",
-      "dateTo": "2025-03-24",
-      "description": "A month-long comedy festival featuring local and international comedians in stand-up, sketch, and improv performances."
-    },
-    {
-      "eventName": "Brisbane Writers Festival",
-      "dateFrom": "2025-09-05",
-      "dateTo": "2025-09-08",
-      "description": "An annual literary festival celebrating books, writing, and ideas with author talks, panel discussions, and workshops."
-    },
-    {
-      "eventName": "Brisbane Asia Pacific Film Festival",
-      "dateFrom": "2025-11-29",
-      "dateTo": "2025-12-08",
-      "description": "Showcasing the best cinema from the Asia Pacific region, including features, documentaries, and short films."
-    }
-  ],
-  "Perth": [
-    {
-      "eventName": "Perth Festival",
-      "dateFrom": "2025-02-07",
-      "dateTo": "2025-03-01",
-      "description": "Australia's longest-running cultural festival, offering a diverse program of music, theatre, dance, literature, and visual arts."
-    },
-    {
-      "eventName": "Fringe World Festival",
-      "dateFrom": "2025-01-16",
-      "dateTo": "2025-02-15",
-      "description": "One of the largest fringe festivals globally, featuring a vast array of performances including comedy, cabaret, theatre, and street arts."
-    },
-    {
-      "eventName": "Sculpture by the Sea",
-      "dateFrom": "2025-03-06",
-      "dateTo": "2025-03-23",
-      "description": "An annual outdoor sculpture exhibition along Cottesloe Beach, showcasing works from Australian and international artists."
-    },
-    {
-      "eventName": "Revelation Perth International Film Festival",
-      "dateFrom": "2025-07-03",
-      "dateTo": "2025-07-13",
-      "description": "A showcase of independent cinema, featuring a diverse selection of films, documentaries, and short films."
-    },
-    {
-      "eventName": "Perth Comedy Festival",
-      "dateFrom": "2025-04-22",
-      "dateTo": "2025-05-19",
-      "description": "A month-long comedy festival featuring local and international comedians in stand-up, sketch, and improv performances."
-    }
-  ],
-  "Adelaide": [
-    {
-      "eventName": "Adelaide Festival",
-      "dateFrom": "2025-02-28",
-      "dateTo": "2025-03-15",
-      "description": "A premier arts festival offering a rich program of theatre, music, dance, and visual arts from renowned international and local artists."
-    },
-    {
-      "eventName": "Adelaide Fringe",
-      "dateFrom": "2025-02-14",
-      "dateTo": "2025-03-15",
-      "description": "The largest open-access arts festival in the Southern Hemisphere, featuring thousands of performances across various genres and venues."
-    },
-    {
-      "eventName": "SALA Festival",
-      "dateFrom": "2025-08-01",
-      "dateTo": "2025-08-31",
-      "description": "South Australia's largest visual arts festival, showcasing the work of local artists in exhibitions, workshops, and events."
-    },
-    {
-      "eventName": "OzAsia Festival",
-      "dateFrom": "2025-09-25",
-      "dateTo": "2025-10-11",
-      "description": "A celebration of Asian arts and culture, featuring performances, exhibitions, and events from across the region."
-    },
-    {
-      "eventName": "Adelaide Film Festival",
-      "dateFrom": "2025-10-16",
-      "dateTo": "2025-10-26",
-      "description": "Showcasing a diverse selection of Australian and international films, including features, documentaries, and shorts."
-    },
-    {
-      "eventName": "Adelaide Writers' Week",
-      "dateFrom": "2025-03-01",
-      "dateTo": "2025-03-06",
-      "description": "An annual literary festival featuring talks, panel discussions, and readings by acclaimed authors and thinkers."
-    }
-  ],
-  "Wellington": [
-    {
-      "eventName": "New Zealand Festival of the Arts",
-      "dateFrom": "2025-02-21",
-      "dateTo": "2025-03-15",
-      "description": "The nation's largest celebration of contemporary arts and culture, featuring a diverse range of performances and exhibitions across various venues in Wellington.",
-      "url": "https://www.festival.nz/"
-    },
-    {
-      "eventName": "Wellington Jazz Festival",
-      "dateFrom": "2025-06-05",
-      "dateTo": "2025-06-09",
-      "description": "A five-day festival showcasing local and international jazz musicians in concerts, workshops, and community events.",
-      "url": "https://www.jazzfestival.co.nz/"
-    },
-    {
-      "eventName": "Wellington on a Plate",
-      "dateFrom": "2025-08-01",
-      "dateTo": "2025-08-16",
-      "description": "A culinary festival celebrating the city's food and beverage industry with special menus, events, and culinary experiences."
-    },
-    {
-      "eventName": "CubaDupa",
-      "dateFrom": "2025-03-28",
-      "dateTo": "2025-03-29",
-      "description": "A vibrant street festival in Wellington's Cuba Street, featuring music, dance, street performers, and food stalls."
-    },
-    {
-      "eventName": "Wellington Pasifika Festival",
-      "dateFrom": "2025-01-18",
-      "dateTo": "2025-01-19",
-      "description": "A celebration of Pacific Island culture with traditional performances, food stalls, and arts and crafts."
-    }
-  ]
-}

tools/find_events.py

@@ -1,64 +0,0 @@
-from datetime import datetime
-from pathlib import Path
-import json
-
-
-def find_events(args: dict) -> dict:
-    search_city = args.get("city", "").lower()
-    search_month = args.get("month", "").capitalize()
-
-    file_path = Path(__file__).resolve().parent / "data" / "find_events_data.json"
-    if not file_path.exists():
-        return {"error": "Data file not found."}
-
-    try:
-        month_number = datetime.strptime(search_month, "%B").month
-    except ValueError:
-        return {"error": "Invalid month provided."}
-
-    # Helper to wrap months into [1..12]
-    def get_adjacent_months(m):
-        prev_m = 12 if m == 1 else (m - 1)
-        next_m = 1 if m == 12 else (m + 1)
-        return [prev_m, m, next_m]
-
-    valid_months = get_adjacent_months(month_number)
-
-    matching_events = []
-    for city_name, events in json.load(open(file_path)).items():
-        if search_city and search_city not in city_name.lower():
-            continue
-
-        for event in events:
-            date_from = datetime.strptime(event["dateFrom"], "%Y-%m-%d")
-            date_to = datetime.strptime(event["dateTo"], "%Y-%m-%d")
-
-            # If the event's start or end month is in our valid months
-            if date_from.month in valid_months or date_to.month in valid_months:
-                # Add metadata explaining how it matches
-                if date_from.month == month_number or date_to.month == month_number:
-                    month_context = "requested month"
-                elif (
-                    date_from.month == valid_months[0]
-                    or date_to.month == valid_months[0]
-                ):
-                    month_context = "previous month"
-                else:
-                    month_context = "next month"
-
-                matching_events.append(
-                    {
-                        "city": city_name,
-                        "eventName": event["eventName"],
-                        "dateFrom": event["dateFrom"],
-                        "dateTo": event["dateTo"],
-                        "description": event["description"],
-                        "month": month_context,
-                    }
-                )
-
-    # Add top-level metadata if you wish
-    return {
-        "note": f"Returning events from {search_month} plus one month either side (i.e., {', '.join(datetime(2025, m, 1).strftime('%B') for m in valid_months)}).",
-        "events": matching_events,
-    }

tools/goal_registry.py

@@ -1,24 +1,28 @@
 from models.tool_definitions import AgentGoal
 from tools.tool_registry import (
-    find_fixtures_tool,
+    search_fixtures_tool,
     search_flights_tool,
+    search_trains_tool,
+    book_train_tool,
     create_invoice_tool,
 )
 
 goal_event_flight_invoice = AgentGoal(
-    tools=[find_fixtures_tool, search_flights_tool, create_invoice_tool],
+    tools=[search_fixtures_tool, search_flights_tool, search_trains_tool, create_invoice_tool],
     description="Help the user gather args for these tools in order: "
-    "1. FindFixtures: Find fixtures for a team in a given month "
+    "1. SearchFixtures: Search for fixtures for a team in a given month"
-    "2. SearchFlights: search for a flight around the event dates "
+    "2. SearchFlights: Search for a flight around the match dates"
-    "3. CreateInvoice: Create a simple invoice for the cost of that flight ",
+    "3. SearchTrains: Search for trains to visit somewhere before or after the match"
+    "4. BookTrain: Book the train tickets"
+    "5. CreateInvoice: Create a simple invoice for the cost of the flights and train tickets",
     starter_prompt="Welcome me, give me a description of what you can do, then ask me for the details you need to do your job",
     example_conversation_history="\n ".join(
         [
             "user: I'd like to travel to a football match",
             "agent: Sure! Let's start by finding an match you'd like to attend. I know about Premier League fixtures in the UK. Could you tell me which team and month you're interested in?",
-            "user: wolverhamptom in may please",
+            "user: Wolves in May please",
-            "agent: Great! Let's find a match for Wolverhampton Wanderers in May.",
+            "agent: Great! Let's find a match for Wolverhampton Wanderers FC in May.",
-            "user_confirmed_tool_run: <user clicks confirm on FindFixtures tool>",
+            "user_confirmed_tool_run: <user clicks confirm on SearchFixtures tool, passing the full team name as an input>",
             'tool_result: results including {"homeTeam": "Wolverhampton Wanderers FC", "awayTeam": "Manchester United", "date": "2025-05-04"}',
             "agent: Found a match! There's an away game against Manchester United on May 4 2025. Would you like to search for flights around this date?",
             "user: Yes, please",
@@ -27,7 +31,19 @@ goal_event_flight_invoice = AgentGoal(
             "agent: Thanks, searching for flights from San Francisco to the Manchester around 2025-05-03 to 2025-05-05.",
             "user_confirmed_tool_run: <user clicks confirm on SearchFlights tool>"
             'tool_result: results including {"flight_number": "CX101", "return_flight_number": "CX102", "price": 850.0}',
-            "agent: Found some flights! The cheapest is CX101 for $850. Would you like to generate an invoice for this flight?",
+            "agent: Found some flights! The cheapest is CX101 for $850. Would you like to go for this one?",
+            "user: Yes, please",
+            "agent: Great! So you're flying to the city the match is in, anywhere else you'd like to visit while you're there?",
+            "user: I'd like to visit my friend in Edinburgh after the match",
+            "agent: Got it! Let's search for trains from Manchester to Edinburgh and back.",
+            "user_confirmed_tool_run: <user clicks confirm on SearchTrains tool>",
+            'tool_result: results including {"train_number": "T1234", "price": 30.0}',
+            "agent: Found these train options! The cheapest is T1234 for $30. Would you like to go for this one?",
+            "user: Yes, please",
+            "agent: Great, lets book the train tickets",
+            "user_confirmed_tool_run: <user clicks confirm on BookTrain tool>",
+            'tool_result: results including {"status": "success"}',
+            "agent: Train tickets booked! Now let's create an invoice for the flights and train tickets",
             "user_confirmed_tool_run: <user clicks confirm on CreateInvoice tool>",
             'tool_result: { "status": "success", "invoice": { "flight_number": "CX101", "amount": 850.0 }, invoiceURL: "https://example.com/invoice" }',
             "agent: Invoice generated! Here's the link: https://example.com/invoice",

tools/search_fixtures.py

@@ -5,7 +5,7 @@ from dotenv import load_dotenv
 
 BASE_URL = 'https://api.football-data.org/v4'
 
-def find_fixtures(args: dict) -> dict:
+def search_fixtures(args: dict) -> dict:
     load_dotenv(override=True)
     api_key = os.getenv("FOOTBALL_DATA_API_KEY", "YOUR_DEFAULT_KEY")
 

tools/search_trains.py

@@ -0,0 +1,65 @@
+import requests
+import os
+from dotenv import load_dotenv
+
+BASE_URL = 'http://localhost:8080/'
+
+def search_trains(args: dict) -> dict:
+    load_dotenv(override=True)
+
+    origin = args.get("origin")
+    destination = args.get("destination")
+    outbound_time = args.get("outbound_time")
+    return_time = args.get("return_time")
+
+    if not origin or not destination or not outbound_time or not return_time:
+        return {"error": "Origin, destination, outbound_time and return_time are required."}
+
+    search_url = f'{BASE_URL}/api/journeys'
+    params = {
+        'from': origin,
+        'to': destination,
+        'outbound_time': outbound_time,
+        'return_time': return_time,
+    }
+
+    response = requests.get(search_url, params=params)
+    if response.status_code != 200:
+        print(response.content)
+        return {"error": "Failed to fetch journey data."}
+
+    journey_data = response.json()
+    return journey_data
+
+def book_train(args: dict) -> dict:
+    load_dotenv(override=True)
+
+    journey_id = args.get("journey_id")
+
+    if not journey_id:
+        return {"error": "Journey ID is required."}
+
+    book_url = f'{BASE_URL}/api/book/{journey_id}'
+    response = requests.post(book_url)
+    if response.status_code != 200:
+        return {"error": "Failed to book ticket."}
+
+    booking_data = response.json()
+    return booking_data
+
+# Example usage
+if __name__ == "__main__":
+    search_args = {
+        "origin": "London Gatwick",
+        "destination": "Manchester",
+        "outbound_time": "2025-03-15T14:00",
+        "return_time": "2025-03-20T14:00"
+    }
+    search_results = search_trains(search_args)
+    print(search_results)
+
+    book_args = {
+        "journey_id": "12345",
+    }
+    booking_results = book_train(book_args)
+    print(booking_results)

tools/tool_registry.py

@@ -1,26 +1,5 @@
 from models.tool_definitions import ToolDefinition, ToolArgument
 
-find_events_tool = ToolDefinition(
-    name="FindEvents",
-    description="Find upcoming events to travel to a given city (e.g., 'Melbourne') and a date or month. "
-    "It knows about events in Oceania only (e.g. major Australian and New Zealand cities). "
-    "It will search 1 month either side of the month provided. "
-    "Returns a list of events. ",
-    arguments=[
-        ToolArgument(
-            name="city",
-            type="string",
-            description="Which city to search for events",
-        ),
-        ToolArgument(
-            name="month",
-            type="string",
-            description="The month to search for events (will search 1 month either side of the month provided)",
-        ),
-    ],
-)
-
-# 2) Define the SearchFlights tool
 search_flights_tool = ToolDefinition(
     name="SearchFlights",
     description="Search for return flights from an origin to a destination within a date range (dateDepart, dateReturn).",
@@ -48,7 +27,45 @@ search_flights_tool = ToolDefinition(
     ],
 )
 
-# 3) Define the CreateInvoice tool
+search_trains_tool = ToolDefinition(
+    name="SearchTrains",
+    description="Search for trains between two stations. Returns a list of trains.",
+    arguments=[
+        ToolArgument(
+            name="origin",
+            type="string",
+            description="The station to depart from",
+        ),
+        ToolArgument(
+            name="destination",
+            type="string",
+            description="The station to arrive at",
+        ),
+        ToolArgument(
+            name="outbound_time",
+            type="ISO8601",
+            description="The date and time to search for outbound trains",
+        ),
+        ToolArgument(
+            name="return_time",
+            type="ISO8601",
+            description="The date and time to search for return trains",
+        ),
+    ],
+)
+
+book_train_tool = ToolDefinition(
+    name="BookTrain",
+    description="Books a train ticket. Returns a booking reference.",
+    arguments=[
+        ToolArgument(
+            name="journey_id",
+            type="string",
+            description="The ID of the journey to book",
+        ),
+    ],
+)
+
 create_invoice_tool = ToolDefinition(
     name="CreateInvoice",
     description="Generate an invoice for the items described for the amount provided. Returns URL to invoice.",
@@ -66,19 +83,19 @@ create_invoice_tool = ToolDefinition(
     ],
 )
 
-find_fixtures_tool = ToolDefinition(
+search_fixtures_tool = ToolDefinition(
-    name="FindFixtures",
+    name="SearchFixtures",
-    description="Find upcoming fixtures for a given team and month",
+    description="Search for upcoming fixtures for a given team and month",
     arguments=[
         ToolArgument(
             name="team",
             type="string",
-            description="The name of the team to search for",
+            description="The full name of the team to search for.",
         ),
         ToolArgument(
             name="month",
             type="string",
-            description="The month to search for fixtures",
+            description="The month to search for fixtures.",
         ),
     ],
 )