feat: Adding retrieve_online_documents endpoint (feast-dev#5002)

* feat: Adding retrieve_online_documents endpoint Signed-off-by: Francisco Javier Arceo <[email protected]> * updated feature server Signed-off-by: Francisco Javier Arceo <[email protected]> --------- Signed-off-by: Francisco Javier Arceo <[email protected]>
tchughesiv · Feb 3, 2025 · 6607d3d · 6607d3d
1 parent a8842b0
commit 6607d3d
Show file tree

Hide file tree

Showing 2 changed files with 42 additions and 8 deletions.
diff --git a/docs/reference/feature-servers/python-feature-server.md b/docs/reference/feature-servers/python-feature-server.md
@@ -226,13 +226,14 @@ feast serve --key /path/to/key.pem --cert /path/to/cert.pem
 
 ## API Endpoints and Permissions
 
-| Endpoint                     | Resource Type                   | Permission                                            | Description                                                              |
-| ---------------------------- |---------------------------------|-------------------------------------------------------| ------------------------------------------------------------------------ |
-| /get-online-features         | FeatureView,OnDemandFeatureView | Read Online                                           | Get online features from the feature store                     |
-| /push                        | FeatureView                     | Write Online, Write Offline, Write Online and Offline | Push features to the feature store (online, offline, or both)  |
-| /write-to-online-store       | FeatureView                     | Write Online                                          | Write features to the online store                             |
-| /materialize                 | FeatureView                     | Write Online                                          | Materialize features within a specified time range             |
-| /materialize-incremental     | FeatureView                     | Write Online                                          | Incrementally materialize features up to a specified timestamp |
+| Endpoint                   | Resource Type                   | Permission                                            | Description                                                    |
+|----------------------------|---------------------------------|-------------------------------------------------------|----------------------------------------------------------------|
+| /get-online-features       | FeatureView,OnDemandFeatureView | Read Online                                           | Get online features from the feature store                     |
+| /retrieve-online-documents | FeatureView                     | Read Online                                           | Retrieve online documents from the feature store for RAG       |
+| /push                      | FeatureView                     | Write Online, Write Offline, Write Online and Offline | Push features to the feature store (online, offline, or both)  |
+| /write-to-online-store     | FeatureView                     | Write Online                                          | Write features to the online store                             |
+| /materialize               | FeatureView                     | Write Online                                          | Materialize features within a specified time range             |
+| /materialize-incremental   | FeatureView                     | Write Online                                          | Incrementally materialize features up to a specified timestamp |
 
 ## How to configure Authentication and Authorization ?
 

diff --git a/sdk/python/feast/feature_server.py b/sdk/python/feast/feature_server.py
@@ -74,6 +74,7 @@ class GetOnlineFeaturesRequest(BaseModel):
     feature_service: Optional[str] = None
     features: Optional[List[str]] = None
     full_feature_names: bool = False
+    query_embedding: Optional[List[float]] = None
 
 
 def _get_features(request: GetOnlineFeaturesRequest, store: "feast.FeatureStore"):
@@ -104,7 +105,6 @@ def _get_features(request: GetOnlineFeaturesRequest, store: "feast.FeatureStore"
                 resource=od_feature_view, actions=[AuthzedAction.READ_ONLINE]
             )
         features = request.features  # type: ignore
-
     return features
 
 
@@ -177,6 +177,39 @@ async def get_online_features(request: GetOnlineFeaturesRequest) -> Dict[str, An
         )
         return response_dict
 
+    @app.post(
+        "/retrieve-online-documents",
+        dependencies=[Depends(inject_user_details)],
+    )
+    async def retrieve_online_documents(
+        request: GetOnlineFeaturesRequest,
+    ) -> Dict[str, Any]:
+        logger.warn(
+            "This endpoint is in alpha and will be moved to /get-online-features when stable."
+        )
+        # Initialize parameters for FeatureStore.retrieve_online_documents_v2(...) call
+        features = await run_in_threadpool(_get_features, request, store)
+
+        read_params = dict(
+            features=features,
+            entity_rows=request.entities,
+            full_feature_names=request.full_feature_names,
+            query=request.query_embedding,
+        )
+
+        response = await run_in_threadpool(
+            lambda: store.retrieve_online_documents_v2(**read_params)  # type: ignore
+        )
+
+        # Convert the Protobuf object to JSON and return it
+        response_dict = await run_in_threadpool(
+            MessageToDict,
+            response.proto,
+            preserving_proto_field_name=True,
+            float_precision=18,
+        )
+        return response_dict
+
     @app.post("/push", dependencies=[Depends(inject_user_details)])
     async def push(request: PushFeaturesRequest) -> None:
         df = pd.DataFrame(request.df)