added tools for searching/booking/add-ons for flights

src/airline_agent/backend/schemas/message.py

@@ -1,5 +1,3 @@
-from __future__ import annotations
-
 from enum import StrEnum
 from typing import Literal
 

src/airline_agent/backend/services/airline_chat.py

@@ -47,6 +47,7 @@
     run_cleanlab_validation_logging_tools,
 )
 from airline_agent.constants import AGENT_INSTRUCTIONS, AGENT_MODEL
+from airline_agent.tools.booking import BookingTools
 from airline_agent.tools.knowledge_base import KnowledgeBase
 
 load_dotenv()
@@ -55,13 +56,14 @@
 logger.setLevel(logging.INFO)
 
 
-def create_agent(kb: KnowledgeBase) -> Agent:
+def create_agent(kb: KnowledgeBase, booking: BookingTools) -> Agent:
     """Create the airline support agent."""
     model = OpenAIChatModel(model_name=AGENT_MODEL, settings=ModelSettings(temperature=0.0))
+
     return Agent(
         model=model,
         instructions=AGENT_INSTRUCTIONS,
-        tools=[kb.get_article, kb.search, kb.list_directory],
+        toolsets=[kb.tools, booking.tools],
     )
 
 
@@ -78,8 +80,9 @@ def get_cleanlab_project() -> Project:
     kb_path=str(pathlib.Path(__file__).parents[4] / "data/kb.json"),
     vector_index_path=str(pathlib.Path(__file__).parents[4] / "data/vector-db"),
 )
+booking = BookingTools()
 project = get_cleanlab_project()
-agent = create_agent(kb)
+agent = create_agent(kb, booking)
 
 thread_to_messages: dict[str, list[ModelMessage]] = {}
 cleanlab_enabled_by_thread: dict[str, bool] = {}

src/airline_agent/cleanlab_utils/conversion_utils.py

@@ -1,18 +1,10 @@
 """Convert pydantic-ai message history and responses to OpenAI Chat Completions format."""
 
-from __future__ import annotations
-
 import base64
-from datetime import UTC, datetime
-from typing import TYPE_CHECKING, Any, Literal, cast
-
-from openai.types.chat import ChatCompletion
+from datetime import datetime
+from typing import Any, Literal, cast
 
-if TYPE_CHECKING:
-    from openai.types.chat import ChatCompletionMessageParam
-    from pydantic_ai.tools import ToolDefinition
-
-from openai.types.chat import ChatCompletionFunctionToolParam
+from openai.types.chat import ChatCompletion, ChatCompletionFunctionToolParam, ChatCompletionMessageParam
 from pydantic_ai.messages import (
     AudioUrl,
     BinaryContent,
@@ -32,8 +24,11 @@
     UserPromptPart,
     VideoUrl,
 )
+from pydantic_ai.tools import ToolDefinition
 from pydantic_ai.usage import RequestUsage
 
+from airline_agent.constants import DEMO_DATETIME
+
 
 def convert_to_openai_messages(message_history: list[ModelMessage]) -> list[ChatCompletionMessageParam]:
     """Convert pydantic-ai message history to OpenAI Chat Completions format."""
@@ -241,7 +236,7 @@ def convert_message_to_chat_completion(message: ChatCompletionMessageParam) -> C
                     "message": choice_message,
                 }
             ],
-            "created": int(datetime.now(UTC).timestamp()),
+            "created": int(DEMO_DATETIME.timestamp()),
             "model": "mock-agent",
             "object": "chat.completion",
             "service_tier": "default",
@@ -284,7 +279,7 @@ def convert_string_to_response_message(
     model_name = None
 
     if timestamp is None:
-        timestamp = datetime.now(UTC)
+        timestamp = DEMO_DATETIME
     text_part = TextPart(content=content)
     usage = RequestUsage(input_tokens=0, output_tokens=0)
     return ModelResponse(

src/airline_agent/cleanlab_utils/validate_utils.py

@@ -1,32 +1,28 @@
-from __future__ import annotations
-
 import logging
 import re
 import warnings
-from typing import TYPE_CHECKING, Any, cast
-
-if TYPE_CHECKING:
-    from cleanlab_codex import Project
-    from codex.types.project_validate_response import ProjectValidateResponse
-    from pydantic_ai.agent import Agent, AgentRunResult
-    from pydantic_ai.tools import ToolDefinition
+from typing import Any, cast
 
+from cleanlab_codex import Project
 from cleanlab_tlm.utils.chat import _ASSISTANT_PREFIX as ASSISTANT_PREFIX
 from cleanlab_tlm.utils.chat import (
     _form_prompt_chat_completions_api as form_prompt_chat_completions_api,
 )
+from codex.types.project_validate_response import ProjectValidateResponse
 from openai.types.chat import (
     ChatCompletionAssistantMessageParam,
     ChatCompletionFunctionToolParam,
     ChatCompletionMessageParam,
 )
+from pydantic_ai.agent import Agent, AgentRunResult
 from pydantic_ai.messages import (
     ModelMessage,
     ModelRequest,
     ModelResponse,
     SystemPromptPart,
     UserPromptPart,
 )
+from pydantic_ai.tools import ToolDefinition
 
 from airline_agent.cleanlab_utils.conversion_utils import (
     convert_message_to_chat_completion,

src/airline_agent/constants.py

@@ -1,7 +1,14 @@
 import logging
+from datetime import date, datetime, timedelta
+from zoneinfo import ZoneInfo
 
 logger = logging.getLogger(__name__)
 
+DEMO_DATE = date(2025, 11, 5)
+DEMO_DATETIME = datetime(2025, 11, 5, 14, 0, 0, tzinfo=ZoneInfo("America/Los_Angeles"))
+FLIGHT_DATA_DATE = DEMO_DATE + timedelta(days=7)
+FLIGHT_DATA_NUM_DAYS = 8
+
 OFFICIAL_DEMO_PROJECT_ID = "3aae1f96-2dda-492f-8c86-17d453d3c298"  # to copy configuration from
 STAGING_DEMO_PROJECT_ID = "6de236e4-c6e7-456c-b248-872236010992"
 RAG_EMBED_MODEL = "text-embedding-3-small"
@@ -10,29 +17,41 @@
 RAG_CHUNK_OVERLAP = 200
 CONTEXT_RETRIEVAL_TOOLS = ["search", "get_article", "list_directory"]
 AGENT_MODEL = "gpt-4o"
-AGENT_INSTRUCTIONS = (
-    """You are an AI customer support agent for Frontier Airlines. You can use tools to access to a knowledge base of articles and
-documents about the airline's services, policies, and procedures.
+AGENT_INSTRUCTIONS = f"""
+You are an AI customer support agent for Frontier Airlines. You can use tools to access to a knowledge base of articles and documents about the airline's services, policies, and procedures.
 
 ## You have access to the following tools:
 - search — find candidate articles by query (keep top-k small, ≤5), returns title/snippet/path.
 - get_article — get the full article by its path.
 - list_directory — list directory structure to make more informed searches.
+- search_flights — search available flights by origin and destination airport codes (IATA) and departure date (YYYY-MM-DD). Always ask for the departure date if the user doesn't provide it.
+- get_fare_details — retrieve fare bundle pricing, included services, and add-ons for a specific flight.
+- book_flights — book one or more flights for the current user. Requires list of flight IDs and fare bundle type (basic, economy, premium, business; defaults to basic). Returns booking confirmation with booking ID and total price.
+- get_booking — retrieve booking details by booking ID.
+- get_my_bookings — retrieve all confirmed bookings for the current user.
+- add_service_to_booking — add an eligible service (bags, seat selection, etc.) to a specific flight within a booking.
+- check_in — complete check-in for a specific flight in a booking.
+- get_flight_timings — get check-in, boarding, and door-close timing windows for a flight.
+- get_flight_status — get the latest status, gates, and delay information for a flight.
 
 ## Tool Use Guidelines:
 - Keep it tight: aim for 1-2 calls per turn (hard cap 4).
 - Answer only from retrieved content.
 - If a missing detail blocks tool use, ask one short clarifying question. If not blocking, proceed and state your assumption.
 - Don't dump raw tool output—summarize clearly.
+- When booking multiple flights (outbound and return), include all flight IDs in a single book_flights call.
 
 ## Response Guidelines:
 - Answer questions based on information you look up in the knowledge base, not based on your own knowledge.
 - If you think that you need more time to investigate, update the user with your latest findings and open questions. You can proceed if the user confirms.
 - Discuss any airline-related topics with the user.
+- When a booking is successfully created, provide the booking ID and confirmation details clearly.
+- If you book flights, provide the booking ID and summarize the flights booked and total price.
 - If the user asks about anything unrelated to the airline, politely inform them that you can only assist with airline-related inquiries.
+
+## Context:
+- Today's date: {DEMO_DATETIME.date().isoformat()}
+- Current time: {DEMO_DATETIME.strftime("%H:%M:%S %Z")}
 """.strip()
-    .replace("\n", " ")
-    .replace("  ", " ")
-)
 
 FALLBACK_RESPONSE = "I'm sorry, but I don't have the information you're looking for. Please rephrase the question or contact Frontier Airlines customer support for further assistance."