Update app.py

app.py

@@ -316,6 +316,10 @@ def generate_answer(message, choice, retrieval_mode):
     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 flights
+    if "flight" in message.lower() and ("JFK" in message or "BHM" in message):
+        response = fetch_google_flights()
+        return response, extract_addresses(response)
 
     prompt_template = QA_CHAIN_PROMPT_1 if choice == "Details" else QA_CHAIN_PROMPT_2
 
@@ -1032,6 +1036,50 @@ def fetch_google_hotels(query="Birmingham Hotel", check_in="2024-08-14", check_o
 
     return hotel_info
 
+from serpapi import GoogleSearch
+
+def fetch_google_flights(departure_id="JFK", arrival_id="BHM", outbound_date="2024-08-14", return_date="2024-08-20"):
+    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()
+    
+    flight_info = ""
+    best_flights = results.get('best_flights', [])
+    
+    for flight in best_flights[:5]:  # Limiting to top 5 flights
+        flight_segments = flight.get('flights', [])
+        if not flight_segments:
+            continue
+        
+        for segment in flight_segments:
+            departure_airport = segment.get('departure_airport', {}).get('name', 'Unknown Departure Airport')
+            departure_time = segment.get('departure_airport', {}).get('time', 'Unknown Time')
+            arrival_airport = segment.get('arrival_airport', {}).get('name', 'Unknown Arrival Airport')
+            arrival_time = segment.get('arrival_airport', {}).get('time', 'Unknown Time')
+            duration = segment.get('duration', 'Unknown Duration')
+            airplane = segment.get('airplane', 'Unknown Airplane')
+            
+            flight_info += f"**Flight Segment:**\n"
+            flight_info += f"  - **Departure:** {departure_airport} at {departure_time}\n"
+            flight_info += f"  - **Arrival:** {arrival_airport} at {arrival_time}\n"
+            flight_info += f"  - **Duration:** {duration} minutes\n"
+            flight_info += f"  - **Airplane:** {airplane}\n"
+        
+        flight_info += "-" * 50 + "\n"
+
+    return flight_info
+
+