Ollama_api_testing

RobuRishabh · RobuRishabh · commit fb69f1b69f3a · 2025-03-03T22:54:23.000-05:00
diff --git a/mock_database.py b/mock_database.py
@@ -0,0 +1,56 @@
+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"
+    }
+]
+
+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
diff --git a/ollama_api.py b/ollama_api.py
@@ -0,0 +1,95 @@
+"""
+Integration with the Ollama API for natural language generation.
+"""
+import os
+import json
+import requests
+from dotenv import load_dotenv
+from typing import Tuple, List
+
+load_dotenv()
+
+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")
+
+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)    
+    """
+    try:
+        response = requests.get(f"{OLLAMA_URL}/api/tags")
+        if response.status_code == 200:
+            return True, "Ollama server is available."
+        return False, "Ollama server is not available. status code: {response.status_code}"
+    except requests.exceptions.ConnectionError:
+        return False, "Ollama server is not available. Connection error."
+    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:
+    """
+    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
+    """
+    if not flights:
+        return "No flights found matching your criteria."
+    
+    response = "Here are the flights that match your criteria:\n"
+    for flight in flights:
+        response += f"Flight {flight['flight_number']} from {flight['origin']} to {flight['destination']} on {flight['time']} by {flight['airline']}\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
+    """
+    # Check is ollama server 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
+        if flights:
+            flight_info = json.dumps(flights, indent = 2)
+            prompt = f"""
+            User Question: {query}
+
+            Available Flights Information:
+            {flight_info}
+            Please provide a concise, natural language response that summarizes these flights, including flight numbers, times, and destinations.
+            """
+        else:
+            prompt = f"""
+            User Question: {query}
+
+            No flights found matching your criteria.
+            """
+        # 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)}"
+    except Exception as e:
+        return generate_fallback_response(query, flights) + f" Unexpected error: {str(e)}"
diff --git a/query_handler.py b/query_handler.py
@@ -0,0 +1,63 @@
+"""
+Query Handler to process the queries from the user and etract the required information from the database.
+"""
+from mock_database import search_flights
+import re
+
+def parse_query(query):
+    """"
+    Parse user query to extract relevant search parameters
+
+    Args: query (str): User query
+    Returns: dict: Query parameters
+    """
+    query = query.lower()
+    search_params = {}
+
+    # Check for origin city
+    origin_indicator = ["from", "originating", "origin", "out of", "departing from", "leaving from"]
+    for indicator in origin_indicator:
+        if indicator in query:
+            parts = query.split(indicator, 1) # maximum split of 1
+            if len(parts) > 1:
+                potential_origin = parts[1].split()[0:2].strip()
+                search_params["origin"] = " ".join(potential_origin).title()
+
+    # Check for destination city
+    dest_indicators = ["to", "arriving at", "going to"]
+    for indicator in dest_indicators:
+        if indicator in query:
+            parts = query.split(indicator, 1)
+            if len(parts) > 1:
+                potential_dest = parts[1].split()[0:2]  # Assuming 2-word city names
+                search_params["destination"] = " ".join(potential_dest).title()
+                
+    # Check for flight number (even if 'flight' isn't explicitly mentioned)
+    flight_number_match = re.search(r'\b[A-Z]{2,4}\d{3,4}\b', query)  # Match patterns like 'NY100' or 'LA200'
+    if flight_number_match:
+        search_params["flight_number"] = flight_number_match.group(0).upper()
+        
+    return search_params
+
+def process_query(query):
+    """
+    Process user query and return relevant flight information
+    
+    Args: query (str): User's question
+    Returns: tuple: (bool success, str message, list matching_flights)
+    """
+    try:
+        search_params = parse_query(query)
+        
+        if not search_params:
+            return False, "I couldn't understand the flight details in your question. Please try asking about specific origins, destinations, or flight numbers.", []
+            
+        matching_flights = search_flights(search_params)
+        
+        if not matching_flights:
+            return False, "No flights found matching your criteria.", []
+            
+        return True, "Found matching flights!", matching_flights
+        
+    except Exception as e:
+        return False, f"An error occurred while processing your query: {str(e)}", []
diff --git a/test_ollama.py b/test_ollama.py
@@ -0,0 +1,16 @@
+from ollama_api import check_ollama_availability, generate_response
+
+# Test Ollama server availability
+is_available, message = check_ollama_availability()
+print(f"Server Availability: {message}")
+
+# Example query and flight data to test generate_response function
+query = "What are the flights from New York to London?"
+flights = [
+    {"flight_number": "NY100", "origin": "New York", "destination": "London", "time": "2025-05-01 08:00", "airline": "Global Airways"},
+    {"flight_number": "NY200", "origin": "New York", "destination": "London", "time": "2025-05-02 12:00", "airline": "Air London"}
+]
+
+# Test generating a response based on the query and flights
+response = generate_response(query, flights)
+print(f"Generated Response: \n{response}")
