UI_Langchain_SemanticSearch_added

RobuRishabh · RobuRishabh · commit 0c6bc59b6b41 · 2025-03-05T01:32:57.000-05:00
diff --git a/app.py b/app.py
@@ -0,0 +1,75 @@
+"""
+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
+st.set_page_config(
+    page_title="Flight Information Assistant",
+    page_icon="✈️"
+)
+
+# ✅ Initialize session state for chat history
+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.
+    """
+    with st.chat_message(role):
+        st.markdown(content)
+
+# ✅ Main UI
+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.")
+
+# ✅ Instructions for the user
+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?
+""")
+
+# ✅ 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)
+
+    try:
+        # ✅ Process user query
+        success, message, flights = process_query(prompt)
+
+        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)
+
+        # ✅ Store and display assistant response
+        st.session_state.messages.append({"role": "assistant", "content": response})
+        display_chat_message("assistant", response)
+
+    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)
diff --git a/mock_database.py b/mock_database.py
@@ -1,3 +1,7 @@
+import ollama
+import numpy as np
+from sklearn.metrics.pairwise import cosine_similarity
+
 flights = [
     {
         "flight_number": "NY100",
@@ -36,21 +40,82 @@
     }
 ]
 
-def search_flights(query_params):
-    """
-    Search for flights based on query parameters
-    """
-    results = []
-    for flight in flights:
-        matches = True
-        for key, value in query_params.items():
-            if key not in flight:
-                matches = False
-                break
-            # Handle the exact matches 
-            if value.lower() != flight[key].lower():
-                matches = False
-                break
-        if matches:
-            results.append(flight)
-    return results
+### ✅ 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
+    try:
+        response = requests.get("http://localhost:11434/api/tags", timeout=3)
+        return response.status_code == 200
+    except requests.exceptions.RequestException:
+        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
+
+    except Exception as e:
+        print(f"❌ Error in search_flights_semantic: {str(e)}")  # Debug print
+        return []
diff --git a/ollama_api.py b/ollama_api.py
@@ -1,95 +1,104 @@
-"""
-Integration with the Ollama API for natural language generation.
-"""
 import os
 import json
 import requests
-from dotenv import load_dotenv
+from langchain_ollama import OllamaLLM
 from typing import Tuple, List
+from dotenv import load_dotenv
 
+# ✅ Load environment variables
 load_dotenv()
 
+print(f"🟢 Loaded OLLAMA_MODEL from .env: {os.getenv('OLLAMA_MODEL')}")  # Debugging check
+
+
+# ✅ Use environment variables dynamically
 OLLAMA_URL = os.getenv("OLLAMA_URL")
-if not OLLAMA_URL:
-    raise ValueError("OLLAMA_URL environment variable is not set.")
-MODEL = os.getenv("OLLAMA_MODEL", "llama2")
+OLLAMA_MODEL = os.getenv("OLLAMA_MODEL")
+if not OLLAMA_MODEL:
+    print("❌ OLLAMA_MODEL is missing! Check your .env file.")
+    exit(1)
+
+# ✅ Initialize Ollama LLM
+ollama_llm = OllamaLLM(model=OLLAMA_MODEL)
 
 def check_ollama_availability() -> Tuple[bool, str]:
     """
-    Check if the Ollama server is available by querying its health endpoint.
-    Returns: Tuple[bool, str]: (success, message)    
+    Check if the locally running Ollama server is available.
+    
+    Returns:
+        Tuple[bool, str]: (is_available, message)
     """
     try:
-        response = requests.get(f"{OLLAMA_URL}/api/tags")
+        response = requests.get(f"{OLLAMA_URL}/api/tags", timeout=3)  # ✅ Use .env URL and set timeout
         if response.status_code == 200:
-            return True, "Ollama server is available."
-        return False, "Ollama server is not available. status code: {response.status_code}"
+            return True, "Ollama server is available"
+        return False, f"Ollama server returned status code: {response.status_code}"
+    except requests.exceptions.Timeout:
+        return False, "⚠️ Ollama server timeout. It may be overloaded or down."
     except requests.exceptions.ConnectionError:
-        return False, "Ollama server is not available. Connection error."
+        return False, "❌ Cannot connect to Ollama server. Ensure it is running."
     except Exception as e:
-        return False, f"An error occurred while checking Ollama availability: {str(e)}"
-    
-def generate_fallback_response(query:str, flights:List[dict]) -> str:
+        return False, f"⚠️ Error checking Ollama server: {str(e)}"
+
+def generate_fallback_response(query: str, flights: List[dict]) -> str:
     """
     Generate a simple response when Ollama is not available.
-    Args: query (str): User's question
-          flights (List[dict]): List of matching flights
-    Returns: str: Fallback response
+    
+    Args:
+        query (str): User's original question.
+        flights (list): Matching flight information.
+    
+    Returns:
+        str: Fallback response.
     """
     if not flights:
-        return "No flights found matching your criteria."
-    
-    response = "Here are the flights that match your criteria:\n"
+        return "I couldn't find any flights matching your criteria. Please try again with different details."
+
+    response = "Here are the flights that match your search:\n\n"
     for flight in flights:
-        response += f"Flight {flight['flight_number']} from {flight['origin']} to {flight['destination']} on {flight['time']} by {flight['airline']}\n"
+        response += (f"Flight {flight['flight_number']} from {flight['origin']} to {flight['destination']}\n"
+                    f"Time: {flight['time']}\n"
+                    f"Airline: {flight['airline']}\n\n")
     return response
- 
+
 def generate_response(query: str, flights: List[dict]) -> str:
     """
-    Generate a natural language response using the Ollama model.
-    Args: query (str): User's question
-            flights (List[dict]): List of matching flights
-    Returns: str: Natural language response
+    Generate a natural language response using the LangChain Ollama integration.
+
+    Args:
+        query (str): User's original question.
+        flights (list): Matching flight information.
+
+    Returns:
+        str: Generated response.
     """
-    # Check is ollama server is available
+    # ✅ First check if Ollama is available
     is_available, message = check_ollama_availability()
     if not is_available:
         return generate_fallback_response(query, flights)
+
     try:
-        # Prepare the prompt for Ollama model
+        # ✅ Prepare flight information for LLM
         if flights:
-            flight_info = json.dumps(flights, indent = 2)
+            flight_info = json.dumps(flights, indent=2)
             prompt = f"""
-            User Question: {query}
+            User Query: {query}
 
-            Available Flights Information:
+            Available Flights:
             {flight_info}
-            Please provide a concise, natural language response that summarizes these flights, including flight numbers, times, and destinations.
+
+            Generate a natural language response summarizing these flights, including flight number, time, and airline details.
             """
         else:
             prompt = f"""
-            User Question: {query}
+            User Query: {query}
 
-            No flights found matching your criteria.
+            No matching flights found. Please provide a response indicating this politely.
             """
-        # Make the api call to Ollama
-        response = requests.post(f"{OLLAMA_URL}/api/generate", 
-                                 json={
-                                    "model": MODEL, 
-                                    "prompt": prompt,
-                                    "stream": False}, 
-                                    timeout=10 # Timeout after 10 seconds
-                                )
-        if response.status_code == 200:
-            try:
-                response_data = response.json()
-                return response_data.get("response", "Sorry, I couldn't generate a response.")
-            except ValueError:
-                return "There was an error processing the response from the server."
-        else:
-            return generate_fallback_response(query, flights)
 
-    except requests.exceptions.RequestException as e:
-        return generate_fallback_response(query, flights) + f" Error: {str(e)}"
+        # ✅ Generate response using LangChain's Ollama LLM
+        response = ollama_llm(prompt)
+        return response.strip()
+
     except Exception as e:
-        return generate_fallback_response(query, flights) + f" Unexpected error: {str(e)}"
+        return generate_fallback_response(query, flights) + f" (⚠️ Error: {str(e)})"
diff --git a/query_handler.py b/query_handler.py
