Update app.py

app.py

@@ -307,16 +307,16 @@ chain_neo4j = (
 def generate_answer(message, choice, retrieval_mode):
     logging.debug(f"generate_answer called with choice: {choice} and retrieval_mode: {retrieval_mode}")
 
+    # Check if the question is about flights
+    if "flight" in message.lower() and ("to" in message.lower() or "from" in message.lower()):
+        response = fetch_flight_info(message)
+        return response, extract_addresses(response)
+
     # Check if the question is about restaurants
     if "restaurant" in message.lower() and "birmingham" in message.lower():
         response = fetch_yelp_restaurants()
         return response, extract_addresses(response)
 
-    # Check if the question is about hotels
-    elif "hotels" in message.lower() and "birmingham" in message.lower():
-        response = fetch_google_hotels()
-        return response, extract_addresses(response)
-
     prompt_template = QA_CHAIN_PROMPT_1 if choice == "Details" else QA_CHAIN_PROMPT_2
 
     if retrieval_mode == "VDB":
@@ -338,6 +338,7 @@ def generate_answer(message, choice, retrieval_mode):
 
 
 
+
 def bot(history, choice, tts_choice, retrieval_mode):
     if not history:
         return history
@@ -1014,6 +1015,48 @@ def fetch_google_hotels():
 
     return response_text
 
+def fetch_flight_info(message):
+    
+    departure_id = "JFK"  # Replace with parsed value from the message
+    arrival_id = "BHM"    # Replace with parsed value from the message
+    outbound_date = "2024-08-14"  # Replace with parsed date from the message
+    return_date = "2024-08-20"    # Replace with parsed date if applicable
+    
+    params = {
+        "engine": "google_flights",
+        "departure_id": departure_id,
+        "arrival_id": arrival_id,
+        "outbound_date": outbound_date,
+        "return_date": return_date,
+        "currency": "USD",
+        "hl": "en",
+        "api_key": os.getenv("SERP_API")
+    }
+
+    search = GoogleSearch(params)
+    results = search.get_dict()
+    flights_results = results.get("flights_results", [])
+
+    response_text = ""
+    
+    for flight in flights_results[:5]:  # Limiting to top 5 flight options
+        price = flight.get("price", "No price")
+        duration = flight.get("duration", "No duration")
+        stops = flight.get("stops", "Direct flight")
+        departure_time = flight.get("departure_time", "No departure time")
+        arrival_time = flight.get("arrival_time", "No arrival time")
+        airlines = flight.get("airlines", ["No airline"])
+
+        # Format the output for each flight
+        response_text += f"Flight: {', '.join(airlines)}\n"
+        response_text += f"Price: {price}\n"
+        response_text += f"Duration: {duration}\n"
+        response_text += f"Stops: {stops}\n"
+        response_text += f"Departure: {departure_time}\n"
+        response_text += f"Arrival: {arrival_time}\n"
+        response_text += "-" * 50 + "\n"
+
+    return response_text