Update app.py

app.py

@@ -1,27 +1,600 @@
-import gradio as gr
-from duckduckgo_search import DDGS
-import spaces
-
-MAX_RESULTS = 5
-
-@spaces.queue  # ✅ CPU-safe, enables Streamable HTTP MCP
-def search_duckduckgo(query: str) -> str:
-    with DDGS() as ddgs:
-        results = ddgs.text(query, max_results=MAX_RESULTS)
-    return "\n".join([f"{r['title']} - {r['href']}" for r in results])
-
-with gr.Blocks() as app:
-    gr.Markdown("# DuckDuckGo Search Tool")
-    gr.Markdown("## Example Video")
-    gr.HTML(
-        """<iframe width="560" height="315" 
-        src="https://www.youtube.com/embed/0EE3sD1MHeg" 
-        frameborder="0" allowfullscreen></iframe>"""
-    )
-
-    query_input = gr.Textbox(label="Enter Search Query")
-    search_btn = gr.Button("Search")
-    output = gr.Textbox(label="Results", interactive=False)
-    search_btn.click(fn=search_duckduckgo, inputs=query_input, outputs=output)
-
-app.launch(mcp_server=True)
+"""
+MCP Server for Indian Railway Seat Availability Checker
+========================================================
+
+This MCP server provides tools for checking train availability between Indian railway stations.
+It uses the RailYatri API to fetch real-time seat availability data.
+
+Installation:
+    pip install fastmcp aiohttp requests
+
+Usage:
+    python train_mcp_server.py
+"""
+
+import asyncio
+import aiohttp
+import requests
+import re
+import logging
+from typing import List, Dict, Any, Union, Optional
+from datetime import datetime
+from fastmcp import FastMCP
+
+# 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
+MAX_RETRIES = 3
+RETRY_DELAY = 1
+
+# Valid quota codes
+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"
+}
+
+# Initialize MCP server
+mcp = FastMCP("Indian Railway Seat Checker")
+
+# --- 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."""
+    if not isinstance(date_str, str):
+        raise ValidationError("Date must be a string")
+    
+    date_str = date_str.strip()
+    
+    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:
+        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."""
+    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}")
+    
+    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']}")
+    
+    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")
+    
+    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}")
+    
+    query["source"] = validate_station_code(query["source"], "source")
+    query["destination"] = validate_station_code(query["destination"], "destination")
+    query["journey_date"] = validate_date_format(query["journey_date"], 'yyyy-mm-dd')
+    
+    if "journey_quota" in query and query["journey_quota"] is not None:
+        query["journey_quota"] = validate_quota(query["journey_quota"])
+
+# --- Core Functions ---
+def get_trains_between_stations(source: str, destination: str, journey_date: str) -> Optional[Dict[str, Any]]:
+    """Fetch train details between two stations."""
+    try:
+        source = validate_station_code(source, "source")
+        destination = validate_station_code(destination, "destination")
+        
+        if not re.match(r'\d{2}-\d{2}-\d{4}', journey_date.strip()):
+            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 Exception as e:
+        logger.error(f"Error 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."""
+    if not isinstance(api_result, dict):
+        return []
+    
+    simplified_trains = []
+
+    def process_train_list(train_list: List[Dict[str, Any]]) -> None:
+        if not isinstance(train_list, list):
+            return
+            
+        for train in train_list:
+            if not isinstance(train, dict):
+                continue
+                
+            try:
+                raw_date = train.get("train_date", "").strip()
+                formatted_date = None
+                
+                if raw_date:
+                    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:
+                        formatted_date = raw_date
+
+                journey_classes = train.get("journey_class", [])
+                if isinstance(journey_classes, str):
+                    journey_classes = [journey_classes]
+                elif not isinstance(journey_classes, list):
+                    journey_classes = []
+                
+                journey_classes = [jc.strip() for jc in journey_classes if jc and str(jc).strip()]
+                
+                if not journey_classes:
+                    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)
+                }
+                
+                if (train_data["train_number"] and 
+                    train_data["source"] and 
+                    train_data["destination"] and
+                    train_data["journey_date"]):
+                    simplified_trains.append(train_data)
+                    
+            except Exception as e:
+                logger.warning(f"Error processing train: {e}")
+                continue
+
+    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])
+    
+    return simplified_trains
+
+async def get_train_availability_async(
+    session: aiohttp.ClientSession,
+    query: Dict[str, Union[str, int]],
+    retry_count: int = 0
+) -> Dict[str, Any]:
+    """Asynchronously fetch seat availability data."""
+    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()
+            response_data = await response.json()
+            
+            if not isinstance(response_data, dict):
+                return {}
+            
+            return response_data
+            
+    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:
+    """Extract 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]]:
+    """Extract train availability details from API data."""
+    if not isinstance(data, dict):
+        return []
+    
+    extracted_data = []
+    
+    try:
+        seat_availability = None
+        
+        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
+        
+        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
+            
+            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
+            
+            availability = entry.get(class_type)
+            if availability is None:
+                for key in ["availability", "status", "seats"]:
+                    if key in entry:
+                        availability = entry[key]
+                        break
+            
+            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
+            
+            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: {e}")
+        return []
+
+    return extracted_data
+
+async def fetch_all_availabilities(queries: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
+    """Fetch train availability for multiple queries."""
+    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}")
+    
+    all_availability_data = []
+    
+    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)
+        except Exception as e:
+            logger.error(f"Error processing response: {e}")
+            continue
+    
+    return all_availability_data
+
+# --- MCP Tool Definitions ---
+
+@mcp.tool()
+async def get_trains(source: str, destination: str, journey_date: str) -> Dict[str, Any]:
+    """
+    Get list of trains between two stations.
+    
+    Args:
+        source: Source station code (e.g., 'HWH' for Howrah, 'NDLS' for New Delhi)
+        destination: Destination station code (e.g., 'PNBE' for Patna, 'BCT' for Mumbai Central)
+        journey_date: Journey date in 'DD-MM-YYYY' format (e.g., '25-09-2025')
+    
+    Returns:
+        Dictionary containing train information with categories:
+        - train_between_stations: Direct trains
+        - alternate_trains: Alternative train options
+        - reserved_trains: Trains with reservation available
+        
+    Example:
+        result = await get_trains("HWH", "PNBE", "25-09-2025")
+    """
+    try:
+        result = get_trains_between_stations(source, destination, journey_date)
+        if result:
+            return {
+                "success": True,
+                "data": result,
+                "message": f"Found trains from {source} to {destination}"
+            }
+        else:
+            return {
+                "success": False,
+                "data": {},
+                "message": "No trains found or API error occurred"
+            }
+    except Exception as e:
+        return {
+            "success": False,
+            "data": {},
+            "message": f"Error: {str(e)}"
+        }
+
+@mcp.tool()
+async def check_seat_availability(
+    source: str, 
+    destination: str, 
+    journey_date: str, 
+    quota: str = "GN"
+) -> Dict[str, Any]:
+    """
+    Check seat availability for all trains between two stations.
+    
+    Args:
+        source: Source station code (e.g., 'HWH' for Howrah)
+        destination: Destination station code (e.g., 'PNBE' for Patna)
+        journey_date: Journey date in 'DD-MM-YYYY' format (e.g., '25-09-2025')
+        quota: Booking quota (default: 'GN' for General)
+               Valid values: GN, TQ (Tatkal), PT (Premium Tatkal), LD (Ladies),
+                           HP (Handicapped), YU (Yuva), SS (Senior Citizen),
+                           DF (Defence), RE (Railway Employee), RQ (RAC)
+    
+    Returns:
+        Dictionary with availability information including:
+        - train_no: Train number
+        - journey_date: Date of journey
+        - class_type: Class of travel (1A, 2A, 3A, SL, etc.)
+        - availability: Seat availability status
+        - total_fare: Fare information
+        
+    Example:
+        result = await check_seat_availability("HWH", "PNBE", "25-09-2025", "GN")
+    """
+    try:
+        trains_data = get_trains_between_stations(source, destination, journey_date)
+        if not trains_data:
+            return {
+                "success": False,
+                "data": [],
+                "message": "No trains found between stations"
+            }
+        
+        simplified_trains = simplify_train_info(trains_data, quota)
+        if not simplified_trains:
+            return {
+                "success": False,
+                "data": [],
+                "message": "No valid trains found for checking availability"
+            }
+        
+        availability_data = await fetch_all_availabilities(simplified_trains)
+        
+        return {
+            "success": True,
+            "data": availability_data,
+            "trains_checked": len(simplified_trains),
+            "availability_records": len(availability_data),
+            "message": f"Checked {len(simplified_trains)} trains, found {len(availability_data)} availability records"
+        }
+        
+    except Exception as e:
+        return {
+            "success": False,
+            "data": [],
+            "message": f"Error: {str(e)}"
+        }
+
+@mcp.tool()
+async def check_specific_train(
+    train_number: str,
+    source: str,
+    destination: str,
+    journey_date: str,
+    journey_class: str,
+    quota: str = "GN"
+) -> Dict[str, Any]:
+    """
+    Check seat availability for a specific train and class.
+    
+    Args:
+        train_number: Train number (e.g., '12301')
+        source: Source station code (e.g., 'HWH')
+        destination: Destination station code (e.g., 'NDLS')
+        journey_date: Journey date in 'YYYY-MM-DD' format
+        journey_class: Class of travel (e.g., '1A', '2A', '3A', 'SL')
+        quota: Booking quota (default: 'GN')
+    
+    Returns:
+        Dictionary with availability information for the specific train
+        
+    Example:
+        result = await check_specific_train("12301", "HWH", "NDLS", "2025-09-25", "3A", "GN")
+    """
+    try:
+        query = {
+            "train_number": train_number,
+            "journey_class": journey_class,
+            "source": source,
+            "destination": destination,
+            "journey_date": journey_date,
+            "journey_quota": quota
+        }
+        
+        validate_query(query)
+        
+        availability_data = await fetch_all_availabilities([query])
+        
+        if availability_data:
+            return {
+                "success": True,
+                "data": availability_data[0] if availability_data else {},
+                "message": f"Availability checked for train {train_number}"
+            }
+        else:
+            return {
+                "success": False,
+                "data": {},
+                "message": "No availability data found"
+            }
+            
+    except Exception as e:
+        return {
+            "success": False,
+            "data": {},
+            "message": f"Error: {str(e)}"
+        }
+
+@mcp.tool()
+async def get_quota_codes() -> Dict[str, str]:
+    """
+    Get all valid quota codes and their descriptions.
+    
+    Returns:
+        Dictionary of quota codes and their full names
+        
+    Example:
+        codes = await get_quota_codes()
+        # Returns: {"GN": "General", "TQ": "Tatkal", ...}
+    """
+    return VALID_QUOTAS
+
+# Run the server
+if __name__ == "__main__":
+    mcp.run(transport="http")