Update app.py

app.py

@@ -95,7 +95,7 @@ def fetch_flight_data(lat_min, lat_max, lon_min, lon_max):
         return {'states': [], 'time': 0}
 
 # Hugging Face model configuration
-HF_API_URL = "https://api-inference.huggingface.co/models/mistralai/Mistral-7B-Instruct-v0.2"
+HF_API_URL = "https://api-inference.huggingface.co/models/google/flan-t5-large"
 HF_TOKEN = os.getenv("HF_TOKEN")
 headers = {"Authorization": f"Bearer {HF_TOKEN}"}
 
@@ -104,20 +104,69 @@ def query_llm(prompt):
         payload = {
             "inputs": prompt,
             "parameters": {
-                "max_new_tokens": 250,
+                "max_length": 250,
                 "temperature": 0.1,
                 "top_p": 0.95,
-                "return_full_text": False
+                "do_sample": False
             }
         }
         
         response = requests.post(HF_API_URL, headers=headers, json=payload)
         response.raise_for_status()
         return response.json()[0]['generated_text']
+    except requests.exceptions.HTTPError as e:
+        if e.response.status_code == 403:
+            st.warning("Language model access is currently restricted. Using direct flight data display instead.")
+        else:
+            st.error(f"Error querying language model: {str(e)}")
+        return None
     except Exception as e:
         st.error(f"Error querying language model: {str(e)}")
         return None
 
+def create_flight_embeddings(geo_df):
+    """Create embeddings for flight data to enable semantic search"""
+    try:
+        from sentence_transformers import SentenceTransformer
+        model = SentenceTransformer('all-MiniLM-L6-v2')
+        
+        # Create text representations of flight data
+        flight_texts = []
+        for _, row in geo_df.iterrows():
+            text = f"Flight {row['callsign']} from {row['origin_country']} "
+            text += f"at altitude {row['baro_altitude']}m, speed {row['velocity']}m/s, "
+            text += f"heading {row['true_track']}°"
+            flight_texts.append(text)
+        
+        # Generate embeddings
+        embeddings = model.encode(flight_texts)
+        return embeddings, flight_texts
+    except Exception as e:
+        st.warning(f"Could not create embeddings: {str(e)}")
+        return None, None
+
+def find_similar_flights(identifier, geo_df, embeddings, flight_texts, threshold=0.7):
+    """Find similar flights using semantic search"""
+    try:
+        from sentence_transformers import SentenceTransformer
+        model = SentenceTransformer('all-MiniLM-L6-v2')
+        
+        # Create query embedding
+        query_embedding = model.encode([identifier])
+        
+        # Calculate similarities
+        from sklearn.metrics.pairwise import cosine_similarity
+        similarities = cosine_similarity(query_embedding, embeddings)[0]
+        
+        # Find similar flights
+        similar_indices = [i for i, sim in enumerate(similarities) if sim > threshold]
+        if similar_indices:
+            return geo_df.iloc[similar_indices]
+        return None
+    except Exception as e:
+        st.warning(f"Error in semantic search: {str(e)}")
+        return None
+
 def query_flight_data(geo_df, question):
     # Preprocess the question to extract key information
     question = question.lower().strip()
@@ -182,10 +231,28 @@ def query_flight_data(geo_df, question):
         except:
             pass
     
+    # If still no match, try semantic search using RAG
+    if flight_data is None or flight_data.empty:
+        try:
+            # Create embeddings for all flights
+            embeddings, flight_texts = create_flight_embeddings(geo_df)
+            if embeddings is not None:
+                # Try to find similar flights
+                similar_flights = find_similar_flights(identifier, geo_df, embeddings, flight_texts)
+                if similar_flights is not None and not similar_flights.empty:
+                    flight_data = similar_flights
+                    st.info(f"Found similar flight(s) to {identifier}")
+        except Exception as e:
+            st.warning(f"Semantic search failed: {str(e)}")
+    
     if flight_data is None or flight_data.empty:
-        return f"Could not find flight information for {identifier}. Please check the flight identifier and try again."
+        # If still no match, show available flights
+        available_flights = geo_df['callsign'].dropna().unique()
+        if len(available_flights) > 0:
+            return f"Could not find flight {identifier}. Available flights: {', '.join(available_flights[:10])}..."
+        return f"Could not find flight {identifier}. No flights currently available in the selected area."
     
-    # Prepare flight data for the LLM
+    # Prepare flight data for display
     flight_info = {}
     for col in flight_data.columns:
         if col in flight_data.columns:
@@ -209,27 +276,27 @@ def query_flight_data(geo_df, question):
     if not flight_info:
         return f"No information available for flight {identifier}."
     
-    # Create a prompt for the LLM
-    prompt = f"""You are a flight information assistant. Answer the following question about flight {identifier} using the provided flight data.
-
-Question: {question}
+    # Try to get LLM response, but fall back to direct display if it fails
+    try:
+        # Create a prompt for the LLM
+        prompt = f"""Answer this question about flight {identifier}: {question}
 
-Flight Data:
+Available flight data:
 {json.dumps(flight_info, indent=2)}
 
-Please provide a clear and concise answer focusing on the specific information requested in the question. If the question asks for information not available in the data, say so clearly."""
-
-    # Get response from LLM
-    llm_response = query_llm(prompt)
+Provide a clear and concise answer focusing on the specific information requested."""
+        
+        llm_response = query_llm(prompt)
+        if llm_response:
+            return llm_response
+    except:
+        pass
     
-    if llm_response:
-        return llm_response
-    else:
-        # Fallback to direct data response if LLM fails
-        response = f"Flight Information for {identifier}:\n"
-        for key, value in flight_info.items():
-            response += f"- {key.replace('_', ' ').title()}: {value}\n"
-        return response
+    # Fallback to direct data display
+    response = f"Flight Information for {identifier}:\n"
+    for key, value in flight_info.items():
+        response += f"- {key.replace('_', ' ').title()}: {value}\n"
+    return response
 
 def flight_tracking(flight_view_level, country, local_time_zone, flight_info, airport, color):
     # Get cached location data