FlightQueryBot_Version1

Author: RobuRishabh

app.py

@@ -1,75 +1,69 @@
 """
-UI using Streamlit
+Flight Information Chatbot UI using Streamlit
 """
 import streamlit as st
 from query_handler import process_query
 from ollama_api import check_ollama_availability, generate_response
 
-# ✅ Set page config
+# Set Streamlit page config
 st.set_page_config(
     page_title="Flight Information Assistant",
-    page_icon="✈️"
+    page_icon="✈️",
 )
 
-# ✅ Initialize session state for chat history
+# Initialize chat history in session state
 if "messages" not in st.session_state:
     st.session_state.messages = []
 
-# ✅ Check Ollama server availability
-ollama_available, ollama_message = check_ollama_availability()
-
 def display_chat_message(role, content):
-    """
-    Display messages in the chat interface.
-    
-    Args:
-        role (str): Role of the message sender (user or assistant).
-        content (str): Message content.
-    """
+    """Display a chat message in Streamlit UI."""
     with st.chat_message(role):
         st.markdown(content)
 
-# ✅ Main UI
+# Display the title
 st.title("✈️ Flight Information Assistant")
 
-# ✅ Display Ollama server availability message
-if not ollama_available:
-    st.warning(f"⚠️ {ollama_message}\nUsing simplified responses until Ollama becomes available.")
+# Check Ollama server availability
+ollama_status, ollama_message = check_ollama_availability()
+if not ollama_status:
+    st.warning("⚠️ Ollama server is unavailable. Responses will be simplified.")
 
-# ✅ Instructions for the user
+# Show instructions
 st.markdown("""
-**Ask me about flights!** Try questions like:
-- What flights are available from New York to London?
-- Show me flight NY100.
-- Are there any flights from Chicago?
+### **Ask me about flights!**  
+Try questions like:  
+- ✈️ What flights are available from New York to London?  
+- 🔍 Show me flight NY100.  
+- 📍 Are there any flights from Chicago?  
 """)
 
-# ✅ Display chat history
+# Display chat history
 for message in st.session_state.messages:
     display_chat_message(message["role"], message["content"])
 
-# ✅ Chat input handling
-if prompt := st.chat_input("Ask about flights..."):
-    # ✅ Store and display user message
-    st.session_state.messages.append({"role": "user", "content": prompt})
-    display_chat_message("user", prompt)
+# Chat input handling
+user_input = st.chat_input("Ask about flights...")
+
+if user_input:
+    # Add user query to chat history
+    st.session_state.messages.append({"role": "user", "content": user_input})
+    display_chat_message("user", user_input)
 
-    try:
-        # ✅ Process user query
-        success, message, flights = process_query(prompt)
+    # Process query
+    with st.spinner("Searching for flights..."):
+        try:
+            success, message, flights = process_query(user_input)
 
-        if not success:
-            st.error(message)  # Show error message in UI
-            response = "⚠️ I couldn't process your request. Try rephrasing your question."
-        else:
-            # ✅ Generate assistant response
-            response = generate_response(prompt, flights)
+            if not success:
+                response = f"⚠️ {message}"
+            else:
+                response = generate_response(user_input, flights)
 
-        # ✅ Store and display assistant response
-        st.session_state.messages.append({"role": "assistant", "content": response})
-        display_chat_message("assistant", response)
+        except ValueError as ve:
+            response = f"❌ Invalid input: {str(ve)}"
+        except Exception as e:
+            response = f"❌ An unexpected error occurred: {str(e)}"
 
-    except Exception as e:
-        error_message = f"❌ An error occurred: {str(e)}"
-        st.session_state.messages.append({"role": "assistant", "content": error_message})
-        display_chat_message("assistant", error_message)
+    # Add assistant response to chat history
+    st.session_state.messages.append({"role": "assistant", "content": response})
+    display_chat_message("assistant", response)

mock_database.py

@@ -1,121 +1,40 @@
-import ollama
-import numpy as np
-from sklearn.metrics.pairwise import cosine_similarity
+import requests
 
-flights = [
-    {
-        "flight_number": "NY100",
-        "origin": "New York",
-        "destination": "London",
-        "time": "2025-05-01 08:00",
-        "airline": "Global Airways"
-    },
-    {
-        "flight_number": "LA200",
-        "origin": "Los Angeles",
-        "destination": "Tokyo",
-        "time": "2025-05-01 10:30",
-        "airline": "Pacific Routes"
-    },
-    {
-        "flight_number": "CH300",
-        "origin": "Chicago",
-        "destination": "Paris",
-        "time": "2025-05-01 15:45",
-        "airline": "Euro Connect"
-    },
-    {
-        "flight_number": "SF400",
-        "origin": "San Francisco",
-        "destination": "Sydney",
-        "time": "2025-05-01 23:15",
-        "airline": "Ocean Pacific"
-    },
-    {
-        "flight_number": "MI500",
-        "origin": "Miami",
-        "destination": "Rio de Janeiro",
-        "time": "2025-05-02 07:30",
-        "airline": "South American Airways"
-    }
-]
-
-### ✅ Check if Ollama is Running Before Sending API Calls
 def check_ollama_availability():
-    """
-    Check if Ollama server is available.
-    Returns: True if available, False otherwise.
-    """
-    import requests
+    """Check if the Ollama server is available."""
     try:
         response = requests.get("http://localhost:11434/api/tags", timeout=3)
         return response.status_code == 200
-    except requests.exceptions.RequestException:
+    except requests.RequestException:
+        print("⚠️ Ollama server is not available.")
         return False
 
-OLLAMA_AVAILABLE = check_ollama_availability()
-
-### ✅ Lazy Embedding Generation with Error Handling
-flight_embeddings = {}
-
-def generate_embedding(text):
-    """
-    Generate text embeddings using Ollama's LLaMA2 API.
-    Only calls Ollama if the server is running.
-    """
-    if not OLLAMA_AVAILABLE:
-        print("⚠️ Ollama server is not available. Using fallback search.")
-        return None
-
-    try:
-        response = ollama.embeddings(model="llama2:latest", prompt=text)
-        return np.array(response["embedding"])
-    except Exception as e:
-        print(f"⚠️ Ollama embedding failed: {e}")
-        return None
-
-def get_flight_embedding(flight_number):
-    """
-    Retrieve or generate an embedding for a flight.
-    """
-    if flight_number in flight_embeddings:
-        return flight_embeddings[flight_number]
-
-    # Get flight details
-    flight = next((f for f in flights if f["flight_number"] == flight_number), None)
-    if not flight:
-        return None
-
-    text = f"{flight['origin']} {flight['destination']} {flight['airline']} {flight['time']}"
-    embedding = generate_embedding(text)
-    if embedding is not None:
-        flight_embeddings[flight_number] = embedding
-    return embedding
 
-def search_flights_semantic(query):
-    """
-    Perform semantic search on flight records using cosine similarity.
-    """
-    print(f"🟢 Searching flights for: {query}")  # Debug print
-
-    try:
-        query_embedding = generate_embedding(query)
-
-        similarities = {}
-        for flight in flights:
-            flight_num = flight["flight_number"]
-            flight_embedding = get_flight_embedding(flight_num)
-
-            if flight_embedding is not None:
-                similarity_score = cosine_similarity([query_embedding], [flight_embedding])[0][0]
-                similarities[flight_num] = similarity_score
-
-        matching_flights = [flight for flight in flights if similarities.get(flight["flight_number"], 0) > 0.6]
-
-        print(f"🟢 Flights Found: {matching_flights}")  # Debug print
-
-        return matching_flights
+# Mock database: list of flights
+flights = [
+    {"flight_number": "NY100", "origin": "New York", "destination": "London", "time": "2025-05-01 08:00", "airline": "Global Airways"},
+    {"flight_number": "LA200", "origin": "Los Angeles", "destination": "Tokyo", "time": "2025-05-01 10:30", "airline": "Pacific Routes"},
+    {"flight_number": "CH300", "origin": "Chicago", "destination": "Paris", "time": "2025-05-01 15:45", "airline": "Euro Connect"},
+    {"flight_number": "SF400", "origin": "San Francisco", "destination": "Sydney", "time": "2025-05-01 23:15", "airline": "Ocean Pacific"},
+    {"flight_number": "MI500", "origin": "Miami", "destination": "Rio de Janeiro", "time": "2025-05-02 07:30", "airline": "South American Airways"}
+]
 
-    except Exception as e:
-        print(f"❌ Error in search_flights_semantic: {str(e)}")  # Debug print
-        return []
+def search_flights(query):
+    """Search for flights based on keywords in the query."""
+    query_lower = query.lower()
+    
+    # Match flights by checking if the query is present in any flight attribute
+    matches = [
+        flight for flight in flights 
+        if query_lower in flight["origin"].lower()
+        or query_lower in flight["destination"].lower()
+        or query_lower in flight["flight_number"].lower()
+        or query_lower in flight["airline"].lower()
+    ]
+    
+    return matches
+
+if __name__ == "__main__":
+    # Example search
+    result = search_flights("Chicago")
+    print(result)