from pydantic import BaseModel
from typing import Optional, List
import json
from openai import OpenAI
from dotenv import load_dotenv
import os
load_dotenv()
# Pydantic models for structured data
class FlightSearchParams(BaseModel):
departure: str
arrival: str
date: Optional[str] = None
class FlightDetails(BaseModel):
flight_number: str
departure: str
arrival: str
departure_time: str
arrival_time: str
price: float
available_seats: int
class ChatbotResponse(BaseModel):
flights: List[FlightDetails]
natural_response: str
class FlightChatbot:
def __init__(self, api_key: str):
self.client = OpenAI(
api_key=api_key,
base_url="https://oai.helicone.ai/v1",
default_headers= {
"Helicone-Auth": f"Bearer {os.getenv('HELICONE_API_KEY')}"
})
self.flights_db = [
{
"flight_number": "BA123",
"departure": "New York",
"arrival": "London",
"departure_time": "2025-01-15T08:30:00",
"arrival_time": "2025-01-15T20:45:00",
"price": 650.00,
"available_seats": 45
},
{
"flight_number": "AA456",
"departure": "London",
"arrival": "New York",
"departure_time": "2025-01-16T10:15:00",
"arrival_time": "2025-01-16T13:30:00",
"price": 720.00,
"available_seats": 12
}
]
def _search_flights(self, departure: str, arrival: str, date: Optional[str] = None) -> List[dict]:
"""Search for flights using the provided parameters."""
matches = []
for flight in self.flights_db:
if (flight["departure"].lower() == departure.lower() and
flight["arrival"].lower() == arrival.lower()):
if date:
flight_date = flight["departure_time"].split("T")[0]
if flight_date == date:
matches.append(flight)
else:
matches.append(flight)
return matches
def process_query(self, user_query: str) -> str:
"""Process a user query and return flight information."""
try:
# First, use function calling to extract parameters
parameter_extraction = self.client.chat.completions.create(
model="gpt-4o-2024-08-06",
messages=[
{
"role": "system",
"content": "You are a flight search assistant. Extract search parameters from user queries."
},
{
"role": "user",
"content": user_query
}
],
tools=[{
"type": "function",
"function": {
"name": "search_flights",
"description": "Search for flights based on departure and arrival cities, and optionally a date",
"parameters": {
"type": "object",
"properties": {
"departure": {
"type": "string",
"description": "Departure city"
},
"arrival": {
"type": "string",
"description": "Arrival city"
},
"date": {
"type": "string",
"description": "Flight date in YYYY-MM-DD format",
"format": "date"
}
},
"required": ["departure", "arrival"]
}
}
}],
tool_choice={"type": "function", "function": {"name": "search_flights"}}
)
# Extract parameters from function call
function_args = json.loads(parameter_extraction.choices[0].message.tool_calls[0].function.arguments)
# Search for flights
found_flights = self._search_flights(
departure=function_args["departure"],
arrival=function_args["arrival"],
date=function_args.get("date")
)
# Use parse helper to generate structured response with natural language
response = self.client.beta.chat.completions.parse(
model="gpt-4o-2024-08-06",
messages=[
{
"role": "system",
"content": """You are a flight search assistant. Generate a response containing:
1. A list of structured flight details
2. A natural language response explaining the search results
For the natural language response:
- Be concise and helpful
- Include key details like flight numbers, times, and prices
- If no flights are found, explain why and suggest alternatives"""
},
{
"role": "user",
"content": f"Original query: {user_query}\nFound flights: {json.dumps(found_flights, indent=2)}"
}
],
response_format=ChatbotResponse
)
return response.choices[0].message
except Exception as e:
error_response = ChatbotResponse(
flights=[],
natural_response=f"I apologize, but I encountered an error processing your request: {str(e)}"
)
return error_response.model_dump_json(indent=2)
def main():
# Initialize chatbot with your API key
chatbot = FlightChatbot(os.getenv('OPENAI_API_KEY'))
# Example queries
example_queries = [
"When is the next flight from New York to London?",
"Find me flights from London to New York on January 16, 2025",
"Are there any flights from Paris to Tokyo tomorrow?"
]
for query in example_queries:
print(f"User Query: {query}")
response = chatbot.process_query(query)
print("\nResponse:")
print(response.refusal or response.parsed)
print("-" * 50 + "\n")
if __name__ == "__main__":
main()