Skip to content

feat: Implement and enhance remote document retrieval functionality #5487

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
franciscojavierarceo merged 6 commits into from
Jul 28, 2025
Merged

feat: Implement and enhance remote document retrieval functionality #5487

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
6 commits
Select commit Hold shift + click to select a range
7c951eb
Remotely retrieve the docs
jyejare Jun 12, 2025
7839fc5
Feature server rectified for documents retrival
jyejare Jun 17, 2025
18bfcf4
Updates to remote retrival for response extraction
jyejare Jun 26, 2025
2043b09
Remote retrive online doc v2
jyejare Jun 26, 2025
8bcb2ae
Unit tests for Remote docuemnts retrival
jyejare Jul 1, 2025
dfafa56
Fixed checks and comments
jyejare Jul 1, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
41 changes: 27 additions & 14 deletions sdk/python/feast/feature_server.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -6,7 +6,7 @@
import traceback
from contextlib import asynccontextmanager
from importlib import resources as importlib_resources
from typing import Any, Dict, List, Optional
from typing import Any, Dict, List, Optional, Union

import pandas as pd
import psutil
Expand Down Expand Up @@ -86,10 +86,18 @@ class MaterializeIncrementalRequest(BaseModel):
class GetOnlineFeaturesRequest(BaseModel):
entities: Dict[str, List[Any]]
feature_service: Optional[str] = None
features: Optional[List[str]] = None
features: List[str] = []
full_feature_names: bool = False
query_embedding: Optional[List[float]] = None


class GetOnlineDocumentsRequest(BaseModel):
feature_service: Optional[str] = None
features: List[str] = []
full_feature_names: bool = False
top_k: Optional[int] = None
query: Optional[List[float]] = None
query_string: Optional[str] = None
api_version: Optional[int] = 1


class ChatMessage(BaseModel):
Expand All @@ -110,7 +118,10 @@ class SaveDocumentRequest(BaseModel):
data: dict


def _get_features(request: GetOnlineFeaturesRequest, store: "feast.FeatureStore"):
def _get_features(
request: Union[GetOnlineFeaturesRequest, GetOnlineDocumentsRequest],
store: "feast.FeatureStore",
):
if request.feature_service:
feature_service = store.get_feature_service(
request.feature_service, allow_cache=True
Expand Down Expand Up @@ -246,24 +257,26 @@ async def get_online_features(request: GetOnlineFeaturesRequest) -> Dict[str, An
dependencies=[Depends(inject_user_details)],
)
async def retrieve_online_documents(
request: GetOnlineFeaturesRequest,
request: GetOnlineDocumentsRequest,
) -> Dict[str, Any]:
logger.warning(
"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,
full_feature_names=request.full_feature_names,
query=request.query_embedding,
query_string=request.query_string,
)
read_params = dict(features=features, query=request.query, top_k=request.top_k)
if request.api_version == 2 and request.query_string is not None:
read_params["query_string"] = request.query_string

response = await run_in_threadpool(
lambda: store.retrieve_online_documents_v2(**read_params) # type: ignore
)
if request.api_version == 2:
response = await run_in_threadpool(
lambda: store.retrieve_online_documents_v2(**read_params) # type: ignore
)
else:
response = await run_in_threadpool(
lambda: store.retrieve_online_documents(**read_params) # type: ignore
)

# Convert the Protobuf object to JSON and return it
response_dict = await run_in_threadpool(
Expand Down
5 changes: 5 additions & 0 deletions sdk/python/feast/feature_store.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2413,6 +2413,11 @@ def _retrieve_from_online_store_v2(
output_len=output_len,
)

utils._populate_result_rows_from_columnar(
online_features_response=online_features_response,
data=entity_key_dict,
)

return OnlineResponse(online_features_response)

def serve(
Expand Down
Loading
Loading