import asyncio import aiohttp import requests import re import pprint import logging from typing import List, Dict, Any, Union, Optional from datetime import datetime import gradio as gr import asyncio import pandas as pd # Configure logging logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s') logger = logging.getLogger(__name__) # --- Configuration Constants --- API_BASE_URL = "https://www.railyatri.in/get-next-days-sa-data" TRAINS_API_URL = "https://trainticketapi.railyatri.in/api/trains-between-station-with-sa.json" DEFAULT_QUOTA = "GN" REQUEST_TIMEOUT = 30 # seconds MAX_RETRIES = 3 RETRY_DELAY = 1 # seconds # Valid quota codes that users can specify VALID_QUOTAS = { "GN": "General", "TQ": "Tatkal", "PT": "Premium Tatkal", "LD": "Ladies", "HP": "Handicapped", "YU": "Yuva", "SS": "Senior Citizen", "DF": "Defence", "RE": "Railway Employee", "RQ": "RAC" } # --- Custom Exceptions --- class TrainAPIError(Exception): """Custom exception for train API related errors.""" pass class ValidationError(Exception): """Custom exception for input validation errors.""" pass # --- Validation Functions --- def validate_station_code(station_code: str, field_name: str) -> str: """Validate and normalize station code.""" if not isinstance(station_code, str): raise ValidationError(f"{field_name} must be a string") station = station_code.strip().upper() if not station or len(station) < 2: raise ValidationError(f"Invalid {field_name}: {station_code}") return station def validate_date_format(date_str: str, input_format: str = None) -> str: """ Validate date and convert to YYYY-M-D format. Args: date_str: Date string to validate input_format: Expected input format ('dd-mm-yyyy' or 'yyyy-mm-dd') Returns: str: Date in YYYY-M-D format """ if not isinstance(date_str, str): raise ValidationError("Date must be a string") date_str = date_str.strip() # Try different date formats formats_to_try = [] if input_format == 'dd-mm-yyyy': formats_to_try = ['%d-%m-%Y', '%d-%m-%y'] elif input_format == 'yyyy-mm-dd': formats_to_try = ['%Y-%m-%d', '%Y-%M-%d'] else: # Auto-detect format formats_to_try = ['%d-%m-%Y', '%Y-%m-%d', '%d-%m-%y', '%Y-%M-%d'] for fmt in formats_to_try: try: dt_obj = datetime.strptime(date_str, fmt) return dt_obj.strftime('%Y-%m-%d') except ValueError: continue raise ValidationError(f"Invalid date format: {date_str}") def validate_quota(quota: str) -> str: """Validate quota code.""" if not isinstance(quota, str): raise ValidationError("Quota must be a string") quota = quota.strip().upper() if quota not in VALID_QUOTAS: valid_codes = ', '.join(VALID_QUOTAS.keys()) raise ValidationError(f"Invalid quota '{quota}'. Valid codes: {valid_codes}") return quota def validate_query(query: Dict[str, Any]) -> None: """Validates a single query dictionary for required fields and data types.""" required_fields = ["train_number", "journey_class", "source", "destination", "journey_date"] for field in required_fields: if field not in query or query[field] is None: raise ValidationError(f"Missing required field: {field}") # Validate train_number train_num = str(query["train_number"]).strip() if not train_num or not train_num.isdigit(): raise ValidationError(f"Invalid train_number: {query['train_number']}") # Validate journey_class journey_class = query["journey_class"] if isinstance(journey_class, str): journey_class = [journey_class] elif not isinstance(journey_class, list): raise ValidationError(f"journey_class must be string or list, got {type(journey_class)}") for j_class in journey_class: if not isinstance(j_class, str) or not j_class.strip(): raise ValidationError(f"Invalid journey_class: {j_class}") # Validate station codes query["source"] = validate_station_code(query["source"], "source") query["destination"] = validate_station_code(query["destination"], "destination") # Validate date query["journey_date"] = validate_date_format(query["journey_date"], 'yyyy-mm-dd') # Validate quota if provided if "journey_quota" in query and query["journey_quota"] is not None: query["journey_quota"] = validate_quota(query["journey_quota"]) def validate_queries(queries: List[Dict[str, Any]]) -> None: """Validates a list of query dictionaries.""" if not isinstance(queries, list): raise ValidationError("queries must be a list") if not queries: raise ValidationError("queries list cannot be empty") for i, query in enumerate(queries): if not isinstance(query, dict): raise ValidationError(f"Query at index {i} must be a dictionary") try: validate_query(query) except ValidationError as e: raise ValidationError(f"Query at index {i}: {str(e)}") # --- Train Discovery Functions --- def get_trains_between_stations(source: str, destination: str, journey_date: str) -> Optional[Dict[str, Any]]: """ Fetch train details between two stations from Railyatri API. Args: source: Source station code (e.g., 'HWH'). destination: Destination station code (e.g., 'PNBE'). journey_date: Journey date in 'dd-mm-yyyy' format (e.g., '22-09-2025'). Returns: Optional[Dict]: JSON response from API if successful, else None. """ try: # Validate inputs source = validate_station_code(source, "source") destination = validate_station_code(destination, "destination") # Keep original date format for this API if not re.match(r'\d{2}-\d{2}-\d{4}', journey_date.strip()): # Try to convert from YYYY-MM-DD to DD-MM-YYYY if needed try: dt = datetime.strptime(journey_date.strip(), '%Y-%m-%d') journey_date = dt.strftime('%d-%m-%Y') except ValueError: raise ValidationError(f"Invalid date format: {journey_date}") params = { "from": source, "to": destination, "dateOfJourney": journey_date.strip() } logger.info(f"Fetching trains from {source} to {destination} on {journey_date}") response = requests.get(TRAINS_API_URL, params=params, timeout=REQUEST_TIMEOUT) response.raise_for_status() result = response.json() logger.info("Successfully fetched train list") return result except requests.exceptions.Timeout: logger.error("Timeout while fetching train list") return None except requests.exceptions.RequestException as e: logger.error(f"Error during train list request: {e}") return None except ValidationError as e: logger.error(f"Validation error: {e}") return None except Exception as e: logger.error(f"Unexpected error while fetching trains: {e}") return None def simplify_train_info(api_result: Dict[str, Any], default_quota: str = DEFAULT_QUOTA) -> List[Dict[str, Any]]: """ Simplify train information from API response. Args: api_result: Raw API response default_quota: Default quota to use for all trains Returns: List[Dict]: Simplified train information """ if not isinstance(api_result, dict): logger.warning("Invalid API result format") return [] simplified_trains = [] def process_train_list(train_list: List[Dict[str, Any]]) -> None: """Process a list of trains and add to simplified_trains.""" if not isinstance(train_list, list): return for train in train_list: if not isinstance(train, dict): continue try: # Handle date conversion raw_date = train.get("train_date", "").strip() formatted_date = None if raw_date: # Try multiple date formats for date_format in ['%d-%m-%Y', '%Y-%m-%d']: try: dt_obj = datetime.strptime(raw_date, date_format) formatted_date = dt_obj.strftime('%Y-%m-%d') break except ValueError: continue if not formatted_date: logger.warning(f"Could not parse date: {raw_date}") formatted_date = raw_date # Get journey classes, ensure it's a list journey_classes = train.get("journey_class", []) if isinstance(journey_classes, str): journey_classes = [journey_classes] elif not isinstance(journey_classes, list): journey_classes = [] # Filter out empty classes journey_classes = [jc.strip() for jc in journey_classes if jc and str(jc).strip()] if not journey_classes: logger.warning(f"No valid journey classes for train {train.get('train_number')}") continue train_data = { "train_number": str(train.get("train_number", "")).strip(), "journey_class": journey_classes, "source": str(train.get("from", "")).strip().upper(), "destination": str(train.get("to", "")).strip().upper(), "journey_date": formatted_date, "journey_quota": validate_quota(default_quota) } # Validate essential fields if (train_data["train_number"] and train_data["source"] and train_data["destination"] and train_data["journey_date"]): simplified_trains.append(train_data) else: logger.warning(f"Incomplete train data: {train_data}") except Exception as e: logger.warning(f"Error processing train: {e}") continue # Process different train categories train_categories = [ "train_between_stations", "alternate_trains", "reserved_trains" ] for category in train_categories: if category in api_result: process_train_list(api_result[category]) logger.info(f"Simplified {len(simplified_trains)} trains") return simplified_trains # --- Core API and Parsing Functions --- async def get_train_availability_async( session: aiohttp.ClientSession, query: Dict[str, Union[str, int]], retry_count: int = 0 ) -> Dict[str, Any]: """ Asynchronously fetches seat availability data from Railyatri for a single request. """ # Normalize parameters params = { "train_number": str(query["train_number"]).strip(), "journey_class": str(query["journey_class"]).strip(), "from": str(query["source"]).strip().upper(), "to": str(query["destination"]).strip().upper(), "journey_date": str(query["journey_date"]).strip(), "journey_quota": str(query.get("journey_quota", DEFAULT_QUOTA)).strip(), } try: timeout = aiohttp.ClientTimeout(total=REQUEST_TIMEOUT) async with session.get(API_BASE_URL, params=params, timeout=timeout) as response: response.raise_for_status() content_type = response.headers.get('content-type', '').lower() if 'application/json' not in content_type: logger.warning(f"Unexpected content type for train {params['train_number']}") response_data = await response.json() if not isinstance(response_data, dict): logger.warning(f"Invalid response format for train {params['train_number']}") return {} return response_data except asyncio.TimeoutError: logger.error(f"Timeout for train {params['train_number']} class {params['journey_class']}") if retry_count < MAX_RETRIES: await asyncio.sleep(RETRY_DELAY * (retry_count + 1)) return await get_train_availability_async(session, query, retry_count + 1) return {} except aiohttp.ClientResponseError as e: if e.status == 429 and retry_count < MAX_RETRIES: await asyncio.sleep(RETRY_DELAY * (retry_count + 1) * 2) return await get_train_availability_async(session, query, retry_count + 1) logger.error(f"HTTP {e.status} for train {params['train_number']}: {e}") return {} except Exception as e: logger.error(f"Error for train {params['train_number']}: {e}") if retry_count < MAX_RETRIES: await asyncio.sleep(RETRY_DELAY * (retry_count + 1)) return await get_train_availability_async(session, query, retry_count + 1) return {} def safe_extract_train_number(ticket_link: str) -> str: """Safely extracts train number from ticket link.""" if not isinstance(ticket_link, str): return "N/A" try: match = re.search(r'train_no=(\d+)', ticket_link) return match.group(1) if match else "N/A" except Exception: return "N/A" def extract_train_availability(data: Dict[str, Any]) -> List[Dict[str, Any]]: """Extracts and transforms train availability details from raw API data.""" if not isinstance(data, dict): return [] extracted_data = [] try: # Find seat availability data with multiple fallback patterns seat_availability = None # Check nested structure first if "result" in data and isinstance(data["result"], dict): for key in ["seat_availibility", "seat_availability"]: if key in data["result"] and isinstance(data["result"][key], list): seat_availability = data["result"][key] break # Check root level if not seat_availability: for key in ["seat_availibility", "seat_availability"]: if key in data and isinstance(data[key], list): seat_availability = data[key] break if not seat_availability: return [] for entry in seat_availability: if not isinstance(entry, dict): continue # Extract fields with fallbacks ticket_link = entry.get("ticket_link", "") train_no = safe_extract_train_number(ticket_link) class_type = None for key in ["class_type", "class", "journey_class"]: if key in entry and entry[key]: class_type = str(entry[key]).strip() break if not class_type: continue # Get availability availability = entry.get(class_type) if availability is None: for key in ["availability", "status", "seats"]: if key in entry: availability = entry[key] break # Get date journey_date = entry.get("Date (DD-MM-YYYY)") if journey_date is None: for key in ["date", "journey_date", "Date"]: if key in entry: journey_date = entry[key] break # Get fare total_fare = entry.get("total_fare") if total_fare is None: for key in ["fare", "price", "cost"]: if key in entry: total_fare = entry[key] break extracted_data.append({ "train_no": train_no, "journey_date": journey_date, "class_type": class_type, "availability": availability, "total_fare": total_fare }) except Exception as e: logger.error(f"Error extracting availability data: {e}") return [] return extracted_data # --- Main Functions --- async def fetch_all_availabilities(queries: List[Dict[str, Any]]) -> List[Dict[str, Any]]: """ Asynchronously fetches and processes train availability for multiple queries. Args: queries: List of query dictionaries Returns: List of extracted availability information """ try: validate_queries(queries) except ValidationError as e: logger.error(f"Input validation failed: {e}") raise logger.info(f"Processing {len(queries)} queries") tasks = [] connector = aiohttp.TCPConnector(limit=10, limit_per_host=5) timeout = aiohttp.ClientTimeout(total=REQUEST_TIMEOUT) try: async with aiohttp.ClientSession( connector=connector, timeout=timeout, headers={'User-Agent': 'Python-TrainAvailability/1.0'} ) as session: for query in queries: journey_classes = query["journey_class"] if isinstance(journey_classes, str): journey_classes = [journey_classes.strip()] elif isinstance(journey_classes, list): journey_classes = [str(jc).strip() for jc in journey_classes if jc] for j_class in journey_classes: if not j_class: continue single_request_query = query.copy() single_request_query["journey_class"] = j_class task = get_train_availability_async(session, single_request_query) tasks.append(task) if not tasks: return [] api_responses = await asyncio.gather(*tasks, return_exceptions=True) except Exception as e: logger.error(f"Error executing requests: {e}") raise TrainAPIError(f"Failed to execute API requests: {e}") # Process results all_availability_data = [] successful_requests = 0 for response in api_responses: if isinstance(response, Exception) or not isinstance(response, dict) or not response: continue try: extracted_info = extract_train_availability(response) if extracted_info: all_availability_data.extend(extracted_info) successful_requests += 1 except Exception as e: logger.error(f"Error processing response: {e}") continue logger.info(f"Processed {successful_requests}/{len(tasks)} requests successfully") logger.info(f"Extracted {len(all_availability_data)} availability records") return all_availability_data @gr.mcp.tool() async def find_available_seats( source: str, destination: str, journey_date: str, quota: str = DEFAULT_QUOTA ) -> List[Dict[str, Any]]: """ Complete workflow to find available seats between stations. Args: source: Source station code (e.g., 'HWH') destination: Destination station code (e.g., 'PNBE') journey_date: Journey date in 'dd-mm-yyyy' format (e.g., '22-09-2025') quota: Journey quota (default: 'GN') Returns: List of available seat information Raises: ValidationError: If input validation fails TrainAPIError: If API requests fail """ try: # Validate inputs source = validate_station_code(source, "source") destination = validate_station_code(destination, "destination") quota = validate_quota(quota) logger.info(f"Finding available seats: {source} → {destination} on {journey_date} (quota: {quota})") # Step 1: Get trains between stations trains_data = get_trains_between_stations(source, destination, journey_date) if not trains_data: logger.warning("No trains data received") return [] # Step 2: Simplify train information simplified_trains = simplify_train_info(trains_data, quota) if not simplified_trains: logger.warning("No valid trains found") return [] logger.info(f"Found {len(simplified_trains)} trains to check") # Step 3: Fetch availability for all trains availability_data = await fetch_all_availabilities(simplified_trains) logger.info(f"Seat search completed: found {len(availability_data)} availability records") return availability_data except ValidationError: raise except TrainAPIError: raise except Exception as e: logger.error(f"Unexpected error in find_available_seats: {e}") raise TrainAPIError(f"Unexpected error: {e}") # --- Gradio Wrapper --- def gradio_find_seats(source: str, destination: str, journey_date: str, quota: str): try: # Run async function in sync wrapper result = asyncio.run(find_available_seats(source, destination, journey_date, quota)) if not result: return "No availability found", None # Convert to DataFrame for nice table display df = pd.DataFrame(result) return "Success", df except Exception as e: return f"Error: {e}", None # --- Gradio UI --- with gr.Blocks() as demo: gr.Markdown("## 🚆 Train Seat Availability Finder") with gr.Row(): source = gr.Textbox(label="Source Station Code", value="HWH") destination = gr.Textbox(label="Destination Station Code", value="PNBE") with gr.Row(): journey_date = gr.Textbox(label="Journey Date (dd-mm-yyyy)", value="22-09-2025") quota = gr.Dropdown(choices=list(VALID_QUOTAS.keys()), value="GN", label="Quota") submit_btn = gr.Button("Find Seats") status = gr.Textbox(label="Status", interactive=False) output_table = gr.Dataframe(label="Availability", interactive=False) submit_btn.click( fn=gradio_find_seats, inputs=[source, destination, journey_date, quota], outputs=[status, output_table] ) # Launch app if __name__ == "__main__": demo.launch(mcp_server= True)