feat: Add Firestore vector embedding support

retrieval_service/datastore/providers/firestore.py

@@ -18,7 +18,10 @@
 
 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
@@ -55,7 +58,7 @@ async def initialize_data(
         policies: list[models.Policy],
     ) -> None:
         async def delete_collections(collection_list: list[AsyncCollectionReference]):
-            # Checks if colelction exists and deletes all documents
+            # Checks if collection exists and deletes all documents
             delete_tasks = []
             for collection_ref in collection_list:
                 collection_exists = collection_ref.limit(1).stream()
@@ -76,6 +79,68 @@ async def delete_collections(collection_list: list[AsyncCollectionReference]):
             [airports_ref, amenities_ref, flights_ref, policies_ref]
         )
 
+        async def delete_indexes(index_list):
+            # Check if the collection exists and deletes all indexes
+            delete_tasks = []
+            for index in index_list:
+                if index:
+                    delete_vector_index = [
+                        "gcloud",
+                        "alpha",
+                        "firestore",
+                        "indexes",
+                        "composite",
+                        "delete",
+                        index,
+                        "--database=(default)",
+                        "--quiet",  # Added to suppress delete warning
+                    ]
+
+                    delete_vector_index_process = await asyncio.create_subprocess_exec(
+                        *delete_vector_index,
+                    )
+                    delete_tasks.append(
+                        asyncio.create_task(delete_vector_index_process.wait())
+                    )
+            await asyncio.gather(*delete_tasks)
+
+        # List indexes and retrieve name field (file-path)
+        list_vector_index = [
+            "gcloud",
+            "alpha",
+            "firestore",
+            "indexes",
+            "composite",
+            "list",
+            "--database=(default)",
+            "--format=value(name)",  # prints name field
+        ]
+
+        list_vector_index_process = await asyncio.create_subprocess_exec(
+            *list_vector_index,
+            stdout=asyncio.subprocess.PIPE,
+        )
+
+        # Capture output and ignore stderr
+        stdout, __ = await list_vector_index_process.communicate()
+
+        # Decode and format output
+        indexes = stdout.decode().strip().split("\n")
+
+        # Check if the indexes already exist; if so, delete indexes
+        # Extract collection and index-id from file path
+        # Assign the index-id to the corresponding collection
+        if indexes == [""]:
+            pass
+        else:
+            collections = {"amenities": "", "policies": ""}
+            for line in indexes:
+                collection, index_id = line.split("/")[-3], line.split("/")[-1]
+                collections[collection] = index_id
+            amenities_ref = collections["amenities"]
+            policies_ref = collections["policies"]
+            await delete_indexes([amenities_ref, policies_ref])
+
         # initialize collections
         create_airports_tasks = []
         for airport in airports:
@@ -105,8 +170,80 @@ async def delete_collections(collection_list: list[AsyncCollectionReference]):
                         "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,
-                        "embedding": amenity.embedding,
+                        # Vector type does not support None value
+                        "embedding": Vector(tuple(amenity.embedding or [])),
                     }
                 )
             )
@@ -122,8 +259,12 @@ async def delete_collections(collection_list: list[AsyncCollectionReference]):
                         "flight_number": flight.flight_number,
                         "departure_airport": flight.departure_airport,
                         "arrival_airport": flight.arrival_airport,
-                        "departure_time": flight.departure_time,
-                        "arrival_time": flight.arrival_time,
+                        "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,
                     }
@@ -142,12 +283,49 @@ async def delete_collections(collection_list: list[AsyncCollectionReference]):
                 .set(
                     {
                         "content": policy.content,
-                        "embedding": policy.embedding,
+                        # Vector type does not accept None value
+                        "embedding": Vector(policy.embedding or []),
                     }
                 )
             )
         await asyncio.gather(*create_policies_tasks)
 
+        # Initialize single-field vector indexes
+        create_amenities_vector_index = [
+            "gcloud",
+            "alpha",
+            "firestore",
+            "indexes",
+            "composite",
+            "create",
+            "--collection-group=amenities",
+            "--query-scope=COLLECTION",
+            '--field-config=field-path=embedding,vector-config={"dimension":768,"flat":"{}"}',
+            "--database=(default)",
+        ]
+        create_amenities_process = await asyncio.create_subprocess_exec(
+            *create_amenities_vector_index,
+        )
+        await create_amenities_process.wait()
+
+        create_policies_vector_index = [
+            "gcloud",
+            "alpha",
+            "firestore",
+            "indexes",
+            "composite",
+            "create",
+            "--collection-group=policies",
+            "--query-scope=COLLECTION",
+            '--field-config=field-path=embedding,vector-config={"dimension":768,"flat":"{}"}',
+            "--database=(default)",
+        ]
+
+        create_policies_process = await asyncio.create_subprocess_exec(
+            *create_policies_vector_index,
+        )
+        await create_policies_process.wait()
+
     async def export_data(
         self,
     ) -> tuple[
@@ -171,6 +349,7 @@ async def export_data(
         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 = []
@@ -183,7 +362,9 @@ async def export_data(
         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) -> Optional[models.Airport]:
@@ -234,12 +415,38 @@ async def get_amenity(self, id: int) -> Optional[models.Amenity]:
         )
         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)
 
     async def amenities_search(
         self, query_embedding: list[float], similarity_threshold: float, top_k: int
     ) -> list[Any]:
-        raise NotImplementedError("Semantic search not yet supported in Firestore.")
+        collection = AsyncQuery(self.__client.collection("amenities"))
+        query_vector = Vector(query_embedding)
+        # Using the same similarity metric to the embedding model's training method
+        # produce the most accurate result
+        distance_measure = DistanceMeasure.DOT_PRODUCT
+        query = collection.find_nearest(
+            vector_field="embedding",
+            query_vector=query_vector,
+            distance_measure=distance_measure,
+            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
 
     async def get_flight(self, flight_id: int) -> Optional[models.Flight]:
         query = self.__client.collection("flights").where(
@@ -326,8 +533,23 @@ async def list_tickets(
 
     async def policies_search(
         self, query_embedding: list[float], similarity_threshold: float, top_k: int
-    ) -> list[str]:
-        raise NotImplementedError("Semantic search not yet supported in Firestore.")
+    ) -> list[Any]:
+        collection = AsyncQuery(self.__client.collection("policies"))
+        query_vector = Vector(query_embedding)
+        distance_measure = DistanceMeasure.DOT_PRODUCT
+        query = collection.find_nearest(
+            vector_field="embedding",
+            query_vector=query_vector,
+            distance_measure=distance_measure,
+            limit=top_k,
+        )
+
+        docs = query.stream()
+        policies = []
+        async for doc in docs:
+            policy_dict = {"id": doc.id, "content": doc.get("content")}
+            policies.append(policy_dict)
+        return policies
 
     async def close(self):
         self.__client.close()

retrieval_service/requirements.txt

@@ -1,7 +1,7 @@
 asyncpg==0.29.0
 fastapi==0.109.2
 google-auth==2.32.0
-google-cloud-firestore==2.14.0
+google-cloud-firestore @ git+https://github.com/googleapis/python-firestore.git@2de16209409c9d9ba41d3444400e6a39ee1b2936
 google-cloud-aiplatform==1.56.0
 google-cloud-spanner==3.46.0
 langchain-core==0.2.8