retrieval_service/datastore/providers/firestore.py

# Copyright 2023 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import asyncio
from datetime import datetime, timedelta
from typing import Any, Literal, Optional

from google.cloud.firestore import AsyncClient  # type: ignore
from google.cloud.firestore_v1.async_collection import AsyncCollectionReference
from google.cloud.firestore_v1.async_query import AsyncQuery
from google.cloud.firestore_v1.base_query import FieldFilter
from google.cloud.firestore_v1.base_vector_query import DistanceMeasure
from google.cloud.firestore_v1.vector import Vector
from pydantic import BaseModel

import models

from .. import datastore

FIRESTORE_IDENTIFIER = "firestore"


class Config(BaseModel, datastore.AbstractConfig):
    kind: Literal["firestore"]
    projectId: Optional[str]


class Client(datastore.Client[Config]):
    __client: AsyncClient

    @datastore.classproperty
    def kind(cls):
        return FIRESTORE_IDENTIFIER

    def __init__(self, client: AsyncClient):
        self.__client = client
        self.__policies_collection = AsyncQuery(self.__client.collection("policies"))
        self.__amenities_collection = AsyncQuery(self.__client.collection("amenities"))

    @classmethod
    async def create(cls, config: Config) -> "Client":
        return cls(AsyncClient(project=config.projectId))

    async def __delete_collections(
        self, collection_list: list[AsyncCollectionReference]
    ):
        # Checks if collection exists and deletes all documents
        delete_tasks = []
        for collection_ref in collection_list:
            collection_exists = collection_ref.limit(1).stream()
            if not collection_exists:
                continue

            docs = collection_ref.stream()
            async for doc in docs:
                delete_tasks.append(asyncio.create_task(doc.reference.delete()))
        await asyncio.gather(*delete_tasks)

    async def parse_index_info(self, line: str) -> tuple[str, str]:
        # Extract collection and index-id from file path
        parts = line.split("/")
        collection_name = parts[-3]
        index_id = parts[-1]
        return collection_name, index_id

    async def __get_indices(self) -> dict[str, str]:
        list_vector_index_process = await asyncio.create_subprocess_exec(
            "gcloud",
            "alpha",
            "firestore",
            "indexes",
            "composite",
            "list",
            "--database=(default)",
            "--format=value(name)",  # prints name field
            stdout=asyncio.subprocess.PIPE,
        )

        # Capture output and ignore stderr
        stdout, __ = await list_vector_index_process.communicate()

        # Decode and format output
        index_lines = stdout.decode().strip().split("\n")

        indices = {}

        # Create a dict with collections and their corresponding vector index.
        for line in index_lines:
            if line:
                collection, index_id = await self.parse_index_info(line)
                indices[collection] = index_id

        return indices

    async def __delete_vector_index(self, indices: list[str]):
        # Check if the collection exists and deletes all indexes
        for index in indices:
            if index:
                delete_vector_index = await asyncio.create_subprocess_exec(
                    "gcloud",
                    "alpha",
                    "firestore",
                    "indexes",
                    "composite",
                    "delete",
                    index,
                    "--database=(default)",
                    "--quiet",  # Added to suppress delete warning
                )
                await delete_vector_index.wait()

    async def __create_vector_index(self, collection_name: str):
        create_vector_index = await asyncio.create_subprocess_exec(
            "gcloud",
            "alpha",
            "firestore",
            "indexes",
            "composite",
            "create",
            f"--collection-group={collection_name}",
            "--query-scope=COLLECTION",
            '--field-config=field-path=embedding,vector-config={"dimension":768,"flat":"{}"}',
            "--database=(default)",
        )
        await create_vector_index.wait()

    async def initialize_data(
        self,
        airports: list[models.Airport],
        amenities: list[models.Amenity],
        flights: list[models.Flight],
        policies: list[models.Policy],
    ) -> None:
        # Check if the collections already exist; if so, delete collections
        airports_ref = self.__client.collection("airports")
        amenities_ref = self.__client.collection("amenities")
        flights_ref = self.__client.collection("flights")
        policies_ref = self.__client.collection("policies")
        await self.__delete_collections(
            [airports_ref, amenities_ref, flights_ref, policies_ref]
        )

        # Retrieve vector indexes and check if the collections already exist; if so, delete collections
        indices = await self.__get_indices()
        amenities_ref = indices.get("amenities", "")
        policies_ref = indices.get("policies", "")
        await self.__delete_vector_index([amenities_ref, policies_ref])

        # Initialize collections
        create_airports_tasks = []
        for airport in airports:
            create_airports_tasks.append(
                self.__client.collection("airports")
                .document(str(airport.id))
                .set(
                    {
                        "iata": airport.iata,
                        "name": airport.name,
                        "city": airport.city,
                        "country": airport.country,
                    }
                )
            )
        await asyncio.gather(*create_airports_tasks)
        create_amenities_tasks = []
        for amenity in amenities:
            create_amenities_tasks.append(
                self.__client.collection("amenities")
                .document(str(amenity.id))
                .set(
                    {
                        "name": amenity.name,
                        "description": amenity.description,
                        "location": amenity.location,
                        "terminal": amenity.terminal,
                        "category": amenity.category,
                        "hour": amenity.hour,
                        # Firebase does not support datetime.time type
                        "sunday_start_hour": (
                            str(amenity.sunday_start_hour)
                            if amenity.sunday_start_hour
                            else None
                        ),
                        "sunday_end_hour": (
                            str(amenity.sunday_end_hour)
                            if amenity.sunday_end_hour
                            else None
                        ),
                        "monday_start_hour": (
                            str(amenity.monday_start_hour)
                            if amenity.monday_start_hour
                            else None
                        ),
                        "monday_end_hour": (
                            str(amenity.monday_end_hour)
                            if amenity.monday_end_hour
                            else None
                        ),
                        "tuesday_start_hour": (
                            str(amenity.tuesday_start_hour)
                            if amenity.tuesday_start_hour
                            else None
                        ),
                        "tuesday_end_hour": (
                            str(amenity.tuesday_end_hour)
                            if amenity.tuesday_end_hour
                            else None
                        ),
                        "wednesday_start_hour": (
                            str(amenity.wednesday_start_hour)
                            if amenity.wednesday_start_hour
                            else None
                        ),
                        "wednesday_end_hour": (
                            str(amenity.wednesday_end_hour)
                            if amenity.wednesday_end_hour
                            else None
                        ),
                        "thursday_start_hour": (
                            str(amenity.thursday_start_hour)
                            if amenity.thursday_start_hour
                            else None
                        ),
                        "thursday_end_hour": (
                            str(amenity.thursday_end_hour)
                            if amenity.thursday_end_hour
                            else None
                        ),
                        "friday_start_hour": (
                            str(amenity.friday_start_hour)
                            if amenity.friday_start_hour
                            else None
                        ),
                        "friday_end_hour": (
                            str(amenity.friday_end_hour)
                            if amenity.friday_end_hour
                            else None
                        ),
                        "saturday_start_hour": (
                            str(amenity.saturday_start_hour)
                            if amenity.saturday_start_hour
                            else None
                        ),
                        "saturday_end_hour": (
                            str(amenity.saturday_end_hour)
                            if amenity.saturday_end_hour
                            else None
                        ),
                        "content": amenity.content,
                        # Vector type does not support None value
                        "embedding": Vector(amenity.embedding or []),
                    }
                )
            )
        await asyncio.gather(*create_amenities_tasks)
        create_flights_tasks = []
        for flight in flights:
            create_flights_tasks.append(
                self.__client.collection("flights")
                .document(str(flight.id))
                .set(
                    {
                        "airline": flight.airline,
                        "flight_number": flight.flight_number,
                        "departure_airport": flight.departure_airport,
                        "arrival_airport": flight.arrival_airport,
                        "departure_time": flight.departure_time.strftime(
                            "%Y-%m-%d %H:%M:%S"
                        ),
                        "arrival_time": flight.arrival_time.strftime(
                            "%Y-%m-%d %H:%M:%S"
                        ),
                        "departure_gate": flight.departure_gate,
                        "arrival_gate": flight.arrival_gate,
                    }
                )
            )
            if len(create_flights_tasks) % 10000 == 0:
                # avoid gRPC batch write timeout error
                await asyncio.gather(*create_flights_tasks)
                create_flights_tasks.clear()
        await asyncio.gather(*create_flights_tasks)
        create_policies_tasks = []
        for policy in policies:
            create_policies_tasks.append(
                self.__client.collection("policies")
                .document(str(policy.id))
                .set(
                    {
                        "content": policy.content,
                        # Vector type does not accept None value
                        "embedding": Vector(policy.embedding or []),
                    }
                )
            )
        await asyncio.gather(*create_policies_tasks)

        # Initialize single-field vector indexes
        await self.__create_vector_index("amenities")
        await self.__create_vector_index("policies")

    async def export_data(
        self,
    ) -> tuple[
        list[models.Airport],
        list[models.Amenity],
        list[models.Flight],
        list[models.Policy],
    ]:
        airport_docs = self.__client.collection("airports").stream()
        amenities_docs = self.__client.collection("amenities").stream()
        flights_docs = self.__client.collection("flights").stream()
        policies_docs = self.__client.collection("policies").stream()

        airports = []
        async for doc in airport_docs:
            airport_dict = doc.to_dict()
            airport_dict["id"] = doc.id
            airports.append(models.Airport.model_validate(airport_dict))

        amenities = []
        async for doc in amenities_docs:
            amenity_dict = doc.to_dict()
            amenity_dict["id"] = doc.id
            amenity_dict["embedding"] = list(amenity_dict["embedding"])
            amenities.append(models.Amenity.model_validate(amenity_dict))

        flights = []
        async for doc in flights_docs:
            flight_dict = doc.to_dict()
            flight_dict["id"] = doc.id
            flights.append(models.Flight.model_validate(flight_dict))

        policies = []
        async for doc in policies_docs:
            policy_dict = doc.to_dict()
            policy_dict["id"] = doc.id
            policy_dict["embedding"] = list(policy_dict["embedding"])
            policies.append(models.Policy.model_validate(policy_dict))

        return airports, amenities, flights, policies

    async def get_airport_by_id(
        self, id: int
    ) -> tuple[Optional[models.Airport], Optional[str]]:
        query = self.__client.collection("airports").where(
            filter=FieldFilter("id", "==", id)
        )
        airport_doc = await query.get()
        airport_dict = airport_doc.to_dict() | {"id": airport_doc.id}
        return models.Airport.model_validate(airport_dict), None

    async def get_airport_by_iata(
        self, iata: str
    ) -> tuple[Optional[models.Airport], Optional[str]]:
        query = self.__client.collection("airports").where(
            filter=FieldFilter("iata", "==", iata)
        )
        airport_doc = await query.get()
        airport_dict = airport_doc.to_dict() | {"id": airport_doc.id}
        return models.Airport.model_validate(airport_dict), None

    async def search_airports(
        self,
        country: Optional[str] = None,
        city: Optional[str] = None,
        name: Optional[str] = None,
    ) -> tuple[list[models.Airport], Optional[str]]:
        query = self.__client.collection("airports")

        if country is not None:
            query = query.where("country", "==", country)

        if city is not None:
            query = query.where("city", "==", city)

        if name is not None:
            query = query.where("name", ">=", name).where("name", "<=", name + "\uf8ff")

        query = query.limit(10)

        docs = query.stream()
        airports = []
        async for doc in docs:
            airport_dict = doc.to_dict() | {"id": doc.id}
            airports.append(models.Airport.model_validate(airport_dict))
        return airports, None

    async def get_amenity(
        self, id: int
    ) -> tuple[Optional[models.Amenity], Optional[str]]:
        query = self.__client.collection("amenities").where(
            filter=FieldFilter("id", "==", id)
        )
        amenity_doc = await query.get()
        amenity_dict = amenity_doc.to_dict() | {"id": amenity_doc.id}
        amenity_dict["embedding"] = list(amenity_dict["embedding"])
        return models.Amenity.model_validate(amenity_dict), None

    async def amenities_search(
        self, query_embedding: list[float], similarity_threshold: float, top_k: int
    ) -> tuple[list[Any], Optional[str]]:
        # Using the same similarity metric to the embedding model's training method
        # produce the most accurate result
        query = self.__amenities_collection.find_nearest(
            vector_field="embedding",
            query_vector=Vector(query_embedding),
            distance_measure=DistanceMeasure.DOT_PRODUCT,
            limit=top_k,
        )

        docs = query.stream()
        amenities = []
        async for doc in docs:
            amenity_dict = {
                "id": doc.id,
                "category": doc.get("category"),
                "description": doc.get("description"),
                "hour": doc.get("hour"),
                "location": doc.get("location"),
                "name": doc.get("name"),
                "terminal": doc.get("terminal"),
            }
            amenities.append(amenity_dict)
        return amenities, None

    async def get_flight(
        self, flight_id: int
    ) -> tuple[Optional[models.Flight], Optional[str]]:
        query = self.__client.collection("flights").where(
            filter=FieldFilter("id", "==", flight_id)
        )
        flight_doc = await query.get()
        flight_dict = flight_doc.to_dict() | {"id": flight_doc.id}
        return models.Flight.model_validate(flight_dict), None

    async def search_flights_by_number(
        self,
        airline: str,
        number: str,
    ) -> tuple[list[models.Flight], Optional[str]]:
        query = (
            self.__client.collection("flights")
            .where(filter=FieldFilter("airline", "==", airline))
            .where(filter=FieldFilter("flight_number", "==", number))
            .limit(10)
        )

        docs = query.stream()
        flights = []
        async for doc in docs:
            flight_dict = doc.to_dict() | {"id": doc.id}
            flights.append(models.Flight.model_validate(flight_dict))
        return flights, None

    async def search_flights_by_airports(
        self,
        date: str,
        departure_airport: Optional[str] = None,
        arrival_airport: Optional[str] = None,
    ) -> tuple[list[models.Flight], Optional[str]]:
        date_obj = datetime.strptime(date, "%Y-%m-%d").date()
        date_timestamp = datetime.combine(date_obj, datetime.min.time())
        query = (
            self.__client.collection("flights")
            .where("departure_time", ">=", date_timestamp)
            .where("departure_time", "<", date_timestamp + timedelta(days=1))
            .limit(10)
        )

        if departure_airport is None:
            query = query.where("departure_airport", "==", departure_airport)
        if arrival_airport is None:
            query = query.where("arrival_airport", "==", arrival_airport)

        docs = query.stream()
        flights = []
        async for doc in docs:
            flight_dict = doc.to_dict() | {"id": doc.id}
            flights.append(models.Flight.model_validate(flight_dict))
        return flights, None

    async def validate_ticket(
        self,
        airline: str,
        flight_number: str,
        departure_airport: str,
        departure_time: str,
    ) -> tuple[Optional[models.Flight], Optional[str]]:
        raise NotImplementedError("Not Implemented")

    async def insert_ticket(
        self,
        user_id: str,
        user_name: str,
        user_email: str,
        airline: str,
        flight_number: str,
        departure_airport: str,
        arrival_airport: str,
        departure_time: str,
        arrival_time: str,
    ):
        raise NotImplementedError("Not Implemented")

    async def list_tickets(
        self,
        user_id: str,
    ) -> tuple[list[Any], Optional[str]]:
        raise NotImplementedError("Not Implemented")

    async def policies_search(
        self, query_embedding: list[float], similarity_threshold: float, top_k: int
    ) -> tuple[list[str], Optional[str]]:
        query = self.__policies_collection.find_nearest(
            vector_field="embedding",
            query_vector=Vector(query_embedding),
            distance_measure=DistanceMeasure.DOT_PRODUCT,
            limit=top_k,
        )

        policies = []
        async for doc in query.stream():
            policies.append(doc.get("content"))
        return policies, None

    async def close(self):
        self.__client.close()