Containerization_done

RobuRishabh · RobuRishabh · commit b32a008b1273 · 2025-03-07T15:54:22.000-05:00
diff --git a/mock_database.py b/mock_database.py
@@ -26,19 +26,38 @@ def check_ollama_availability():
     {"flight_number": "MI500", "origin": "Miami", "destination": "Rio de Janeiro", "time": "2025-05-02 07:30", "airline": "South American Airways"}
 ]
 
-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
+def search_flights(origin=None, destination=None, flight_number=None, airline=None):
+    """
+    Search for flights based on exact matches for origin, destination, flight number, or airline.
+    Ensures that at least one valid filter is applied.
+    """
+    print(f"🔍 Searching for: Origin={origin}, Destination={destination}, Flight Number={flight_number}, Airline={airline}")
+
+    # ✅ If flight number is provided, prioritize searching by flight number only
+    if flight_number:
+        matches = [flight for flight in flights if flight["flight_number"].lower() == flight_number.lower()]
+        print(f"🔎 Flight number search results: {matches}")
+        return matches
+
+    # ✅ If no flight number, apply standard search
+    if not any([origin, destination, airline]):
+        print("⚠️ No valid search parameters provided. Returning an empty list.")
+        return []
+
+    # Normalize input (convert to lowercase)
+    origin = origin.lower() if origin else None
+    destination = destination.lower() if destination and destination != "city name" else None  # Ignore "City Name"
+    airline = airline.lower() if airline else None
+
+    # Filter flights with exact matching for provided fields
     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()
+        flight for flight in flights
+        if (not origin or flight["origin"].lower() == origin)
+        and (not destination or flight["destination"].lower() == destination)
+        and (not airline or flight["airline"].lower() == airline)
     ]
-    
+
+    print(f"🔎 Found flights: {matches}")
     return matches
 
 if __name__ == "__main__":
diff --git a/ollama_api.py b/ollama_api.py
@@ -1,23 +1,15 @@
 import os
 import json
 import requests
-from langchain_ollama import OllamaLLM
+from langchain_ollama import OllamaLLM  # Correct import
 from typing import Tuple, List
 from dotenv import load_dotenv
 
-# Load environment variables
 load_dotenv()
-
-# Set default Ollama settings
 OLLAMA_URL = os.getenv("OLLAMA_URL", "http://localhost:11434")
-OLLAMA_MODEL = os.getenv("OLLAMA_MODEL", "llama2:latest")  # Use default if not set
-
-if not OLLAMA_MODEL:
-    print("❌ OLLAMA_MODEL is missing! Using default model: 'llama2:latest'.")
+OLLAMA_MODEL = os.getenv("OLLAMA_MODEL", "llama2:latest")
 
-# Initialize Ollama LLM
 def initialize_ollama():
-    """Initialize the Ollama LLM model safely."""
     try:
         ollama_llm = OllamaLLM(model=OLLAMA_MODEL, base_url=OLLAMA_URL)
         print(f"🟢 Successfully initialized Ollama LLM with model: {OLLAMA_MODEL}")
@@ -29,75 +21,43 @@ def initialize_ollama():
 ollama_llm = initialize_ollama()
 
 def check_ollama_availability() -> Tuple[bool, str]:
-    """Check if the locally running Ollama server is available."""
     try:
         response = requests.get(f"{OLLAMA_URL}/api/tags", timeout=3)
-        if response.status_code == 200:
-            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, "❌ Cannot connect to Ollama server. Ensure it is running."
-    except Exception as e:
-        return False, f"⚠️ Error checking Ollama server: {str(e)}"
+        return response.status_code == 200, f"🟢 Ollama server is available at {OLLAMA_URL}"
+    except requests.RequestException as e:
+        return False, f"⚠️ Ollama server is not available at {OLLAMA_URL}: {str(e)}"
 
 def generate_fallback_response(query: str, flights: List[dict]) -> str:
-    """Generate a simple response when Ollama is not available."""
     if not flights:
-        return "I couldn't find any flights matching your criteria. Please try again with different details."
-
+        return "I couldn't find any flights matching your criteria. Please try again."
     response = "Here are the flights that match your search:\n\n"
     for flight in flights:
         response += (
-            "✈️ Flight {flight_number} from {origin} to {destination}\n"
-            "⏰ Time: {time}\n"
-            "🏢 Airline: {airline}\n\n"
-        ).format(
-            flight_number=flight.get("flight_number", "Unknown"),
-            origin=flight.get("origin", "Unknown"),
-            destination=flight.get("destination", "Unknown"),
-            time=flight.get("time", "N/A"),
-            airline=flight.get("airline", "N/A")
+            f"✈️ Flight {flight.get('flight_number', 'Unknown')} from {flight.get('origin', 'Unknown')} to {flight.get('destination', 'Unknown')}\n"
+            f"⏰ Time: {flight.get('time', 'N/A')}\n"
+            f"🏢 Airline: {flight.get('airline', 'N/A')}\n\n"
         )
     return response.strip()
 
 def generate_response(query: str, flights: List[dict]) -> str:
-    """Generate a natural language response using the LangChain Ollama integration."""
     is_available, message = check_ollama_availability()
-    if not is_available:
-        print(f"⚠️ {message}")
-        return generate_fallback_response(query, flights)
-
-    if not ollama_llm:
-        print("⚠️ Ollama model is not initialized. Using fallback response.")
+    if not is_available or not ollama_llm:
+        print(f"⚠️ {message if not is_available else 'Ollama model not initialized'}")
         return generate_fallback_response(query, flights)
 
     try:
-        if flights:
-            flight_info = json.dumps(flights, indent=2)
-            prompt = f"""
-            User Query: {query}
-
-            Available Flights:
-            {flight_info}
-
-            Generate a natural language response summarizing these flights, including flight number, time, and airline details.
-            """
-        else:
-            prompt = f"""
-            User Query: {query}
-
-            No matching flights found. Please provide a response indicating this politely.
-            """
-
+        flight_info = json.dumps(flights, indent=2) if flights else "No matching flights found."
+        prompt = f"""
+        User Query: {query}
+        Available Flights: {flight_info}
+        Generate a natural language response summarizing these flights, including flight number, time, and airline details if available, or politely indicate no flights were found.
+        """
         print("🟢 Sending prompt to Ollama for response generation...")
         response = ollama_llm.invoke(prompt)
         return response.strip() if response else generate_fallback_response(query, flights)
-
     except Exception as e:
         print(f"⚠️ Ollama LLM generation failed: {str(e)}")
-        return generate_fallback_response(query, flights) + f" (⚠️ Error: {str(e)})"
+        return generate_fallback_response(query, flights)
 
 # Test
 if __name__ == "__main__":
diff --git a/query_handler.py b/query_handler.py
@@ -1,90 +1,159 @@
-import ollama
 import json
 import os
 import re
 from dotenv import load_dotenv
+from langchain_ollama import OllamaLLM
 from mock_database import search_flights, check_ollama_availability
 
 # Load environment variables
 load_dotenv()
-OLLAMA_MODEL = os.getenv("OLLAMA_MODEL", "llama2:latest")  # Load model dynamically
-
+OLLAMA_MODEL = os.getenv("OLLAMA_MODEL", "llama2:latest")
 OLLAMA_AVAILABLE = check_ollama_availability()
 
-def extract_entities_llama2(query):
+# Initialize Ollama LLM
+def initialize_ollama():
+    """Initialize the Ollama LLM model safely."""
+    try:
+        ollama_llm = OllamaLLM(model=OLLAMA_MODEL)
+        print(f"🟢 Successfully initialized Ollama LLM with model: {OLLAMA_MODEL}")
+        return ollama_llm
+    except Exception as e:
+        print(f"❌ Failed to initialize Ollama LLM: {str(e)}")
+        return None
+
+ollama_llm = initialize_ollama()
+
+CITY_MAPPING = {
+    "ny": "New York",
+    "la": "Los Angeles",
+    "ch": "Chicago",
+    "sf": "San Francisco",
+    "mi": "Miami"
+}
+
+def extract_entities_ollama(query):
     """
-    Uses Ollama to extract structured flight details from a query.
-    Ensures the response is always valid JSON.
+    Uses Ollama to extract structured flight details from a query and ensures correct data mapping.
+    If Ollama fails to extract an entity, fallback to a keyword-based search.
     """
-    if not OLLAMA_AVAILABLE:
+    if not OLLAMA_AVAILABLE or not ollama_llm:
         print("⚠️ Ollama server is unavailable. Using basic keyword search.")
         return {}
 
     print(f"🟢 Using Ollama model: {OLLAMA_MODEL}")
-    prompt = f'''
-    Extract flight details from this user query and respond **only** in valid JSON.
-    Do not include any extra text, explanations, or markdown.
-    
+    prompt = f"""
+    Extract flight details from the following user query and return only valid JSON.
+    Do not include any explanations, additional text, or markdown.
+
     Query: "{query}"
-    
-    Example valid responses:
-    {{"origin": "New York", "destination": "London", "flight_number": "AA123", "time": "10:00 AM", "date": "2025-05-01", "airline": "Global Airways"}}
-    {{"origin": "Chicago", "destination": null, "flight_number": null, "time": null, "date": null, "airline": null}}
+
+    The response should be in this exact JSON format:
+    {{
+      "origin": "City Name",
+      "destination": "City Name",
+      "flight_number": "Flight Number",
+      "date": "YYYY-MM-DD",
+      "airline": "Airline Name"
+    }}
 
     If a value is missing, set it to `null`.
-    
-    '''
+    """
 
     try:
-        print(f"🟢 Sending request to Ollama with model: {OLLAMA_MODEL}")
-        response = ollama.chat(model=OLLAMA_MODEL.strip(), messages=[{"role": "user", "content": prompt}])
-        content = response["message"]["content"].strip()
+        print(f"🟢 Sending request to Ollama for entity extraction...")
+        response = ollama_llm.invoke(prompt)
 
-        # 🔍 Extract the valid JSON from the response
-        json_match = re.search(r"\{.*\}", content, re.DOTALL)
+        # Extract valid JSON from response
+        json_match = re.search(r"\{.*\}", response, re.DOTALL)
         if json_match:
-            json_str = json_match.group(0)  # Extract JSON part
-            extracted = json.loads(json_str)  # Convert to dictionary
+            json_str = json_match.group(0)
+            extracted = json.loads(json_str)
 
-            # 🛠️ Clean extracted data (remove placeholders)
-            extracted_clean = {k: v for k, v in extracted.items() if v and v.lower() not in ["null", "none"]}
+            # ✅ Clean extracted values
+            if extracted.get("destination") == "City Name":
+                extracted["destination"] = None  # Ignore placeholder values
+            if extracted.get("origin") == "NYC":
+                extracted["origin"] = "New York"  # Fix Ollama misinterpretation
 
-            print(f"🟢 Extracted Entities: {extracted_clean}")
+            # ✅ If no flight number is extracted, try regex
+            if not extracted.get("flight_number"):
+                extracted["flight_number"] = extract_flight_number(query)
+
+            extracted_clean = {k: v for k, v in extracted.items() if v}  # Remove None values
+            print(f"🟢 Extracted Entities from Ollama: {extracted_clean}")
             return extracted_clean
 
         else:
-            print(f"⚠️ No valid JSON found in response. Raw response: {content}")
-            return {}
+            print(f"⚠️ No valid JSON found in response. Falling back to keyword search.")
+            return extract_entities_from_keywords(query)
 
     except json.JSONDecodeError as jde:
-        print(f"⚠️ JSONDecodeError: {jde}. Raw response: {content}")
-        return {}
+        print(f"⚠️ JSONDecodeError: {jde}. Falling back to keyword search.")
+        return extract_entities_from_keywords(query)
     except Exception as e:
-        print(f"⚠️ Error during entity extraction: {e}")
-        return {}
+        print(f"⚠️ Error during entity extraction: {e}. Falling back to keyword search.")
+        return extract_entities_from_keywords(query)
+
+def extract_flight_number(query):
+    """
+    Extracts a flight number from a query using regex.
+    Flight numbers typically have a two-letter airline code followed by 3-4 digits (e.g., "NY100").
+    """
+    match = re.search(r"\b[A-Z]{2}\d{3,4}\b", query)
+    return match.group(0) if match else None
+
+
+def extract_entities_from_keywords(query):
+    """
+    Fallback function to extract entities from a query using simple keyword matching.
+    """
+    keywords = query.lower()
+
+    cities = ["new york", "los angeles", "chicago", "san francisco", "miami", "paris", "tokyo", "london", "rio de janeiro", "sydney"]
+    airlines = ["global airways", "pacific routes", "euro connect", "ocean pacific", "south american airways"]
+
+    origin = next((city for city in cities if city in keywords), None)
+    destination = None  # Can't always determine destination reliably from keywords
+    flight_number_match = re.search(r"\b[A-Z]{2}\d{3,4}\b", query)  # Match flight numbers like "NY100"
+    flight_number = flight_number_match.group(0) if flight_number_match else None
+    airline = next((air for air in airlines if air in keywords), None)
+
+    extracted = {
+        "origin": origin,
+        "destination": destination,
+        "flight_number": flight_number,
+        "airline": airline
+    }
+
+    extracted_clean = {k: v for k, v in extracted.items() if v}  # Remove None values
+    print(f"🟢 Extracted Entities from Keywords: {extracted_clean}")
+    return extracted_clean
 
 
 def process_query(query):
     """
     Process user query and return relevant flight information.
-    Uses keyword search if Ollama is unavailable.
+    Uses Ollama for entity extraction instead of Transformers.
     """
     try:
         print(f"🟢 Processing query: {query}")
 
-        # Extract structured entities (if Ollama is available)
-        search_params = extract_entities_llama2(query) if OLLAMA_AVAILABLE else {}
+        # Extract structured entities using Ollama
+        search_params = extract_entities_ollama(query)
 
-        # Perform a keyword-based search instead of semantic search
-        if search_params:
-            print(f"🟢 Searching with extracted parameters: {search_params}")
-            matching_flights = search_flights(search_params.get("origin", "") or search_params.get("destination", ""))
-        else:
-            print("⚠️ No structured entities found. Performing direct keyword search.")
-            matching_flights = search_flights(query)
+        # Ensure extracted values are correct before searching
+        origin = search_params.get("origin")
+        destination = search_params.get("destination")
+        flight_number = search_params.get("flight_number")
+        airline = search_params.get("airline")
+
+        print(f"🟢 Searching with extracted parameters: {search_params}")
+
+        # Use extracted details for searching
+        matching_flights = search_flights(origin, destination, flight_number, airline)
 
         if not matching_flights:
-            return False, "No flights found matching your criteria.", []
+            return False, "⚠️ No flights found matching your criteria. Please try again with different details.", []
 
         return True, "Here are the flights that match your criteria:", matching_flights
 
@@ -95,7 +164,9 @@ def process_query(query):
         print(f"❌ Unexpected error in process_query: {str(e)}")
         return False, f"An error occurred while processing your query: {str(e)}", []
 
+
+
 # Test
 if __name__ == "__main__":
-    success, message, flights = process_query("Show me flight NY100")
+    success, message, flights = process_query("Show me flights from New York to London on May 1st")
     print(f"Success: {success}, Message: {message}, Flights: {flights}")
